Coverage for gco_mcp/tools/mission.py: 100.00%

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 Iterable, Mapping

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

100

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

105

106 # ------------------------------------------------------------------ #

107 # Session helpers

108 # ------------------------------------------------------------------ #

109

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.

112

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)

120

121 def _strip_private_fields_iterations(iterations: Iterable[object]) -> list[Any]:

122 """Strip private keys from each iteration's ``criteria_evaluation`` shape.

123

124 Thin alias over

125 :func:`mission.validation.strip_private_fields_iterations`.

126 """

127 return mission_validation.strip_private_fields_iterations(iterations)

128

129 # ------------------------------------------------------------------ #

130 # Engine wiring

131 # ------------------------------------------------------------------ #

132

133 # The engine factory itself lives in :mod:`mcp.mission._engine_factory`

134 # so the CLI can reuse the same wiring without crossing the

135 # ``GCO_ENABLE_MISSION`` gate. We keep a thin alias here so call

136 # sites in this module stay readable without having to spell out

137 # the long import path.

138 from mission._engine_factory import build_mission_engine as _build_engine

139

140 # ------------------------------------------------------------------ #

141 # mission_start

142 # ------------------------------------------------------------------ #

143

144 @mcp.tool(tags={"low-risk", "mission"})

145 @audit_logged

146 async def mission_start(

147 directive: str,

148 criteria: list[dict[str, Any]],

149 budget: dict[str, Any],

150 tool_allowlist: list[str] | None = None,

151 checkpoint_cadence: dict[str, Any] | None = None,

152 stagnation_threshold: int = 3,

153 use_sampling: bool | None = None,

154 allow_scripted_strategies: bool = False,

155 allow_all_tools: bool = False,

156 ) -> str:

157 """[gated by GCO_ENABLE_MISSION] Start a new Mission session.

158

159 Args:

160 directive: Natural-language goal description.

161 criteria: List of success criterion dicts (``metric_threshold``,

162 ``event``, or ``predicate`` kinds).

163 budget: Budget controls dict with ``max_iterations`` and

164 ``max_wall_clock_seconds``. Cost guardrails live

165 out-of-band via AWS Budgets and Cost Anomaly Detection.

166 tool_allowlist: List of tool names the session may invoke.

167 Optional; omit it when ``allow_all_tools`` is set.

168 checkpoint_cadence: Optional cadence dict (default

169 ``{"kind": "every_iteration"}``).

170 stagnation_threshold: Iterations of no progress before

171 terminating (default 3).

172 use_sampling: Three-state opt-in. ``None`` auto-detects,

173 ``True`` opts in explicitly, ``False`` opts out.

174 Sampling always runs server-side through Bedrock — MCP

175 client sampling left the protocol with FastMCP 4.

176 allow_scripted_strategies: When True, the session permits

177 scripted strategies (validated via the sandbox AST).

178 allow_all_tools: When True (default ``False``), resolve the

179 session's allowlist to every currently-registered tool

180 minus the ``mission_*`` control tools, instead of an

181 explicit ``tool_allowlist``. Mutually exclusive with a

182 non-empty ``tool_allowlist``.

183

184 Returns a JSON string with the new ``session_id`` and the

185 resolved sampling state, or an error envelope.

186 """

187 try:

188 directive_clean = mission_validation.validate_directive(directive)

189 criteria_clean = mission_validation.validate_criteria(criteria)

190 registered_tools = await _registered_tools_dict()

191 registered_tags = await _registered_tool_tags()

192 control_tools = {n for n, tags in registered_tags.items() if "mission" in tags}

193 allowlist_clean = mission_validation.resolve_effective_allowlist(

194 allow_all_tools=allow_all_tools,

195 explicit_allowlist=tool_allowlist,

196 registered_tools=registered_tools,

197 control_tools=control_tools,

198 )

199 budget_clean = mission_validation.validate_budget(

200 budget, allowlist_clean, registered_tags

201 )

