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