Coverage for gco_mcp/mission/_engine_factory.py: 90.53%
137 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""Shared :class:`MissionEngine` factory used by the MCP tool and the CLI.
3Both the ``gco_mcp/tools/mission.py`` MCP tool surface and the
4``cli/commands/mission_cmd.py`` Click subcommands need to build a
5:class:`mcp.mission.engine.MissionEngine` with production-wired
6dependencies — a real tool dispatcher that routes through the live
7FastMCP registry, a sampling callable that runs the
8``Strategy_Revision`` prompt against the resolved backend, and an
9optional sandbox runner for scripted strategies. The wiring used to
10live only inside the MCP tool module, which left the CLI with a stub
11dispatcher and ``sampling_callable=None``. That made
12``gco mission run`` and ``gco mission iterate`` useful for smoke-
13testing the engine bookkeeping but unable to converge on goals that
14depend on actual tool-result content.
16This module hosts the shared factory. The MCP tool surface and the
17CLI both call :func:`build_engine_dependencies` to obtain the same
18``(tool_dispatcher, sampling_callable, sandbox_runner)`` triple, then
19hand them to :class:`mcp.mission.engine.MissionEngine`. The CLI also
20uses :func:`make_stub_dispatcher` to opt into the canned-response
21behaviour explicitly through ``--dry-run`` for the smoke-test use
22case the original stub was designed for.
24Why a separate module rather than living inside the MCP tool? The
25MCP tool's body is gated by ``GCO_ENABLE_MISSION``; the CLI must be
26able to import the factory regardless of the flag, because the
27flag-gating happens in the Click group, not at import time. Splitting
28the factory out keeps the MCP tool body lean and lets the CLI reach
29the same wiring without crossing a feature-flag boundary.
30"""
32from __future__ import annotations
34import json
35import sys
36from collections.abc import Awaitable, Callable, Mapping
37from datetime import UTC, datetime
38from pathlib import Path
39from typing import TYPE_CHECKING, Any, cast
41# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
42# Generated at (UTC): 2026-07-18T01:03:40Z
43# Flowchart(s) generated from this file:
44# * ``build_engine_dependencies`` -> ``diagrams/code_diagrams/gco_mcp/mission/_engine_factory.build_engine_dependencies.html``
45# (PNG: ``diagrams/code_diagrams/gco_mcp/mission/_engine_factory.build_engine_dependencies.png``)
46# Regenerate with ``python diagrams/code_diagrams/generate.py``.
47# <pyflowchart-code-diagram> END
50# The mission package and the FastMCP server module both live under
51# ``gco_mcp/``; the path-injection pattern matches the rest of the MCP
52# surface so ``import server`` and ``import mission.*`` resolve
53# without making the ``mcp`` directory a package.
54sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
56from mission import sampling as mission_sampling # noqa: E402
57from mission import state as mission_state # noqa: E402
58from mission.engine import MissionEngine, SandboxRunner, ToolDispatcher # noqa: E402
60if TYPE_CHECKING: # pragma: no cover - import only for type checkers
61 from mission.types import SessionState, ToolCallRecord
64__all__ = [
65 "EngineDependencies",
66 "build_engine_dependencies",
67 "build_mission_engine",
68 "fetch_registered_tool_metadata",
69 "make_stub_dispatcher",
70 "remaining_wall_clock_seconds",
71]
74# ---------------------------------------------------------------------------
75# Public types
76# ---------------------------------------------------------------------------
79class EngineDependencies:
80 """Dependency triple consumed by :class:`MissionEngine`.
82 Holds the callables :class:`MissionEngine` needs at construction
83 time so callers can build them once and pass them through. A simple
84 namespace class rather than a NamedTuple so fields can be optional
85 without the boilerplate.
86 """
88 __slots__ = ("final_lessons_callable", "sampling_callable", "sandbox_runner", "tool_dispatcher")
90 def __init__(
91 self,
92 *,
93 tool_dispatcher: ToolDispatcher,
94 sampling_callable: Callable[..., Awaitable[Any]] | None,
95 sandbox_runner: SandboxRunner | None,
96 final_lessons_callable: Callable[..., Awaitable[Any]] | None = None,
97 ) -> None:
98 self.tool_dispatcher = tool_dispatcher
99 self.sampling_callable = sampling_callable
100 self.sandbox_runner = sandbox_runner
101 self.final_lessons_callable = final_lessons_callable
104# ---------------------------------------------------------------------------
105# Helpers shared between the MCP tool and the CLI
106# ---------------------------------------------------------------------------
109def remaining_wall_clock_seconds(session: Mapping[str, Any]) -> float | None:
110 """Return remaining wall-clock seconds for ``session``, or ``None``.
112 Mirrors the behaviour the MCP tool surface used to keep private:
113 sessions that haven't started yet report the full cap; sessions
114 whose ``max_wall_clock_seconds`` is the ``-1`` "uncapped" sentinel
115 return ``None`` so the sampling prompt's budget context renders
116 ``"remaining_wall_clock_seconds": null``; running sessions return
117 the difference between cap and elapsed time clamped at zero.
118 """
119 cap = session.get("budget", {}).get("max_wall_clock_seconds")
120 if cap is None or cap == -1:
121 return None
122 started_raw = session.get("started_at")
123 if not started_raw:
124 return float(cap)
125 try:
126 started = datetime.fromisoformat(started_raw)
127 except TypeError, ValueError:
128 return float(cap)
129 elapsed = (datetime.now(UTC) - started).total_seconds()
130 return max(0.0, float(cap) - elapsed)
133async def fetch_registered_tool_metadata() -> tuple[dict[str, Any], dict[str, str]]:
134 """Return (registered_tools, tool_docstrings) from the live FastMCP registry.
136 Reads through ``server.mcp._list_tools()``, the same low-level path
137 the MCP tool surface used to walk inline. Returns two parallel
138 dicts so the sampler closure can be built without the caller
139 needing to call ``mcp.*`` introspection twice.
141 A blanket ``Exception`` swallow yields ``({}, {})`` so the engine
142 factory still produces a usable dispatcher when the registry is
143 not yet populated (CLI path before ``register_all_tools`` ran,
144 test harness with a stub mcp instance, etc.).
145 """
146 try:
147 from server import mcp # noqa: PLC0415 - lazy
148 except Exception:
149 return {}, {}
150 try:
151 tools = await mcp._list_tools()
152 except Exception:
153 return {}, {}
154 registered = {t.name: t for t in tools}
155 docstrings = {name: (getattr(t, "description", "") or "") for name, t in registered.items()}
156 return registered, docstrings
159# ---------------------------------------------------------------------------
160# Tool dispatcher
161# ---------------------------------------------------------------------------
164async def _live_dispatch_tool(
165 tool_name: str,
166 args: dict[str, Any],
167 ctx_inner: Any | None,
168) -> Any:
169 """Dispatch ``tool_name`` against the live FastMCP registry.
171 Looks the tool up via ``server.mcp.get_tool`` and invokes it with
172 ``args``. The raw FastMCP ``ToolResult`` Pydantic model is not
173 JSON-serialisable, so the helper unwraps it:
175 * Prefer ``structured_content`` when present — every FastMCP tool
176 with a typed return surfaces a JSON-able dict here.
177 * Fall back to the first content block's ``text`` field;
178 best-effort JSON-parse so structured string-returning tools
179 round-trip as dicts.
180 * Anything else returns ``None`` so the engine records a benign
181 placeholder rather than a non-serialisable object.
183 ``RuntimeError`` propagates for unknown tool names so the engine's
184 per-call try/except records a ``failed`` outcome rather than
185 silently invoking nothing.
186 """
187 # Reuse the active request context when one exists so tools that
188 # introspect ``get_context()`` see the right one. Fall back to
189 # ``ctx_inner`` when no request is active (CLI path / unit-test
190 # path).
191 context: Any | None
192 try:
193 from fastmcp.server.dependencies import get_context # noqa: PLC0415
195 try:
196 context = get_context()
197 except Exception:
198 context = ctx_inner
199 except Exception:
200 context = ctx_inner
201 del context # FastMCP uses contextvars internally
203 from server import mcp # noqa: PLC0415
205 tool_obj = await mcp.get_tool(tool_name)
206 if tool_obj is None:
207 raise RuntimeError(f"tool {tool_name!r} not registered")
208 result = await tool_obj.run(args)
210 structured = getattr(result, "structured_content", None)
211 if isinstance(structured, dict):
212 return structured
213 content_blocks = getattr(result, "content", None) or []
214 if content_blocks:
215 first = content_blocks[0]
216 text_payload = getattr(first, "text", None)
217 if isinstance(text_payload, str): 217 ↛ 222line 217 didn't jump to line 222 because the condition on line 217 was always true
218 try:
219 return json.loads(text_payload)
220 except TypeError, ValueError:
221 return text_payload
222 return None
225def make_stub_dispatcher() -> ToolDispatcher:
226 """Return a tool dispatcher that returns canned-ok responses.
228 Reserved for ``--dry-run`` smoke testing. Returns
229 ``{"_status": "ok", "_stub": True, ...}`` for every call so the
230 engine bookkeeping converges without invoking any real tool —
231 useful for exercising the loop without spending Bedrock or AWS
232 credits, and for unit tests that don't want a live registry.
234 Production code paths use :func:`build_engine_dependencies` so
235 this stub never fires unless the operator opts in explicitly.
236 """
238 async def _dispatch(tool_name: str, args: dict[str, Any], ctx: Any) -> dict[str, Any]:
239 return {
240 "_status": "ok",
241 "_stub": True,
242 "tool_name": tool_name,
243 "args": dict(args),
244 }
246 return _dispatch
249# ---------------------------------------------------------------------------
250# Sandbox runner
251# ---------------------------------------------------------------------------
254def _build_sandbox_runner(session: Mapping[str, Any]) -> SandboxRunner | None:
255 """Wire a real sandbox runner when the session permits scripted strategies."""
256 if not session.get("allow_scripted_strategies"):
257 return None
258 try:
259 from mission.sandbox import MissionSandbox # noqa: PLC0415
260 except ImportError:
261 return None
263 sandbox = MissionSandbox(
264 list(session.get("tool_allowlist") or []),
265 cast("SessionState", session),
266 )
268 async def _sandbox_runner(
269 script: str,
270 ctx_arg: Any,
271 dispatcher: ToolDispatcher,
272 ) -> tuple[dict[str, Any], list[ToolCallRecord]]:
273 obs, calls = await sandbox.run(script, ctx_arg, dispatcher)
274 return obs, cast("list[ToolCallRecord]", calls)
276 return _sandbox_runner
279# ---------------------------------------------------------------------------
280# Sampling callable
281# ---------------------------------------------------------------------------
284def _build_sampling_callable(
285 session: Mapping[str, Any],
286 ctx: Any | None,
287 *,
288 registered_tools: dict[str, Any],
289 tool_docstrings: dict[str, str],
290) -> Callable[..., Awaitable[Any]] | None:
291 """Wire the Strategy_Revision sampler when the session opted in."""
292 if not session.get("use_sampling"):
293 return None
294 if session.get("sampling_backend_resolved") == "none":
295 return None
297 backend_obj = mission_sampling.select_sampling_backend(
298 ctx,
299 model_id=session.get("bedrock_model_id"),
300 prefs=session.get("sampling_model_preferences"),
301 )
302 if backend_obj is None:
303 return None
305 # Slow-moving live signals (per-region queue depth, GPU utilisation,
306 # deployed-region list, reservation counts). Cached on the closure
307 # so one ``build_engine_dependencies`` call — which spans the
308 # multi-iteration drive of one CLI invocation or one MCP request —
309 # only pays the AWS round-trip once. Outer-list trick keeps the
310 # cache mutable through the inner closure without ``nonlocal``.
311 from mission._environment import gather_session_environment # noqa: PLC0415
313 env_cache: list[Mapping[str, Any] | None] = []
315 async def _sampler(*, session: dict[str, Any], ctx: Any | None) -> Any:
316 iterations = session.get("iterations") or []
317 latest = iterations[-1] if iterations else None
318 if latest is None:
319 return None
320 budget = session.get("budget") or {}
321 # ``max_iterations=-1`` is the "uncapped" sentinel; the prompt
322 # expects an informational remaining-iterations count, so we
323 # report zero in that mode (the model is told nothing about
324 # the iteration axis when there's no cap to count down from).
325 # Finite caps subtract the count of recorded iterations and
326 # clamp at zero.
327 cap = int(budget.get("max_iterations", 0))
328 remaining_iters = 0 if cap == -1 else max(0, cap - len(iterations))
329 if not env_cache:
330 try:
331 env_cache.append(gather_session_environment(session))
332 except Exception: # noqa: BLE001
333 env_cache.append(None)
334 env_ctx = env_cache[0]
335 return await mission_sampling.maybe_sample_strategy_revision(
336 backend=backend_obj,
337 session=cast("SessionState", session),
338 iteration=latest,
339 allowlist=list(session.get("tool_allowlist") or []),
340 registered_tools=registered_tools,
341 tool_docstrings=tool_docstrings,
342 remaining_iterations=remaining_iters,
343 remaining_wall_clock_secs=remaining_wall_clock_seconds(session),
344 allow_scripts=bool(session.get("allow_scripted_strategies", False)),
345 environment_context=env_ctx,
346 )
348 return _sampler
351# ---------------------------------------------------------------------------
352# Final lessons callable
353# ---------------------------------------------------------------------------
356def _build_final_lessons_callable(
357 session: Mapping[str, Any],
358 ctx: Any | None,
359 tool_docstrings: dict[str, str],
360) -> Callable[..., Awaitable[Any]] | None:
361 """Wire the Final_Report lessons overlay when sampling is enabled.
363 When the session opted into sampling and a backend resolves, the
364 engine calls this after a terminal verdict to produce model-derived
365 ``lessons`` and ``recommended_followups`` for the Final_Report.
366 Without it, the report uses deterministic templates.
367 """
368 if not session.get("use_sampling"):
369 return None
370 if session.get("sampling_backend_resolved") == "none": 370 ↛ 371line 370 didn't jump to line 371 because the condition on line 370 was never true
371 return None
373 backend_obj = mission_sampling.select_sampling_backend(
374 ctx,
375 model_id=session.get("bedrock_model_id"),
376 prefs=session.get("sampling_model_preferences"),
377 )
378 if backend_obj is None: 378 ↛ 379line 378 didn't jump to line 379 because the condition on line 378 was never true
379 return None
381 async def _final_lessons(*, session: dict[str, Any], ctx: Any | None) -> Any:
382 return await mission_sampling.maybe_sample_final_lessons(
383 backend=backend_obj,
384 session=cast("SessionState", session),
385 tool_docstrings=tool_docstrings,
386 )
388 return _final_lessons
391# ---------------------------------------------------------------------------
392# Public factory
393# ---------------------------------------------------------------------------
396async def build_engine_dependencies(
397 session: Mapping[str, Any],
398 ctx: Any | None,
399 *,
400 use_stub_dispatcher: bool = False,
401) -> EngineDependencies:
402 """Build the :class:`MissionEngine` dependency triple for ``session``.
404 Looks up the live FastMCP registry to populate the registered-tools
405 map and per-tool docstring cache, then assembles the production
406 wiring: a live tool dispatcher (or the canned-stub when
407 ``use_stub_dispatcher`` is True), the Strategy_Revision sampler
408 when sampling resolved to a real backend, and a sandbox runner
409 when the session permits scripted strategies.
411 The CLI passes ``use_stub_dispatcher=True`` only on the
412 ``--dry-run`` path; the MCP tool surface always uses the live
413 dispatcher.
414 """
415 if use_stub_dispatcher:
416 # The stub dispatcher means real tools never run, which means
417 # there's nothing to inform the sampling prompt; downgrade to
418 # the deterministic propose path for symmetry. The session's
419 # ``use_sampling`` flag stays as the operator set it so the
420 # criteria-scaffold path still runs through Bedrock.
421 return EngineDependencies(
422 tool_dispatcher=make_stub_dispatcher(),
423 sampling_callable=None,
424 sandbox_runner=_build_sandbox_runner(session),
425 )
427 registered_tools, tool_docstrings = await fetch_registered_tool_metadata()
428 sampling_callable = _build_sampling_callable(
429 session,
430 ctx,
431 registered_tools=registered_tools,
432 tool_docstrings=tool_docstrings,
433 )
434 sandbox_runner = _build_sandbox_runner(session)
435 final_lessons = _build_final_lessons_callable(session, ctx, tool_docstrings)
436 return EngineDependencies(
437 tool_dispatcher=_live_dispatch_tool,
438 sampling_callable=sampling_callable,
439 sandbox_runner=sandbox_runner,
440 final_lessons_callable=final_lessons,
441 )
444async def build_mission_engine(
445 session: Mapping[str, Any],
446 ctx: Any | None,
447 *,
448 use_stub_dispatcher: bool = False,
449) -> MissionEngine:
450 """Build a :class:`MissionEngine` instance ready to drive ``session``.
452 Convenience wrapper over :func:`build_engine_dependencies` that
453 also resolves the persistence backend through
454 :func:`mcp.mission.state.get_backend`. Most callers should use
455 this; the lower-level :func:`build_engine_dependencies` is
456 available for code that needs to instantiate the engine itself
457 (e.g. with a custom backend).
458 """
459 deps = await build_engine_dependencies(session, ctx, use_stub_dispatcher=use_stub_dispatcher)
460 backend = mission_state.get_backend()
461 return MissionEngine(
462 backend=backend,
463 tool_dispatcher=deps.tool_dispatcher,
464 sampling_callable=deps.sampling_callable,
465 sandbox_runner=deps.sandbox_runner,
466 final_lessons_callable=deps.final_lessons_callable,
467 )