202 cadence_clean = mission_validation.validate_cadence(

203 checkpoint_cadence

204 if checkpoint_cadence is not None

205 else {"kind": "every_iteration"}

206 )

207 except MissionValidationError as err:

208 return json.dumps({"code": err.code, "details": err.details})

209

210 use_sampling_resolved, backend_resolved = mission_sampling.resolve_sampling_state(

211 use_sampling

212 )

213

214 session_id = f"mission-{secrets.token_hex(8)}"

215 session: dict[str, Any] = {

216 "version": SCHEMA_VERSION,

217 "session_id": session_id,

218 "directive_text": directive_clean,

219 "criteria": criteria_clean,

220 "budget": budget_clean,

221 "tool_allowlist": allowlist_clean,

222 "checkpoint_cadence": cadence_clean,

223 "stagnation_threshold": stagnation_threshold,

224 "use_sampling": use_sampling_resolved,

225 "sampling_backend_resolved": backend_resolved,

226 "allow_scripted_strategies": bool(allow_scripted_strategies),

227 "status": "pending",

228 "created_at": datetime.now(UTC).isoformat(),

229 "iterations": [],

230 "no_progress_counter": 0,

231 }

232

233 backend = mission_state.get_backend()

234 # Strip the validator's cached ``_parsed_ast`` AST nodes from

235 # predicate criteria before persistence — ``ast.Expression``

236 # is not JSON-serialisable and the FilesystemBackend writes

237 # via ``json.dump``. The engine re-parses on demand from the

238 # ``expression`` string when it next loads the session, so

239 # stripping here is lossless. Mirrors ``_strip_private_criteria``

240 # in ``cli/commands/mission_cmd.py``.

241 backend.save_session(cast("SessionState", _strip_private_fields(session)))

242

243 return json.dumps(

244 {

245 "session_id": session_id,

246 "status": "pending",

247 "use_sampling": use_sampling_resolved,

248 "sampling_backend_resolved": backend_resolved,

249 }

250 )

251

252 # ------------------------------------------------------------------ #

253 # mission_status

254 # ------------------------------------------------------------------ #

255

256 @mcp.tool(tags={"safe", "mission"})

257 @audit_logged

258 async def mission_status(session_id: str) -> str:

259 """[gated by GCO_ENABLE_MISSION] Get the full state of a Mission session.

260

261 Args:

262 session_id: The session identifier returned by

263 :func:`mission_start`.

264

265 Returns the full session JSON or an error envelope when the

266 session is unknown.

267 """

268 backend = mission_state.get_backend()

269 session = backend.load_session(session_id)

270 if session is None:

271 return json.dumps(

272 {

273 "code": "session_not_found",

274 "details": {"session_id": session_id},

275 }

276 )

277 cleaned = _strip_private_fields(session)

278 return json.dumps(cleaned, default=str)

279

280 # ------------------------------------------------------------------ #

281 # mission_iterate

282 # ------------------------------------------------------------------ #

283

284 @mcp.tool(tags={"low-risk", "mission"})

285 @audit_logged

286 async def mission_iterate(

287 session_id: str,

288 max_iterations_this_call: int = 1,

289 ) -> str:

290 """[gated by GCO_ENABLE_MISSION] Run iteration(s) on a Mission session.

291

292 Args:

293 session_id: The session to iterate.

294 max_iterations_this_call: How many iterations to run before

295 returning (default 1). The loop exits early on a

296 terminal verdict (``complete`` or ``terminate``).

297

298 Returns a JSON object with a ``session_id`` and an

299 ``iterations`` list of iteration summaries (verdict, reason,

300 iteration index). On engine errors, returns an error envelope

301 with whatever summaries had accumulated before the failure.

302 """

303 if max_iterations_this_call < 1:

304 return json.dumps(

305 {

306 "code": "invalid_argument",

307 "details": {"reason": "max_iterations_this_call must be >= 1"},

308 }

309 )

310

311 ctx = _try_get_context()

312 backend = mission_state.get_backend()

