Coverage for gco_mcp / tools / mission.py: 100.00%
188 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Mission goal-directed iteration loop tools.
3The whole module body is gated by :data:`feature_flags.FLAG_MISSION` so
4the ten ``mission_*`` tool decorators only fire when
5``GCO_ENABLE_MISSION=true``. With the flag unset, this module imports
6cleanly and FastMCP never sees the tools.
8[gated by GCO_ENABLE_MISSION]
9"""
11from __future__ import annotations
13import contextlib
14import json
15import secrets
16import sys
17from collections.abc import Mapping, Sequence
18from datetime import UTC, datetime
19from pathlib import Path
20from typing import Any, cast
22from audit import audit_logged
23from feature_flags import FLAG_MISSION, is_enabled
24from server import mcp
26# Mission package lives under ``gco_mcp/mission/``; the path-injection
27# pattern matches the rest of the MCP module surface so ``import
28# mission.*`` resolves without making the ``mcp`` directory a package.
29sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
32def _try_get_context() -> Any | None:
33 """Return the active FastMCP Context if inside a request, else ``None``.
35 Mirrors :func:`mcp.tools.jobs._ctx_warning`: wraps the optional
36 ``fastmcp.server.dependencies.get_context`` import so the helper
37 works in unit tests that don't go through an MCP request — those
38 raise ``RuntimeError`` from ``get_context()``, which we swallow.
39 """
40 try:
41 from fastmcp.server.dependencies import get_context
43 return get_context()
44 except Exception:
45 return None
48# Module body is entirely gated by the feature flag. When the flag is
49# unset, none of the tool decorators below fire and FastMCP never sees
50# the registrations.
51if is_enabled(FLAG_MISSION):
52 from mission import (
53 sampling as mission_sampling,
54 )
55 from mission import (
56 state as mission_state,
57 )
58 from mission import (
59 validation as mission_validation,
60 )
61 from mission.decide import decide_verdict
62 from mission.engine import MissionEngineError
63 from mission.types import SCHEMA_VERSION, TERMINAL_STATES, SessionState
64 from mission.validation import MissionValidationError
66 # ------------------------------------------------------------------ #
67 # Registry introspection helpers
68 # ------------------------------------------------------------------ #
70 async def _registered_tools_dict() -> dict[str, Any]:
71 """Return a name -> Tool object mapping for every registered tool.
73 Uses ``mcp._list_tools()`` because that bypasses the catalog-
74 replacement transforms (BM25 / Code Mode) and gives us the
75 underlying registry. Tolerates FastMCP API drift by falling
76 back to an empty mapping on any exception — the validators
77 downstream interpret an empty dict as "nothing registered",
78 which is benign for sessions that don't lean on tag-based
79 cost gating.
80 """
81 try:
82 tools = await mcp._list_tools()
83 except Exception:
84 return {}
85 return {t.name: t for t in tools}
87 async def _registered_tool_tags() -> dict[str, set[str]]:
88 """Return name -> tag-set for every registered tool.
90 Same defensive shape as :func:`_registered_tools_dict`. Tools
91 with no declared ``tags`` attribute contribute an empty set so
92 the budget validator treats them as non-cost-incurring.
93 """
94 registered = await _registered_tools_dict()
95 out: dict[str, set[str]] = {}
96 for name, tool in registered.items():
97 tags = getattr(tool, "tags", None)
98 out[name] = set(tags) if tags else set()
99 return out
101 async def _tool_docstrings_dict() -> dict[str, str]:
102 """Return name -> docstring/description mapping for sampling prompts."""
103 registered = await _registered_tools_dict()
104 return {name: (getattr(t, "description", "") or "") for name, t in registered.items()}
106 # ------------------------------------------------------------------ #
107 # Session helpers
108 # ------------------------------------------------------------------ #
110 def _strip_private_fields(session: Mapping[str, Any]) -> dict[str, Any]:
111 """Return a JSON-safe copy of ``session`` with private criterion keys dropped.
113 Thin alias over :func:`mission.validation.strip_private_fields`.
114 Kept as a module-private name so call sites in this file
115 (``mission_start``, ``mission_complete``, etc.) read at a
116 glance without having to qualify the canonical helper through
117 the long ``mission.validation`` path.
118 """
119 return mission_validation.strip_private_fields(session)
121 def _strip_private_fields_iterations(
122 iterations: Sequence[Mapping[str, Any]],
123 ) -> list[dict[str, Any]]:
124 """Strip private keys from each iteration's ``criteria_evaluation`` shape.
126 Thin alias over
127 :func:`mission.validation.strip_private_fields_iterations`.
128 """
129 return mission_validation.strip_private_fields_iterations(iterations)
131 # ------------------------------------------------------------------ #
132 # Engine wiring
133 # ------------------------------------------------------------------ #
135 # The engine factory itself lives in :mod:`mcp.mission._engine_factory`
136 # so the CLI can reuse the same wiring without crossing the
137 # ``GCO_ENABLE_MISSION`` gate. We keep a thin alias here so call
138 # sites in this module stay readable without having to spell out
139 # the long import path.
140 from mission._engine_factory import build_mission_engine as _build_engine # noqa: PLC0415
142 # ------------------------------------------------------------------ #
143 # mission_start
144 # ------------------------------------------------------------------ #
146 @mcp.tool(tags={"low-risk", "mission"})
147 @audit_logged
148 async def mission_start(
149 directive: str,
150 criteria: list[dict[str, Any]],
151 budget: dict[str, Any],
152 tool_allowlist: list[str] | None = None,
153 checkpoint_cadence: dict[str, Any] | None = None,
154 stagnation_threshold: int = 3,
155 use_sampling: bool | None = None,
156 allow_scripted_strategies: bool = False,
157 allow_all_tools: bool = False,
158 ) -> str:
159 """[gated by GCO_ENABLE_MISSION] Start a new Mission session.
161 Args:
162 directive: Natural-language goal description.
163 criteria: List of success criterion dicts (``metric_threshold``,
164 ``event``, or ``predicate`` kinds).
165 budget: Budget controls dict with ``max_iterations`` and
166 ``max_wall_clock_seconds``. Cost guardrails live
167 out-of-band via AWS Budgets and Cost Anomaly Detection.
168 tool_allowlist: List of tool names the session may invoke.
169 Optional; omit it when ``allow_all_tools`` is set.
170 checkpoint_cadence: Optional cadence dict (default
171 ``{"kind": "every_iteration"}``).
172 stagnation_threshold: Iterations of no progress before
173 terminating (default 3).
174 use_sampling: Three-state opt-in. ``None`` auto-detects,
175 ``True`` opts in explicitly, ``False`` opts out.
176 Sampling always runs server-side through Bedrock — MCP
177 client sampling left the protocol with FastMCP 4.
178 allow_scripted_strategies: When True, the session permits
179 scripted strategies (validated via the sandbox AST).
180 allow_all_tools: When True (default ``False``), resolve the
181 session's allowlist to every currently-registered tool
182 minus the ``mission_*`` control tools, instead of an
183 explicit ``tool_allowlist``. Mutually exclusive with a
184 non-empty ``tool_allowlist``.
186 Returns a JSON string with the new ``session_id`` and the
187 resolved sampling state, or an error envelope.
188 """
189 try:
190 directive_clean = mission_validation.validate_directive(directive)
191 criteria_clean = mission_validation.validate_criteria(criteria)
192 registered_tools = await _registered_tools_dict()
193 registered_tags = await _registered_tool_tags()
194 control_tools = {n for n, tags in registered_tags.items() if "mission" in tags}
195 allowlist_clean = mission_validation.resolve_effective_allowlist(
196 allow_all_tools=allow_all_tools,
197 explicit_allowlist=tool_allowlist,
198 registered_tools=registered_tools,
199 control_tools=control_tools,
200 )
201 budget_clean = mission_validation.validate_budget(
202 budget, allowlist_clean, registered_tags
203 )
204 cadence_clean = mission_validation.validate_cadence(
205 checkpoint_cadence
206 if checkpoint_cadence is not None
207 else {"kind": "every_iteration"}
208 )
209 except MissionValidationError as err:
210 return json.dumps({"code": err.code, "details": err.details})
212 use_sampling_resolved, backend_resolved = mission_sampling.resolve_sampling_state(
213 use_sampling
214 )
216 session_id = f"mission-{secrets.token_hex(8)}"
217 session: dict[str, Any] = {
218 "version": SCHEMA_VERSION,
219 "session_id": session_id,
220 "directive_text": directive_clean,
221 "criteria": criteria_clean,
222 "budget": budget_clean,
223 "tool_allowlist": allowlist_clean,
224 "checkpoint_cadence": cadence_clean,
225 "stagnation_threshold": stagnation_threshold,
226 "use_sampling": use_sampling_resolved,
227 "sampling_backend_resolved": backend_resolved,
228 "allow_scripted_strategies": bool(allow_scripted_strategies),
229 "status": "pending",
230 "created_at": datetime.now(UTC).isoformat(),
231 "iterations": [],
232 "no_progress_counter": 0,
233 }
235 backend = mission_state.get_backend()
236 # Strip the validator's cached ``_parsed_ast`` AST nodes from
237 # predicate criteria before persistence — ``ast.Expression``
238 # is not JSON-serialisable and the FilesystemBackend writes
239 # via ``json.dump``. The engine re-parses on demand from the
240 # ``expression`` string when it next loads the session, so
241 # stripping here is lossless. Mirrors ``_strip_private_criteria``
242 # in ``cli/commands/mission_cmd.py``.
243 backend.save_session(cast("SessionState", _strip_private_fields(session)))
245 return json.dumps(
246 {
247 "session_id": session_id,
248 "status": "pending",
249 "use_sampling": use_sampling_resolved,
250 "sampling_backend_resolved": backend_resolved,
251 }
252 )
254 # ------------------------------------------------------------------ #
255 # mission_status
256 # ------------------------------------------------------------------ #
258 @mcp.tool(tags={"safe", "mission"})
259 @audit_logged
260 async def mission_status(session_id: str) -> str:
261 """[gated by GCO_ENABLE_MISSION] Get the full state of a Mission session.
263 Args:
264 session_id: The session identifier returned by
265 :func:`mission_start`.
267 Returns the full session JSON or an error envelope when the
268 session is unknown.
269 """
270 backend = mission_state.get_backend()
271 session = backend.load_session(session_id)
272 if session is None:
273 return json.dumps(
274 {
275 "code": "session_not_found",
276 "details": {"session_id": session_id},
277 }
278 )
279 cleaned = _strip_private_fields(session)
280 return json.dumps(cleaned, default=str)
282 # ------------------------------------------------------------------ #
283 # mission_iterate
284 # ------------------------------------------------------------------ #
286 @mcp.tool(tags={"low-risk", "mission"})
287 @audit_logged
288 async def mission_iterate(
289 session_id: str,
290 max_iterations_this_call: int = 1,
291 ) -> str:
292 """[gated by GCO_ENABLE_MISSION] Run iteration(s) on a Mission session.
294 Args:
295 session_id: The session to iterate.
296 max_iterations_this_call: How many iterations to run before
297 returning (default 1). The loop exits early on a
298 terminal verdict (``complete`` or ``terminate``).
300 Returns a JSON object with a ``session_id`` and an
301 ``iterations`` list of iteration summaries (verdict, reason,
302 iteration index). On engine errors, returns an error envelope
303 with whatever summaries had accumulated before the failure.
304 """
305 if max_iterations_this_call < 1:
306 return json.dumps(
307 {
308 "code": "invalid_argument",
309 "details": {"reason": "max_iterations_this_call must be >= 1"},
310 }
311 )
313 ctx = _try_get_context()
314 backend = mission_state.get_backend()
316 session = backend.load_session(session_id)
317 if session is None:
318 return json.dumps(
319 {
320 "code": "session_not_found",
321 "details": {"session_id": session_id},
322 }
323 )
325 engine = await _build_engine(session, ctx)
327 summaries: list[dict[str, Any]] = []
328 for _ in range(max_iterations_this_call):
329 try:
330 record = await engine.run_iteration(session_id, ctx=ctx)
331 except MissionEngineError as err:
332 return json.dumps(
333 {
334 "code": err.code,
335 "details": {"session_id": session_id},
336 "iterations": summaries,
337 }
338 )
339 summaries.append(
340 {
341 "iteration_index": record["iteration_index"],
342 "verdict": record["verdict"],
343 "verdict_reason": record["verdict_reason"],
344 }
345 )
346 if record["verdict"] in ("complete", "terminate"):
347 break
349 return json.dumps({"session_id": session_id, "iterations": summaries})
351 # ------------------------------------------------------------------ #
352 # mission_checkpoint
353 # ------------------------------------------------------------------ #
355 @mcp.tool(tags={"safe", "mission"})
356 @audit_logged
357 async def mission_checkpoint(session_id: str) -> str:
358 """[gated by GCO_ENABLE_MISSION] Re-run the verdict cascade on the latest iteration.
360 Args:
361 session_id: The session whose latest iteration should be
362 re-evaluated.
364 Returns a JSON object with the freshly-computed verdict and
365 reason. Does not run the propose / execute / observe phases —
366 only the deterministic decide cascade. Returns an error envelope
367 when the session is missing or has no iterations yet.
368 """
369 backend = mission_state.get_backend()
370 session = backend.load_session(session_id)
371 if session is None:
372 return json.dumps(
373 {
374 "code": "session_not_found",
375 "details": {"session_id": session_id},
376 }
377 )
378 iterations = session.get("iterations") or []
379 if not iterations:
380 return json.dumps(
381 {
382 "code": "no_iterations",
383 "details": {"session_id": session_id},
384 }
385 )
387 latest = iterations[-1]
388 verdict, reason = decide_verdict(session, latest, datetime.now(UTC))
389 return json.dumps(
390 {
391 "session_id": session_id,
392 "iteration_index": latest["iteration_index"],
393 "verdict": verdict,
394 "verdict_reason": reason,
395 }
396 )
398 # ------------------------------------------------------------------ #
399 # mission_complete
400 # ------------------------------------------------------------------ #
402 @mcp.tool(tags={"low-risk", "mission"})
403 @audit_logged
404 async def mission_complete(session_id: str, reason: str = "forced_complete") -> str:
405 """[gated by GCO_ENABLE_MISSION] Force a Mission session into completed status.
407 Args:
408 session_id: The session to complete.
409 reason: Free-form reason recorded alongside the synthetic
410 final verdict (default ``forced_complete``).
412 Stamps a synthetic ``complete`` final verdict and an
413 ``ended_at`` timestamp. Refuses sessions already in a terminal
414 state.
415 """
416 del reason # currently informational; the synthetic verdict is fixed
417 backend = mission_state.get_backend()
418 session = backend.load_session(session_id)
419 if session is None:
420 return json.dumps(
421 {
422 "code": "session_not_found",
423 "details": {"session_id": session_id},
424 }
425 )
426 if session["status"] in TERMINAL_STATES:
427 return json.dumps(
428 {
429 "code": "session_terminal",
430 "details": {
431 "session_id": session_id,
432 "status": session["status"],
433 },
434 }
435 )
436 session["status"] = "completed"
437 session["final_verdict"] = "complete"
438 session["ended_at"] = datetime.now(UTC).isoformat()
439 # Defensive strip — sessions loaded from the backend were
440 # already private-field-clean, but a future change that
441 # re-attaches ``_parsed_ast`` somewhere in the flow shouldn't
442 # silently break persistence. ``_strip_private_fields`` is
443 # cheap and idempotent on already-clean inputs.
444 backend.save_session(cast("SessionState", _strip_private_fields(session)))
445 return json.dumps({"session_id": session_id, "status": "completed"})
447 # ------------------------------------------------------------------ #
448 # mission_abort
449 # ------------------------------------------------------------------ #
451 @mcp.tool(tags={"low-risk", "mission"})
452 @audit_logged
453 async def mission_abort(session_id: str, pause: bool = False) -> str:
454 """[gated by GCO_ENABLE_MISSION] Pause or terminate a Mission session.
456 Args:
457 session_id: The session to transition.
458 pause: When True, transition to ``paused``. When False
459 (default), transition to ``terminated`` with a
460 synthetic ``terminate`` final verdict.
462 Refuses sessions already in a terminal state. Best-effort
463 emits a ``ctx.warning`` when the active request context is
464 available so the operator sees the side-effect.
465 """
466 backend = mission_state.get_backend()
467 session = backend.load_session(session_id)
468 if session is None:
469 return json.dumps(
470 {
471 "code": "session_not_found",
472 "details": {"session_id": session_id},
473 }
474 )
475 if session["status"] in TERMINAL_STATES:
476 return json.dumps(
477 {
478 "code": "session_terminal",
479 "details": {
480 "session_id": session_id,
481 "status": session["status"],
482 },
483 }
484 )
485 if pause:
486 session["status"] = "paused"
487 else:
488 session["status"] = "terminated"
489 session["final_verdict"] = "terminate"
490 session["ended_at"] = datetime.now(UTC).isoformat()
491 ctx = _try_get_context()
492 if ctx is not None:
493 with contextlib.suppress(Exception):
494 await ctx.warning(f"Mission session {session_id} terminated by operator.")
495 # Defensive strip — see mission_complete for the rationale.
496 backend.save_session(cast("SessionState", _strip_private_fields(session)))
497 return json.dumps({"session_id": session_id, "status": session["status"]})
499 # ------------------------------------------------------------------ #
500 # mission_resume
501 # ------------------------------------------------------------------ #
503 @mcp.tool(tags={"low-risk", "mission"})
504 @audit_logged
505 async def mission_resume(session_id: str) -> str:
506 """[gated by GCO_ENABLE_MISSION] Resume a paused Mission session.
508 Args:
509 session_id: The session to resume.
511 Transitions ``paused -> running``. Returns an error envelope
512 when the session is missing or not in ``paused`` state.
513 """
514 backend = mission_state.get_backend()
515 session = backend.load_session(session_id)
516 if session is None:
517 return json.dumps(
518 {
519 "code": "session_not_found",
520 "details": {"session_id": session_id},
521 }
522 )
523 if session["status"] != "paused":
524 return json.dumps(
525 {
526 "code": "not_paused",
527 "details": {
528 "session_id": session_id,
529 "status": session["status"],
530 },
531 }
532 )
533 session["status"] = "running"
534 # Defensive strip — see mission_complete for the rationale.
535 backend.save_session(cast("SessionState", _strip_private_fields(session)))
536 return json.dumps({"session_id": session_id, "status": "running"})
538 # ------------------------------------------------------------------ #
539 # mission_history
540 # ------------------------------------------------------------------ #
542 @mcp.tool(tags={"safe", "mission"})
543 @audit_logged
544 async def mission_history(session_id: str, format: str = "summary") -> str:
545 """[gated by GCO_ENABLE_MISSION] Get iteration history for a Mission session.
547 Args:
548 session_id: The session whose history to retrieve.
549 format: ``"summary"`` (default) returns a compact list of
550 ``{iteration_index, verdict, verdict_reason,
551 started_at, ended_at}`` dicts; ``"full"`` returns the
552 complete iteration record dicts.
554 Returns a JSON object with an ``iterations`` list, or an error
555 envelope when the session is unknown.
556 """
557 backend = mission_state.get_backend()
558 session = backend.load_session(session_id)
559 if session is None:
560 return json.dumps(
561 {
562 "code": "session_not_found",
563 "details": {"session_id": session_id},
564 }
565 )
566 iterations = session.get("iterations") or []
567 if format == "full":
568 return json.dumps(
569 {"iterations": _strip_private_fields_iterations(iterations)},
570 default=str,
571 )
572 summaries = [
573 {
574 "iteration_index": it.get("iteration_index"),
575 "verdict": it.get("verdict"),
576 "verdict_reason": it.get("verdict_reason"),
577 "started_at": it.get("started_at"),
578 "ended_at": it.get("ended_at"),
579 }
580 for it in iterations
581 ]
582 return json.dumps({"iterations": summaries})
584 # ------------------------------------------------------------------ #
585 # mission_list
586 # ------------------------------------------------------------------ #
588 @mcp.tool(tags={"safe", "mission"})
589 @audit_logged
590 async def mission_list(status: str | None = None) -> str:
591 """[gated by GCO_ENABLE_MISSION] List Mission sessions.
593 Args:
594 status: Optional filter. Recognised values are ``running``,
595 ``completed``, ``terminated``, ``failed``, ``paused``,
596 ``pending``. Omit to list every known session.
598 Returns a JSON object with a ``sessions`` list of summary dicts
599 (``session_id``, ``status``, ``created_at``,
600 ``iteration_count``).
601 """
602 backend = mission_state.get_backend()
603 filter_dict = {"status": status} if status else None
604 sessions = backend.list_sessions(filter_dict)
605 return json.dumps({"sessions": sessions})
607 # ------------------------------------------------------------------ #
608 # mission_memory_search
609 # ------------------------------------------------------------------ #
611 @mcp.tool(tags={"safe", "mission"})
612 @audit_logged
613 async def mission_memory_search(
614 directive: str,
615 top_k: int = 3,
616 final_verdict: str | None = None,
617 ) -> str:
618 """[gated by GCO_ENABLE_MISSION] Search mission memory for similar past missions.
620 Embeds ``directive`` and queries the ``{project}-mission-memory``
621 DynamoDB vector index for the closest completed missions — the
622 same institutional memory the engine consults on sampling
623 sessions. Requires the mission-memory add-on to be deployed
624 (``mission_memory.enabled`` in cdk.json, on by default).
626 Args:
627 directive: Natural-language mission goal to search by.
628 top_k: Number of similar missions to return (default 3,
629 mirroring the ``mission_memory.top_k`` config default).
630 final_verdict: Optional inline filter on the stored
631 terminal verdict — ``"complete"`` or ``"terminate"``.
633 Returns a JSON object with a ``results`` list (each entry
634 carries ``session_id``, ``directive``, ``lessons``,
635 ``recommended_followups``, ``final_verdict``, ``verdict_reason``,
636 ``iteration_count``, ``completed_at``, and a similarity
637 ``score``), or an error envelope: ``mission_memory_unavailable``
638 when the table/index is absent or still backfilling,
639 ``mission_memory_search_failed`` for anything else.
640 """
641 from mission.memory import ( # noqa: PLC0415
642 MissionMemoryStore,
643 MissionMemoryUnavailableError,
644 )
646 try:
647 results = MissionMemoryStore().search_similar(
648 directive, top_k=top_k, final_verdict=final_verdict
649 )
650 except MissionMemoryUnavailableError as err:
651 return json.dumps(
652 {
653 "code": "mission_memory_unavailable",
654 "details": {"message": str(err)},
655 }
656 )
657 except Exception as err: # noqa: BLE001 — tool surface must envelope, not raise
658 return json.dumps(
659 {
660 "code": "mission_memory_search_failed",
661 "details": {"message": str(err)},
662 }
663 )
664 return json.dumps({"results": results})