Coverage for gco_mcp / mission / criteria_scaffold.py: 100.00%
385 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""Helpers for ``gco mission scaffold-criteria``.
3The CLI subcommand turns a natural-language directive into a JSON
4array of Criterion objects that ``mission.validation.validate_criteria``
5accepts. Two paths are exposed:
7* :func:`generate_deterministic_criteria` — pure, no I/O. Keyword-matches
8 the directive against a small template table to pick a kind and shape
9 the criterion. The default fallback is a single ``predicate`` with
10 ``expression: "True"`` so the operator notices and edits before use.
11 Always emits at most ``max_criteria`` entries.
12* :func:`generate_sampled_criteria` — async, drives a resolved
13 :class:`SamplingBackend` to produce JSON. The response is parsed,
14 validated through ``validate_criteria``, and on rejection is retried
15 up to ``retries`` times with a feedback prompt mentioning the
16 rejection ``reason``. After the retry budget is exhausted, the helper
17 raises :class:`ScaffoldSamplingError` so the caller can fall back to
18 the deterministic path.
19* :func:`build_scaffold_prompt` — render the prompt the sampling
20 backend sees. Pure; lives here so tests can pin the exact text.
22The module is import-light: no FastMCP, no boto3, no MCP server. It
23imports the validators (and through them the predicate AST validator)
24and the sampling Protocol type, but nothing that touches a transport.
25The CLI wires the two paths together; this module keeps them
26decoupled so each can be tested in isolation.
27"""
29from __future__ import annotations
31import ast
32import itertools
33import json
34import re
35from collections.abc import Mapping, Sequence
36from dataclasses import dataclass
37from typing import TYPE_CHECKING, Any
39from gco.bedrock import BedrockFTUFormNotAcceptedError
41from . import validation as _validation
42from .predicate import PredicateRejected, parse_predicate
43from .validation import MissionValidationError
45# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
46# Generated at (UTC): 2026-09-18T02:11:36Z
47# Generated from Git commit: b8faa9689385cea16155a285a7f70cf6d488e512
48# Flowchart(s) generated from this file:
49# * ``generate_sampled_criteria`` -> ``diagrams/code_diagrams/gco_mcp/mission/criteria_scaffold.generate_sampled_criteria.html``
50# (PNG: ``diagrams/code_diagrams/gco_mcp/mission/criteria_scaffold.generate_sampled_criteria.png``)
51# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
52# <pyflowchart-code-diagram> END
55if TYPE_CHECKING: # pragma: no cover - type-checker only
56 from .sampling import SamplingBackend
59__all__ = [
60 "DEFAULT_MAX_CRITERIA",
61 "DEFAULT_RETRIES",
62 "ScaffoldSamplingError",
63 "build_scaffold_prompt",
64 "generate_deterministic_criteria",
65 "generate_sampled_criteria",
66]
69# ---------------------------------------------------------------------------
70# Tunables
71# ---------------------------------------------------------------------------
73#: Default cap on the number of criteria scaffolded per call.
74DEFAULT_MAX_CRITERIA: int = 5
76#: Default retry count for the sampling path. Each retry re-prompts
77#: the model with a feedback message containing the rejection reason.
78DEFAULT_RETRIES: int = 3
80# ---------------------------------------------------------------------------
81# Keyword templates for the deterministic fallback
82# ---------------------------------------------------------------------------
84# Each entry is (regex, builder). The first match wins; builders
85# return a single Criterion dict that ``validate_criteria`` accepts.
86# The regex is matched case-insensitively against the directive.
87# Order matters: more specific patterns appear first.
89# "Lower is better" metrics (loss, error rate, latency, cost).
90_LOWER_IS_BETTER_RE = re.compile(r"\b(loss|error|latency|cost)\b", re.IGNORECASE)
92# "Higher is better" metrics (accuracy, throughput, recall, F1).
93_HIGHER_IS_BETTER_RE = re.compile(r"\b(accuracy|throughput|f1|recall|precision)\b", re.IGNORECASE)
95# Search-flavoured directives.
96_SEARCH_RE = re.compile(r"\b(find|search|discover|locate|lookup)\b", re.IGNORECASE)
98# Event-style directives.
99_EVENT_RE = re.compile(
100 r"\b(succeed|succeeded|complete|completed|finish|finished|emit)\b",
101 re.IGNORECASE,
102)
105def _slugify(value: str, fallback: str = "criterion") -> str:
106 """Turn a directive snippet into a stable criterion_id-friendly slug.
108 Lowercase, ASCII letters / digits / underscores only. Empty input
109 falls back to ``fallback``. Non-empty results are capped at 32
110 chars so the audit log entries don't get unwieldy.
111 """
112 cleaned = re.sub(r"[^A-Za-z0-9]+", "_", value.strip().lower()).strip("_")
113 if not cleaned:
114 return fallback
115 return cleaned[:32]
118@dataclass(frozen=True)
119class _DirectiveMatch:
120 """Internal: a directive's matched template plus the captured token."""
122 kind: str
123 captured: str # the matched keyword; informs slug + metric name
126def _classify_directive(directive: str) -> _DirectiveMatch | None:
127 """Pick the matching template for ``directive``, or ``None`` for default.
129 The first match wins so more specific patterns can take precedence
130 over the generic "search" template by listing first. Returns
131 ``None`` when nothing matches; the caller then emits the
132 placeholder predicate fallback.
133 """
134 if (m := _LOWER_IS_BETTER_RE.search(directive)) is not None:
135 return _DirectiveMatch(kind="metric_threshold_lower", captured=m.group(1).lower())
136 if (m := _HIGHER_IS_BETTER_RE.search(directive)) is not None:
137 return _DirectiveMatch(kind="metric_threshold_higher", captured=m.group(1).lower())
138 if _SEARCH_RE.search(directive) is not None:
139 return _DirectiveMatch(kind="predicate_search", captured="search")
140 if _EVENT_RE.search(directive) is not None:
141 return _DirectiveMatch(kind="event", captured="job_succeeded")
142 return None
145def _build_metric_threshold(directive: str, captured: str, op: str) -> dict[str, Any]:
146 """Build a ``metric_threshold`` criterion for the given keyword.
148 The metric name uses ``val_<keyword>`` so it lines up with the
149 common validation-loss / val-accuracy convention; the target is a
150 placeholder the operator should override (0.1 for lower-is-better
151 metrics, 0.9 for higher-is-better metrics).
153 The dot-path is prefixed with ``metrics.`` because the engine's
154 Observe_Phase merges the dispatcher's top-level ``metrics`` dict
155 into the Observation under the ``metrics`` key, and the
156 ``_evaluate_metric_threshold`` resolver walks the path against the
157 Observation root. A bare ``val_loss`` (no prefix) would land on
158 every iteration as ``inconclusive: metric_path_missing`` because
159 the Observation's top level carries ``tool_results``, ``metrics``,
160 ``events`` — not loose metric values. See
161 :data:`tests.test_mission_e2e_train_to_loss` for the canonical
162 end-to-end shape this prefix lines up with.
163 """
164 slug = _slugify(captured, fallback="metric")
165 target = 0.1 if op in ("<", "<=") else 0.9
166 metric_name = f"val_{captured}" if captured in ("loss", "accuracy") else captured
167 return {
168 "criterion_id": f"{slug}_target",
169 "kind": "metric_threshold",
170 "required": True,
171 "metric": f"metrics.{metric_name}",
172 "op": op,
173 "target": target,
174 }
177def _build_predicate_search() -> dict[str, Any]:
178 """The canonical search predicate: the iteration produced any results.
180 Uses subscript form (``obs["tool_results"]``) rather than
181 ``obs.get(...)`` because the predicate AST validator rejects
182 method calls on ``obs`` — only the eight pure stdlib callables
183 are allowed. Subscript notation is the documented surface for
184 reading from the Observation.
185 """
186 return {
187 "criterion_id": "results_present",
188 "kind": "predicate",
189 "required": True,
190 "expression": "len(obs['tool_results']) > 0",
191 }
194def _build_tool_call_succeeded(tool_name: str) -> dict[str, Any]:
195 """Build a ``tool_call_succeeded`` criterion targeting ``tool_name``.
197 The slug is derived from the tool name so two ``tool_call_succeeded``
198 entries in the same list don't collide on ``criterion_id``. The
199 default ``min_count`` of 1 is left implicit on the criterion shape
200 so the operator can edit it after scaffolding without first
201 deleting an explicit value.
202 """
203 slug = _slugify(tool_name, fallback="tool")
204 return {
205 "criterion_id": f"{slug}_called",
206 "kind": "tool_call_succeeded",
207 "required": True,
208 "tool_name": tool_name,
209 }
212def _build_event(captured: str) -> dict[str, Any]:
213 """Build an ``event`` criterion using the captured keyword as the name."""
214 return {
215 "criterion_id": "expected_event",
216 "kind": "event",
217 "required": True,
218 "event_name": captured,
219 }
222def _build_default_placeholder() -> dict[str, Any]:
223 """Return the deterministic placeholder predicate.
225 The expression is the literal ``True`` so the criterion is always
226 met — this is intentional. The TODO note in the description is the
227 cue for the operator to edit the file before running. Mission's
228 validators accept the criterion as-is so the scaffolded output is
229 always usable, but a session run with this criterion unmodified
230 completes on iteration 0.
231 """
232 return {
233 "criterion_id": "todo_placeholder",
234 "kind": "predicate",
235 "required": True,
236 "expression": "True",
237 # Non-required pass-through key (not on the validator's
238 # required-keys list) so we don't trip schema validation.
239 # It surfaces in the JSON for the operator to read.
240 "description": "TODO: replace this placeholder with a real success condition.",
241 }
244def generate_deterministic_criteria(
245 directive: str,
246 *,
247 allowlist: list[str] | None = None,
248 max_criteria: int = DEFAULT_MAX_CRITERIA,
249) -> list[dict[str, Any]]:
250 """Build a criteria list deterministically from a directive.
252 Always returns a list that ``validate_criteria`` accepts. The
253 keyword-template lookup is naive on purpose — the fallback is
254 *guidance for the operator*, not a substitute for thinking about
255 the goal. The placeholder predicate is the explicit signal that
256 no template matched.
258 Args:
259 directive: The natural-language goal.
260 allowlist: Optional list of tool names. When the directive is
261 a search-flavoured goal *and* an allowlist is supplied,
262 the generator emits one ``tool_call_succeeded`` criterion
263 per allowlisted tool (capped at ``max_criteria``) instead
264 of the loose ``len(obs['tool_results']) > 0`` predicate.
265 That gives the operator concrete per-tool success
266 signals out of the box and keeps the criterion server-
267 evaluated rather than going through the predicate AST
268 sandbox. Falls back to the predicate when no allowlist
269 is supplied so existing callers keep their shape.
270 max_criteria: Cap on the number of entries returned. Always
271 at least 1; values less than 1 are clamped.
273 Returns:
274 A list of one or more Criterion dicts. The list always
275 validates through :func:`mission.validation.validate_criteria`.
276 """
277 if max_criteria < 1:
278 max_criteria = 1
279 match = _classify_directive(directive)
280 if match is None:
281 return [_build_default_placeholder()]
282 if match.kind == "metric_threshold_lower":
283 return [_build_metric_threshold(directive, match.captured, "<=")]
284 if match.kind == "metric_threshold_higher":
285 return [_build_metric_threshold(directive, match.captured, ">=")]
286 if match.kind == "predicate_search":
287 # Prefer per-tool ``tool_call_succeeded`` criteria when the
288 # operator told us what tools they intend to allowlist —
289 # those are server-evaluated and require zero predicate
290 # syntax. Fall back to the loose predicate when no
291 # allowlist is available so the no-allowlist call shape
292 # stays exactly as it was.
293 if allowlist:
294 tool_names = list(allowlist)[:max_criteria]
295 return [_build_tool_call_succeeded(name) for name in tool_names]
296 return [_build_predicate_search()]
297 if match.kind == "event":
298 return [_build_event(match.captured)]
299 # Defensive fallback — keeps mypy happy with the exhaustive return.
300 return [_build_default_placeholder()] # pragma: no cover
303# ---------------------------------------------------------------------------
304# Sampling path
305# ---------------------------------------------------------------------------
308class ScaffoldSamplingError(Exception):
309 """Raised when every sampling attempt was rejected.
311 The caller (the CLI) catches this and falls back to the
312 deterministic path. The ``last_reason`` attribute carries the
313 rejection token from the final retry so the CLI can surface it
314 in a one-line warning.
315 """
317 def __init__(self, last_reason: str, message: str | None = None) -> None:
318 self.last_reason: str = last_reason
319 super().__init__(message or last_reason)
322def build_scaffold_prompt(
323 directive: str,
324 *,
325 allowlist: list[str] | None = None,
326 max_criteria: int = DEFAULT_MAX_CRITERIA,
327 feedback: str | None = None,
328) -> str:
329 """Render the prompt the sampling backend sees.
331 The prompt asks for a strict JSON-array response, one entry per
332 Criterion. The shape is described inline so the model doesn't
333 need to fetch a schema document. The ``feedback`` argument carries
334 the rejection reason from a prior attempt — when present, it is
335 appended as a "feedback" block telling the model why the previous
336 response was rejected.
337 """
338 allowlist_block = "(none specified)" if not allowlist else ", ".join(allowlist)
339 folded_directive = directive.casefold().replace("-", "_")
340 context_guardrails: list[str] = []
341 if allowlist:
342 context_guardrails.append(
343 "The only valid tool_name values are: " + ", ".join(allowlist) + "."
344 )
345 else:
346 context_guardrails.append("Do not emit tool_call_succeeded criteria.")
347 if "loss" not in folded_directive:
348 context_guardrails.append("Do not emit or reference val_loss.")
349 if "goal_reached" not in folded_directive and "goal reached" not in folded_directive:
350 context_guardrails.append("Do not emit or reference goal_reached.")
351 sections: list[str] = []
352 sections.append(
353 "You are drafting Success_Criteria for a Mission goal-directed "
354 "iteration loop. The operator's directive and the tool "
355 "allowlist follow. Produce a JSON array of criterion objects "
356 "the operator can hand to `gco mission start --criteria-file`."
357 )
358 sections.append("")
359 sections.append("=== Directive ===")
360 sections.append(directive)
361 sections.append("")
362 sections.append("=== Tool allowlist ===")
363 sections.append(allowlist_block)
364 sections.append("")
365 sections.append(f"=== Cap: at most {max_criteria} criterion entries ===")
366 sections.append("")
367 sections.append("=== Observation shape (read by predicates and metric paths) ===")
368 sections.append(
369 "Each iteration's Observation is a dict with these fields:\n"
370 ' - "tool_results": list[dict] — every tool the iteration\n'
371 " called returns one entry. Each entry is whatever the\n"
372 " tool itself returned, plus a top-level ``_status`` flag.\n"
373 ' - "metrics": dict[str, Any] — numeric / scalar values\n'
374 " surfaced by tools that emit them. The dot-path for a\n"
375 " metric_threshold criterion against ``val_loss`` is\n"
376 ' ``"metrics.val_loss"`` (NOT ``"val_loss"``); the engine\n'
377 " walks the path against the Observation root and a bare\n"
378 " name will land as ``inconclusive: metric_path_missing``\n"
379 " on every iteration.\n"
380 ' - "events": list[dict] — emitted events, each with an\n'
381 " ``event_name`` key.\n"
382 ' - "errors" (optional): list[dict] — errors any tool raised.\n'
383 ' - "phase_started_at" / "phase_ended_at": ISO-8601 strings.'
384 )
385 sections.append("")
386 sections.append("=== Output schema ===")
387 sections.append(
388 "Return a single JSON array. Each entry is an object with "
389 "these required keys:\n"
390 ' - "criterion_id": unique non-empty string\n'
391 ' - "kind": one of "metric_threshold" / "event" / '
392 '"predicate" / "tool_call_succeeded"\n'
393 ' - "required": JSON boolean\n'
394 "Plus the kind-specific keys:\n"
395 ' metric_threshold -> "metric" (DOT-PATH into the\n'
396 " Observation, e.g.\n"
397 " ``metrics.val_loss``,\n"
398 " ``tool_results.0.score``), "
399 '"op" (one of <, <=, >, >=, ==, !=), "target" (number)\n'
400 ' event -> "event_name" (non-empty string;\n'
401 ' matched against entries in obs["events"])\n'
402 ' tool_call_succeeded -> "tool_name" (non-empty string;\n'
403 " matched against entries in\n"
404 ' ``obs["tool_results"]`` whose\n'
405 ' ``_status`` equals ``"ok"``).\n'
406 ' Optional: "min_count" (positive\n'
407 " int, default 1).\n"
408 " PREFER this kind over a predicate\n"
409 ' when the goal is "this tool ran\n'
410 ' and succeeded" — it is server-\n'
411 " evaluated and never goes through\n"
412 " the predicate AST sandbox.\n"
413 ' predicate -> "expression" (a Python expression\n'
414 " evaluated against `obs` — see\n"
415 " the predicate vocabulary section\n"
416 " below for the exact surface)"
417 )
418 sections.append("")
419 sections.append("=== Predicate vocabulary ===")
420 sections.append(
421 "Predicate expressions run inside a tight AST sandbox. The\n"
422 "allowed surface:\n"
423 "\n"
424 "Names: ``obs`` (the Observation dict).\n"
425 "Top-level callables (twelve, all pure stdlib):\n"
426 " ``len``, ``min``, ``max``, ``sum``, ``abs``,\n"
427 " ``any``, ``all``, ``sorted``,\n"
428 " ``str``, ``int``, ``float``, ``bool`` (type coercions).\n"
429 "Read-only method calls on any value (eight, all pure):\n"
430 " ``.get(key[, default])``, ``.keys()``, ``.values()``,\n"
431 " ``.items()``, ``.lower()``, ``.upper()``, ``.strip()``,\n"
432 " ``.startswith(prefix[, start[, end]])``\n"
433 "Operators: arithmetic, comparisons (<, <=, >, >=, ==, !=,\n"
434 " is, is not, in, not in), boolean (and, or, not), ternary\n"
435 " (a if b else c).\n"
436 "Containers: list/tuple/dict/set literals, list / set / dict\n"
437 " / generator comprehensions (the comprehension target may\n"
438 " not shadow ``obs`` or any callable name).\n"
439 "Subscripts: ``obs['key']``, ``obs['k']['nested']``,\n"
440 " ``obs['list'][0]``, etc.\n"
441 "Attribute access: ONLY single-level on ``obs`` (e.g. ``obs.events``\n"
442 " for read-only access; subscript form is preferred). Nested\n"
443 " walks like ``obs.a.b`` are rejected — use ``obs['a']['b']``.\n"
444 "\n"
445 "Method calls outside the eight pure-accessor names are\n"
446 "rejected (no ``.append``, ``.update``, ``.pop``, ``.count``,\n"
447 "``.split``, etc.). Calls to non-allowlisted names\n"
448 "(``list``, ``dict``, ``getattr``, ``isinstance``, ...)\n"
449 "are rejected."
450 )
451 sections.append("")
452 sections.append("=== Predicate examples (do NOT use rejected forms) ===")
453 sections.append(
454 "ACCEPTED predicate expressions:\n"
455 " len(obs['tool_results']) > 0\n"
456 " obs['metrics']['val_loss'] < 0.1\n"
457 " any(e['event_name'] == 'goal_reached' for e in obs['events'])\n"
458 " any(r.get('_status') == 'ok' for r in obs['tool_results'])\n"
459 " all(r.get('_status') == 'ok' for r in obs['tool_results'])\n"
460 " any(r.get('_status') == 'ok' and r.get('tool_name') == 'find_docs'\n"
461 " for r in obs['tool_results'])\n"
462 " any('inference' in str(r).lower() for r in obs['tool_results'])\n"
463 " len(obs.get('errors', [])) == 0\n"
464 " any(k == 'val_loss' for k in obs['metrics'].keys())\n"
465 " any(k.startswith('val_') for k in obs['metrics'].keys())\n"
466 "\n"
467 "REJECTED predicate expressions (will fail validation):\n"
468 " obs.metrics.val_loss < 0.1 # nested attribute walk; use obs['metrics']['val_loss']\n"
469 " obs['tool_results'].count('ok') # ``.count`` is not on the method allowlist\n"
470 " obs['tool_results'].append(1) # ``.append`` mutates and is not allowed\n"
471 " any(r.split(',') for r in obs['tool_results']) # ``.split`` not on method allowlist\n"
472 " getattr(obs, 'tool_results') # ``getattr`` not on callable allowlist\n"
473 " obs['x'].y.z # attribute walk after subscript"
474 )
475 sections.append("")
476 sections.append("=== Final checklist ===")
477 sections.append(
478 "Use only conditions directly relevant to the operator's directive. "
479 "Do not copy example-only metric, event, or tool names unless the "
480 "directive calls for them. Every object must contain criterion_id, "
481 "kind, and an explicit required boolean, and all criterion_id values "
482 "must be unique. An event criterion must use event_name and must never "
483 "use expression. " + " ".join(context_guardrails)
484 )
485 sections.append("")
486 sections.append("Output only the JSON array. No prose, no markdown fences.")
487 if feedback:
488 sections.append("")
489 sections.append("=== Feedback on previous attempt ===")
490 sections.append(feedback)
491 return "\n".join(sections)
494def _strip_markdown_fence(text: str, *, require_closed: bool = False) -> str:
495 """Remove one Markdown fence, optionally requiring a closing delimiter."""
496 stripped = text.strip()
497 if not stripped.startswith("```"):
498 return stripped
499 closed = len(stripped) >= 6 and stripped.endswith("```")
500 if not closed:
501 if require_closed:
502 raise ValueError("unclosed Markdown fence in model response")
503 first_newline = stripped.find("\n")
504 return stripped[first_newline + 1 :].strip() if first_newline != -1 else stripped
505 body = stripped[3:-3]
506 first_newline = body.find("\n")
507 if first_newline != -1:
508 body = body[first_newline + 1 :]
509 return body.strip()
512def _parse_response(text: str) -> list[dict[str, Any]]:
513 """Extract a JSON array from a model response.
515 Models occasionally wrap JSON in Markdown fences. Reasoning models can
516 also expose prose followed by a standalone ``</think>`` terminator and a
517 final JSON document. For that exact envelope, only a complete strict-JSON
518 suffix is accepted; scanning rationale for an earlier draft array would be
519 ambiguous. Responses without the marker retain the established first-array
520 extraction behavior. Raises ``ValueError`` when no array is recoverable.
521 """
522 think_pattern = r"(?m)^[ \t]*</think>[ \t]*$"
523 stripped = _strip_markdown_fence(
524 text,
525 require_closed=re.search(think_pattern, text.strip()) is not None,
526 )
527 think_ends = list(re.finditer(think_pattern, stripped))
528 if think_ends:
529 stripped = _strip_markdown_fence(
530 stripped[think_ends[-1].end() :],
531 require_closed=True,
532 )
533 if not stripped:
534 raise ValueError("no JSON document follows </think> terminator")
535 parsed = json.loads(stripped)
536 else:
537 # Find the first '[' and last ']' so ordinary responses can carry
538 # leading/trailing prose without accepting arbitrary non-JSON syntax.
539 start = stripped.find("[")
540 end = stripped.rfind("]")
541 if start == -1 or end == -1 or end < start:
542 raise ValueError("no JSON array found in response")
543 parsed = json.loads(stripped[start : end + 1])
544 if not isinstance(parsed, list):
545 raise ValueError("JSON payload is not a list")
546 out: list[dict[str, Any]] = []
547 for entry in parsed:
548 if not isinstance(entry, dict):
549 raise ValueError("array entry is not an object")
550 out.append(entry)
551 return out
554def _normalize_kind_name(criterion: dict[str, Any]) -> dict[str, Any]:
555 """Rewrite obvious ``kind`` typos to the canonical names.
557 Models occasionally emit a near-miss for the criterion ``kind``
558 field — pluralising (``tool_calls_succeeded`` instead of the
559 canonical ``tool_call_succeeded``), abbreviating
560 (``threshold`` instead of ``metric_threshold``), or hyphenating
561 (``tool-call-succeeded`` instead of underscore form). The
562 structural validator rejects these with ``kind_invalid`` and the
563 retry-with-feedback path can recover, but the typos are
564 mechanical: a closed alias map covers every captured emission
565 we have seen across Bedrock models.
567 The map is intentionally narrow — we only canonicalise a name
568 when it is unambiguously a typo for one of the four valid
569 kinds, never a name a future kind extension might claim. Returns
570 the input unchanged when the kind is already canonical, missing,
571 or not a string. Returns a shallow copy when a rewrite fires so
572 the input dict is never mutated.
573 """
574 kind = criterion.get("kind")
575 if not isinstance(kind, str):
576 return criterion
577 canonical = _KIND_ALIASES.get(kind)
578 if canonical is None:
579 return criterion
580 if canonical == kind:
581 return criterion
582 out = dict(criterion)
583 out["kind"] = canonical
584 return out
587# Closed alias map for ``_normalize_kind_name``. Every entry here was
588# observed in the captured fixture corpus under
589# ``tests/fixtures/scaffold_responses/`` — adding a new entry is the
590# right move only when a captured model emits a near-miss the
591# rejection-feedback retry doesn't recover on the next attempt.
592_KIND_ALIASES: dict[str, str] = {
593 # Llama 4 Scout pluralises the kind name in its first emission.
594 "tool_calls_succeeded": "tool_call_succeeded",
595 # Hyphenated forms occasionally surface from JSON-schema-trained
596 # smaller models that map ``snake_case`` onto ``kebab-case``.
597 "tool-call-succeeded": "tool_call_succeeded",
598 "metric-threshold": "metric_threshold",
599}
602def _normalize_metric_path(criterion: dict[str, Any]) -> dict[str, Any]:
603 """Auto-prefix bare metric names with ``metrics.`` for ``metric_threshold``.
605 The engine's metric path resolver walks the dot-path against the
606 Observation root, where canonical metric values live under the
607 ``metrics`` sub-dict. A bare ``"val_loss"`` lands as
608 ``inconclusive: metric_path_missing`` on every iteration.
610 Models trained on generic metric semantics tend to emit bare
611 names anyway. Rather than reject the response and burn a retry,
612 this normaliser injects the ``metrics.`` prefix when:
614 1. ``kind == "metric_threshold"``,
615 2. ``metric`` is a non-empty string,
616 3. The string contains no ``.`` separator (so already-qualified
617 paths like ``tool_results.0.score`` or
618 ``metrics.something.nested`` pass through verbatim).
620 Returns a shallow copy so the input is never mutated. The strip is
621 idempotent on already-prefixed values: ``"metrics.foo"`` has a
622 ``.`` so it falls through unchanged.
623 """
624 if criterion.get("kind") != "metric_threshold":
625 return criterion
626 metric = criterion.get("metric")
627 if not isinstance(metric, str) or not metric:
628 return criterion
629 if "." in metric:
630 return criterion
631 out = dict(criterion)
632 out["metric"] = f"metrics.{metric}"
633 return out
636class _AttributeToSubscriptRewriter(ast.NodeTransformer):
637 """Rewrite ``obs.<attr>`` chains as ``obs['<attr>']`` chains.
639 The predicate validator accepts a single-level attribute read on
640 ``obs`` (``obs.tool_results``) but rejects nested attribute walks
641 (``obs.metrics.val_loss``) and method-style calls
642 (``obs.x.any()``, ``obs.get('x')``). Models routinely emit those
643 shapes because they are the obvious Pythonic idioms. This
644 transformer rewrites the *attribute-walk* shapes mechanically;
645 method-call shapes that need creative rewriting are left alone so
646 the standard retry-with-feedback loop can teach the model.
648 The walk only rewrites attribute reads whose innermost base is the
649 ``Name('obs')`` — every other attribute access (e.g. on a list
650 element returned from a comprehension, on a number) is left
651 untouched so the validator's other guards still apply.
652 """
654 def visit_Attribute(self, node: ast.Attribute) -> ast.AST: # ast hook name
655 # Recurse into the value first so a nested attribute walk gets
656 # rewritten bottom-up: ``obs.metrics.val_loss`` -> visit
657 # ``obs.metrics`` first (which becomes ``obs['metrics']``)
658 # then wrap the result in ``[...]['val_loss']``.
659 self.generic_visit(node)
660 # Only rewrite when the rewritten base is one of:
661 # * Name('obs') — the simple ``obs.x`` case
662 # * Subscript whose ultimate base is Name('obs') — the
663 # already-rewritten ``obs['metrics']`` case
664 # Anything else (attribute on a Call, on a list literal, on a
665 # comprehension target) is left as-is so the validator's
666 # rejections still fire on shapes the autofix shouldn't try to
667 # silently rescue.
668 base = node.value
669 innermost = base
670 while isinstance(innermost, ast.Subscript):
671 innermost = innermost.value
672 if not (isinstance(innermost, ast.Name) and innermost.id == "obs"):
673 return node
674 return ast.Subscript(
675 value=base,
676 slice=ast.Constant(value=node.attr),
677 ctx=node.ctx,
678 )
681def _autofix_predicate(criterion: dict[str, Any]) -> dict[str, Any]:
682 """Best-effort rewrite of attribute-walk predicates into subscript form.
684 Keeps the crit dict unchanged when:
686 * ``kind != "predicate"``
687 * the expression is missing or non-string
688 * the expression already parses cleanly through
689 :func:`mission.predicate.parse_predicate`
690 * source has a syntax error (the validator will reject it with the
691 original code anyway)
692 * the rewritten expression *still* fails validation (so the
693 retry-with-feedback path runs against the original source the
694 model emitted, not a partially-rewritten one)
696 Returns a shallow copy with the rewritten ``expression`` only when
697 the rewrite produced a predicate that clears the validator. This
698 mirrors :func:`_normalize_metric_path` — never mutates input,
699 always returns a JSON-safe dict.
700 """
701 if criterion.get("kind") != "predicate":
702 return criterion
703 expression = criterion.get("expression")
704 if not isinstance(expression, str) or not expression:
705 return criterion
706 # Cheap fast path: if the source is already valid, don't pay the
707 # cost of an AST round-trip on the happy case.
708 try:
709 parse_predicate(expression)
710 return criterion
711 except PredicateRejected:
712 pass
714 try:
715 tree = ast.parse(expression, mode="eval")
716 except SyntaxError:
717 return criterion
719 rewritten_tree = _AttributeToSubscriptRewriter().visit(tree)
720 ast.fix_missing_locations(rewritten_tree)
721 try:
722 rewritten_src = ast.unparse(rewritten_tree)
723 except Exception: # unparse failure leaves us no better off
724 return criterion
726 # Re-validate the rewrite. If the rewrite still doesn't validate
727 # (e.g. a method call like ``obs.x.any()`` produced
728 # ``obs['x'].any()`` which is still a method-call-on-subscript),
729 # fall back to the original so the retry-with-feedback loop sees
730 # the model's actual emission.
731 try:
732 parse_predicate(rewritten_src)
733 except PredicateRejected:
734 return criterion
736 out = dict(criterion)
737 out["expression"] = rewritten_src
738 return out
741def _materialize_required_default(criterion: dict[str, Any]) -> dict[str, Any]:
742 """Materialize the fail-closed default when a model omits ``required``.
744 Explicit values, including invalid non-booleans, remain untouched so the
745 validator can reject them. Only an absent key becomes ``True``; this cannot
746 make Mission completion less strict.
747 """
748 if "required" in criterion:
749 return criterion
750 out = dict(criterion)
751 out["required"] = True
752 return out
755def _normalize_criterion_ids(criteria: list[dict[str, Any]]) -> list[dict[str, Any]]:
756 """Trim valid IDs and suffix model-generated collisions deterministically.
758 Every condition is retained. Missing, non-string, or blank IDs remain
759 untouched for validation, while later duplicates receive an unused numeric
760 suffix. Original suffixed IDs are reserved up front so a repair never
761 steals a name that a later criterion already owns.
762 """
763 original_ids: list[str | None] = []
764 for criterion in criteria:
765 criterion_id = criterion.get("criterion_id")
766 if isinstance(criterion_id, str) and criterion_id.strip():
767 original_ids.append(criterion_id.strip())
768 else:
769 original_ids.append(None)
771 reserved = {criterion_id for criterion_id in original_ids if criterion_id is not None}
772 seen: set[str] = set()
773 normalized: list[dict[str, Any]] = []
774 for criterion, base in zip(criteria, original_ids, strict=True):
775 if base is None:
776 normalized.append(criterion)
777 continue
778 candidate = base
779 if candidate in seen:
780 suffix = 2
781 while f"{base}_{suffix}" in reserved:
782 suffix += 1
783 candidate = f"{base}_{suffix}"
784 reserved.add(candidate)
785 seen.add(candidate)
786 if criterion.get("criterion_id") == candidate:
787 normalized.append(criterion)
788 continue
789 out = dict(criterion)
790 out["criterion_id"] = candidate
791 normalized.append(out)
792 return normalized
795def _normalize_sampled_criteria(criteria: list[dict[str, Any]]) -> list[dict[str, Any]]:
796 """Apply the conservative model-output normalizations in one stable order."""
797 normalized = [_normalize_kind_name(criterion) for criterion in criteria]
798 normalized = [_normalize_metric_path(criterion) for criterion in normalized]
799 normalized = [_autofix_predicate(criterion) for criterion in normalized]
800 normalized = [_materialize_required_default(criterion) for criterion in normalized]
801 return _normalize_criterion_ids(normalized)
804def _is_tool_name_accessor(node: ast.AST) -> bool:
805 """Return whether an expression reads a result's ``tool_name`` field."""
806 if isinstance(node, ast.Subscript):
807 return isinstance(node.slice, ast.Constant) and node.slice.value == "tool_name"
808 return bool(
809 isinstance(node, ast.Call)
810 and isinstance(node.func, ast.Attribute)
811 and node.func.attr == "get"
812 and node.args
813 and isinstance(node.args[0], ast.Constant)
814 and node.args[0].value == "tool_name"
815 )
818def _literal_string_options(node: ast.AST) -> tuple[set[str], bool]:
819 """Return literal string choices and whether the whole operand is proven."""
820 if isinstance(node, ast.Constant) and isinstance(node.value, str):
821 return {node.value}, True
822 if isinstance(node, (ast.List, ast.Tuple, ast.Set)):
823 if not node.elts:
824 return set(), False
825 values: set[str] = set()
826 for element in node.elts:
827 if not isinstance(element, ast.Constant) or not isinstance(element.value, str):
828 return set(), False
829 values.add(element.value)
830 return values, True
831 return set(), False
834def _predicate_tool_references(criterion: Mapping[str, Any]) -> tuple[set[str], bool]:
835 """Return explicit tool IDs and whether any tool-name read is unanalyzable.
837 A tool-name accessor is accepted only when it is a direct comparison
838 operand against one or more literal strings. Wrappers, computed keys,
839 starred ``get`` arguments, dictionary enumeration, and other dynamic reads
840 from ``tool_results`` fail closed because the analyzer cannot prove that
841 their resulting tool IDs stay inside the caller's allowlist.
842 """
843 expression = criterion.get("expression")
844 if not isinstance(expression, str):
845 return set(), False
847 parsed = criterion.get("_parsed_ast")
848 if not isinstance(parsed, ast.Expression):
849 try:
850 parsed = parse_predicate(expression)
851 except PredicateRejected:
852 return set(), "tool_name" in expression
854 nodes = tuple(ast.walk(parsed))
855 string_literals = [
856 node.value
857 for node in nodes
858 if isinstance(node, ast.Constant) and isinstance(node.value, str)
859 ]
860 compact_source = re.sub(r"[^a-z0-9_]", "", expression.casefold())
861 mentions_tool_name = "tool_name" in compact_source or "tool_name" in "".join(
862 value.casefold() for value in string_literals
863 )
864 references_tool_results = any(
865 value.casefold() == "tool_results" for value in string_literals
866 ) or any(
867 isinstance(node, ast.Attribute)
868 and isinstance(node.value, ast.Name)
869 and node.value.id == "obs"
870 and node.attr == "tool_results"
871 for node in nodes
872 )
873 has_dynamic_observation_access = any(
874 (
875 isinstance(node, ast.Subscript)
876 and isinstance(node.value, ast.Name)
877 and node.value.id == "obs"
878 and not (
879 isinstance(node.slice, ast.Constant) and isinstance(node.slice.value, (str, int))
880 )
881 )
882 or (
883 isinstance(node, ast.Call)
884 and isinstance(node.func, ast.Attribute)
885 and isinstance(node.func.value, ast.Name)
886 and node.func.value.id == "obs"
887 and node.func.attr == "get"
888 and (
889 not node.args
890 or not isinstance(node.args[0], ast.Constant)
891 or not isinstance(node.args[0].value, str)
892 )
893 )
894 for node in nodes
895 )
897 has_dynamic_tool_result_access = False
898 if references_tool_results:
899 for node in nodes:
900 if (
901 isinstance(node, ast.Call)
902 and isinstance(node.func, ast.Attribute)
903 and (
904 node.func.attr == "items"
905 or (
906 node.func.attr == "get"
907 and (
908 not node.args
909 or not isinstance(node.args[0], ast.Constant)
910 or not isinstance(node.args[0].value, str)
911 )
912 )
913 )
914 ) or (
915 isinstance(node, ast.Subscript)
916 and not (
917 isinstance(node.slice, ast.Constant)
918 and isinstance(node.slice.value, (str, int))
919 )
920 ):
921 has_dynamic_tool_result_access = True
923 accessor_ids = {id(node) for node in nodes if _is_tool_name_accessor(node)}
924 if not accessor_ids:
925 return (
926 set(),
927 mentions_tool_name or has_dynamic_observation_access or has_dynamic_tool_result_access,
928 )
930 names: set[str] = set()
931 matched_accessor_ids: set[int] = set()
932 for node in nodes:
933 if not isinstance(node, ast.Compare):
934 continue
935 operands = [node.left, *node.comparators]
936 if len(node.ops) > 1 and any(id(operand) in accessor_ids for operand in operands):
937 # One accessor participates in multiple comparison legs. Treat the
938 # chain as unanalyzable rather than letting one literal leg mask a
939 # computed or off-allowlist sibling.
940 continue
941 for left, right in itertools.pairwise(operands):
942 if id(left) in accessor_ids:
943 literals, fully_literal = _literal_string_options(right)
944 if fully_literal:
945 names.update(literals)
946 matched_accessor_ids.add(id(left))
947 if id(right) in accessor_ids:
948 literals, fully_literal = _literal_string_options(left)
949 if fully_literal:
950 names.update(literals)
951 matched_accessor_ids.add(id(right))
952 return (
953 names,
954 has_dynamic_observation_access
955 or has_dynamic_tool_result_access
956 or matched_accessor_ids != accessor_ids,
957 )
960def _validate_sampled_criteria_context(
961 criteria: Sequence[Mapping[str, Any]],
962 *,
963 directive: str,
964 allowlist: list[str] | None,
965) -> None:
966 """Reject model criteria that cannot be satisfied in the supplied context.
968 Structural validation alone cannot know which tools Mission may execute or
969 whether schema examples leaked into an unrelated directive. This sampled-
970 output-only gate keeps direct and predicate tool references inside the
971 caller's allowlist and rejects the two exact example signals when the
972 directive does not mention them. It never rewrites model intent.
973 """
974 allowed_tools = set(allowlist or ())
975 folded_directive = directive.casefold().replace("-", "_")
976 mentions_loss = "loss" in folded_directive
977 mentions_goal = "goal_reached" in folded_directive or "goal reached" in folded_directive
979 for criterion in criteria:
980 criterion_id = criterion.get("criterion_id")
981 referenced_tools: set[str] = set()
982 has_unanalyzable_tool_reference = False
983 if criterion.get("kind") == "tool_call_succeeded":
984 tool_name = criterion.get("tool_name")
985 if isinstance(tool_name, str):
986 referenced_tools.add(tool_name)
987 elif criterion.get("kind") == "predicate":
988 referenced_tools, has_unanalyzable_tool_reference = _predicate_tool_references(
989 criterion
990 )
991 if has_unanalyzable_tool_reference:
992 raise MissionValidationError(
993 "validation_error",
994 details={
995 "field": "criteria",
996 "criterion_id": criterion_id,
997 "reason": "tool_name_reference_not_statically_allowlisted",
998 },
999 )
1000 disallowed = sorted(referenced_tools - allowed_tools)
1001 if disallowed:
1002 raise MissionValidationError(
1003 "validation_error",
1004 details={
1005 "field": "criteria",
1006 "criterion_id": criterion_id,
1007 "reason": "tool_name_not_allowlisted",
1008 "tool_names": disallowed,
1009 },
1010 )
1012 metric = criterion.get("metric")
1013 expression = criterion.get("expression")
1014 event_name = criterion.get("event_name")
1015 if not mentions_loss and (
1016 metric == "metrics.val_loss"
1017 or (isinstance(expression, str) and "val_loss" in expression.casefold())
1018 ):
1019 raise MissionValidationError(
1020 "validation_error",
1021 details={
1022 "field": "criteria",
1023 "criterion_id": criterion_id,
1024 "reason": "criterion_not_relevant_to_directive",
1025 "signal": "val_loss",
1026 },
1027 )
1028 if not mentions_goal and (
1029 event_name == "goal_reached"
1030 or (isinstance(expression, str) and "goal_reached" in expression.casefold())
1031 ):
1032 raise MissionValidationError(
1033 "validation_error",
1034 details={
1035 "field": "criteria",
1036 "criterion_id": criterion_id,
1037 "reason": "criterion_not_relevant_to_directive",
1038 "signal": "goal_reached",
1039 },
1040 )
1043async def generate_sampled_criteria(
1044 backend: SamplingBackend,
1045 directive: str,
1046 *,
1047 allowlist: list[str] | None = None,
1048 max_criteria: int = DEFAULT_MAX_CRITERIA,
1049 retries: int = DEFAULT_RETRIES,
1050) -> list[dict[str, Any]]:
1051 """Drive a sampling backend to produce a validated criteria list.
1053 Builds the prompt, calls ``backend.sample(prompt_str)``, parses
1054 the JSON, validates through :func:`validate_criteria`, and on
1055 rejection retries up to ``retries`` times with feedback. Returns
1056 the validated list (with private ``_parsed_ast`` keys stripped so
1057 the result is JSON-safe). Raises :class:`ScaffoldSamplingError`
1058 when every attempt was rejected, and propagates
1059 :class:`gco.bedrock.BedrockFTUFormNotAcceptedError` unwrapped so a
1060 missing Anthropic first-time-use form is reported rather than
1061 silently downgraded to deterministic criteria.
1063 The backend is duck-typed against the ``SamplingBackend`` protocol
1064 on purpose — tests can substitute a stub object whose ``sample``
1065 method returns canned strings without bringing in a transport.
1066 """
1067 feedback: str | None = None
1068 last_reason = "no_attempts"
1069 # We do retries + 1 total attempts — the first attempt is "free",
1070 # then each retry is one extra try.
1071 for attempt in range(retries + 1):
1072 prompt_str = build_scaffold_prompt(
1073 directive,
1074 allowlist=allowlist,
1075 max_criteria=max_criteria,
1076 feedback=feedback,
1077 )
1078 try:
1079 raw = await _call_backend(backend, prompt_str)
1080 except BedrockFTUFormNotAcceptedError:
1081 # A missing Anthropic FTU form is a permanent misconfiguration, not
1082 # a transport fault. Let it escape the transport-agnostic catch
1083 # below so the caller reports it instead of quietly scaffolding
1084 # deterministic criteria.
1085 raise
1086 except Exception as exc: # transport-agnostic catch
1087 # Transport-layer failures are not retriable from the
1088 # scaffolder's point of view — the backend itself decides
1089 # whether to recover. Surface as a sampling error so the
1090 # CLI falls back deterministically.
1091 raise ScaffoldSamplingError(
1092 "transport_error",
1093 message=f"sampling backend raised {type(exc).__name__}: {exc}",
1094 ) from exc
1095 try:
1096 parsed = _parse_response(raw)
1097 except (ValueError, json.JSONDecodeError) as exc:
1098 last_reason = "json_parse"
1099 feedback = (
1100 "Your previous response could not be parsed as a JSON "
1101 "array. Return a single JSON array, no prose, no "
1102 f"markdown fences. ({exc})"
1103 )
1104 continue
1105 # Cap before normalization so model output can never exceed the
1106 # operator-selected criterion budget. The shared helper then applies
1107 # only conservative, meaning-preserving model-output repairs before
1108 # strict structural validation.
1109 if len(parsed) > max_criteria:
1110 parsed = parsed[:max_criteria]
1111 parsed = _normalize_sampled_criteria(parsed)
1112 try:
1113 validated = _validation.validate_criteria(parsed)
1114 _validate_sampled_criteria_context(
1115 validated,
1116 directive=directive,
1117 allowlist=allowlist,
1118 )
1119 except MissionValidationError as exc:
1120 details = exc.details or {}
1121 last_reason = str(details.get("reason") or exc.code)
1122 feedback = (
1123 "Your previous response was rejected by the validator. "
1124 f"Rejection reason: {last_reason}. Details: {details!r}. "
1125 "Re-emit a corrected JSON array."
1126 )
1127 continue
1128 # Strip private cached AST keys so the JSON written to disk is
1129 # round-trippable. ``_parsed_ast`` is attached to predicate
1130 # entries by ``validate_criteria``.
1131 del attempt
1132 return [
1133 {k: v for k, v in entry.items() if not str(k).startswith("_")} for entry in validated
1134 ]
1135 raise ScaffoldSamplingError(last_reason)
1138async def _call_backend(backend: SamplingBackend, prompt_str: str) -> str:
1139 """Adapt the protocol's ``sample(SamplingPrompt)`` call to a string prompt.
1141 The :class:`SamplingBackend` protocol takes a structured
1142 :class:`SamplingPrompt`. The criteria-scaffold use case is a
1143 one-off prompt rather than a full Mission round-trip, so we
1144 construct a minimal ``SamplingPrompt`` whose render produces
1145 exactly ``prompt_str``. Backends that need extra context
1146 (Bedrock's region, MCP's model preferences) read from their bound
1147 state and ignore the prompt's surrounding fields.
1149 Tests can substitute a stub backend whose ``sample`` returns a
1150 canned string; those tests pass the stub directly to
1151 :func:`generate_sampled_criteria` and bypass the protocol entirely.
1152 """
1153 # Lazy import to avoid the import cycle: sampling imports validation
1154 # which would otherwise import this module.
1155 from .sampling import SamplingPrompt
1157 # Wrap the prompt string in a dataclass that renders to itself.
1158 # The full SamplingPrompt has many required fields; the scaffolder
1159 # uses a thin adapter that overrides ``assemble`` so the existing
1160 # backend implementations call ``assemble()`` and get the prompt.
1161 prompt_obj = _PromptAdapter(prompt_str)
1162 # Backends accept any object with an ``assemble`` method —
1163 # BedrockSamplingBackend calls ``prompt.assemble()`` to get the
1164 # rendered string.
1165 del SamplingPrompt # imported only for documentation linkage
1166 return await backend.sample(prompt_obj) # type: ignore[arg-type]
1169class _PromptAdapter:
1170 """Minimal duck-typed stand-in for :class:`SamplingPrompt`.
1172 Both backends call ``prompt.assemble()`` to render the prompt
1173 string. This adapter satisfies that single contract so the
1174 scaffolder can route a free-form prompt through the same backend
1175 surface the engine uses, without constructing a full
1176 SamplingPrompt with iteration history that does not exist for a
1177 one-off scaffolding call.
1178 """
1180 def __init__(self, text: str) -> None:
1181 self._text: str = text
1183 def assemble(self) -> str:
1184 return self._text