Coverage for gco_mcp / mission / swarm_scaffold.py: 100.00%
127 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Swarm_Plan generation: sampled decomposition, deterministic fallback.
3A Swarm_Plan is a list of **spawn requests** — plain JSON-safe dicts in
4exactly the shape :func:`mission.swarm.validate_spawn` accepts — so a
5plan can be printed for review, saved to disk, and fed straight to the
6runner's spawn seam. Every plan this module returns has already been
7admission-validated end to end (fleet cap, iteration pool, finite child
8budgets, allowlist exclusions, and mutating-tool overlap are enforced
9*across* the plan by simulating the registry the spawns would build), so
10a returned plan cannot be rejected at spawn time against the same config
11and registered-tool set.
13Two producers, mirroring the criteria scaffolder one level up:
15* :func:`generate_sampled_plan` — asks a sampling backend for a JSON
16 array of child specs, validates every entry, and feeds rejection
17 reasons back into a bounded retry loop
18 (:func:`mission.criteria_scaffold.generate_sampled_criteria` is the
19 precedent). Exhaustion raises :class:`SwarmScaffoldError`; callers
20 fall back to the deterministic path, warning once, exactly like
21 ``gco mission run`` does.
22* :func:`generate_deterministic_plan` — always available, no sampling,
23 no AWS: a single child mirroring the swarm directive with criteria
24 from the deterministic criteria scaffold. Degenerate on purpose — a
25 swarm of one is safe and correct when nothing smarter is available,
26 and it is the CI path.
28The advisory-only boundary holds: sampling proposes plans, the pure
29validators admit them, and nothing in this module touches a verdict.
30:func:`sample_revised_directive` supplies the optional
31``on_failure_with_revision`` directive text for respawns — again advisory
32text only; the respawn decision lives in the deterministic restart table.
33"""
35from __future__ import annotations
37import json
38from collections.abc import Collection, Mapping
39from typing import Any, Final
41from gco.bedrock import BedrockFTUFormNotAcceptedError
43from . import criteria_scaffold
44from . import swarm as swarm_rules
45from .types import SessionState, SwarmConfig
46from .validation import MissionValidationError, validate_directive
48__all__ = [
49 "DEFAULT_CHILD_MAX_ITERATIONS",
50 "DEFAULT_CHILD_MAX_WALL_CLOCK_SECONDS",
51 "DEFAULT_PLAN_RETRIES",
52 "SwarmScaffoldError",
53 "build_plan_prompt",
54 "generate_deterministic_plan",
55 "generate_sampled_plan",
56 "sample_revised_directive",
57 "validate_plan",
58]
60DEFAULT_CHILD_MAX_ITERATIONS: Final[int] = 5
61"""Default per-child iteration budget when the caller supplies none."""
63DEFAULT_CHILD_MAX_WALL_CLOCK_SECONDS: Final[int] = 300
64"""Default per-child wall-clock budget when the caller supplies none."""
66DEFAULT_PLAN_RETRIES: Final[int] = 3
67"""Sampled-path retry budget before deterministic fallback."""
69_PROMPT_TOOL_LIMIT: Final[int] = 40
70"""Cap on catalog entries rendered into the plan prompt."""
72_PROMPT_DOCSTRING_CHARS: Final[int] = 200
73"""Per-tool docstring budget in the plan prompt."""
76class SwarmScaffoldError(Exception):
77 """Every sampled plan attempt was rejected.
79 ``last_reason`` carries the final rejection token (a validator
80 ``details.reason``, ``json_parse``, or ``transport_error``) so the
81 caller's fallback warning names the cause.
82 """
84 def __init__(self, last_reason: str, message: str | None = None) -> None:
85 self.last_reason = last_reason
86 super().__init__(message if message is not None else last_reason)
89def _spec_to_request(spec: swarm_rules.SpawnSpec) -> dict[str, Any]:
90 """Serialize a validated SpawnSpec back into a JSON-safe spawn request.
92 Criteria drop the validator's cached ``_parsed_ast`` so the request
93 round-trips through ``json.dumps``; re-validation re-parses on
94 demand. The request re-admits by construction against the same
95 config and registered-tool inputs it was validated with.
96 """
97 criteria = [
98 {k: v for k, v in criterion.items() if not str(k).startswith("_")}
99 for criterion in spec["criteria"]
100 ]
101 return {
102 "slot": spec["slot"],
103 "directive": spec["directive"],
104 "criteria": criteria,
105 "budget": dict(spec["budget"]),
106 "tool_allowlist": list(spec["tool_allowlist"]),
107 "cadence": dict(spec["checkpoint_cadence"]),
108 "restart_policy": spec["restart_policy"],
109 "max_respawns": spec["max_respawns"],
110 "use_sampling": spec["use_sampling"],
111 }
114def validate_plan(
115 entries: list[dict[str, Any]],
116 *,
117 config: SwarmConfig,
118 registered_tools: dict[str, Any],
119 registered_tags: Mapping[str, Collection[str]],
120 flag_lookup: dict[str, str] | None = None,
121) -> list[dict[str, Any]]:
122 """Admission-validate a whole plan, simulating the registry it builds.
124 Entries are validated in order against a simulated child registry so
125 the fleet cap, the iteration pool, and the mutating-tool overlap rule
126 apply across the plan exactly as they will at spawn time. Raises
127 :class:`~mission.validation.MissionValidationError` on the first
128 failing entry (its ``details`` gain a ``plan_index``); returns the
129 normalized, JSON-safe request list on success.
130 """
131 if not isinstance(entries, list) or not entries:
132 raise MissionValidationError(
133 "validation_error",
134 details={"field": "plan", "reason": "empty_or_not_a_list"},
135 )
136 simulated: list[Any] = []
137 sibling_allowlists: dict[str, list[str]] = {}
138 requests: list[dict[str, Any]] = []
139 for index, entry in enumerate(entries):
140 if not isinstance(entry, dict):
141 raise MissionValidationError(
142 "validation_error",
143 details={"field": "plan", "reason": "entry_not_a_dict", "plan_index": index},
144 )
145 try:
146 spec = swarm_rules.validate_spawn(
147 parent_role="orchestrator",
148 config=config,
149 children=simulated,
150 request=entry,
151 registered_tools=registered_tools,
152 registered_tags=registered_tags,
153 sibling_allowlists=sibling_allowlists,
154 flag_lookup=flag_lookup,
155 )
156 except MissionValidationError as err:
157 details = dict(err.details or {})
158 details["plan_index"] = index
159 raise MissionValidationError(err.code, details=details) from err
160 simulated.append(
161 swarm_rules.new_registry_entry(spec, f"plan-{index}", "1970-01-01T00:00:00+00:00")
162 )
163 sibling_allowlists[spec["slot"]] = list(spec["tool_allowlist"])
164 requests.append(_spec_to_request(spec))
165 return requests
168def generate_deterministic_plan(
169 directive: str,
170 *,
171 config: SwarmConfig,
172 registered_tools: dict[str, Any],
173 registered_tags: Mapping[str, Collection[str]],
174 tool_allowlist: list[str] | None = None,
175 allow_all_tools: bool = False,
176 flag_lookup: dict[str, str] | None = None,
177) -> list[dict[str, Any]]:
178 """The always-available fallback: one child mirroring the directive.
180 Criteria come from the deterministic criteria scaffold (so a
181 search-flavoured directive with an allowlist gets concrete
182 ``tool_call_succeeded`` entries), the budget is the module default
183 bounded by the pool, and the restart policy is ``never``. The single
184 entry runs through :func:`validate_plan` before returning, so the
185 fallback can never emit a plan the spawn seam would reject.
186 """
187 iterations = min(DEFAULT_CHILD_MAX_ITERATIONS, config["child_iteration_pool"])
188 entry: dict[str, Any] = {
189 "slot": "worker-1",
190 "directive": directive,
191 "criteria": criteria_scaffold.generate_deterministic_criteria(
192 directive, allowlist=tool_allowlist
193 ),
194 "budget": {
195 "max_iterations": iterations,
196 "max_wall_clock_seconds": DEFAULT_CHILD_MAX_WALL_CLOCK_SECONDS,
197 },
198 "restart_policy": "never",
199 "use_sampling": False,
200 }
201 if allow_all_tools:
202 entry["allow_all_tools"] = True
203 else:
204 entry["tool_allowlist"] = list(tool_allowlist or [])
205 return validate_plan(
206 [entry],
207 config=config,
208 registered_tools=registered_tools,
209 registered_tags=registered_tags,
210 flag_lookup=flag_lookup,
211 )
214def build_plan_prompt(
215 directive: str,
216 *,
217 config: SwarmConfig,
218 registered_tools: dict[str, Any],
219 tool_docstrings: Mapping[str, str] | None = None,
220 max_children: int | None = None,
221 tool_allowlist: list[str] | None = None,
222 feedback: str | None = None,
223) -> str:
224 """Assemble the deterministic decomposition prompt.
226 Deterministic given its inputs — tool names sorted, docstrings
227 truncated to a fixed budget, the catalog capped — matching the
228 byte-identity discipline of the wider sampling prompt builders.
230 ``tool_allowlist`` narrows the advertised catalog to the operator's
231 permitted set. The narrowing is advisory here and enforced by
232 :func:`generate_sampled_plan` narrowing the registry it validates
233 against; showing the model only what it may use keeps the two in
234 agreement instead of inviting rejections.
235 """
236 cap = max_children if max_children is not None else config["max_children"]
237 cap = max(1, min(cap, config["max_children"]))
238 permitted = set(registered_tools)
239 if tool_allowlist:
240 permitted &= set(tool_allowlist)
241 names = sorted(permitted)[:_PROMPT_TOOL_LIMIT]
242 docs = tool_docstrings or {}
243 catalog_lines = []
244 for name in names:
245 doc = str(docs.get(name, "")).strip().splitlines()
246 first = doc[0][:_PROMPT_DOCSTRING_CHARS] if doc else ""
247 catalog_lines.append(f"- {name}: {first}" if first else f"- {name}")
248 sections = [
249 "You are decomposing an operator goal into a fleet of supervised,",
250 "budgeted worker sessions (a swarm plan). Respond with a single JSON",
251 "array — no prose, no markdown fences. Each element is one child spec:",
252 "",
253 '{"slot": "<unique-name>", "directive": "<child goal>",',
254 ' "criteria": [<mission criteria objects>],',
255 ' "budget": {"max_iterations": <int >= 1>,',
256 ' "max_wall_clock_seconds": <int >= 1>},',
257 ' "tool_allowlist": ["<registered tool name>", ...],',
258 ' "restart_policy": "never" | "on_failure" | "on_failure_with_revision"}',
259 "",
260 "Rules:",
261 f"- At most {cap} children.",
262 f"- The sum of max_iterations across children must not exceed "
263 f"{config['child_iteration_pool']} (the shared iteration pool).",
264 "- Budgets are finite: -1 is rejected on children.",
265 "- Slot names: 1-64 chars, alphanumeric plus . _ - only.",
266 "- Only tools from the catalog below may appear in an allowlist.",
267 "- Two children must not share a non-read-only tool.",
268 "",
269 "Every criterion object requires all three of these keys:",
270 ' - "criterion_id": unique non-empty string (unique within the child)',
271 ' - "kind": one of "metric_threshold" / "metric_trend" / "event" /',
272 ' "tool_call_succeeded" / "predicate"',
273 ' - "required": JSON boolean',
274 "Plus the kind-specific keys:",
275 ' metric_threshold -> "metric" (dot-path into the Observation,',
276 ' e.g. metrics.results_count), "op" (one of',
277 ' <, <=, >, >=, ==, !=), "target" (number)',
278 ' metric_trend -> "metric" (dot-path), "direction" (one of',
279 " increasing, decreasing, non_increasing,",
280 " non_decreasing)",
281 ' event -> "event_name" (non-empty string)',
282 ' tool_call_succeeded -> "tool_name" (non-empty string, from the',
283 " child's own allowlist). PREFER this kind",
284 ' when the goal is "this tool ran and',
285 ' succeeded" — it is server-evaluated.',
286 ' predicate -> "expression" (Python expression over `obs`)',
287 "",
288 "CRITICAL — every criterion must be decidable from the Observation.",
289 "A criterion whose metric path is absent evaluates *inconclusive*, and",
290 "ANY inconclusive criterion — required or not — blocks completion for",
291 'the child\'s entire budget. "required": false does NOT make a',
292 "criterion safe to guess at. Most tools emit no metrics at all, so do",
293 "not invent metric paths: use metric_threshold / metric_trend only for",
294 "metrics you are confident the child's tools actually emit, and prefer",
295 "tool_call_succeeded otherwise. A single well-chosen criterion beats a",
296 "speculative second one that can never be decided.",
297 "",
298 "Worked example of one complete child spec:",
299 '{"slot": "docs-worker", "directive": "Find the inference docs.",',
300 ' "criteria": [{"criterion_id": "docs_found",',
301 ' "kind": "tool_call_succeeded",',
302 ' "required": true, "tool_name": "find_docs"}],',
303 ' "budget": {"max_iterations": 5, "max_wall_clock_seconds": 300},',
304 ' "tool_allowlist": ["find_docs"], "restart_policy": "never"}',
305 "",
306 "=== Operator directive ===",
307 directive,
308 "",
309 "=== Registered tool catalog ===",
310 *catalog_lines,
311 ]
312 if feedback:
313 sections.extend(["", "=== Validator feedback on your previous attempt ===", feedback])
314 return "\n".join(sections)
317class _PromptAdapter:
318 """Thin ``SamplingPrompt`` look-alike over a pre-assembled string.
320 Both shipped backends call ``prompt.assemble()`` and nothing else —
321 the same duck-typing seam the criteria scaffolder and the
322 semantic-progress judge exploit.
323 """
325 def __init__(self, text: str) -> None:
326 self._text = text
328 def assemble(self) -> str:
329 return self._text
332async def generate_sampled_plan(
333 backend: Any,
334 directive: str,
335 *,
336 config: SwarmConfig,
337 registered_tools: dict[str, Any],
338 registered_tags: Mapping[str, Collection[str]],
339 tool_docstrings: Mapping[str, str] | None = None,
340 max_children: int | None = None,
341 tool_allowlist: list[str] | None = None,
342 retries: int = DEFAULT_PLAN_RETRIES,
343 flag_lookup: dict[str, str] | None = None,
344) -> list[dict[str, Any]]:
345 """Drive a sampling backend to produce an admission-validated plan.
347 Same loop discipline as ``generate_sampled_criteria``: attempt,
348 parse, validate through :func:`validate_plan`, feed the precise
349 rejection back, retry up to ``retries`` extra times, and raise
350 :class:`SwarmScaffoldError` on exhaustion. Transport failures are
351 not retriable here — the backend owns its own recovery — and
352 surface as ``transport_error`` immediately.
354 ``tool_allowlist`` is the operator's permitted tool set. When given,
355 the registry this function prompts with *and validates against* is
356 narrowed to it, so a sampled child can never carry a tool the
357 operator did not permit: the allowlist is the primary blast-radius
358 control over a fleet, and honouring it only on the deterministic
359 path would mean a successful sample silently widened it.
360 """
361 if tool_allowlist:
362 permitted = set(tool_allowlist)
363 registered_tools = {
364 name: tool for name, tool in registered_tools.items() if name in permitted
365 }
366 registered_tags = {
367 name: tags for name, tags in registered_tags.items() if name in permitted
368 }
369 feedback: str | None = None
370 last_reason = "no_attempts"
371 for _attempt in range(retries + 1):
372 prompt = build_plan_prompt(
373 directive,
374 config=config,
375 registered_tools=registered_tools,
376 tool_docstrings=tool_docstrings,
377 max_children=max_children,
378 tool_allowlist=tool_allowlist,
379 feedback=feedback,
380 )
381 try:
382 raw = str(await backend.sample(_PromptAdapter(prompt)))
383 except BedrockFTUFormNotAcceptedError:
384 # A missing Anthropic first-time-use form is a permanent
385 # account misconfiguration; report it rather than silently
386 # downgrading to the deterministic plan (criteria-scaffold
387 # precedent).
388 raise
389 except Exception as exc: # noqa: BLE001 — transport-agnostic catch
390 raise SwarmScaffoldError(
391 "transport_error",
392 message=f"sampling backend raised {type(exc).__name__}: {exc}",
393 ) from exc
394 try:
395 parsed = _parse_json_array(raw)
396 except ValueError as exc:
397 last_reason = "json_parse"
398 feedback = (
399 "Your previous response could not be parsed as a JSON array. "
400 f"Return a single JSON array, no prose, no markdown fences. ({exc})"
401 )
402 continue
403 try:
404 return validate_plan(
405 parsed,
406 config=config,
407 registered_tools=registered_tools,
408 registered_tags=registered_tags,
409 flag_lookup=flag_lookup,
410 )
411 except MissionValidationError as exc:
412 details = exc.details or {}
413 last_reason = str(details.get("reason") or exc.code)
414 feedback = (
415 "Your previous response was rejected by the spawn validator. "
416 f"Rejection reason: {last_reason}. Details: {details!r}. "
417 "Re-emit a corrected JSON array."
418 )
419 continue
420 raise SwarmScaffoldError(last_reason)
423def _parse_json_array(raw: str) -> list[dict[str, Any]]:
424 """Parse the model response as a JSON array, tolerating code fences."""
425 text = raw.strip()
426 if text.startswith("```"):
427 lines = [line for line in text.splitlines() if not line.strip().startswith("```")]
428 text = "\n".join(lines).strip()
429 parsed = json.loads(text)
430 if not isinstance(parsed, list):
431 raise ValueError("expected a JSON array")
432 return parsed
435async def sample_revised_directive(
436 backend: Any,
437 failed_session: SessionState,
438 *,
439 lessons: list[str] | None = None,
440) -> str | None:
441 """Ask the backend for a revised directive after a child failure.
443 Advisory text supply for ``on_failure_with_revision`` respawns —
444 the deterministic restart table already made the respawn decision.
445 Returns validated directive text, or ``None`` on any failure or
446 unusable response so the caller falls back to the original
447 directive. Never raises.
448 """
449 original = str(failed_session.get("directive_text", "")).strip()
450 if not original:
451 return None
452 lesson_lines = [f"- {lesson}" for lesson in (lessons or []) if str(lesson).strip()]
453 sections = [
454 "A supervised worker session failed to meet its criteria and is being",
455 "respawned. Revise its directive so the next attempt is more likely to",
456 "succeed. Respond with exactly one line of plain text — the revised",
457 "directive. No JSON, no prose around it.",
458 "",
459 "=== Original directive ===",
460 original,
461 ]
462 if lesson_lines:
463 sections.extend(["", "=== Lessons from the failed attempt ===", *lesson_lines])
464 try:
465 raw = str(await backend.sample(_PromptAdapter("\n".join(sections))))
466 candidate = raw.strip().splitlines()[0].strip() if raw.strip() else ""
467 return validate_directive(candidate)
468 except Exception: # noqa: BLE001 — advisory path degrades, never raises
469 return None