Coverage for gco_mcp / mission / _engine_factory.py: 100.00%
169 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +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, Sequence
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-09-03T18:56:22Z
43# Generated from Git commit: 37fd4384775eeebf18fea3e5e085cef9645077be
44# Flowchart(s) generated from this file:
45# * ``build_engine_dependencies`` -> ``diagrams/code_diagrams/gco_mcp/mission/_engine_factory.build_engine_dependencies.html``
46# (PNG: ``diagrams/code_diagrams/gco_mcp/mission/_engine_factory.build_engine_dependencies.png``)
47# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
48# <pyflowchart-code-diagram> END
51# The mission package and the FastMCP server module both live under
52# ``gco_mcp/``; the path-injection pattern matches the rest of the MCP
53# surface so ``import server`` and ``import mission.*`` resolve
54# without making the ``mcp`` directory a package.
55sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
57from mission import sampling as mission_sampling # noqa: E402
58from mission import state as mission_state # noqa: E402
59from mission.engine import ( # noqa: E402
60 MissionEngine,
61 ObservationAugmenter,
62 SandboxRunner,
63 ToolDispatcher,
64)
66if TYPE_CHECKING: # pragma: no cover - import only for type checkers
67 from mission.types import SessionState, ToolCallRecord
70__all__ = [
71 "EngineDependencies",
72 "MissionToolResultError",
73 "build_engine_dependencies",
74 "build_mission_engine",
75 "fetch_registered_tool_metadata",
76 "make_stub_dispatcher",
77 "remaining_wall_clock_seconds",
78]
81# ---------------------------------------------------------------------------
82# Public types
83# ---------------------------------------------------------------------------
86class EngineDependencies:
87 """Dependency triple consumed by :class:`MissionEngine`.
89 Holds the callables :class:`MissionEngine` needs at construction
90 time so callers can build them once and pass them through. A simple
91 namespace class rather than a NamedTuple so fields can be optional
92 without the boilerplate.
93 """
95 __slots__ = (
96 "final_lessons_callable",
97 "memory_store",
98 "observation_augmenters",
99 "sampling_callable",
100 "sandbox_runner",
101 "tool_dispatcher",
102 )
104 def __init__(
105 self,
106 *,
107 tool_dispatcher: ToolDispatcher,
108 sampling_callable: Callable[..., Awaitable[Any]] | None,
109 sandbox_runner: SandboxRunner | None,
110 final_lessons_callable: Callable[..., Awaitable[Any]] | None = None,
111 memory_store: Any | None = None,
112 observation_augmenters: Sequence[ObservationAugmenter] | None = None,
113 ) -> None:
114 self.tool_dispatcher = tool_dispatcher
115 self.sampling_callable = sampling_callable
116 self.sandbox_runner = sandbox_runner
117 self.final_lessons_callable = final_lessons_callable
118 self.memory_store = memory_store
119 self.observation_augmenters = observation_augmenters
122# ---------------------------------------------------------------------------
123# Helpers shared between the MCP tool and the CLI
124# ---------------------------------------------------------------------------
127def remaining_wall_clock_seconds(session: Mapping[str, Any]) -> float | None:
128 """Return remaining wall-clock seconds for ``session``, or ``None``.
130 Mirrors the behaviour the MCP tool surface used to keep private:
131 sessions that haven't started yet report the full cap; sessions
132 whose ``max_wall_clock_seconds`` is the ``-1`` "uncapped" sentinel
133 return ``None`` so the sampling prompt's budget context renders
134 ``"remaining_wall_clock_seconds": null``; running sessions return
135 the difference between cap and elapsed time clamped at zero.
136 """
137 cap = session.get("budget", {}).get("max_wall_clock_seconds")
138 if cap is None or cap == -1:
139 return None
140 started_raw = session.get("started_at")
141 if not started_raw:
142 return float(cap)
143 try:
144 started = datetime.fromisoformat(started_raw)
145 except TypeError, ValueError:
146 return float(cap)
147 elapsed = (datetime.now(UTC) - started).total_seconds()
148 return max(0.0, float(cap) - elapsed)
151async def fetch_registered_tool_metadata() -> tuple[dict[str, Any], dict[str, str]]:
152 """Return (registered_tools, tool_docstrings) from the live FastMCP registry.
154 Reads through ``server.mcp._list_tools()``, the same low-level path
155 the MCP tool surface used to walk inline. Returns two parallel
156 dicts so the sampler closure can be built without the caller
157 needing to call ``mcp.*`` introspection twice.
159 A blanket ``Exception`` swallow yields ``({}, {})`` so the engine
160 factory still produces a usable dispatcher when the registry is
161 not yet populated (CLI path before ``register_all_tools`` ran,
162 test harness with a stub mcp instance, etc.).
163 """
164 try:
165 from server import mcp # noqa: PLC0415 - lazy
166 except Exception:
167 return {}, {}
168 try:
169 tools = await mcp._list_tools()
170 except Exception:
171 return {}, {}
172 registered = {t.name: t for t in tools}
173 docstrings = {name: (getattr(t, "description", "") or "") for name, t in registered.items()}
174 return registered, docstrings
177# ---------------------------------------------------------------------------
178# Tool dispatcher
179# ---------------------------------------------------------------------------
182class MissionToolResultError(RuntimeError):
183 """A live FastMCP tool returned a typed transport-level error result."""
185 def __init__(self, tool_name: str, details: Any) -> None:
186 self.tool_name = tool_name
187 self.details = details
188 if details is None:
189 summary = "no error details"
190 elif isinstance(details, str):
191 summary = details
192 else:
193 summary = json.dumps(details, sort_keys=True, default=str)
194 super().__init__(f"tool {tool_name!r} returned an error result: {summary}")
197def _tool_result_payload(result: Any) -> Any:
198 """Unwrap a FastMCP result into a JSON-serialisable observation payload."""
199 structured = getattr(result, "structured_content", None)
200 if isinstance(structured, dict):
201 return structured
202 content_blocks = getattr(result, "content", None) or []
203 if content_blocks:
204 first = content_blocks[0]
205 text_payload = getattr(first, "text", None)
206 if isinstance(text_payload, str):
207 try:
208 return json.loads(text_payload)
209 except TypeError, ValueError:
210 return text_payload
211 return None
214async def _live_dispatch_tool(
215 tool_name: str,
216 args: dict[str, Any],
217 ctx_inner: Any | None,
218) -> Any:
219 """Dispatch ``tool_name`` against the live FastMCP registry.
221 Looks the tool up via ``server.mcp.get_tool`` and invokes it with
222 ``args``. The raw FastMCP ``ToolResult`` Pydantic model is not
223 JSON-serialisable, so the helper unwraps it:
225 * Prefer ``structured_content`` when present — every FastMCP tool
226 with a typed return surfaces a JSON-able dict here.
227 * Fall back to the first content block's ``text`` field;
228 best-effort JSON-parse so structured string-returning tools
229 round-trip as dicts.
230 * A result with FastMCP's typed ``is_error`` flag raises
231 :class:`MissionToolResultError`, so the engine records the call as
232 ``failed`` instead of treating an error body as a successful observation.
233 * Anything else returns ``None`` so the engine records a benign
234 placeholder rather than a non-serialisable object.
236 ``RuntimeError`` propagates for unknown tool names so the engine's
237 per-call try/except records a ``failed`` outcome rather than
238 silently invoking nothing.
239 """
240 # Reuse the active request context when one exists so tools that
241 # introspect ``get_context()`` see the right one. Fall back to
242 # ``ctx_inner`` when no request is active (CLI path / unit-test
243 # path).
244 context: Any | None
245 try:
246 from fastmcp.server.dependencies import get_context # noqa: PLC0415
248 try:
249 context = get_context()
250 except Exception:
251 context = ctx_inner
252 except Exception:
253 context = ctx_inner
254 del context # FastMCP uses contextvars internally
256 from server import mcp # noqa: PLC0415
258 tool_obj = await mcp.get_tool(tool_name)
259 if tool_obj is None:
260 raise RuntimeError(f"tool {tool_name!r} not registered")
261 result = await tool_obj.run(args)
262 payload = _tool_result_payload(result)
263 if getattr(result, "is_error", False) is True:
264 raise MissionToolResultError(tool_name, payload)
265 return payload
268def make_stub_dispatcher() -> ToolDispatcher:
269 """Return a tool dispatcher that returns canned-ok responses.
271 Reserved for ``--dry-run`` smoke testing. Returns
272 ``{"_status": "ok", "_stub": True, ...}`` for every call so the
273 engine bookkeeping converges without invoking any real tool —
274 useful for exercising the loop without spending Bedrock or AWS
275 credits, and for unit tests that don't want a live registry.
277 Production code paths use :func:`build_engine_dependencies` so
278 this stub never fires unless the operator opts in explicitly.
279 """
281 async def _dispatch(tool_name: str, args: dict[str, Any], ctx: Any) -> dict[str, Any]:
282 return {
283 "_status": "ok",
284 "_stub": True,
285 "tool_name": tool_name,
286 "args": dict(args),
287 }
289 return _dispatch
292# ---------------------------------------------------------------------------
293# Sandbox runner
294# ---------------------------------------------------------------------------
297def _build_sandbox_runner(session: Mapping[str, Any]) -> SandboxRunner | None:
298 """Wire a real sandbox runner when the session permits scripted strategies."""
299 if not session.get("allow_scripted_strategies"):
300 return None
301 try:
302 from mission.sandbox import MissionSandbox # noqa: PLC0415
303 except ImportError:
304 return None
306 sandbox = MissionSandbox(
307 list(session.get("tool_allowlist") or []),
308 cast("SessionState", session),
309 )
311 async def _sandbox_runner(
312 script: str,
313 ctx_arg: Any,
314 dispatcher: ToolDispatcher,
315 ) -> tuple[dict[str, Any], list[ToolCallRecord]]:
316 obs, calls = await sandbox.run(script, ctx_arg, dispatcher)
317 return obs, cast("list[ToolCallRecord]", calls)
319 return _sandbox_runner
322# ---------------------------------------------------------------------------
323# Mission-memory store
324# ---------------------------------------------------------------------------
327def _build_memory_store() -> Any | None:
328 """Construct the mission-memory store for production wiring.
330 One seam for both memory paths — the engine's best-effort terminal
331 write and the sampler closure's prior-missions retrieval — so the
332 test-suite conftest can neutralise real AWS reach by patching this
333 single function. Construction itself is free (table/index names
334 resolve lazily from SSM on first use), but it is still guarded: any
335 unexpected failure degrades to "no memory", never to a failed
336 engine build.
337 """
338 try:
339 from mission.memory import MissionMemoryStore # noqa: PLC0415
341 return MissionMemoryStore()
342 except Exception: # noqa: BLE001
343 return None
346# ---------------------------------------------------------------------------
347# Sampling callable
348# ---------------------------------------------------------------------------
351def _build_sampling_callable(
352 session: Mapping[str, Any],
353 ctx: Any | None,
354 *,
355 registered_tools: dict[str, Any],
356 tool_docstrings: dict[str, str],
357) -> Callable[..., Awaitable[Any]] | None:
358 """Wire the Strategy_Revision sampler when the session opted in."""
359 if not session.get("use_sampling"):
360 return None
361 if session.get("sampling_backend_resolved") == "none":
362 return None
364 backend_obj = mission_sampling.select_sampling_backend(
365 model_id=session.get("bedrock_model_id"),
366 )
368 # Slow-moving live signals (per-region queue depth, GPU utilisation,
369 # deployed-region list, reservation counts). Cached on the closure
370 # so one ``build_engine_dependencies`` call — which spans the
371 # multi-iteration drive of one CLI invocation or one MCP request —
372 # only pays the AWS round-trip once. Outer-list trick keeps the
373 # cache mutable through the inner closure without ``nonlocal``.
374 from mission._environment import gather_session_environment # noqa: PLC0415
376 env_cache: list[Mapping[str, Any] | None] = []
377 # Prior similar missions from the memory vector index. Same
378 # one-slot cache trick as ``env_cache``: retrieval costs one
379 # embedding call plus one SearchVectors round-trip, and the
380 # directive never changes mid-session, so pay it once per engine
381 # wiring. Retrieval is inherently gated on ``use_sampling`` —
382 # this closure only exists for sampling sessions — which keeps the
383 # deterministic Propose path free of network calls, exactly what
384 # the determinism suite pins down. Best-effort: any failure
385 # (absent table, backfilling index, Bedrock down, no credentials)
386 # degrades to "no prior context". An empty result list also maps
387 # to ``None`` so the prompt section only renders when there is
388 # something to say.
389 memory_cache: list[list[Mapping[str, Any]] | None] = []
391 async def _sampler(*, session: dict[str, Any], ctx: Any | None) -> Any:
392 iterations = session.get("iterations") or []
393 latest = iterations[-1] if iterations else None
394 if latest is None:
395 return None
396 budget = session.get("budget") or {}
397 # ``max_iterations=-1`` is the "uncapped" sentinel; the prompt
398 # expects an informational remaining-iterations count, so we
399 # report zero in that mode (the model is told nothing about
400 # the iteration axis when there's no cap to count down from).
401 # Finite caps subtract the count of recorded iterations and
402 # clamp at zero.
403 cap = int(budget.get("max_iterations", 0))
404 remaining_iters = 0 if cap == -1 else max(0, cap - len(iterations))
405 if not env_cache:
406 try:
407 env_cache.append(gather_session_environment(session))
408 except Exception: # noqa: BLE001
409 env_cache.append(None)
410 env_ctx = env_cache[0]
411 if not memory_cache:
412 try:
413 store = _build_memory_store()
414 results = (
415 store.search_similar(str(session.get("directive_text") or ""))
416 if store is not None
417 else None
418 )
419 memory_cache.append(results or None)
420 except Exception: # noqa: BLE001
421 memory_cache.append(None)
422 prior_missions = memory_cache[0]
423 return await mission_sampling.maybe_sample_strategy_revision(
424 backend=backend_obj,
425 session=cast("SessionState", session),
426 iteration=latest,
427 allowlist=list(session.get("tool_allowlist") or []),
428 registered_tools=registered_tools,
429 tool_docstrings=tool_docstrings,
430 remaining_iterations=remaining_iters,
431 remaining_wall_clock_secs=remaining_wall_clock_seconds(session),
432 allow_scripts=bool(session.get("allow_scripted_strategies", False)),
433 environment_context=env_ctx,
434 prior_missions=prior_missions,
435 )
437 return _sampler
440# ---------------------------------------------------------------------------
441# Final lessons callable
442# ---------------------------------------------------------------------------
445def _build_final_lessons_callable(
446 session: Mapping[str, Any],
447 ctx: Any | None,
448 tool_docstrings: dict[str, str],
449) -> Callable[..., Awaitable[Any]] | None:
450 """Wire the Final_Report lessons overlay when sampling is enabled.
452 When the session opted into sampling and a backend resolves, the
453 engine calls this after a terminal verdict to produce model-derived
454 ``lessons`` and ``recommended_followups`` for the Final_Report.
455 Without it, the report uses deterministic templates.
456 """
457 if not session.get("use_sampling"):
458 return None
459 if session.get("sampling_backend_resolved") == "none":
460 return None
462 backend_obj = mission_sampling.select_sampling_backend(
463 model_id=session.get("bedrock_model_id"),
464 )
466 async def _final_lessons(*, session: dict[str, Any]) -> Any:
467 return await mission_sampling.maybe_sample_final_lessons(
468 backend=backend_obj,
469 session=cast("SessionState", session),
470 tool_docstrings=tool_docstrings,
471 )
473 return _final_lessons
476# ---------------------------------------------------------------------------
477# Public factory
478# ---------------------------------------------------------------------------
481async def build_engine_dependencies(
482 session: Mapping[str, Any],
483 ctx: Any | None,
484 *,
485 use_stub_dispatcher: bool = False,
486 extra_tool_metadata: tuple[dict[str, Any], dict[str, str]] | None = None,
487) -> EngineDependencies:
488 """Build the :class:`MissionEngine` dependency triple for ``session``.
490 Looks up the live FastMCP registry to populate the registered-tools
491 map and per-tool docstring cache, then assembles the production
492 wiring: a live tool dispatcher (or the canned-stub when
493 ``use_stub_dispatcher`` is True), the Strategy_Revision sampler
494 when sampling resolved to a real backend, and a sandbox runner
495 when the session permits scripted strategies.
497 The CLI passes ``use_stub_dispatcher=True`` only on the
498 ``--dry-run`` path; the MCP tool surface always uses the live
499 dispatcher.
501 ``extra_tool_metadata`` is an optional ``(tools_map, docstrings)``
502 pair merged over the live registry snapshot before the sampler is
503 built. The swarm layer uses it to teach an orchestrator's
504 Strategy_Revision sampler the in-process supervisor tools — which
505 are deliberately never FastMCP-registered — so spawn proposals
506 validate against the catalog like any other call. It never touches
507 dispatch: the dispatcher wrapper routes those names in-process.
508 """
509 if use_stub_dispatcher:
510 # The stub dispatcher means real tools never run, which means
511 # there's nothing to inform the sampling prompt; downgrade to
512 # the deterministic propose path for symmetry. The session's
513 # ``use_sampling`` flag stays as the operator set it so the
514 # criteria-scaffold path still runs through Bedrock.
515 return EngineDependencies(
516 tool_dispatcher=make_stub_dispatcher(),
517 sampling_callable=None,
518 sandbox_runner=_build_sandbox_runner(session),
519 )
521 registered_tools, tool_docstrings = await fetch_registered_tool_metadata()
522 if extra_tool_metadata is not None:
523 extra_tools, extra_docs = extra_tool_metadata
524 registered_tools = {**registered_tools, **extra_tools}
525 tool_docstrings = {**tool_docstrings, **extra_docs}
526 sampling_callable = _build_sampling_callable(
527 session,
528 ctx,
529 registered_tools=registered_tools,
530 tool_docstrings=tool_docstrings,
531 )
532 sandbox_runner = _build_sandbox_runner(session)
533 final_lessons = _build_final_lessons_callable(session, ctx, tool_docstrings)
534 return EngineDependencies(
535 tool_dispatcher=_live_dispatch_tool,
536 sampling_callable=sampling_callable,
537 sandbox_runner=sandbox_runner,
538 final_lessons_callable=final_lessons,
539 # The terminal-verdict memory write is best-effort inside the
540 # engine, so the store is wired unconditionally on the live
541 # path (the stub-dispatcher / --dry-run branch above stays
542 # memory-free: throwaway smoke sessions must not become
543 # institutional memory).
544 memory_store=_build_memory_store(),
545 )
548async def build_mission_engine(
549 session: Mapping[str, Any],
550 ctx: Any | None,
551 *,
552 use_stub_dispatcher: bool = False,
553) -> MissionEngine:
554 """Build a :class:`MissionEngine` instance ready to drive ``session``.
556 Convenience wrapper over :func:`build_engine_dependencies` that
557 also resolves the persistence backend through
558 :func:`mcp.mission.state.get_backend`. Most callers should use
559 this; the lower-level :func:`build_engine_dependencies` is
560 available for code that needs to instantiate the engine itself
561 (e.g. with a custom backend).
562 """
563 deps = await build_engine_dependencies(session, ctx, use_stub_dispatcher=use_stub_dispatcher)
564 backend = mission_state.get_backend()
565 return MissionEngine(
566 backend=backend,
567 tool_dispatcher=deps.tool_dispatcher,
568 sampling_callable=deps.sampling_callable,
569 sandbox_runner=deps.sandbox_runner,
570 final_lessons_callable=deps.final_lessons_callable,
571 memory_store=deps.memory_store,
572 observation_augmenters=deps.observation_augmenters,
573 )