Coverage for gco_mcp / tools / swarm.py: 100.00%
160 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Swarm supervision tools: one orchestrator Mission driving child Missions.
3The whole module body is gated by :data:`feature_flags.FLAG_SWARM` so the
4six ``swarm_*`` tool decorators only fire when ``GCO_ENABLE_SWARM=true``.
5With the flag unset, this module imports cleanly and FastMCP never sees
6the tools.
8The tools are thin wrappers over the mission-package swarm machinery:
9validation and pool rules in ``mission/swarm.py``, the concurrent child
10runner in ``mission/swarm_runner.py``, and plan scaffolding in
11``mission/swarm_scaffold.py``. The three in-process supervisor tools
12(``mission_spawn`` / ``children_status`` / ``child_abort``) are **not**
13registered here — they exist only inside an orchestrator engine's
14dispatcher, which is the recursion guard.
16[gated by GCO_ENABLE_SWARM]
17"""
19from __future__ import annotations
21import json
22import secrets
23import sys
24from collections.abc import Awaitable, Callable, Mapping
25from datetime import UTC, datetime
26from pathlib import Path
27from typing import Any, cast
29from audit import audit_logged
30from feature_flags import FLAG_SWARM, is_enabled
31from server import mcp
33# Mission package lives under ``gco_mcp/mission/``; the path-injection
34# pattern matches the rest of the MCP module surface.
35sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
38def _try_get_context() -> Any | None:
39 """Return the active FastMCP Context if inside a request, else ``None``."""
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_SWARM):
52 from mission import sampling as mission_sampling
53 from mission import state as mission_state
54 from mission import swarm as swarm_rules
55 from mission import swarm_scaffold
56 from mission._engine_factory import (
57 EngineDependencies,
58 build_engine_dependencies,
59 )
60 from mission.swarm_runner import (
61 SwarmRunner,
62 SwarmRunnerBusyError,
63 abort_swarm,
64 build_children_snapshot,
65 build_fleet_rollup,
66 list_swarms,
67 )
68 from mission.types import TERMINAL_STATES, SessionState
69 from mission.validation import (
70 MissionValidationError,
71 resolve_effective_allowlist,
72 validate_budget,
73 validate_cadence,
74 validate_criteria,
75 validate_directive,
76 )
78 # ------------------------------------------------------------------ #
79 # Registry introspection helpers (mirroring tools/mission.py)
80 # ------------------------------------------------------------------ #
82 async def _registered_tools_dict() -> dict[str, Any]:
83 """Live registered-tool map from the shared FastMCP instance."""
84 tools = await mcp._list_tools()
85 return {tool.name: tool for tool in tools}
87 async def _registered_tool_tags() -> dict[str, set[str]]:
88 """Live tool-name → tag-set map for risk-tier checks."""
89 tools = await mcp._list_tools()
90 return {tool.name: set(getattr(tool, "tags", None) or ()) for tool in tools}
92 async def _tool_docstrings_dict() -> dict[str, str]:
93 """Live tool-name → description map for the plan prompt."""
94 tools = await mcp._list_tools()
95 return {tool.name: str(getattr(tool, "description", "") or "") for tool in tools}
97 def _error(err: MissionValidationError) -> str:
98 return json.dumps({"code": err.code, "details": err.details})
100 def _strip_private_criteria(criteria: list[Any]) -> list[Any]:
101 return [
102 {k: v for k, v in c.items() if not str(k).startswith("_")} if isinstance(c, dict) else c
103 for c in criteria
104 ]
106 def _load_orchestrator(session_id: str) -> tuple[Any, SessionState | None, str | None]:
107 """Load and role-check; returns (backend, session, error_json)."""
108 backend = mission_state.get_backend()
109 session = backend.load_session(session_id)
110 if session is None:
111 return (
112 backend,
113 None,
114 json.dumps({"code": "session_not_found", "details": {"session_id": session_id}}),
115 )
116 if session.get("role") != "orchestrator" or "swarm" not in session:
117 return (
118 backend,
119 None,
120 json.dumps(
121 {
122 "code": "validation_error",
123 "details": {
124 "field": "role",
125 "reason": "not_an_orchestrator",
126 "session_id": session_id,
127 },
128 }
129 ),
130 )
131 return (backend, session, None)
133 class _SchemaShim:
134 """Duck-typed Pydantic-model stand-in over a plain JSON schema dict."""
136 def __init__(self, schema: dict[str, Any]) -> None:
137 self._schema = schema
139 def model_json_schema(self) -> dict[str, Any]:
140 return self._schema
142 class _SupervisorToolStub:
143 """Tool-shaped catalog entry for a never-registered supervisor tool.
145 Gives the Strategy_Revision sampler a name, description, and args
146 schema to validate spawn proposals against, without the tool ever
147 touching the FastMCP registry — dispatch still routes in-process
148 through the runner's wrapper, and every spawn re-validates.
149 """
151 def __init__(self, name: str) -> None:
152 self.name = name
153 self.description = swarm_rules.SUPERVISOR_TOOL_DOCSTRINGS[name]
154 self.tags = {"swarm", "supervisor"}
155 self.input_schema = _SchemaShim(swarm_rules.SUPERVISOR_TOOL_SCHEMAS[name])
157 def _supervisor_tool_metadata() -> tuple[dict[str, Any], dict[str, str]]:
158 tools = {name: _SupervisorToolStub(name) for name in swarm_rules.SUPERVISOR_TOOLS}
159 return tools, dict(swarm_rules.SUPERVISOR_TOOL_DOCSTRINGS)
161 def _deps_builder_for(ctx: Any | None) -> Callable[[Mapping[str, Any]], Awaitable[Any]]:
162 async def build(session: Mapping[str, Any]) -> EngineDependencies:
163 extra = _supervisor_tool_metadata() if session.get("role") == "orchestrator" else None
164 return await build_engine_dependencies(session, ctx, extra_tool_metadata=extra)
166 return build
168 # ------------------------------------------------------------------ #
169 # swarm_start
170 # ------------------------------------------------------------------ #
172 @mcp.tool(tags={"low-risk", "swarm"})
173 @audit_logged
174 async def swarm_start(
175 directive: str,
176 criteria: list[dict[str, Any]],
177 budget: dict[str, Any],
178 swarm: dict[str, Any],
179 tool_allowlist: list[str] | None = None,
180 allow_all_tools: bool = False,
181 checkpoint_cadence: dict[str, Any] | None = None,
182 stagnation_threshold: int = 3,
183 use_sampling: bool | None = None,
184 ) -> str:
185 """[gated by GCO_ENABLE_SWARM] Start a new swarm (orchestrator) session.
187 Args:
188 directive: The swarm-level goal, natural language.
189 criteria: Orchestrator success criteria. These evaluate over
190 the fleet snapshot — aggregate metrics like
191 ``metrics.children_completed`` and predicates over
192 ``obs['children']``.
193 budget: Orchestrator loop caps (``max_iterations``,
194 ``max_wall_clock_seconds``; Mission semantics, ``-1``
195 sentinel allowed on one axis).
196 swarm: The swarm rails: ``max_children`` and
197 ``child_iteration_pool`` (required, strictly positive),
198 ``max_concurrent_children`` (default 3),
199 ``allow_overlapping_mutating_tools`` (default false).
200 tool_allowlist: Optional extra tools for the orchestrator's
201 own Execute phase. The three in-process supervisor tools
202 are always present; ``children_status`` leads the list so
203 the deterministic strategy polls the fleet.
204 allow_all_tools: Resolve the extra allowlist to every
205 registered tool (minus loop-management names). Mutually
206 exclusive with ``tool_allowlist``.
207 checkpoint_cadence: Optional cadence dict (default
208 ``{"kind": "every_iteration"}``).
209 stagnation_threshold: Evaluated iterations with no criterion
210 transition before the cascade terminates (default 3) —
211 also the swarm's pool-exhaustion exit.
212 use_sampling: Three-state opt-in for the orchestrator's
213 advisory sampler. ``None`` auto-detects.
215 Returns a JSON string with the new ``session_id``, or a
216 structured error envelope.
217 """
218 try:
219 directive_clean = validate_directive(directive)
220 criteria_clean = validate_criteria(criteria)
221 swarm_clean = swarm_rules.validate_swarm_config(swarm)
222 registered_tools = await _registered_tools_dict()
223 registered_tags = await _registered_tool_tags()
224 budget_clean = validate_budget(budget, [], registered_tags)
225 cadence_clean = validate_cadence(
226 checkpoint_cadence
227 if checkpoint_cadence is not None
228 else {"kind": "every_iteration"}
229 )
230 extra: list[str] = []
231 if allow_all_tools or tool_allowlist:
232 extra = resolve_effective_allowlist(
233 allow_all_tools=allow_all_tools,
234 explicit_allowlist=tool_allowlist,
235 registered_tools=registered_tools,
236 )
237 except MissionValidationError as err:
238 return _error(err)
240 use_resolved, backend_resolved = mission_sampling.resolve_sampling_state(use_sampling)
241 session_id = f"mission-{secrets.token_hex(8)}"
242 session = swarm_rules.build_orchestrator_session(
243 session_id=session_id,
244 directive=directive_clean,
245 criteria=criteria_clean,
246 budget=budget_clean,
247 swarm_config=swarm_clean,
248 cadence=cadence_clean,
249 extra_allowlist=extra,
250 stagnation_threshold=stagnation_threshold,
251 use_sampling=use_resolved,
252 sampling_backend_resolved=backend_resolved,
253 created_at=datetime.now(UTC).isoformat(),
254 )
255 mission_state.get_backend().save_session(cast("SessionState", session))
256 return json.dumps(
257 {
258 "session_id": session_id,
259 "status": "pending",
260 "use_sampling": use_resolved,
261 "sampling_backend_resolved": backend_resolved,
262 "swarm": swarm_clean,
263 }
264 )
266 # ------------------------------------------------------------------ #
267 # swarm_iterate
268 # ------------------------------------------------------------------ #
270 @mcp.tool(tags={"low-risk", "swarm"})
271 @audit_logged
272 async def swarm_iterate(
273 session_id: str,
274 max_orchestrator_iterations: int | None = None,
275 ) -> str:
276 """[gated by GCO_ENABLE_SWARM] Drive a swarm's fleet forward.
278 Long-running: builds the child runner, schedules every live
279 child, and iterates the orchestrator — to its terminal verdict
280 by default, or detaching after ``max_orchestrator_iterations``
281 with the fleet left resumable. Exactly one live runner may
282 drive a swarm at a time; a second call while one runs returns a
283 ``swarm_runner_active`` envelope naming the holding process.
284 """
285 backend, session, error = _load_orchestrator(session_id)
286 if error is not None:
287 return error
288 assert session is not None
289 if session["status"] in TERMINAL_STATES:
290 return json.dumps(
291 {
292 "code": "session_terminal",
293 "details": {"session_id": session_id, "status": session["status"]},
294 }
295 )
296 ctx = _try_get_context()
297 registered_tools = await _registered_tools_dict()
298 registered_tags = await _registered_tool_tags()
299 try:
300 runner = SwarmRunner(
301 backend=backend,
302 orchestrator_id=session_id,
303 deps_builder=_deps_builder_for(ctx),
304 registered_tools=registered_tools,
305 registered_tags=registered_tags,
306 )
307 final = await runner.run_to_completion(
308 max_orchestrator_iterations=max_orchestrator_iterations
309 )
310 except SwarmRunnerBusyError as busy:
311 return json.dumps(
312 {
313 "code": "swarm_runner_active",
314 "details": {
315 "session_id": session_id,
316 "holder_pid": busy.holder_pid,
317 },
318 }
319 )
320 except MissionValidationError as err:
321 return _error(err)
322 snapshot = build_children_snapshot(
323 final["swarm"], final.get("children", []), backend.load_session
324 )
325 return json.dumps(
326 {
327 "session_id": session_id,
328 "status": final["status"],
329 "final_verdict": final.get("final_verdict"),
330 "iterations_run": len(final.get("iterations", [])),
331 "children": snapshot["metrics"],
332 }
333 )
335 # ------------------------------------------------------------------ #
336 # swarm_status
337 # ------------------------------------------------------------------ #
339 @mcp.tool(tags={"safe", "swarm"})
340 @audit_logged
341 async def swarm_status(session_id: str) -> str:
342 """[gated by GCO_ENABLE_SWARM] One-call fleet rollup for a swarm.
344 Returns the orchestrator summary, the swarm rails, the iteration
345 pool balance, the slot-ordered child table, and a findings list
346 (orphaned runner heartbeat, unreadable children, exhausted pool)
347 in the ``fleet_status`` one-document style.
348 """
349 backend, session, error = _load_orchestrator(session_id)
350 if error is not None:
351 return error
352 assert session is not None
353 return json.dumps(build_fleet_rollup(backend, session))
355 # ------------------------------------------------------------------ #
356 # swarm_abort
357 # ------------------------------------------------------------------ #
359 @mcp.tool(tags={"low-risk", "swarm"})
360 @audit_logged
361 async def swarm_abort(session_id: str) -> str:
362 """[gated by GCO_ENABLE_SWARM] Terminate a swarm and its children.
364 Transitions the orchestrator to ``terminated`` and aborts every
365 non-terminal child through the standard abort transition,
366 settling each slot's pool reservation. Works with no live
367 runner; a live runner observes the terminal orchestrator at its
368 next boundary and stands down.
369 """
370 backend, session, error = _load_orchestrator(session_id)
371 if error is not None:
372 return error
373 assert session is not None
374 if session["status"] in TERMINAL_STATES:
375 return json.dumps(
376 {
377 "code": "session_terminal",
378 "details": {"session_id": session_id, "status": session["status"]},
379 }
380 )
381 return json.dumps(abort_swarm(backend, session))
383 # ------------------------------------------------------------------ #
384 # swarm_list
385 # ------------------------------------------------------------------ #
387 @mcp.tool(tags={"safe", "swarm"})
388 @audit_logged
389 async def swarm_list(status: str | None = None) -> str:
390 """[gated by GCO_ENABLE_SWARM] List swarm (orchestrator) sessions.
392 Optionally filtered by lifecycle ``status``. Child and standalone
393 sessions never appear here — use ``mission_list`` for those.
394 """
395 backend = mission_state.get_backend()
396 return json.dumps({"swarms": list_swarms(backend, status=status)})
398 # ------------------------------------------------------------------ #
399 # swarm_plan
400 # ------------------------------------------------------------------ #
402 @mcp.tool(tags={"safe", "swarm"})
403 @audit_logged
404 async def swarm_plan(
405 directive: str,
406 swarm: dict[str, Any],
407 tool_allowlist: list[str] | None = None,
408 allow_all_tools: bool = False,
409 max_children: int | None = None,
410 use_sampling: bool | None = None,
411 retries: int = 3,
412 ) -> str:
413 """[gated by GCO_ENABLE_SWARM] Draft a validated swarm plan.
415 Decomposes the directive into admission-validated spawn requests
416 — the sampled path with retry-and-feedback when a sampling
417 backend resolves, always falling back to the deterministic
418 single-worker plan. The returned plan feeds ``mission_spawn``
419 requests (or ``gco swarm run``) verbatim; review before running.
420 """
421 try:
422 directive_clean = validate_directive(directive)
423 config = swarm_rules.validate_swarm_config(swarm)
424 except MissionValidationError as err:
425 return _error(err)
426 registered_tools = await _registered_tools_dict()
427 registered_tags = await _registered_tool_tags()
428 docstrings = await _tool_docstrings_dict()
429 use_resolved, backend_name = mission_sampling.resolve_sampling_state(use_sampling)
430 plan: list[dict[str, Any]] | None = None
431 fallback_reason: str | None = None
432 if use_resolved:
433 backend_obj = mission_sampling.select_sampling_backend(None)
434 if backend_obj is not None:
435 try:
436 plan = await swarm_scaffold.generate_sampled_plan(
437 backend_obj,
438 directive_clean,
439 config=config,
440 registered_tools=registered_tools,
441 registered_tags=registered_tags,
442 tool_docstrings=docstrings,
443 max_children=max_children,
444 tool_allowlist=(None if allow_all_tools else tool_allowlist or None),
445 retries=retries,
446 )
447 except swarm_scaffold.SwarmScaffoldError as err:
448 fallback_reason = err.last_reason
449 else:
450 fallback_reason = "sampling_backend_unavailable"
451 if plan is None:
452 try:
453 plan = swarm_scaffold.generate_deterministic_plan(
454 directive_clean,
455 config=config,
456 registered_tools=registered_tools,
457 registered_tags=registered_tags,
458 tool_allowlist=tool_allowlist,
459 allow_all_tools=allow_all_tools,
460 )
461 except MissionValidationError as err:
462 return _error(err)
463 return json.dumps(
464 {
465 "plan": plan,
466 "sampling_path": fallback_reason is None and use_resolved,
467 "sampling_backend_resolved": backend_name,
468 "fallback_reason": fallback_reason,
469 }
470 )