313

314 session = backend.load_session(session_id)

315 if session is None:

316 return json.dumps(

317 {

318 "code": "session_not_found",

319 "details": {"session_id": session_id},

320 }

321 )

322

323 engine = await _build_engine(session, ctx)

324

325 summaries: list[dict[str, Any]] = []

326 for _ in range(max_iterations_this_call):

327 try:

328 record = await engine.run_iteration(session_id, ctx=ctx)

329 except MissionEngineError as err:

330 return json.dumps(

331 {

332 "code": err.code,

333 "details": {"session_id": session_id},

334 "iterations": summaries,

335 }

336 )

337 summaries.append(

338 {

339 "iteration_index": record["iteration_index"],

340 "verdict": record["verdict"],

341 "verdict_reason": record["verdict_reason"],

342 }

343 )

344 if record["verdict"] in ("complete", "terminate"):

345 break

346

347 return json.dumps({"session_id": session_id, "iterations": summaries})

348

349 # ------------------------------------------------------------------ #

350 # mission_checkpoint

351 # ------------------------------------------------------------------ #

352

353 @mcp.tool(tags={"safe", "mission"})

354 @audit_logged

355 async def mission_checkpoint(session_id: str) -> str:

356 """[gated by GCO_ENABLE_MISSION] Re-run the verdict cascade on the latest iteration.

357

358 Args:

359 session_id: The session whose latest iteration should be

360 re-evaluated.

361

362 Returns a JSON object with the freshly-computed verdict and

363 reason. Does not run the propose / execute / observe phases —

364 only the deterministic decide cascade. Returns an error envelope

365 when the session is missing or has no iterations yet.

366 """

367 backend = mission_state.get_backend()

368 session = backend.load_session(session_id)

369 if session is None:

370 return json.dumps(

371 {

372 "code": "session_not_found",

373 "details": {"session_id": session_id},

374 }

375 )

376 iterations = session.get("iterations") or []

377 if not iterations:

378 return json.dumps(

379 {

380 "code": "no_iterations",

381 "details": {"session_id": session_id},

382 }

383 )

384

385 latest = iterations[-1]

386 verdict, reason = decide_verdict(session, latest, datetime.now(UTC))

387 return json.dumps(

388 {

389 "session_id": session_id,

390 "iteration_index": latest["iteration_index"],

391 "verdict": verdict,

392 "verdict_reason": reason,

393 }

394 )

395

396 # ------------------------------------------------------------------ #

397 # mission_complete

398 # ------------------------------------------------------------------ #

399

400 @mcp.tool(tags={"low-risk", "mission"})

401 @audit_logged

402 async def mission_complete(session_id: str, reason: str = "forced_complete") -> str:

403 """[gated by GCO_ENABLE_MISSION] Force a Mission session into completed status.

404

405 Args:

406 session_id: The session to complete.

407 reason: Free-form reason recorded alongside the synthetic

408 final verdict (default ``forced_complete``).

409

410 Stamps a synthetic ``complete`` final verdict and an

411 ``ended_at`` timestamp. Refuses sessions already in a terminal

412 state.

413 """

414 del reason # currently informational; the synthetic verdict is fixed

415 backend = mission_state.get_backend()

416 session = backend.load_session(session_id)

417 if session is None:

418 return json.dumps(

419 {

420 "code": "session_not_found",

421 "details": {"session_id": session_id},

422 }

423 )

424 if session["status"] in TERMINAL_STATES:

425 return json.dumps(

426 {

427 "code": "session_terminal",

428 "details": {

429 "session_id": session_id,

430 "status": session["status"],

431 },

432 }

433 )

434 session["status"] = "completed"

435 session["final_verdict"] = "complete"

436 session["ended_at"] = datetime.now(UTC).isoformat()

437 # Defensive strip — sessions loaded from the backend were

438 # already private-field-clean, but a future change that

439 # re-attaches ``_parsed_ast`` somewhere in the flow shouldn't

