Coverage for gco_mcp/mission/sampling.py: 91.42%
478 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""Mission sampling — prompt builders for the advisory LLM path.
3The Mission engine routes optional model-driven advice
4(Strategy_Revision rationales / next-strategy proposals on ``adjust``,
5Final_Report ``lessons`` / ``recommended_followups`` on ``complete`` and
6``terminate``) through a small, transport-agnostic plumbing pipe that
7starts here. This module is the **prompt-assembly half** of that pipe:
8pure Python, sync, no MCP / boto3 / fastmcp imports. Backends, capability
9detection, response validation, and orchestration helpers land in sibling
10sections of this file in subsequent commits.
12The two render methods on :class:`SamplingPrompt` produce a
13deterministic ``str`` payload from the bare data the caller passes in:
15* :meth:`SamplingPrompt.assemble` — the Strategy_Revision prompt. Includes
16 the directive, the Success_Criteria with current per-criterion status,
17 the resolved Tool_Allowlist with each tool's docstring, an explicit
18 budget context block, the last five Iteration summaries (Observation
19 fields larger than :data:`OBSERVATION_FIELD_BYTE_CAP` truncated to
20 :data:`OBSERVATION_FIELD_TRUNCATE_TO` bytes plus the marker
21 :data:`TRUNCATION_MARKER`, with the original byte lengths recorded
22 under the ``_original_bytes`` map), and the JSON Schema instruction
23 block built around :data:`STRATEGY_REVISION_SCHEMA`.
24* :meth:`SamplingPrompt.assemble_final_lessons` — the Final_Report
25 prompt. Reuses the directive / criteria assembly but emits
26 :data:`FINAL_LESSONS_SCHEMA` instead of the strategy-revision schema,
27 and replaces the iteration-by-iteration Observation summaries with a
28 short ``verdict`` / ``verdict_reason`` summary list because the
29 Final_Report path does not need raw Observation history.
31Both render methods cap total output at :data:`PROMPT_BYTE_BUDGET`
32bytes (UTF-8). When the assembled prompt exceeds the cap, the oldest
33Iteration summary is dropped and the prompt re-rendered, repeating
34until the prompt fits. Truncation and dropping are deterministic — the
35same inputs always produce a byte-identical output. This is the
36property the tests under
37``tests/test_mission_sampling.py`` pin down.
38"""
40from __future__ import annotations
42import asyncio
43import json
44import os
45from collections.abc import Mapping, Sequence
46from dataclasses import dataclass, field
47from typing import TYPE_CHECKING, Any, Literal, Protocol, cast, runtime_checkable
49from gco.bedrock import (
50 BEDROCK_READ_TIMEOUT_SECONDS,
51 BedrockResponseTruncatedError,
52 build_bedrock_converse_options,
53 extract_bedrock_converse_text,
54 get_default_bedrock_model_id,
55 raise_if_bedrock_ftu_form_error,
56)
58from . import validation as _validation
59from .types import Criterion, CriterionResult, IterationRecord, Observation, Strategy
60from .validation import MissionValidationError
62# <pyflowchart-code-diagram> BEGIN - auto-inserted, do not edit
63# Generated at (UTC): 2026-07-18T01:03:40Z
64# Flowchart(s) generated from this file:
65# * ``maybe_sample_strategy_revision`` -> ``diagrams/code_diagrams/gco_mcp/mission/sampling.maybe_sample_strategy_revision.html``
66# (PNG: ``diagrams/code_diagrams/gco_mcp/mission/sampling.maybe_sample_strategy_revision.png``)
67# Regenerate with ``python diagrams/code_diagrams/generate.py``.
68# <pyflowchart-code-diagram> END
71if TYPE_CHECKING: # pragma: no cover - import-time only
72 # Runtime access is provided lazily by ``__getattr__`` below.
73 DEFAULT_BEDROCK_MODEL_ID: str
75 # ``fastmcp.Context`` is the concrete type expected by
76 # :class:`MCPSamplingBackend`. Kept behind ``TYPE_CHECKING`` so the
77 # runtime import surface stays pure-stdlib; the backend itself
78 # accepts any object that exposes a compatible ``sample`` method.
79 from fastmcp import Context as _FastMCPContext # noqa: F401
81__all__ = [
82 "BEDROCK_READ_TIMEOUT_SECONDS",
83 "BEDROCK_TEMPERATURE",
84 "DEFAULT_BEDROCK_MODEL_ID",
85 "DEFAULT_BEDROCK_REGION",
86 "ENV_BEDROCK_MODEL_ID",
87 "ENV_BEDROCK_REGION",
88 "ENVIRONMENT_CONTEXT_BYTE_CAP",
89 "FINAL_LESSONS_SCHEMA",
90 "BedrockSamplingBackend",
91 "MCPSamplingBackend",
92 "MissionValidationError",
93 "OBSERVATION_FIELD_BYTE_CAP",
94 "OBSERVATION_FIELD_TRUNCATE_TO",
95 "PROMPT_BYTE_BUDGET",
96 "RECENT_ITERATIONS_LIMIT",
97 "STRATEGY_REVISION_SCHEMA",
98 "STRATEGY_SHAPE_SCHEMA",
99 "SamplingBackend",
100 "SamplingFallback",
101 "SamplingPrompt",
102 "SamplingTransportError",
103 "SamplingUsed",
104 "TRUNCATION_MARKER",
105 "maybe_sample_final_lessons",
106 "maybe_sample_strategy_revision",
107 "resolve_sampling_state",
108 "select_sampling_backend",
109 "validate_strategy_against_catalog",
110]
113def __getattr__(name: str) -> Any:
114 """Resolve the historical Mission default only when explicitly accessed."""
115 if name == "DEFAULT_BEDROCK_MODEL_ID":
116 return get_default_bedrock_model_id()
117 raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
120def __dir__() -> list[str]:
121 """Advertise the lazy compatibility alias to introspection tools."""
122 return sorted({*globals(), "DEFAULT_BEDROCK_MODEL_ID"})
125# ---------------------------------------------------------------------------
126# Tunables (named so tests can reference them without hard-coding magic)
127# ---------------------------------------------------------------------------
129#: Per-Observation-field byte cap. Fields whose JSON-serialised UTF-8
130#: byte length exceeds this value are truncated. A field whose byte
131#: length is exactly equal to this value is **not** truncated — the
132#: comparison uses strict greater-than to keep the boundary stable.
133OBSERVATION_FIELD_BYTE_CAP: int = 4096
135#: Target byte length after truncation. The truncated string is the
136#: first ``OBSERVATION_FIELD_TRUNCATE_TO`` bytes of the JSON-serialised
137#: form, decoded with ``errors="ignore"`` so a multi-byte boundary in
138#: the middle of a UTF-8 codepoint cannot raise, with
139#: :data:`TRUNCATION_MARKER` appended.
140OBSERVATION_FIELD_TRUNCATE_TO: int = 2048
142#: Marker appended to every truncated field so the reader can see at a
143#: glance the field was clipped.
144TRUNCATION_MARKER: str = "... [truncated]"
146#: Total prompt byte budget. The render methods drop the oldest
147#: Iteration summary one at a time until ``len(prompt.encode("utf-8"))
148#: <= PROMPT_BYTE_BUDGET``.
149PROMPT_BYTE_BUDGET: int = 32768
151#: Maximum number of Iteration summaries to include even if the byte
152#: budget is plentiful. The caller is expected to pass at most this
153#: many already; the builder slices defensively.
154RECENT_ITERATIONS_LIMIT: int = 5
156#: Per-Environment-context byte cap. The optional environment context
157#: block (``=== Environment context ===``) is its own truncation
158#: domain so the section can never grow without bound and push the
159#: rest of the prompt over :data:`PROMPT_BYTE_BUDGET`. The cap mirrors
160#: :data:`OBSERVATION_FIELD_BYTE_CAP` because both surfaces hold the
161#: same flavour of structured live signal (cluster + queue snapshots
162#: in this case) and the same truncation marker convention applies.
163ENVIRONMENT_CONTEXT_BYTE_CAP: int = 4096
166# ---------------------------------------------------------------------------
167# Environment context summarisation
168# ---------------------------------------------------------------------------
171def _summarise_environment_context(env: Mapping[str, Any]) -> dict[str, Any]:
172 """Return a JSON-safe, byte-capped summary of the environment context.
174 The block is rendered into the Strategy_Revision prompt under
175 ``=== Environment context ===``. It carries small, slow-moving
176 live signals — per-region queue depths, GPU utilisation, deployed
177 region list, reservation counts — that the model would otherwise
178 have to spend tool calls to discover.
180 Two guarantees on the output:
182 1. The serialised form fits inside :data:`ENVIRONMENT_CONTEXT_BYTE_CAP`
183 UTF-8 bytes. When the input does not, top-level fields are
184 evaluated in sorted-key order, dropped one at a time from the
185 largest contributor down, and the dropped key list is recorded
186 under ``"_dropped_fields"`` so the operator can spot which
187 inputs got pruned.
188 2. Top-level keys are emitted in sorted order so two callers
189 passing semantically-identical dicts produce a byte-identical
190 block — the same property the determinism tests pin down for
191 Observation summaries.
192 """
193 # Defensive copy + sort so insertion order doesn't leak.
194 ordered: dict[str, Any] = {key: env[key] for key in sorted(env.keys())}
195 serialised = _dumps(ordered)
196 if _utf8_len(serialised) <= ENVIRONMENT_CONTEXT_BYTE_CAP:
197 return ordered
199 # Drop largest top-level field first, repeating until under cap.
200 # Records dropped keys so the operator (and the audit pipeline)
201 # can see what got pruned without having to diff against the
202 # gather helper's output.
203 dropped: list[str] = []
204 working = dict(ordered)
205 while _utf8_len(_dumps(working)) > ENVIRONMENT_CONTEXT_BYTE_CAP and working:
206 biggest_key = max(working, key=lambda k: _utf8_len(_dumps(working[k])))
207 dropped.append(biggest_key)
208 del working[biggest_key]
210 if dropped: 210 ↛ 214line 210 didn't jump to line 214 because the condition on line 210 was always true
211 # Sort the dropped list so its position in the prompt is stable
212 # regardless of which key happened to be biggest first.
213 working["_dropped_fields"] = sorted(dropped)
214 return working
217# ---------------------------------------------------------------------------
218# Bedrock backend tunables
219# ---------------------------------------------------------------------------
221#: Default Bedrock model identifier, loaded lazily from
222#: ``cdk.json`` ``context.bedrock.default_model_id`` through the lightweight
223#: :mod:`gco.bedrock` resolver. The module-level compatibility attribute keeps
224#: existing Mission integrations stable without coupling unrelated imports to
225#: Bedrock configuration resolution.
226#:
227#: Operators with regulatory or model-governance requirements can override per
228#: call via ``GCO_MISSION_BEDROCK_MODEL_ID`` or ``--bedrock-model-id``; see
229#: docs/CUSTOMIZATION.md ("Bedrock Model Selection").
231#: Default Bedrock region. The capacity advisor pins ``us-east-1`` for
232#: the same reason: cross-region inference profiles routinely surface
233#: in ``us-east-1`` first and our installations have it whitelisted.
234DEFAULT_BEDROCK_REGION: str = "us-east-1"
236#: Env var that overrides :data:`DEFAULT_BEDROCK_MODEL_ID` at runtime.
237ENV_BEDROCK_MODEL_ID: str = "GCO_MISSION_BEDROCK_MODEL_ID"
239#: Env var that overrides :data:`DEFAULT_BEDROCK_REGION` at runtime.
240ENV_BEDROCK_REGION: str = "GCO_MISSION_BEDROCK_REGION"
242#: Sampling temperature for non-default Bedrock model overrides. The
243#: canonical Nova 2 default uses high reasoning, for which AWS requires
244#: temperature to be unset; :func:`build_bedrock_converse_options` removes it.
245BEDROCK_TEMPERATURE: float = 0.2
248# ---------------------------------------------------------------------------
249# JSON Schemas — embedded as module-level constants
250# ---------------------------------------------------------------------------
252# The Strategy shape mirrors the ``Strategy`` TypedDict from
253# ``gco_mcp/mission/types.py``: every key is optional in isolation and the
254# validator enforces the mutual-exclusivity invariant (exactly one of
255# ``tool_calls`` or ``script`` populated). The schema below mirrors
256# that with a ``oneOf`` clause. The model-side validator in subsequent
257# commits performs the same check on parsed responses; this schema is
258# the textual instruction the prompt embeds for the model.
259STRATEGY_SHAPE_SCHEMA: dict[str, Any] = {
260 "type": "object",
261 "additionalProperties": False,
262 "properties": {
263 "tool_calls": {
264 "type": "array",
265 "minItems": 1,
266 "items": {
267 "type": "object",
268 "additionalProperties": True,
269 "required": ["tool_name", "args"],
270 "properties": {
271 "tool_name": {"type": "string", "minLength": 1},
272 "args": {"type": "object"},
273 },
274 },
275 },
276 "script": {"type": "string", "minLength": 1},
277 "expected_observation_keys": {
278 "type": "array",
279 "items": {"type": "string"},
280 },
281 "rationale": {"type": "string"},
282 },
283 "oneOf": [
284 {"required": ["tool_calls"]},
285 {"required": ["script"]},
286 ],
287}
289#: JSON Schema for the model's response when called for a
290#: Strategy_Revision. The model must return exactly these three keys.
291STRATEGY_REVISION_SCHEMA: dict[str, Any] = {
292 "$schema": "http://json-schema.org/draft-07/schema#",
293 "title": "Mission strategy revision",
294 "type": "object",
295 "additionalProperties": False,
296 "required": ["revision_rationale", "next_strategy", "confidence"],
297 "properties": {
298 "revision_rationale": {"type": "string", "minLength": 1},
299 "next_strategy": STRATEGY_SHAPE_SCHEMA,
300 "confidence": {
301 "type": "number",
302 "minimum": 0.0,
303 "maximum": 1.0,
304 },
305 },
306}
308#: JSON Schema for the model's response when called from the
309#: Final_Report writer.
310FINAL_LESSONS_SCHEMA: dict[str, Any] = {
311 "$schema": "http://json-schema.org/draft-07/schema#",
312 "title": "Mission final lessons",
313 "type": "object",
314 "additionalProperties": False,
315 "required": ["lessons", "recommended_followups"],
316 "properties": {
317 "lessons": {
318 "type": "array",
319 "minItems": 1,
320 "items": {"type": "string", "minLength": 1},
321 },
322 "recommended_followups": {
323 "type": "array",
324 "items": {"type": "string", "minLength": 1},
325 },
326 },
327}
330# ---------------------------------------------------------------------------
331# JSON helpers — every dump in this module routes through ``_dumps``
332# so the byte-counting and the rendered prompt agree on the encoding.
333# ---------------------------------------------------------------------------
336def _dumps(value: Any, *, indent: int | None = None) -> str:
337 """Deterministic JSON encoder used everywhere in this module.
339 ``sort_keys=True`` is the source of determinism — Python dicts are
340 insertion-ordered, but the Hypothesis strategies that drive the
341 determinism tests build dicts via ``fixed_dictionaries`` whose
342 insertion order is implementation-defined, so sorting is the only
343 way to get byte-identical output across two draws of the same
344 abstract dict shape. ``ensure_ascii=False`` keeps non-ASCII text
345 intact so the byte-budget bookkeeping matches what the LLM sees.
346 """
347 return json.dumps(
348 value,
349 sort_keys=True,
350 ensure_ascii=False,
351 indent=indent,
352 separators=(",", ": ") if indent is not None else (",", ":"),
353 )
356def _utf8_len(s: str) -> int:
357 """UTF-8 byte length of ``s`` — the only "size" the budget cares about."""
358 return len(s.encode("utf-8"))
361def _truncate_serialised(serialised: str) -> str:
362 """Slice a serialised value down to ``OBSERVATION_FIELD_TRUNCATE_TO``
363 bytes plus :data:`TRUNCATION_MARKER`. Decode-safe.
365 The slicing is byte-level rather than codepoint-level because the
366 cap itself is a byte budget. Using ``errors="ignore"`` strips any
367 partial codepoint at the boundary so the result is always valid
368 UTF-8 — at the cost of dropping at most three bytes' worth of an
369 incomplete codepoint, which is acceptable for an advisory summary.
370 """
371 truncated_bytes = serialised.encode("utf-8")[:OBSERVATION_FIELD_TRUNCATE_TO]
372 truncated_str = truncated_bytes.decode("utf-8", errors="ignore")
373 return truncated_str + TRUNCATION_MARKER
376# ---------------------------------------------------------------------------
377# Observation summarisation
378# ---------------------------------------------------------------------------
381def _summarise_observation(obs: Mapping[str, Any] | Observation) -> dict[str, Any]:
382 """Return a JSON-safe summary of an Observation with oversized fields
383 truncated and the original byte lengths recorded.
385 The summary mirrors the Observation's top-level keys. For each key
386 whose JSON-serialised value exceeds :data:`OBSERVATION_FIELD_BYTE_CAP`
387 bytes, the value is replaced by the byte-clamped + marker string and
388 the original byte length is recorded under
389 ``summary["_original_bytes"][<key>]``. Fields at or below the cap pass
390 through unchanged.
392 The ``_original_bytes`` private key is omitted entirely when no field
393 was truncated so the summary stays clean for the common case.
394 """
395 obs_map: Mapping[str, Any] = cast("Mapping[str, Any]", obs)
396 summary: dict[str, Any] = {}
397 original_bytes: dict[str, int] = {}
398 # Sorting the keys guarantees the rendered prompt is byte-identical
399 # even when the caller's dict was built in a different insertion
400 # order than another caller's identical-shape dict.
401 for key in sorted(obs_map.keys()):
402 if key == "_original_bytes": 402 ↛ 406line 402 didn't jump to line 406 because the condition on line 402 was never true
403 # A defensively-guarded passthrough: a previous summarisation
404 # round (e.g., a re-render after dropping iterations) must
405 # not double-count the marker map.
406 continue
407 value = obs_map[key]
408 serialised = _dumps(value)
409 n_bytes = _utf8_len(serialised)
410 if n_bytes > OBSERVATION_FIELD_BYTE_CAP:
411 summary[key] = _truncate_serialised(serialised)
412 original_bytes[key] = n_bytes
413 else:
414 summary[key] = value
415 if original_bytes:
416 summary["_original_bytes"] = original_bytes
417 return summary
420def _summarise_iteration(iteration: Mapping[str, Any] | IterationRecord) -> dict[str, Any]:
421 """Build the per-iteration summary that feeds the Strategy_Revision prompt.
423 The summary keeps just the fields a downstream model needs to
424 reason about: the iteration index, the strategy that was tried,
425 the verdict + reason, and the size-capped Observation. Phase
426 timestamps and the criteria-evaluation list are intentionally
427 omitted because a) they are deterministic functions of fields the
428 model already sees in the criteria-status block, and b) keeping
429 them out shrinks the per-iteration footprint so the byte budget
430 holds with five iterations more often.
431 """
432 obs = iteration.get("observation") or {}
433 return {
434 "iteration_index": iteration.get("iteration_index"),
435 "strategy": iteration.get("strategy") or {},
436 "verdict": iteration.get("verdict"),
437 "verdict_reason": iteration.get("verdict_reason"),
438 "observation_summary": _summarise_observation(obs),
439 }
442def _summarise_iteration_for_lessons(
443 iteration: Mapping[str, Any] | IterationRecord,
444) -> dict[str, Any]:
445 """Final_Report summary — verdict + reason only, no Observation.
447 The Final_Report path needs to reason about *what happened* across
448 the run, not the per-iteration tool output. Dropping the Observation
449 keeps the prompt small enough that the byte budget never bites in
450 practice for sessions of any reasonable length.
451 """
452 return {
453 "iteration_index": iteration.get("iteration_index"),
454 "verdict": iteration.get("verdict"),
455 "verdict_reason": iteration.get("verdict_reason"),
456 }
459# ---------------------------------------------------------------------------
460# Criteria status pairing
461# ---------------------------------------------------------------------------
464def _pair_criteria_with_status(
465 criteria: Sequence[Criterion],
466 statuses: Sequence[CriterionResult],
467) -> list[dict[str, Any]]:
468 """Return ``criteria`` annotated with their most recent status entry.
470 Each entry in the result mirrors the Criterion definition (the kind,
471 the required-flag, the kind-specific payload keys) and adds a
472 nested ``status`` block populated from the matching ``CriterionResult``
473 by ``criterion_id``. Criteria with no matching status entry get
474 ``status`` set to ``{"status": "inconclusive", "evidence": null}``
475 so the model always sees a stable shape.
477 The ``_parsed_ast`` private key on a ``predicate`` criterion is
478 stripped — it is a Python ``ast.Expression`` object that is not
479 JSON-serialisable and that the model has no use for.
480 """
481 by_id: dict[str, CriterionResult] = {}
482 for s in statuses:
483 cid = s.get("criterion_id")
484 if cid is None: 484 ↛ 485line 484 didn't jump to line 485 because the condition on line 484 was never true
485 continue
486 # If the caller passes duplicates (older then newer), prefer the
487 # last entry — that's the most-recent-wins convention the engine
488 # uses everywhere else.
489 by_id[cid] = s
491 out: list[dict[str, Any]] = []
492 for c in criteria:
493 cid = c.get("criterion_id")
494 # Strip private cached AST and surface only the prompt-relevant fields.
495 public = {k: v for k, v in c.items() if not k.startswith("_")}
496 match = by_id.get(cid) if cid is not None else None
497 if match is None:
498 public["status"] = {
499 "status": "inconclusive",
500 "evidence": None,
501 }
502 else:
503 public["status"] = {
504 "status": match.get("status"),
505 "evidence": match.get("evidence"),
506 "evaluated_at": match.get("evaluated_at"),
507 }
508 out.append(public)
509 return out
512# ---------------------------------------------------------------------------
513# Tool allowlist rendering
514# ---------------------------------------------------------------------------
517def _render_tool_allowlist(
518 allowlist: Sequence[str],
519 docstrings: Mapping[str, str],
520 schemas: Mapping[str, Any] | None = None,
521) -> list[dict[str, Any]]:
522 """Pair every allowlisted tool name with its docstring and input schema.
524 Tools without a registered docstring get an empty string — the
525 prompt remains valid; the model just sees a tool name with no
526 inline description. Tools without a schema get ``null`` so the
527 model knows no args are required. Names are emitted in the
528 caller's allowlist order so the prompt is identical for two
529 callers that pass the same list.
530 """
531 rendered: list[dict[str, Any]] = []
532 for name in allowlist:
533 entry: dict[str, Any] = {
534 "tool_name": name,
535 "docstring": str(docstrings.get(name, "")),
536 }
537 if schemas:
538 schema = schemas.get(name)
539 if schema is not None: 539 ↛ 541line 539 didn't jump to line 541 because the condition on line 539 was always true
540 entry["input_schema"] = schema
541 rendered.append(entry)
542 return rendered
545# ---------------------------------------------------------------------------
546# Budget context rendering
547# ---------------------------------------------------------------------------
550def _render_budget_context(
551 *,
552 remaining_iterations: int,
553 remaining_wall_clock_secs: float | None,
554 allow_scripts: bool,
555) -> dict[str, Any]:
556 """Render the budget context block. Stable shape regardless of inputs.
558 ``None`` for the wall-clock cap is rendered verbatim as JSON
559 ``null`` so the model can disambiguate "unbounded" from "0".
560 """
561 return {
562 "remaining_iterations": int(remaining_iterations),
563 "remaining_wall_clock_seconds": (
564 float(remaining_wall_clock_secs) if remaining_wall_clock_secs is not None else None
565 ),
566 "allow_scripted_strategies": bool(allow_scripts),
567 }
570# ---------------------------------------------------------------------------
571# SamplingPrompt — the public class
572# ---------------------------------------------------------------------------
575@dataclass(frozen=True)
576class SamplingPrompt:
577 """Bundle the bare data needed to assemble a sampling prompt string.
579 The dataclass is ``frozen=True`` so callers cannot mutate the inputs
580 between an :meth:`assemble` call and an :meth:`assemble_final_lessons`
581 call — both methods produce deterministic outputs from the same
582 bound state, which is the property the determinism tests pin down.
584 All inputs are required positionally or by keyword; defaults are
585 only provided where the design spec defines a default.
586 """
588 directive: str
589 success_criteria: Sequence[Criterion]
590 criteria_status: Sequence[CriterionResult]
591 recent_iterations: Sequence[IterationRecord]
592 tool_allowlist: Sequence[str]
593 tool_docstrings: Mapping[str, str]
594 remaining_iterations: int
595 remaining_wall_clock_secs: float | None
596 allow_scripts: bool = field(default=False)
597 #: Per-tool JSON Schema for the input parameters. Keyed by tool
598 #: name; values are the JSON-serialisable schema dict (or ``None``
599 #: for tools that take no args). Included in the prompt so the
600 #: Strategy_Revision model can propose valid ``args`` dicts.
601 tool_schemas: Mapping[str, Any] = field(default_factory=dict)
602 #: Optional snapshot of slow-moving live signals (per-region queue
603 #: depth, GPU utilisation, deployed-region list, reservation
604 #: counts, etc.) gathered once at session start and reused on
605 #: every iteration's prompt. ``None`` (the default) suppresses the
606 #: ``=== Environment context ===`` section entirely so the prompt
607 #: stays byte-identical to the pre-environment-context shape —
608 #: that's what every existing determinism test pins down.
609 environment_context: Mapping[str, Any] | None = field(default=None)
611 # ---- Strategy_Revision rendering --------------------------------------
613 def assemble(self) -> str:
614 """Return the Strategy_Revision prompt string.
616 The output is capped at :data:`PROMPT_BYTE_BUDGET` UTF-8 bytes.
617 When the freshly-assembled prompt exceeds the cap, the oldest
618 Iteration summary is dropped and the prompt re-rendered. The
619 loop terminates because each drop monotonically shrinks the
620 prompt and there is a non-iteration baseline that fits well
621 under the cap on its own (the directive, criteria, allowlist,
622 budget block, and schema instruction together are ~6-10 KB
623 for any reasonable session shape).
624 """
625 # Defensive slice — the caller is asked to pass at most five,
626 # but if they pass more, take the most recent five.
627 iterations: list[IterationRecord] = list(self.recent_iterations[-RECENT_ITERATIONS_LIMIT:])
629 while True:
630 text = self._render(
631 schema=STRATEGY_REVISION_SCHEMA,
632 iterations=[_summarise_iteration(it) for it in iterations],
633 schema_purpose="strategy_revision",
634 )
635 if _utf8_len(text) <= PROMPT_BYTE_BUDGET or not iterations:
636 return text
637 # Drop the oldest iteration and try again.
638 iterations = iterations[1:]
640 # ---- Final_Report rendering -------------------------------------------
642 def assemble_final_lessons(self) -> str:
643 """Return the Final_Report ``lessons`` prompt string.
645 The shape parallels :meth:`assemble` but emits
646 :data:`FINAL_LESSONS_SCHEMA` and uses iteration **verdict
647 summaries only** instead of full Observation summaries. The same
648 :data:`PROMPT_BYTE_BUDGET` byte cap applies; the same
649 oldest-first drop policy kicks in if the cap is exceeded.
650 """
651 iterations: list[IterationRecord] = list(self.recent_iterations)
653 while True:
654 text = self._render(
655 schema=FINAL_LESSONS_SCHEMA,
656 iterations=[_summarise_iteration_for_lessons(it) for it in iterations],
657 schema_purpose="final_lessons",
658 )
659 if _utf8_len(text) <= PROMPT_BYTE_BUDGET or not iterations: 659 ↛ 661line 659 didn't jump to line 661 because the condition on line 659 was always true
660 return text
661 iterations = iterations[1:]
663 # ---- Internal renderer ------------------------------------------------
665 def _render(
666 self,
667 *,
668 schema: dict[str, Any],
669 iterations: Sequence[Mapping[str, Any]],
670 schema_purpose: str,
671 ) -> str:
672 """Format the full prompt from the section blocks.
674 The text layout is fixed — every section is delimited by a
675 ``=== <name> ===`` header so the model can latch onto a
676 predictable structure. Section bodies are JSON wherever the
677 content is structured; the directive itself is rendered as
678 free text because that is how the operator wrote it.
679 """
680 criteria_block = _pair_criteria_with_status(self.success_criteria, self.criteria_status)
681 tool_block = _render_tool_allowlist(
682 self.tool_allowlist, self.tool_docstrings, self.tool_schemas
683 )
684 budget_block = _render_budget_context(
685 remaining_iterations=self.remaining_iterations,
686 remaining_wall_clock_secs=self.remaining_wall_clock_secs,
687 allow_scripts=self.allow_scripts,
688 )
690 if schema_purpose == "strategy_revision":
691 preamble = (
692 "You are advising a Mission goal-directed iteration loop. "
693 "Propose the next Strategy that moves the Mission toward "
694 "satisfying its Success_Criteria. The Verdict label, "
695 "budget enforcement, and Criteria evaluation are all "
696 "computed server-side and are unaffected by your output. "
697 "Your role is advisory: the rationale and next_strategy "
698 "you produce are validated against the Tool_Allowlist and "
699 "the remaining budget before being adopted.\n\n"
700 "IMPORTANT: You may propose MULTIPLE tool calls in a "
701 "single iteration by including multiple entries in the "
702 "tool_calls array. This is especially useful when the "
703 "unmet criteria require results from different tools — "
704 "calling them all in one iteration lets the evaluator "
705 "see all results together. Use the input_schema in the "
706 "Tool allowlist section to construct valid args for each "
707 "tool call."
708 )
709 recent_header = "Recent iterations (oldest first)"
710 else:
711 preamble = (
712 "You are advising a Mission goal-directed iteration loop "
713 "that has just reached a terminal Verdict. Produce the "
714 "lessons learned and the recommended follow-ups for the "
715 "operator. Your output is merged into the Final_Report; "
716 "the Verdict label, budget bookkeeping, and Criteria "
717 "evaluation that produced the terminal state are "
718 "deterministic server-side outputs and are not under "
719 "review."
720 )
721 recent_header = "Iteration verdict summary (oldest first)"
723 sections: list[str] = []
724 sections.append(preamble)
725 sections.append("")
726 sections.append("=== Mission directive ===")
727 sections.append(self.directive)
728 sections.append("")
729 sections.append("=== Success criteria with current status ===")
730 sections.append(_dumps(criteria_block, indent=2))
731 sections.append("")
732 sections.append("=== Tool allowlist ===")
733 sections.append(_dumps(tool_block, indent=2))
734 sections.append("")
735 sections.append("=== Budget context ===")
736 sections.append(_dumps(budget_block, indent=2))
737 sections.append("")
738 if self.environment_context is not None:
739 # Truncated + key-sorted — see :func:`_summarise_environment_context`.
740 # Emitting a header even for an empty dict means a session
741 # that opted in but had a probe failure still surfaces
742 # "we tried" so the operator can act on the gap.
743 env_summary = _summarise_environment_context(self.environment_context)
744 sections.append("=== Environment context (slow-moving live signals) ===")
745 sections.append(_dumps(env_summary, indent=2))
746 sections.append("")
747 sections.append(f"=== {recent_header} ===")
748 sections.append(_dumps(list(iterations), indent=2))
749 sections.append("")
750 sections.append("=== Output schema ===")
751 sections.append(
752 "Respond with a single JSON object that validates against "
753 "the JSON Schema below. Do not include any prose outside "
754 "the JSON object."
755 )
756 sections.append(_dumps(schema, indent=2))
758 return "\n".join(sections)
761# ---------------------------------------------------------------------------
762# Backend protocol and transport-error type
763# ---------------------------------------------------------------------------
766@runtime_checkable
767class SamplingBackend(Protocol):
768 """Transport-agnostic surface for the advisory LLM call.
770 Implementations bind a concrete transport (e.g., the MCP
771 ``Context.sample`` capability or ``bedrock-runtime:Converse``) and
772 expose a single async ``sample`` entry point. The protocol is
773 ``runtime_checkable`` so call sites — and tests — can use
774 ``isinstance(backend, SamplingBackend)`` to gate dispatch on a
775 duck-typed backend instance.
777 Attributes:
778 backend_name: Stable identifier the audit pipeline emits in the
779 ``sampling_backend`` field. Constrained to the two
780 transports the system supports today.
781 model_id: The concrete model identifier the backend will route
782 the prompt to. Echoed in audit events so replay can
783 reproduce the exact request.
784 """
786 backend_name: Literal["mcp", "bedrock"]
787 model_id: str
789 async def sample(self, prompt: SamplingPrompt) -> str:
790 """Render ``prompt`` through the bound transport and return the
791 raw model output text. Implementations raise
792 :class:`SamplingTransportError` (with a transport-tagged
793 ``code``) on any transport-layer failure so the engine's
794 fallback policy can branch on a single, well-typed exception.
795 """
796 ...
799class SamplingTransportError(Exception):
800 """Transport-layer failure raised by a :class:`SamplingBackend`.
802 The mandatory ``code`` attribute tags the failure class so the
803 engine's deterministic-fallback path can branch on a stable string
804 without parsing the message. The convention is
805 ``"<backend>_<error_class>"`` for backend-specific failures and a
806 short, snake-cased label for backend-agnostic failures.
808 Documented codes (used elsewhere in the Mission stack):
810 * ``"bedrock_AccessDeniedException"`` — IAM denied
811 ``bedrock:InvokeModel`` for the resolved model.
812 * ``"mcp_unavailable"`` — the MCP transport reported no sampling
813 capability or the in-flight call was cancelled.
814 * ``"bedrock_malformed_response"`` — Converse returned a payload
815 that did not have the expected ``output.message.content[0].text``
816 shape.
817 * ``"bedrock_truncated_response"`` — the answer was cut off by an
818 output-token limit (``stopReason == "max_tokens"``), so its text
819 cannot be trusted to be complete.
820 * ``"bedrock_no_credentials"`` — the local ``boto3`` session could
821 not resolve credentials.
823 Args:
824 code: Mandatory failure tag (see examples above).
825 message: Optional human-readable detail. When present, it is
826 joined to ``code`` with ``": "`` for the string
827 representation; when absent, ``str(self)`` is just the
828 ``code``.
829 """
831 def __init__(self, code: str, message: str | None = None) -> None:
832 self.code: str = code
833 self.message: str | None = message
834 # Forward the most useful single-line representation to
835 # ``Exception.__init__`` so ``logging`` / ``traceback`` modules
836 # show the same string ``str(self)`` produces below.
837 if message is None:
838 super().__init__(code)
839 else:
840 super().__init__(f"{code}: {message}")
842 def __str__(self) -> str:
843 if self.message is None:
844 return self.code
845 return f"{self.code}: {self.message}"
848# ---------------------------------------------------------------------------
849# MCPSamplingBackend — routes the prompt through a FastMCP Context
850# ---------------------------------------------------------------------------
853class MCPSamplingBackend:
854 """Sampling backend that calls ``ctx.sample`` on a FastMCP-style Context.
856 The constructor accepts any object exposing an awaitable ``sample``
857 method — typically ``fastmcp.Context``. The type is intentionally
858 duck-typed so ``fastmcp`` is not a runtime import requirement of
859 this module.
861 Two FastMCP API quirks are absorbed here:
863 1. ``ctx.sample`` returns either a bare string or an object with a
864 ``.text`` attribute, depending on the FastMCP version. Both
865 shapes are accepted; anything else surfaces as
866 :class:`SamplingTransportError` with code
867 ``"mcp_unexpected_response_type"``.
868 2. The keyword carrying ``ModelPreferences`` has been spelled both
869 ``modelPreferences`` (camelCase, MCP wire format) and
870 ``model_preferences`` (snake_case, older Python binding). The
871 backend tries the camelCase form first and falls back to the
872 snake_case form on a ``TypeError`` so it works against either
873 FastMCP release without pinning a version.
875 Any other exception from the transport is re-raised as
876 :class:`SamplingTransportError` with code
877 ``"mcp_<ExceptionClassName>"`` and the original exception preserved
878 via ``__cause__``.
879 """
881 backend_name: Literal["mcp", "bedrock"] = "mcp"
883 def __init__(
884 self,
885 ctx: Any,
886 model_id: str | None = None,
887 prefs: dict[str, Any] | None = None,
888 ) -> None:
889 """Bind the Context, the optional model id, and the preferences dict.
891 Args:
892 ctx: A FastMCP ``Context`` (or any duck-compatible object
893 exposing an awaitable ``sample(text, **kwargs)``
894 method). Stored verbatim.
895 model_id: Optional concrete model identifier. Echoed in
896 audit events. Defaults to the empty string when not
897 provided so the attribute is always present.
898 prefs: Optional FastMCP ``ModelPreferences`` payload. Passed
899 straight through to ``ctx.sample`` when not ``None``;
900 omitted from the call entirely when ``None``. The
901 payload is not validated here — that is the transport's
902 job.
903 """
904 self._ctx = ctx
905 self.model_id: str = model_id if model_id is not None else ""
906 self._prefs = prefs
908 async def sample(self, prompt: SamplingPrompt) -> str:
909 """Render ``prompt`` through ``ctx.sample`` and return the raw text.
911 Raises:
912 SamplingTransportError: On any transport-level failure or
913 when the transport returns an unexpected response shape.
914 The original exception is chained via ``__cause__``.
915 """
916 text = prompt.assemble()
917 try:
918 if self._prefs is None:
919 result = await self._ctx.sample(text)
920 else:
921 # Compatibility shim: try MCP-spec camelCase first; on a
922 # signature mismatch, fall back to the snake_case form.
923 # The ``TypeError`` here is a binding-version concern,
924 # not a transport failure, so it is NOT translated to a
925 # SamplingTransportError.
926 try:
927 result = await self._ctx.sample(text, modelPreferences=self._prefs)
928 except TypeError:
929 result = await self._ctx.sample(text, model_preferences=self._prefs)
931 if hasattr(result, "text"):
932 return str(result.text)
933 if isinstance(result, str):
934 return result
935 raise SamplingTransportError(
936 "mcp_unexpected_response_type",
937 message=f"got {type(result).__name__}",
938 )
939 except SamplingTransportError:
940 # Already tagged by us — let it propagate untouched so the
941 # ``code`` attribute is preserved.
942 raise
943 except Exception as err: # noqa: BLE001 - intentional broad catch
944 raise SamplingTransportError(f"mcp_{type(err).__name__}") from err
947# ---------------------------------------------------------------------------
948# BedrockSamplingBackend — routes the prompt through bedrock-runtime:Converse
949# ---------------------------------------------------------------------------
952class BedrockSamplingBackend:
953 """Sampling backend that calls ``bedrock-runtime:Converse``.
955 The backend resolves its model id and region at construction time
956 from (in order of precedence) the explicit constructor argument,
957 the matching environment variable
958 (:data:`ENV_BEDROCK_MODEL_ID` / :data:`ENV_BEDROCK_REGION`), and
959 finally the shared ``cdk.json`` default
960 (:data:`DEFAULT_BEDROCK_MODEL_ID` /
961 :data:`DEFAULT_BEDROCK_REGION`). The ``boto3`` client itself is
962 constructed lazily on the first :meth:`sample` call so that
963 ``import mission.sampling`` does not pull ``boto3`` into the
964 import graph and so that test code can swap the import in via
965 ``unittest.mock.patch`` without paying for a real session at
966 construction time.
968 Failure modes:
970 * Missing or partial AWS credentials at client-construction time
971 surface as :class:`SamplingTransportError` with code
972 ``"bedrock_no_credentials"``; the original exception is chained
973 via ``__cause__``.
974 * A ``botocore.exceptions.ClientError`` from the ``Converse`` call
975 surfaces as :class:`SamplingTransportError` with code
976 ``"bedrock_<ErrorCode>"`` where ``<ErrorCode>`` is read from the
977 error envelope (defaulting to ``"Unknown"`` when the envelope is
978 malformed). The one exception is the Anthropic first-time-use
979 gate, which raises
980 :class:`gco.bedrock.BedrockFTUFormNotAcceptedError` instead of a
981 transport error so it is never absorbed by a deterministic
982 fallback. See ``docs/CUSTOMIZATION.md`` (Bedrock Model Selection).
983 * A response without a non-empty ``text`` block under
984 ``output.message.content`` — including reasoning-only and empty
985 ``content`` lists — surfaces as :class:`SamplingTransportError`
986 with code ``"bedrock_malformed_response"``.
987 * A response cut off by an output-token limit
988 (``stopReason == "max_tokens"``) surfaces as
989 :class:`SamplingTransportError` with code
990 ``"bedrock_truncated_response"``.
991 """
993 backend_name: Literal["mcp", "bedrock"] = "bedrock"
995 def __init__(
996 self,
997 model_id: str | None = None,
998 region: str | None = None,
999 ) -> None:
1000 """Resolve the model id and region; defer client construction.
1002 Args:
1003 model_id: Optional explicit model id. When ``None``, falls
1004 back to the :data:`ENV_BEDROCK_MODEL_ID` environment
1005 variable, then to the shared ``cdk.json`` default exposed as
1006 :data:`DEFAULT_BEDROCK_MODEL_ID`.
1007 region: Optional explicit region. When ``None``, falls back
1008 to :data:`ENV_BEDROCK_REGION`, then to
1009 :data:`DEFAULT_BEDROCK_REGION`.
1010 """
1011 if model_id is not None:
1012 self.model_id = model_id
1013 self._uses_default_model = False
1014 elif ENV_BEDROCK_MODEL_ID in os.environ:
1015 # Preserve the existing explicit-environment semantics, including
1016 # an intentionally empty value, without evaluating the fallback.
1017 self.model_id = os.environ[ENV_BEDROCK_MODEL_ID]
1018 self._uses_default_model = False
1019 else:
1020 self.model_id = get_default_bedrock_model_id()
1021 self._uses_default_model = True
1022 self._region: str = (
1023 region
1024 if region is not None
1025 else os.environ.get(ENV_BEDROCK_REGION, DEFAULT_BEDROCK_REGION)
1026 )
1027 # The boto3 client is built on first ``sample`` call. ``None``
1028 # here is the sentinel for "not yet constructed".
1029 self._client: Any = None
1031 @classmethod
1032 def from_canonical_default(
1033 cls,
1034 region: str | None = None,
1035 ) -> BedrockSamplingBackend:
1036 """Build a backend that deliberately applies canonical reasoning.
1038 Unlike ``cls(model_id=None)``, this bypasses the model environment
1039 override. Fixture capture uses it to reproduce the checked-in default
1040 exactly, while ordinary explicit model IDs retain override semantics.
1041 """
1042 backend = cls(model_id=get_default_bedrock_model_id(), region=region)
1043 backend._uses_default_model = True
1044 return backend
1046 def _get_client(self) -> Any:
1047 """Return the cached ``bedrock-runtime`` client, building it on first use.
1049 ``boto3`` and ``botocore.exceptions`` are imported here rather
1050 than at module top-level so that pure-Python consumers of this
1051 module (the prompt builder, the protocol, the error type) do
1052 not pay for the ``boto3`` import. This also lets tests patch
1053 ``mission.sampling.boto3`` after import.
1054 """
1055 if self._client is not None:
1056 return self._client
1057 # Local import — keeps the module's import surface boto3-free.
1058 import boto3
1059 from botocore.config import Config
1060 from botocore.exceptions import (
1061 NoCredentialsError,
1062 PartialCredentialsError,
1063 )
1065 try:
1066 self._client = boto3.Session().client(
1067 "bedrock-runtime",
1068 region_name=self._region,
1069 config=Config(read_timeout=BEDROCK_READ_TIMEOUT_SECONDS),
1070 )
1071 except (NoCredentialsError, PartialCredentialsError) as err:
1072 raise SamplingTransportError("bedrock_no_credentials") from err
1073 return self._client
1075 async def sample(self, prompt: SamplingPrompt) -> str:
1076 """Render ``prompt`` through ``Converse`` and return the response text.
1078 Raises:
1079 SamplingTransportError: On any transport-level failure.
1080 * ``bedrock_no_credentials`` — credentials could not be
1081 resolved by ``boto3`` at client-construction time.
1082 * ``bedrock_<ErrorCode>`` — the ``Converse`` call raised
1083 a ``ClientError``; ``<ErrorCode>`` is the AWS error
1084 code from the envelope.
1085 * ``bedrock_malformed_response`` — the response did not
1086 contain a non-empty final text content block.
1087 * ``bedrock_truncated_response`` — the response was cut
1088 off by an output-token limit and cannot be trusted to
1089 be complete.
1090 gco.bedrock.BedrockFTUFormNotAcceptedError: The account has
1091 not submitted Anthropic's one-time first-time-use case
1092 form. Raised instead of a transport error so callers
1093 cannot silently fall back past a permanent, one-line-fix
1094 misconfiguration.
1095 """
1096 # Local import — see ``_get_client`` for the rationale.
1097 from botocore.exceptions import ClientError
1099 client = self._get_client()
1101 text = prompt.assemble()
1102 converse_options = build_bedrock_converse_options(
1103 self.model_id,
1104 # Deliberately no maxTokens: the Converse default is the model's
1105 # own maximum output length, so a rationale can never be cut off
1106 # by a GCO-imposed cap. A cap is opt-in — pass maxTokens here to
1107 # restore one.
1108 inference_config={"temperature": BEDROCK_TEMPERATURE},
1109 apply_default_reasoning=self._uses_default_model,
1110 )
1111 try:
1112 response = await asyncio.to_thread(
1113 client.converse,
1114 modelId=self.model_id,
1115 messages=[{"role": "user", "content": [{"text": text}]}],
1116 **converse_options,
1117 )
1118 except ClientError as err:
1119 # A missing Anthropic FTU form is a permanent account-scoped
1120 # misconfiguration, not a transport fault: escalate it instead of
1121 # letting the deterministic-fallback path absorb it silently.
1122 raise_if_bedrock_ftu_form_error(err)
1123 # ``e.response`` is documented to be present on ClientError
1124 # but the envelope shape can vary; defend against missing
1125 # keys so the audit pipeline always sees a tagged code.
1126 envelope = getattr(err, "response", None) or {}
1127 error_block = envelope.get("Error", {}) if isinstance(envelope, dict) else {}
1128 code = (
1129 error_block.get("Code", "Unknown") if isinstance(error_block, dict) else "Unknown"
1130 )
1131 raise SamplingTransportError(f"bedrock_{code}") from err
1133 # Capture token usage from the Converse response for the audit
1134 # trail. The ``usage`` block is present on every successful
1135 # Converse response and carries ``inputTokens`` and
1136 # ``outputTokens``. Store on the instance so callers can read
1137 # it after each sample() call without changing the protocol.
1138 usage = response.get("usage") or {}
1139 self.last_input_tokens: int | None = usage.get("inputTokens")
1140 self.last_output_tokens: int | None = usage.get("outputTokens")
1142 try:
1143 return extract_bedrock_converse_text(response)
1144 except BedrockResponseTruncatedError as err:
1145 # A cut-off rationale is unusable; let the deterministic-fallback
1146 # path absorb it like any other transport-shaped fault.
1147 raise SamplingTransportError("bedrock_truncated_response") from err
1148 except (KeyError, IndexError, TypeError) as err:
1149 raise SamplingTransportError("bedrock_malformed_response") from err
1152# ---------------------------------------------------------------------------
1153# Backend resolver
1154# ---------------------------------------------------------------------------
1157def _ctx_has_sampling_capability(ctx: Any) -> bool:
1158 """Return ``True`` when the MCP context advertises sampling support.
1160 FastMCP exposes the negotiated client capabilities under a couple
1161 of attribute paths depending on its release: the modern
1162 ``ctx.session_capabilities.sampling`` and the older
1163 ``ctx.fastmcp.client_capabilities.sampling``. Either path may be
1164 missing or set to ``None`` on a context that has not finished
1165 capability negotiation. The probe walks both paths defensively
1166 using ``getattr`` so a missing attribute never raises.
1167 """
1168 caps = getattr(ctx, "session_capabilities", None)
1169 if caps is None:
1170 # Older FastMCP versions surface capabilities via fastmcp.client_capabilities.
1171 fastmcp_attr = getattr(ctx, "fastmcp", None)
1172 caps = (
1173 getattr(fastmcp_attr, "client_capabilities", None) if fastmcp_attr is not None else None
1174 )
1175 if caps is None:
1176 return False
1177 sampling = getattr(caps, "sampling", None)
1178 return bool(sampling)
1181def select_sampling_backend(
1182 ctx: Any | None,
1183 model_id: str | None,
1184 prefs: dict[str, Any] | None,
1185) -> SamplingBackend | None:
1186 """Resolve which sampling backend to use for the given call site.
1188 The resolver picks one of three outcomes:
1190 * MCP path — ``ctx`` is non-``None`` and the context advertises
1191 sampling capability. Returns an :class:`MCPSamplingBackend`
1192 bound to ``ctx`` with ``model_id`` and ``prefs`` forwarded.
1193 * CLI path — ``ctx`` is ``None``. Returns a
1194 :class:`BedrockSamplingBackend` constructed with ``model_id``;
1195 credential resolution is deferred to the first ``sample`` call.
1196 * Deterministic-fallback only — ``ctx`` is non-``None`` but the
1197 context does not advertise sampling capability. Returns ``None``.
1199 Args:
1200 ctx: A FastMCP-style ``Context`` for the MCP path, or ``None``
1201 for the CLI path. Anything else passes through the same
1202 duck-typed capability probe used by the MCP backend.
1203 model_id: Optional concrete model identifier. Forwarded to
1204 either backend constructor verbatim.
1205 prefs: Optional FastMCP ``ModelPreferences`` payload. Forwarded
1206 only to :class:`MCPSamplingBackend`; the Bedrock backend
1207 uses its own pinned inference config.
1209 Returns:
1210 A :class:`SamplingBackend` instance, or ``None`` when the only
1211 available path is the deterministic fallback.
1212 """
1213 if ctx is None:
1214 return BedrockSamplingBackend(model_id)
1215 if _ctx_has_sampling_capability(ctx):
1216 return MCPSamplingBackend(ctx, model_id, prefs)
1217 return None
1220# ---------------------------------------------------------------------------
1221# Strategy-against-catalog validator
1222# ---------------------------------------------------------------------------
1225def _resolve_input_schema(tool: Any) -> Any:
1226 """Return the registered Pydantic input model for a Tool, or None.
1228 FastMCP exposes the model under ``input_schema`` in newer releases
1229 and ``inputSchema`` in older ones. Tolerate both. Tools that genuinely
1230 take no args (or test catalog mocks that omit the attribute) yield
1231 ``None``, in which case the caller skips per-call args validation.
1232 """
1233 schema = getattr(tool, "input_schema", None)
1234 if schema is None:
1235 schema = getattr(tool, "inputSchema", None)
1236 return schema
1239def _extract_tool_json_schemas(
1240 allowlist: Sequence[str],
1241 registered_tools: Mapping[str, Any],
1242) -> dict[str, Any]:
1243 """Extract JSON Schema dicts for each allowlisted tool's input parameters.
1245 Calls ``.model_json_schema()`` on the Pydantic model exposed by
1246 ``_resolve_input_schema``. Falls back gracefully: tools without a
1247 schema, tools whose schema isn't a Pydantic model, and any
1248 exception during schema extraction all yield ``None`` for that
1249 tool (omitted from the output dict). The caller renders the
1250 result into the Strategy_Revision prompt so the model can propose
1251 valid ``args`` dicts.
1252 """
1253 schemas: dict[str, Any] = {}
1254 for name in allowlist:
1255 tool = registered_tools.get(name)
1256 if tool is None: 1256 ↛ 1257line 1256 didn't jump to line 1257 because the condition on line 1256 was never true
1257 continue
1258 model = _resolve_input_schema(tool)
1259 if model is None: 1259 ↛ 1260line 1259 didn't jump to line 1260 because the condition on line 1259 was never true
1260 continue
1261 try:
1262 # Pydantic v2 models expose model_json_schema() as a classmethod.
1263 json_schema = model.model_json_schema()
1264 schemas[name] = json_schema
1265 except Exception:
1266 # Non-Pydantic schema, or a mock that doesn't support it.
1267 continue
1268 return schemas
1271def validate_strategy_against_catalog(
1272 strategy: Strategy,
1273 allowlist: list[str],
1274 registered_tools: dict[str, Any],
1275 allow_scripts: bool,
1276) -> None:
1277 """Validate a Strategy against the live tool catalog.
1279 Returns ``None`` on accept; raises :class:`MissionValidationError`
1280 with a structured ``details.reason`` enum on reject. The function
1281 layers catalog-aware checks on top of the structural validation in
1282 :func:`mission.validation.validate_strategy`:
1284 1. Mutual-exclusivity (exactly one of ``tool_calls`` / ``script``).
1285 2. Per-call ``tool_name`` is in ``allowlist``.
1286 3. Per-call ``args`` validates against the registered Pydantic model
1287 exposed under ``Tool.input_schema`` (or the older
1288 ``Tool.inputSchema``); calls whose tool has neither attribute or
1289 a ``None`` schema skip args validation.
1290 4. For scripted strategies, ``allow_scripts`` is True and the
1291 script's AST passes :func:`mission.sandbox.validate_script_ast`.
1293 Args:
1294 strategy: The Strategy dict to validate.
1295 allowlist: The session's resolved Tool_Allowlist.
1296 registered_tools: Mapping from tool name to a registered tool
1297 object (typed ``Any`` so the module imports without
1298 FastMCP). Read-only — only ``input_schema`` /
1299 ``inputSchema`` is consulted.
1300 allow_scripts: Session-level flag gating scripted strategies.
1301 """
1302 # 1. Structural validation: mutual exclusivity, script-allow gating,
1303 # and AST validation for scripts. Reuses the existing validator
1304 # so error shapes for those rejection classes stay aligned with
1305 # the rest of the input pipeline.
1306 _validation.validate_strategy(cast("dict[str, Any]", strategy), allowlist, allow_scripts)
1308 # The structural validator has already accepted exactly one of the
1309 # two shapes. Branch on which one is present.
1310 if "tool_calls" in strategy:
1311 tool_calls = strategy["tool_calls"]
1312 # Empty list is rejected by validate_strategy; this is a defence
1313 # in depth for callers that might bypass that path.
1314 if not tool_calls: 1314 ↛ 1315line 1314 didn't jump to line 1315 because the condition on line 1314 was never true
1315 raise MissionValidationError(
1316 "validation_error",
1317 details={
1318 "field": "strategy",
1319 "subfield": "tool_calls",
1320 "reason": "tool_calls_empty",
1321 },
1322 )
1324 # 2. Per-call name-in-allowlist check.
1325 for call in tool_calls:
1326 name = call.get("tool_name")
1327 if name not in allowlist:
1328 raise MissionValidationError(
1329 "validation_error",
1330 details={
1331 "field": "strategy",
1332 "subfield": "tool_calls",
1333 "tool_name": name,
1334 "reason": "tool_not_allowlisted",
1335 "allowlist": list(allowlist),
1336 },
1337 )
1339 # 3. Per-call args validation against the tool's Pydantic model.
1340 for call in tool_calls:
1341 name = call["tool_name"]
1342 tool = registered_tools.get(name)
1343 if tool is None: 1343 ↛ 1347line 1343 didn't jump to line 1347 because the condition on line 1343 was never true
1344 # Catalog could have a name in the allowlist that is not
1345 # currently registered (gating, dynamic load). Mirror the
1346 # unknown-tool shape used elsewhere.
1347 raise MissionValidationError(
1348 "validation_error",
1349 details={
1350 "field": "strategy",
1351 "subfield": "tool_calls",
1352 "tool_name": name,
1353 "reason": "tool_not_registered",
1354 },
1355 )
1356 schema = _resolve_input_schema(tool)
1357 if schema is None:
1358 # Either the tool genuinely takes no args, or the test
1359 # catalog omitted a model. Skip args validation rather
1360 # than reject — the design treats missing schema as
1361 # "trust the dispatcher".
1362 continue
1363 args = call.get("args", {})
1364 if not isinstance(args, dict): 1364 ↛ 1365line 1364 didn't jump to line 1365 because the condition on line 1364 was never true
1365 raise MissionValidationError(
1366 "validation_error",
1367 details={
1368 "field": "strategy",
1369 "subfield": "tool_calls",
1370 "tool_name": name,
1371 "reason": "tool_args_invalid",
1372 "errors": [
1373 {
1374 "type": "args_not_a_dict",
1375 "actual_type": type(args).__name__,
1376 }
1377 ],
1378 },
1379 )
1380 try:
1381 schema.model_validate(args)
1382 except Exception as exc: # noqa: BLE001 - pydantic ValidationError + similar
1383 # Pydantic v2 ValidationError exposes ``.errors()`` as a
1384 # list of structured dicts. Tolerate any other exception
1385 # type (e.g. older Pydantic, custom validators) by
1386 # falling back to ``str(exc)``.
1387 errors_method = getattr(exc, "errors", None)
1388 if callable(errors_method): 1388 ↛ 1394line 1388 didn't jump to line 1394 because the condition on line 1388 was always true
1389 try:
1390 errors_payload: Any = errors_method()
1391 except Exception: # noqa: BLE001 - defensive
1392 errors_payload = [{"type": "unknown", "msg": str(exc)}]
1393 else:
1394 errors_payload = [{"type": "unknown", "msg": str(exc)}]
1395 raise MissionValidationError(
1396 "validation_error",
1397 details={
1398 "field": "strategy",
1399 "subfield": "tool_calls",
1400 "tool_name": name,
1401 "reason": "tool_args_invalid",
1402 "errors": errors_payload,
1403 },
1404 ) from exc
1406 # 4. Cost estimation against remaining budget. Removed —
1407 # cost guardrails live out-of-band via AWS Budgets / Cost
1408 # Anomaly Detection rather than in the Mission cascade.
1409 # Scripted strategies: validate_strategy already ran allow_scripts
1410 # gating and the AST validator. No catalog-aware checks are layered
1411 # on top here — the script-side enforcement happens at execute time
1412 # via the in-script tool callable wrappers.
1415# ---------------------------------------------------------------------------
1416# Orchestration helpers — bind a backend to a SessionState and return either
1417# a used result or a deterministic fallback.
1418# ---------------------------------------------------------------------------
1420# Local imports kept inside this section so the prompt-builder /
1421# backend half above stays free of audit / decide dependencies.
1422from . import audit as _mission_audit # noqa: E402
1423from . import decide as _decide # noqa: E402
1425# Type alias used by the helpers below. ``SessionState`` is a TypedDict
1426# whose runtime value is just ``dict``; the alias keeps the signatures
1427# expressive without forcing the import to leak through ``__all__``.
1428from .types import SessionState as _SessionState # noqa: E402
1431@dataclass(frozen=True)
1432class SamplingUsed:
1433 """A successful sampling call's accepted output."""
1435 output_text: str
1436 """Raw model output (the text returned by the bound backend)."""
1438 parsed: dict[str, Any]
1439 """Parsed JSON payload that has cleared the schema and catalog checks."""
1441 backend_name: Literal["mcp", "bedrock"]
1442 """Stable backend identifier — echoes the bound backend's tag."""
1444 model_id: str
1445 """The concrete model id the backend routed the prompt to."""
1448@dataclass(frozen=True)
1449class SamplingFallback:
1450 """A rejected or unavailable sampling call's deterministic substitute.
1452 Returned when the bound backend was ``None``, the transport raised,
1453 the model output failed to parse / validate, or any catalog or
1454 budget check rejected the proposed strategy. The ``rationale`` is
1455 a pure function of the bound :class:`SessionState` and the most
1456 recent :class:`IterationRecord`, so the engine can replay or
1457 reproduce a fallback exactly from persisted state.
1458 """
1460 rationale: str
1461 """Deterministic fallback text. Empty string for ``final_lessons``
1462 — the final-report writer fills in its own deterministic text in
1463 that case."""
1465 reason: str
1466 """Stable token tagging *why* the fallback fired. Examples:
1467 ``"transport_error"``, ``"json_parse"``, ``"schema_mismatch"``,
1468 ``"tool_not_allowlisted"``, ``"tool_args_invalid"``,
1469 ``"over_budget"``, ``"script_rejected"``,
1470 ``"no_backend_resolved"``, ``"disabled"``."""
1472 backend_name: Literal["mcp", "bedrock", "none"]
1473 """The bound backend's tag, or ``"none"`` when no backend was
1474 resolved at the call site."""
1476 model_id: str | None
1477 """The bound backend's model id, or ``None`` when no backend was
1478 resolved."""
1481# ---------------------------------------------------------------------------
1482# JSON / schema helpers (private to the orchestration layer)
1483# ---------------------------------------------------------------------------
1486def _extract_json_object(text: str) -> dict[str, Any]:
1487 """Parse the first JSON object embedded in ``text``.
1489 Models routinely wrap JSON in prose. The implementation slices from
1490 the first ``{`` to the last ``}`` and feeds the result to
1491 :func:`json.loads`. When no balanced braces are present, or the
1492 sliced substring is not valid JSON, the function raises
1493 :class:`json.JSONDecodeError` so the caller can branch on a single
1494 well-typed exception.
1495 """
1496 start = text.find("{")
1497 end = text.rfind("}")
1498 if start == -1 or end == -1 or end < start:
1499 # No braces at all → treat as a parse error so the calling
1500 # branch surfaces ``reason="json_parse"``.
1501 raise json.JSONDecodeError("no JSON object found", text, 0)
1502 candidate = text[start : end + 1]
1503 parsed = json.loads(candidate)
1504 if not isinstance(parsed, dict): 1504 ↛ 1507line 1504 didn't jump to line 1507 because the condition on line 1504 was never true
1505 # The sliced substring parsed but is not an object — surface as
1506 # a parse error too, since downstream code requires a dict.
1507 raise json.JSONDecodeError("top-level JSON value is not an object", candidate, 0)
1508 return parsed
1511def _validate_revision_schema(parsed: dict[str, Any]) -> None:
1512 """Reject a parsed payload that is not a valid Strategy_Revision.
1514 Required keys: ``revision_rationale`` (non-empty str),
1515 ``next_strategy`` (dict), ``confidence`` (number in [0, 1]).
1516 """
1517 rationale = parsed.get("revision_rationale")
1518 if not isinstance(rationale, str) or not rationale:
1519 raise ValueError("schema_mismatch: revision_rationale must be non-empty str")
1520 next_strategy = parsed.get("next_strategy")
1521 if not isinstance(next_strategy, dict): 1521 ↛ 1522line 1521 didn't jump to line 1522 because the condition on line 1521 was never true
1522 raise ValueError("schema_mismatch: next_strategy must be a dict")
1523 confidence = parsed.get("confidence")
1524 # ``bool`` is excluded explicitly — it is a subclass of ``int`` in
1525 # Python and would otherwise sneak past the numeric check.
1526 if isinstance(confidence, bool) or not isinstance(confidence, (int, float)): 1526 ↛ 1527line 1526 didn't jump to line 1527 because the condition on line 1526 was never true
1527 raise ValueError("schema_mismatch: confidence must be a number")
1528 if not (0.0 <= float(confidence) <= 1.0): 1528 ↛ 1529line 1528 didn't jump to line 1529 because the condition on line 1528 was never true
1529 raise ValueError("schema_mismatch: confidence must be in [0, 1]")
1532def _validate_lessons_schema(parsed: dict[str, Any]) -> None:
1533 """Reject a parsed payload that is not a valid final-lessons dict.
1535 Required keys: ``lessons`` (non-empty list of non-empty str),
1536 ``recommended_followups`` (list of str — may be empty).
1537 """
1538 lessons = parsed.get("lessons")
1539 if not isinstance(lessons, list) or not lessons: 1539 ↛ 1540line 1539 didn't jump to line 1540 because the condition on line 1539 was never true
1540 raise ValueError("schema_mismatch: lessons must be a non-empty list")
1541 for item in lessons:
1542 if not isinstance(item, str) or not item: 1542 ↛ 1543line 1542 didn't jump to line 1543 because the condition on line 1542 was never true
1543 raise ValueError("schema_mismatch: each lesson must be a non-empty str")
1544 followups = parsed.get("recommended_followups")
1545 if not isinstance(followups, list): 1545 ↛ 1546line 1545 didn't jump to line 1546 because the condition on line 1545 was never true
1546 raise ValueError("schema_mismatch: recommended_followups must be a list")
1547 for item in followups:
1548 if not isinstance(item, str): 1548 ↛ 1549line 1548 didn't jump to line 1549 because the condition on line 1548 was never true
1549 raise ValueError("schema_mismatch: each follow-up must be a str")
1552# ---------------------------------------------------------------------------
1553# maybe_sample_strategy_revision
1554# ---------------------------------------------------------------------------
1557async def maybe_sample_strategy_revision(
1558 *,
1559 backend: SamplingBackend | None,
1560 session: _SessionState,
1561 iteration: IterationRecord,
1562 allowlist: list[str],
1563 registered_tools: dict[str, Any],
1564 tool_docstrings: dict[str, str],
1565 remaining_iterations: int,
1566 remaining_wall_clock_secs: float | None,
1567 allow_scripts: bool,
1568 environment_context: Mapping[str, Any] | None = None,
1569) -> SamplingUsed | SamplingFallback:
1570 """Consult the advisory LLM for a Strategy_Revision, or fall back.
1572 Returns a :class:`SamplingUsed` when the bound backend produces a
1573 JSON object that clears schema validation and the catalog checks.
1574 Returns a :class:`SamplingFallback` carrying the deterministic
1575 rationale from
1576 :func:`mission.decide.build_revision_rationale_template` on every
1577 rejection class. Emits exactly one
1578 :func:`mission.audit.emit_sampling_event` per call.
1579 """
1580 session_id = session["session_id"]
1581 iteration_index = iteration["iteration_index"]
1582 template = _decide.build_revision_rationale_template(session, iteration)
1584 # ---- No backend resolved: short-circuit. ------------------------------
1585 if backend is None:
1586 _mission_audit.emit_sampling_event(
1587 session_id,
1588 iteration_index,
1589 sampling_purpose="strategy_revision",
1590 sampling_status="disabled",
1591 sampling_backend="none",
1592 )
1593 return SamplingFallback(
1594 rationale=template,
1595 reason="no_backend_resolved",
1596 backend_name="none",
1597 model_id=None,
1598 )
1600 backend_name = backend.backend_name
1601 model_id = backend.model_id
1603 # ---- Build the prompt. ------------------------------------------------
1604 # The in-progress iteration that triggered ``adjust`` is already in
1605 # ``session["iterations"][-1]``, so the most-recent-five window is a
1606 # plain slice; ``RECENT_ITERATIONS_LIMIT`` is enforced inside the
1607 # prompt builder as a defensive cap.
1608 recent_iterations = list(session["iterations"][-RECENT_ITERATIONS_LIMIT:])
1609 tool_schemas = _extract_tool_json_schemas(allowlist, registered_tools)
1610 prompt = SamplingPrompt(
1611 directive=session["directive_text"],
1612 success_criteria=session["criteria"],
1613 criteria_status=iteration["criteria_evaluation"],
1614 recent_iterations=recent_iterations,
1615 tool_allowlist=allowlist,
1616 tool_docstrings=tool_docstrings,
1617 remaining_iterations=remaining_iterations,
1618 remaining_wall_clock_secs=remaining_wall_clock_secs,
1619 allow_scripts=allow_scripts,
1620 tool_schemas=tool_schemas,
1621 environment_context=environment_context,
1622 )
1624 # ---- Transport: backend.sample. --------------------------------------
1625 try:
1626 output_text = await backend.sample(prompt)
1627 except SamplingTransportError as err:
1628 _mission_audit.emit_sampling_event(
1629 session_id,
1630 iteration_index,
1631 sampling_purpose="strategy_revision",
1632 sampling_status="rejected",
1633 sampling_backend=backend_name,
1634 sampling_model_id=model_id or None,
1635 validation_error=err.code,
1636 )
1637 return SamplingFallback(
1638 rationale=template,
1639 reason="transport_error",
1640 backend_name=backend_name,
1641 model_id=model_id,
1642 )
1644 # ---- Parse the output as JSON. ---------------------------------------
1645 try:
1646 parsed = _extract_json_object(output_text)
1647 except json.JSONDecodeError:
1648 _mission_audit.emit_sampling_event(
1649 session_id,
1650 iteration_index,
1651 sampling_purpose="strategy_revision",
1652 sampling_status="rejected",
1653 sampling_backend=backend_name,
1654 sampling_model_id=model_id or None,
1655 validation_error="json_parse",
1656 )
1657 return SamplingFallback(
1658 rationale=template,
1659 reason="json_parse",
1660 backend_name=backend_name,
1661 model_id=model_id,
1662 )
1664 # ---- Schema validation. ----------------------------------------------
1665 try:
1666 _validate_revision_schema(parsed)
1667 except ValueError:
1668 _mission_audit.emit_sampling_event(
1669 session_id,
1670 iteration_index,
1671 sampling_purpose="strategy_revision",
1672 sampling_status="rejected",
1673 sampling_backend=backend_name,
1674 sampling_model_id=model_id or None,
1675 validation_error="schema_mismatch",
1676 )
1677 return SamplingFallback(
1678 rationale=template,
1679 reason="schema_mismatch",
1680 backend_name=backend_name,
1681 model_id=model_id,
1682 )
1684 # ---- Catalog validation on the proposed next_strategy. ---------------
1685 try:
1686 validate_strategy_against_catalog(
1687 parsed["next_strategy"],
1688 allowlist,
1689 registered_tools,
1690 allow_scripts,
1691 )
1692 except MissionValidationError as err:
1693 # ``err.details["reason"]`` carries the structured rejection
1694 # token (e.g. ``"tool_not_allowlisted"``,
1695 # ``"tool_args_invalid"``). Fall back to a generic label when
1696 # the validator emits a rejection without a ``reason`` key.
1697 details = err.details or {}
1698 reason = details.get("reason", "validation_error")
1699 _mission_audit.emit_sampling_event(
1700 session_id,
1701 iteration_index,
1702 sampling_purpose="strategy_revision",
1703 sampling_status="rejected",
1704 sampling_backend=backend_name,
1705 sampling_model_id=model_id or None,
1706 validation_error=str(reason),
1707 )
1708 return SamplingFallback(
1709 rationale=template,
1710 reason=str(reason),
1711 backend_name=backend_name,
1712 model_id=model_id,
1713 )
1715 # ---- Success path. ---------------------------------------------------
1716 # Extract token usage from the backend if available (Bedrock backend
1717 # stores it as a side-channel after each sample() call).
1718 _input_tokens = getattr(backend, "last_input_tokens", None)
1719 _output_tokens = getattr(backend, "last_output_tokens", None)
1720 _mission_audit.emit_sampling_event(
1721 session_id,
1722 iteration_index,
1723 sampling_purpose="strategy_revision",
1724 sampling_status="used",
1725 sampling_backend=backend_name,
1726 sampling_model_id=model_id or None,
1727 model_output_bytes=len(output_text.encode("utf-8")),
1728 input_tokens=_input_tokens,
1729 output_tokens=_output_tokens,
1730 )
1731 return SamplingUsed(
1732 output_text=output_text,
1733 parsed=parsed,
1734 backend_name=backend_name,
1735 model_id=model_id,
1736 )
1739# ---------------------------------------------------------------------------
1740# maybe_sample_final_lessons
1741# ---------------------------------------------------------------------------
1744async def maybe_sample_final_lessons(
1745 *,
1746 backend: SamplingBackend | None,
1747 session: _SessionState,
1748 remaining_iterations: int = 0,
1749 remaining_wall_clock_secs: float | None = None,
1750 allow_scripts: bool = False,
1751 tool_docstrings: dict[str, str] | None = None,
1752 environment_context: Mapping[str, Any] | None = None,
1753) -> SamplingUsed | SamplingFallback:
1754 """Consult the advisory LLM for final lessons, or fall back.
1756 Returns a :class:`SamplingUsed` when the bound backend produces a
1757 JSON object that clears the lessons schema. Returns a
1758 :class:`SamplingFallback` with an *empty* rationale on every
1759 rejection class — the final-report writer is responsible for the
1760 deterministic-text path when sampling does not produce usable
1761 output. Emits exactly one
1762 :func:`mission.audit.emit_sampling_event` per call, with
1763 ``iteration_index_or_purpose=None`` since the call is out-of-loop.
1764 """
1765 session_id = session["session_id"]
1767 # ---- No backend resolved: short-circuit. ------------------------------
1768 if backend is None:
1769 _mission_audit.emit_sampling_event(
1770 session_id,
1771 None,
1772 sampling_purpose="final_lessons",
1773 sampling_status="disabled",
1774 sampling_backend="none",
1775 )
1776 return SamplingFallback(
1777 rationale="",
1778 reason="no_backend_resolved",
1779 backend_name="none",
1780 model_id=None,
1781 )
1783 backend_name = backend.backend_name
1784 model_id = backend.model_id
1786 # ---- Build the prompt. ------------------------------------------------
1787 # Pass *all* iterations; the prompt builder trims / drops as needed
1788 # to fit the byte budget.
1789 prompt = SamplingPrompt(
1790 directive=session["directive_text"],
1791 success_criteria=session["criteria"],
1792 # The lessons prompt has no per-iteration criteria status; the
1793 # builder still expects the field, so reuse the most recent
1794 # iteration's evaluation when available, else an empty list.
1795 criteria_status=(
1796 list(session["iterations"][-1]["criteria_evaluation"]) if session["iterations"] else []
1797 ),
1798 recent_iterations=list(session["iterations"]),
1799 tool_allowlist=session.get("tool_allowlist", []),
1800 tool_docstrings=tool_docstrings or {},
1801 remaining_iterations=remaining_iterations,
1802 remaining_wall_clock_secs=remaining_wall_clock_secs,
1803 allow_scripts=allow_scripts,
1804 environment_context=environment_context,
1805 )
1807 # ---- Transport: backend.sample (uses lessons assembler). -------------
1808 try:
1809 # We render the lessons-specific prompt here so the byte-cap
1810 # bookkeeping uses the right schema header. The backend's own
1811 # ``sample`` calls ``prompt.assemble()`` under the hood for the
1812 # Strategy_Revision flow, but for lessons we assemble here and
1813 # invoke a thin shim through the backend.
1814 rendered = prompt.assemble_final_lessons()
1815 output_text = await _sample_with_assembled_text(backend, rendered)
1816 except SamplingTransportError as err:
1817 _mission_audit.emit_sampling_event(
1818 session_id,
1819 None,
1820 sampling_purpose="final_lessons",
1821 sampling_status="rejected",
1822 sampling_backend=backend_name,
1823 sampling_model_id=model_id or None,
1824 validation_error=err.code,
1825 )
1826 return SamplingFallback(
1827 rationale="",
1828 reason="transport_error",
1829 backend_name=backend_name,
1830 model_id=model_id,
1831 )
1833 # ---- Parse the output as JSON. ---------------------------------------
1834 try:
1835 parsed = _extract_json_object(output_text)
1836 except json.JSONDecodeError:
1837 _mission_audit.emit_sampling_event(
1838 session_id,
1839 None,
1840 sampling_purpose="final_lessons",
1841 sampling_status="rejected",
1842 sampling_backend=backend_name,
1843 sampling_model_id=model_id or None,
1844 validation_error="json_parse",
1845 )
1846 return SamplingFallback(
1847 rationale="",
1848 reason="json_parse",
1849 backend_name=backend_name,
1850 model_id=model_id,
1851 )
1853 # ---- Schema validation. ----------------------------------------------
1854 try:
1855 _validate_lessons_schema(parsed)
1856 except ValueError:
1857 _mission_audit.emit_sampling_event(
1858 session_id,
1859 None,
1860 sampling_purpose="final_lessons",
1861 sampling_status="rejected",
1862 sampling_backend=backend_name,
1863 sampling_model_id=model_id or None,
1864 validation_error="schema_mismatch",
1865 )
1866 return SamplingFallback(
1867 rationale="",
1868 reason="schema_mismatch",
1869 backend_name=backend_name,
1870 model_id=model_id,
1871 )
1873 # ---- Success path. ---------------------------------------------------
1874 _input_tokens = getattr(backend, "last_input_tokens", None)
1875 _output_tokens = getattr(backend, "last_output_tokens", None)
1876 _mission_audit.emit_sampling_event(
1877 session_id,
1878 None,
1879 sampling_purpose="final_lessons",
1880 sampling_status="used",
1881 sampling_backend=backend_name,
1882 sampling_model_id=model_id or None,
1883 model_output_bytes=len(output_text.encode("utf-8")),
1884 input_tokens=_input_tokens,
1885 output_tokens=_output_tokens,
1886 )
1887 return SamplingUsed(
1888 output_text=output_text,
1889 parsed=parsed,
1890 backend_name=backend_name,
1891 model_id=model_id,
1892 )
1895async def _sample_with_assembled_text(backend: SamplingBackend, rendered: str) -> str:
1896 """Route a pre-assembled prompt string through a backend.
1898 Both shipped backends accept a :class:`SamplingPrompt` and call
1899 ``assemble`` themselves to render the strategy-revision shape. For
1900 the final-lessons path we render the lessons-shaped prompt here
1901 and need to deliver that exact text to the transport. The shim
1902 builds a tiny prompt-shaped wrapper whose :meth:`assemble` returns
1903 the pre-rendered text and forwards it to the backend.
1904 """
1905 pre_rendered = rendered
1907 class _PreRendered:
1908 """Thin :class:`SamplingPrompt` look-alike with a fixed assemble()."""
1910 def assemble(self) -> str:
1911 return pre_rendered
1913 # The two shipped backends only call ``prompt.assemble()`` so the
1914 # duck-typed wrapper above is enough to drive either of them.
1915 return await backend.sample(_PreRendered()) # type: ignore[arg-type]
1918# ---------------------------------------------------------------------------
1919# Session-start sampling-state resolver
1920# ---------------------------------------------------------------------------
1923def _bedrock_credentials_available() -> bool:
1924 """Lightweight probe: do local AWS credentials resolve?
1926 Instantiates a ``boto3.Session()`` and asks for ``get_credentials()``
1927 without making any network call. ``boto3`` is imported inside the
1928 function so the module's top-level import surface stays free of
1929 SDK dependencies — and so a host that has no ``boto3`` installed
1930 (or any other unexpected import-time failure) cleanly degrades to
1931 "no credentials available" rather than crashing the helper.
1932 """
1933 try:
1934 import boto3
1936 session = boto3.Session()
1937 creds = session.get_credentials()
1938 return creds is not None
1939 except Exception:
1940 return False
1943def resolve_sampling_state(
1944 ctx: Any | None,
1945 use_sampling_param: bool | None,
1946) -> tuple[bool, Literal["mcp", "bedrock", "none"]]:
1947 """Decide whether sampling is enabled for a session and which backend resolves.
1949 Resolution precedence (first match wins):
1951 1. ``use_sampling_param is False`` — caller explicitly disabled
1952 sampling, so the result is ``(False, "none")`` regardless of
1953 any capability the environment advertises.
1954 2. ``ctx`` is non-``None`` and advertises MCP sampling capability —
1955 the caller is on the MCP path and the host can perform sampling,
1956 so the result is ``(True, "mcp")``.
1957 3. ``ctx is None`` (CLI path) — probe local AWS credentials. When
1958 they resolve, return ``(True, "bedrock")``. When they do not,
1959 return ``(True, "none")`` if the caller opted in explicitly with
1960 ``use_sampling_param is True`` (so the caller can decide whether
1961 to error or proceed deterministic-only), and ``(False, "none")``
1962 otherwise.
1963 4. ``ctx`` is non-``None`` but advertises no sampling capability —
1964 same explicit-opt-in handling as the no-credentials CLI branch:
1965 ``(True, "none")`` when the caller opted in explicitly,
1966 ``(False, "none")`` otherwise.
1968 Args:
1969 ctx: A FastMCP-style Context for the MCP path, or ``None`` for
1970 the CLI path. The capability probe is duck-typed so any
1971 object exposing the same attributes that
1972 :func:`_ctx_has_sampling_capability` checks works here.
1973 use_sampling_param: Three-state opt-in flag. ``None`` means the
1974 caller did not specify and the helper should auto-detect.
1975 ``False`` short-circuits to a disabled state. ``True`` means
1976 the caller explicitly opted in; the backend is auto-detected
1977 and ``"none"`` is allowed when no concrete backend resolves.
1979 Returns:
1980 A ``(use_sampling, backend)`` tuple. The caller persists both
1981 values on its ``SessionState`` so the audit pipeline can stamp
1982 every later sampling event with the resolved backend.
1983 """
1984 # 1. Explicit opt-out wins outright.
1985 if use_sampling_param is False:
1986 return (False, "none")
1988 # 2. MCP path with capability advertised.
1989 if ctx is not None and _ctx_has_sampling_capability(ctx):
1990 return (True, "mcp")
1992 # 3. CLI path — probe local AWS credentials.
1993 if ctx is None:
1994 if _bedrock_credentials_available():
1995 return (True, "bedrock")
1996 if use_sampling_param is True:
1997 return (True, "none")
1998 return (False, "none")
2000 # 4. ctx present but no MCP capability — only honour an explicit True.
2001 if use_sampling_param is True:
2002 return (True, "none")
2003 return (False, "none")