440 # silently break persistence. ``_strip_private_fields`` is

441 # cheap and idempotent on already-clean inputs.

442 backend.save_session(cast("SessionState", _strip_private_fields(session)))

443 return json.dumps({"session_id": session_id, "status": "completed"})

444

445 # ------------------------------------------------------------------ #

446 # mission_abort

447 # ------------------------------------------------------------------ #

448

449 @mcp.tool(tags={"low-risk", "mission"})

450 @audit_logged

451 async def mission_abort(session_id: str, pause: bool = False) -> str:

452 """[gated by GCO_ENABLE_MISSION] Pause or terminate a Mission session.

453

454 Args:

455 session_id: The session to transition.

456 pause: When True, transition to ``paused``. When False

457 (default), transition to ``terminated`` with a

458 synthetic ``terminate`` final verdict.

459

460 Refuses sessions already in a terminal state. Best-effort

461 emits a ``ctx.warning`` when the active request context is

462 available so the operator sees the side-effect.

463 """

464 backend = mission_state.get_backend()

465 session = backend.load_session(session_id)

466 if session is None:

467 return json.dumps(

468 {

469 "code": "session_not_found",

470 "details": {"session_id": session_id},

471 }

472 )

473 if session["status"] in TERMINAL_STATES:

474 return json.dumps(

475 {

476 "code": "session_terminal",

477 "details": {

478 "session_id": session_id,

479 "status": session["status"],

480 },

481 }

482 )

483 if pause:

484 session["status"] = "paused"

485 else:

486 session["status"] = "terminated"

487 session["final_verdict"] = "terminate"

488 session["ended_at"] = datetime.now(UTC).isoformat()

489 ctx = _try_get_context()

490 if ctx is not None:

491 with contextlib.suppress(Exception):

492 await ctx.warning(f"Mission session {session_id} terminated by operator.")

493 # Defensive strip — see mission_complete for the rationale.

494 backend.save_session(cast("SessionState", _strip_private_fields(session)))

495 return json.dumps({"session_id": session_id, "status": session["status"]})

496

497 # ------------------------------------------------------------------ #

498 # mission_resume

499 # ------------------------------------------------------------------ #

500

501 @mcp.tool(tags={"low-risk", "mission"})

502 @audit_logged

503 async def mission_resume(session_id: str) -> str:

504 """[gated by GCO_ENABLE_MISSION] Resume a paused Mission session.

505

506 Args:

507 session_id: The session to resume.

508

509 Transitions ``paused -> running``. Returns an error envelope

510 when the session is missing or not in ``paused`` state.

511 """

512 backend = mission_state.get_backend()

513 session = backend.load_session(session_id)

514 if session is None:

515 return json.dumps(

516 {

517 "code": "session_not_found",

518 "details": {"session_id": session_id},

519 }

520 )

521 if session["status"] != "paused":

522 return json.dumps(

523 {

524 "code": "not_paused",

525 "details": {

526 "session_id": session_id,

527 "status": session["status"],

528 },

529 }

530 )

531 session["status"] = "running"

532 # Defensive strip — see mission_complete for the rationale.

533 backend.save_session(cast("SessionState", _strip_private_fields(session)))

534 return json.dumps({"session_id": session_id, "status": "running"})

535

536 # ------------------------------------------------------------------ #

537 # mission_history

538 # ------------------------------------------------------------------ #

539

540 @mcp.tool(tags={"safe", "mission"})

541 @audit_logged

542 async def mission_history(session_id: str, format: str = "summary") -> str:

543 """[gated by GCO_ENABLE_MISSION] Get iteration history for a Mission session.

544

545 Args:

546 session_id: The session whose history to retrieve.

547 format: ``"summary"`` (default) returns a compact list of

548 ``{iteration_index, verdict, verdict_reason,

549 started_at, ended_at}`` dicts; ``"full"`` returns the

550 complete iteration record dicts.

551

552 Returns a JSON object with an ``iterations`` list, or an error

553 envelope when the session is unknown.

554 """

555 backend = mission_state.get_backend()

556 session = backend.load_session(session_id)

557 if session is None:

558 return json.dumps(

559 {

560 "code": "session_not_found",

561 "details": {"session_id": session_id},

562 }

563 )

564 iterations = session.get("iterations") or []

565 if format == "full":

566 return json.dumps(

567 {"iterations": _strip_private_fields_iterations(iterations)},

568 default=str,

569 )

570 summaries = [

571 {

572 "iteration_index": it.get("iteration_index"),

573 "verdict": it.get("verdict"),

574 "verdict_reason": it.get("verdict_reason"),

575 "started_at": it.get("started_at"),

576 "ended_at": it.get("ended_at"),

577 }

578 for it in iterations

579 ]

580 return json.dumps({"iterations": summaries})

581

582 # ------------------------------------------------------------------ #

583 # mission_list

584 # ------------------------------------------------------------------ #

585

586 @mcp.tool(tags={"safe", "mission"})

587 @audit_logged

588 async def mission_list(status: str | None = None) -> str:

589 """[gated by GCO_ENABLE_MISSION] List Mission sessions.

590

591 Args:

592 status: Optional filter. Recognised values are ``running``,

593 ``completed``, ``terminated``, ``failed``, ``paused``,

594 ``pending``. Omit to list every known session.

595

596 Returns a JSON object with a ``sessions`` list of summary dicts

597 (``session_id``, ``status``, ``created_at``,

598 ``iteration_count``).

599 """

600 backend = mission_state.get_backend()

601 filter_dict = {"status": status} if status else None

602 sessions = backend.list_sessions(filter_dict)

603 return json.dumps({"sessions": sessions})

604

605 # ------------------------------------------------------------------ #

606 # mission_memory_search

607 # ------------------------------------------------------------------ #

608

609 @mcp.tool(tags={"safe", "mission"})

610 @audit_logged

611 async def mission_memory_search(

612 directive: str,

613 top_k: int = 3,

614 final_verdict: str | None = None,

615 ) -> str:

616 """[gated by GCO_ENABLE_MISSION] Search mission memory for similar past missions.

617

618 Embeds ``directive`` and queries the ``{project}-mission-memory``

619 DynamoDB vector index for the closest completed missions — the

620 same institutional memory the engine consults on sampling

621 sessions. Requires the mission-memory add-on to be deployed

622 (``mission_memory.enabled`` in cdk.json, on by default).

623

624 Args:

625 directive: Natural-language mission goal to search by.

626 top_k: Number of similar missions to return (default 3,

627 mirroring the ``mission_memory.top_k`` config default).

628 final_verdict: Optional inline filter on the stored

629 terminal verdict — ``"complete"`` or ``"terminate"``.

630

631 Returns a JSON object with a ``results`` list (each entry

632 carries ``session_id``, ``directive``, ``lessons``,

633 ``recommended_followups``, ``final_verdict``, ``verdict_reason``,

634 ``iteration_count``, ``completed_at``, and a similarity

635 ``score``), or an error envelope: ``mission_memory_unavailable``

636 when the table/index is absent or still backfilling,

637 ``mission_memory_search_failed`` for anything else.

638 """

639 from mission.memory import (

640 MissionMemoryStore,

641 MissionMemoryUnavailableError,

642 )

643

644 try:

645 results = MissionMemoryStore().search_similar(

646 directive, top_k=top_k, final_verdict=final_verdict

647 )

648 except MissionMemoryUnavailableError as err:

649 return json.dumps(

650 {

651 "code": "mission_memory_unavailable",

652 "details": {"message": str(err)},

653 }

654 )

655 except Exception as err: # tool surface must envelope, not raise

656 return json.dumps(

657 {

658 "code": "mission_memory_search_failed",

659 "details": {"message": str(err)},

660 }

661 )

662 return json.dumps({"results": results})

Coverage for gco_mcp / tools / mission.py: 100.00%

188 statements