Coverage for gco_mcp / mission / sampling.py: 100.00%
448 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +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 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_mission_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-09-09T17:36:47Z
64# Generated from Git commit: d03cb5dc20f9b805636c85ce7af957eebb94c28e
65# Flowchart(s) generated from this file:
66# * ``maybe_sample_strategy_revision`` -> ``diagrams/code_diagrams/gco_mcp/mission/sampling.maybe_sample_strategy_revision.html``
67# (PNG: ``diagrams/code_diagrams/gco_mcp/mission/sampling.maybe_sample_strategy_revision.png``)
68# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> python diagrams/generate.py --code-only``.
69# <pyflowchart-code-diagram> END
72__all__ = [
73 "BEDROCK_READ_TIMEOUT_SECONDS",
74 "BEDROCK_TEMPERATURE",
75 "DEFAULT_BEDROCK_REGION",
76 "ENV_BEDROCK_MODEL_ID",
77 "ENV_BEDROCK_REGION",
78 "ENVIRONMENT_CONTEXT_BYTE_CAP",
79 "FINAL_LESSONS_SCHEMA",
80 "BedrockSamplingBackend",
81 "MissionValidationError",
82 "OBSERVATION_FIELD_BYTE_CAP",
83 "OBSERVATION_FIELD_TRUNCATE_TO",
84 "PRIOR_MISSIONS_BYTE_CAP",
85 "PROMPT_BYTE_BUDGET",
86 "RECENT_ITERATIONS_LIMIT",
87 "STRATEGY_REVISION_SCHEMA",
88 "STRATEGY_SHAPE_SCHEMA",
89 "SamplingBackend",
90 "SamplingFallback",
91 "SamplingPrompt",
92 "SamplingTransportError",
93 "SamplingUsed",
94 "TRUNCATION_MARKER",
95 "maybe_sample_final_lessons",
96 "maybe_sample_strategy_revision",
97 "resolve_sampling_state",
98 "select_sampling_backend",
99 "validate_strategy_against_catalog",
100]
103# ---------------------------------------------------------------------------
104# Tunables (named so tests can reference them without hard-coding magic)
105# ---------------------------------------------------------------------------
107#: Per-Observation-field byte cap. Fields whose JSON-serialised UTF-8
108#: byte length exceeds this value are truncated. A field whose byte
109#: length is exactly equal to this value is **not** truncated — the
110#: comparison uses strict greater-than to keep the boundary stable.
111OBSERVATION_FIELD_BYTE_CAP: int = 4096
113#: Target byte length after truncation. The truncated string is the
114#: first ``OBSERVATION_FIELD_TRUNCATE_TO`` bytes of the JSON-serialised
115#: form, decoded with ``errors="ignore"`` so a multi-byte boundary in
116#: the middle of a UTF-8 codepoint cannot raise, with
117#: :data:`TRUNCATION_MARKER` appended.
118OBSERVATION_FIELD_TRUNCATE_TO: int = 2048
120#: Marker appended to every truncated field so the reader can see at a
121#: glance the field was clipped.
122TRUNCATION_MARKER: str = "... [truncated]"
124#: Total prompt byte budget. The render methods drop the oldest
125#: Iteration summary one at a time until ``len(prompt.encode("utf-8"))
126#: <= PROMPT_BYTE_BUDGET``.
127PROMPT_BYTE_BUDGET: int = 32768
129#: Maximum number of Iteration summaries to include even if the byte
130#: budget is plentiful. The caller is expected to pass at most this
131#: many already; the builder slices defensively.
132RECENT_ITERATIONS_LIMIT: int = 5
134#: Per-Environment-context byte cap. The optional environment context
135#: block (``=== Environment context ===``) is its own truncation
136#: domain so the section can never grow without bound and push the
137#: rest of the prompt over :data:`PROMPT_BYTE_BUDGET`. The cap mirrors
138#: :data:`OBSERVATION_FIELD_BYTE_CAP` because both surfaces hold the
139#: same flavour of structured live signal (cluster + queue snapshots
140#: in this case) and the same truncation marker convention applies.
141ENVIRONMENT_CONTEXT_BYTE_CAP: int = 4096
143#: Byte cap for the optional prior-missions block (``=== Prior similar
144#: missions ===``). Its own truncation domain for the same reason as
145#: :data:`ENVIRONMENT_CONTEXT_BYTE_CAP`: retrieved lessons are
146#: free-text of unbounded length and must never crowd the rest of the
147#: prompt out of :data:`PROMPT_BYTE_BUDGET`.
148PRIOR_MISSIONS_BYTE_CAP: int = 4096
150#: The memory-item fields the prior-missions block passes through to
151#: the prompt — the vector index's ``INCLUDE`` projection plus the key
152#: and the similarity score. Anything else a future projection might
153#: surface is dropped so the block's shape stays stable.
154_PRIOR_MISSION_FIELDS: frozenset[str] = frozenset(
155 {
156 "session_id",
157 "directive",
158 "lessons",
159 "recommended_followups",
160 "final_verdict",
161 "verdict_reason",
162 "iteration_count",
163 "completed_at",
164 "score",
165 }
166)
169# ---------------------------------------------------------------------------
170# Environment context summarisation
171# ---------------------------------------------------------------------------
174def _summarise_environment_context(env: Mapping[str, Any]) -> dict[str, Any]:
175 """Return a JSON-safe, byte-capped summary of the environment context.
177 The block is rendered into the Strategy_Revision prompt under
178 ``=== Environment context ===``. It carries small, slow-moving
179 live signals — per-region queue depths, GPU utilisation, deployed
180 region list, reservation counts — that the model would otherwise
181 have to spend tool calls to discover.
183 Two guarantees on the output:
185 1. The serialised form fits inside :data:`ENVIRONMENT_CONTEXT_BYTE_CAP`
186 UTF-8 bytes. When the input does not, top-level fields are
187 evaluated in sorted-key order, dropped one at a time from the
188 largest contributor down, and the dropped key list is recorded
189 under ``"_dropped_fields"`` so the operator can spot which
190 inputs got pruned.
191 2. Top-level keys are emitted in sorted order so two callers
192 passing semantically-identical dicts produce a byte-identical
193 block — the same property the determinism tests pin down for
194 Observation summaries.
195 """
196 # Defensive copy + sort so insertion order doesn't leak.
197 ordered: dict[str, Any] = {key: env[key] for key in sorted(env.keys())}
198 serialised = _dumps(ordered)
199 if _utf8_len(serialised) <= ENVIRONMENT_CONTEXT_BYTE_CAP:
200 return ordered
202 # Drop largest top-level field first, repeating until under cap.
203 # Records dropped keys so the operator (and the audit pipeline)
204 # can see what got pruned without having to diff against the
205 # gather helper's output.
206 dropped: list[str] = []
207 working = dict(ordered)
208 while _utf8_len(_dumps(working)) > ENVIRONMENT_CONTEXT_BYTE_CAP and working:
209 biggest_key = max(working, key=lambda k: _utf8_len(_dumps(working[k])))
210 dropped.append(biggest_key)
211 del working[biggest_key]
213 if dropped:
214 # Sort the dropped list so its position in the prompt is stable
215 # regardless of which key happened to be biggest first.
216 working["_dropped_fields"] = sorted(dropped)
217 return working
220# ---------------------------------------------------------------------------
221# Prior-missions summarisation
222# ---------------------------------------------------------------------------
225def _summarise_prior_missions(
226 missions: Sequence[Mapping[str, Any]],
227) -> list[dict[str, Any]]:
228 """Return a JSON-safe, byte-capped summary of retrieved prior missions.
230 Rendered into the prompt under ``=== Prior similar missions ===``.
231 The input is the :meth:`mcp.mission.memory.MissionMemoryStore.search_similar`
232 result list, ordered most-similar-first.
234 Three guarantees on the output:
236 1. Only the fields in :data:`_PRIOR_MISSION_FIELDS` pass through,
237 emitted in sorted-key order — so two semantically-identical
238 inputs produce a byte-identical block (the determinism property
239 every prompt section pins down), and a recreated index with a
240 wider projection cannot change the block's shape.
241 2. Each mission's ``lessons`` field is truncated to
242 :data:`OBSERVATION_FIELD_TRUNCATE_TO` bytes with
243 :data:`TRUNCATION_MARKER` when it exceeds
244 :data:`OBSERVATION_FIELD_BYTE_CAP` — one verbose write-up must
245 not evict every other retrieved mission.
246 3. The serialised list fits in :data:`PRIOR_MISSIONS_BYTE_CAP`
247 UTF-8 bytes. When it does not, the *least similar* mission (the
248 list tail) is dropped first, repeating until under cap.
249 """
250 summarised: list[dict[str, Any]] = []
251 for mission in missions:
252 entry = {key: mission[key] for key in sorted(_PRIOR_MISSION_FIELDS) if key in mission}
253 lessons = entry.get("lessons")
254 if isinstance(lessons, str) and _utf8_len(lessons) > OBSERVATION_FIELD_BYTE_CAP:
255 entry["lessons"] = _truncate_serialised(lessons)
256 summarised.append(entry)
258 while _utf8_len(_dumps(summarised)) > PRIOR_MISSIONS_BYTE_CAP and summarised:
259 summarised.pop()
260 return summarised
263# ---------------------------------------------------------------------------
264# Bedrock backend tunables
265# ---------------------------------------------------------------------------
267#: The default Bedrock model identifier is read on demand from ``cdk.json``
268#: ``context.bedrock.mission_default_model_id`` through the lightweight
269#: :func:`gco.bedrock.get_default_mission_model_id` resolver, so unrelated
270#: imports never couple to Bedrock configuration resolution.
271#:
272#: Operators with regulatory or model-governance requirements can override per
273#: call via ``GCO_MISSION_BEDROCK_MODEL_ID`` or ``--bedrock-model-id``; see
274#: docs/CUSTOMIZATION.md ("Bedrock Model Selection").
276#: Default Bedrock region. The capacity advisor pins ``us-east-1`` for
277#: the same reason: cross-region inference profiles routinely surface
278#: in ``us-east-1`` first and our installations have it whitelisted.
279DEFAULT_BEDROCK_REGION: str = "us-east-1"
281#: Env var that overrides the canonical Mission model default at runtime.
282ENV_BEDROCK_MODEL_ID: str = "GCO_MISSION_BEDROCK_MODEL_ID"
284#: Env var that overrides :data:`DEFAULT_BEDROCK_REGION` at runtime.
285ENV_BEDROCK_REGION: str = "GCO_MISSION_BEDROCK_REGION"
287#: Sampling temperature requested for Bedrock models that accept one. The
288#: canonical Claude Opus 5 default does not: Opus 4.7 onward deprecated
289#: ``temperature``, ``topP``, and ``topK``, so
290#: :func:`build_bedrock_converse_options` drops this field for every restricted
291#: Claude line — default or explicit override — and for OpenAI and xAI
292#: profiles. Models outside those families keep it.
293BEDROCK_TEMPERATURE: float = 0.2
296# ---------------------------------------------------------------------------
297# JSON Schemas — embedded as module-level constants
298# ---------------------------------------------------------------------------
300# The Strategy shape mirrors the ``Strategy`` TypedDict from
301# ``gco_mcp/mission/types.py``: every key is optional in isolation and the
302# validator enforces the mutual-exclusivity invariant (exactly one of
303# ``tool_calls`` or ``script`` populated). The schema below mirrors
304# that with a ``oneOf`` clause. The model-side validator in subsequent
305# commits performs the same check on parsed responses; this schema is
306# the textual instruction the prompt embeds for the model.
307STRATEGY_SHAPE_SCHEMA: dict[str, Any] = {
308 "type": "object",
309 "additionalProperties": False,
310 "properties": {
311 "tool_calls": {
312 "type": "array",
313 "minItems": 1,
314 "items": {
315 "type": "object",
316 "additionalProperties": True,
317 "required": ["tool_name", "args"],
318 "properties": {
319 "tool_name": {"type": "string", "minLength": 1},
320 "args": {"type": "object"},
321 },
322 },
323 },
324 "script": {"type": "string", "minLength": 1},
325 "expected_observation_keys": {
326 "type": "array",
327 "items": {"type": "string"},
328 },
329 "rationale": {"type": "string"},
330 },
331 "oneOf": [
332 {"required": ["tool_calls"]},
333 {"required": ["script"]},
334 ],
335}
337#: JSON Schema for the model's response when called for a
338#: Strategy_Revision. The model must return exactly these three keys.
339STRATEGY_REVISION_SCHEMA: dict[str, Any] = {
340 "$schema": "http://json-schema.org/draft-07/schema#",
341 "title": "Mission strategy revision",
342 "type": "object",
343 "additionalProperties": False,
344 "required": ["revision_rationale", "next_strategy", "confidence"],
345 "properties": {
346 "revision_rationale": {"type": "string", "minLength": 1},
347 "next_strategy": STRATEGY_SHAPE_SCHEMA,
348 "confidence": {
349 "type": "number",
350 "minimum": 0.0,
351 "maximum": 1.0,
352 },
353 },
354}
356#: JSON Schema for the model's response when called from the
357#: Final_Report writer.
358FINAL_LESSONS_SCHEMA: dict[str, Any] = {
359 "$schema": "http://json-schema.org/draft-07/schema#",
360 "title": "Mission final lessons",
361 "type": "object",
362 "additionalProperties": False,
363 "required": ["lessons", "recommended_followups"],
364 "properties": {
365 "lessons": {
366 "type": "array",
367 "minItems": 1,
368 "items": {"type": "string", "minLength": 1},
369 },
370 "recommended_followups": {
371 "type": "array",
372 "items": {"type": "string", "minLength": 1},
373 },
374 },
375}
378# ---------------------------------------------------------------------------
379# JSON helpers — every dump in this module routes through ``_dumps``
380# so the byte-counting and the rendered prompt agree on the encoding.
381# ---------------------------------------------------------------------------
384def _dumps(value: Any, *, indent: int | None = None) -> str:
385 """Deterministic JSON encoder used everywhere in this module.
387 ``sort_keys=True`` is the source of determinism — Python dicts are
388 insertion-ordered, but the Hypothesis strategies that drive the
389 determinism tests build dicts via ``fixed_dictionaries`` whose
390 insertion order is implementation-defined, so sorting is the only
391 way to get byte-identical output across two draws of the same
392 abstract dict shape. ``ensure_ascii=False`` keeps non-ASCII text
393 intact so the byte-budget bookkeeping matches what the LLM sees.
394 """
395 return json.dumps(
396 value,
397 sort_keys=True,
398 ensure_ascii=False,
399 indent=indent,
400 separators=(",", ": ") if indent is not None else (",", ":"),
401 )
404def _utf8_len(s: str) -> int:
405 """UTF-8 byte length of ``s`` — the only "size" the budget cares about."""
406 return len(s.encode("utf-8"))
409def _truncate_serialised(serialised: str) -> str:
410 """Slice a serialised value down to ``OBSERVATION_FIELD_TRUNCATE_TO``
411 bytes plus :data:`TRUNCATION_MARKER`. Decode-safe.
413 The slicing is byte-level rather than codepoint-level because the
414 cap itself is a byte budget. Using ``errors="ignore"`` strips any
415 partial codepoint at the boundary so the result is always valid
416 UTF-8 — at the cost of dropping at most three bytes' worth of an
417 incomplete codepoint, which is acceptable for an advisory summary.
418 """
419 truncated_bytes = serialised.encode("utf-8")[:OBSERVATION_FIELD_TRUNCATE_TO]
420 truncated_str = truncated_bytes.decode("utf-8", errors="ignore")
421 return truncated_str + TRUNCATION_MARKER
424# ---------------------------------------------------------------------------
425# Observation summarisation
426# ---------------------------------------------------------------------------
429def _summarise_observation(obs: Mapping[str, Any] | Observation) -> dict[str, Any]:
430 """Return a JSON-safe summary of an Observation with oversized fields
431 truncated and the original byte lengths recorded.
433 The summary mirrors the Observation's top-level keys. For each key
434 whose JSON-serialised value exceeds :data:`OBSERVATION_FIELD_BYTE_CAP`
435 bytes, the value is replaced by the byte-clamped + marker string and
436 the original byte length is recorded under
437 ``summary["_original_bytes"][<key>]``. Fields at or below the cap pass
438 through unchanged.
440 The ``_original_bytes`` private key is omitted entirely when no field
441 was truncated so the summary stays clean for the common case.
442 """
443 obs_map: Mapping[str, Any] = cast("Mapping[str, Any]", obs)
444 summary: dict[str, Any] = {}
445 original_bytes: dict[str, int] = {}
446 # Sorting the keys guarantees the rendered prompt is byte-identical
447 # even when the caller's dict was built in a different insertion
448 # order than another caller's identical-shape dict.
449 for key in sorted(obs_map.keys()):
450 if key == "_original_bytes":
451 # A defensively-guarded passthrough: a previous summarisation
452 # round (e.g., a re-render after dropping iterations) must
453 # not double-count the marker map.
454 continue
455 value = obs_map[key]
456 serialised = _dumps(value)
457 n_bytes = _utf8_len(serialised)
458 if n_bytes > OBSERVATION_FIELD_BYTE_CAP:
459 summary[key] = _truncate_serialised(serialised)
460 original_bytes[key] = n_bytes
461 else:
462 summary[key] = value
463 if original_bytes:
464 summary["_original_bytes"] = original_bytes
465 return summary
468def _summarise_iteration(iteration: Mapping[str, Any] | IterationRecord) -> dict[str, Any]:
469 """Build the per-iteration summary that feeds the Strategy_Revision prompt.
471 The summary keeps just the fields a downstream model needs to
472 reason about: the iteration index, the strategy that was tried,
473 the verdict + reason, and the size-capped Observation. Phase
474 timestamps and the criteria-evaluation list are intentionally
475 omitted because a) they are deterministic functions of fields the
476 model already sees in the criteria-status block, and b) keeping
477 them out shrinks the per-iteration footprint so the byte budget
478 holds with five iterations more often.
479 """
480 obs = iteration.get("observation") or {}
481 return {
482 "iteration_index": iteration.get("iteration_index"),
483 "strategy": iteration.get("strategy") or {},
484 "verdict": iteration.get("verdict"),
485 "verdict_reason": iteration.get("verdict_reason"),
486 "observation_summary": _summarise_observation(obs),
487 }
490def _summarise_iteration_for_lessons(
491 iteration: Mapping[str, Any] | IterationRecord,
492) -> dict[str, Any]:
493 """Final_Report summary — verdict + reason only, no Observation.
495 The Final_Report path needs to reason about *what happened* across
496 the run, not the per-iteration tool output. Dropping the Observation
497 keeps the prompt small enough that the byte budget never bites in
498 practice for sessions of any reasonable length.
499 """
500 return {
501 "iteration_index": iteration.get("iteration_index"),
502 "verdict": iteration.get("verdict"),
503 "verdict_reason": iteration.get("verdict_reason"),
504 }
507# ---------------------------------------------------------------------------
508# Criteria status pairing
509# ---------------------------------------------------------------------------
512def _pair_criteria_with_status(
513 criteria: Sequence[Criterion],
514 statuses: Sequence[CriterionResult],
515) -> list[dict[str, Any]]:
516 """Return ``criteria`` annotated with their most recent status entry.
518 Each entry in the result mirrors the Criterion definition (the kind,
519 the required-flag, the kind-specific payload keys) and adds a
520 nested ``status`` block populated from the matching ``CriterionResult``
521 by ``criterion_id``. Criteria with no matching status entry get
522 ``status`` set to ``{"status": "inconclusive", "evidence": null}``
523 so the model always sees a stable shape.
525 The ``_parsed_ast`` private key on a ``predicate`` criterion is
526 stripped — it is a Python ``ast.Expression`` object that is not
527 JSON-serialisable and that the model has no use for.
528 """
529 by_id: dict[str, CriterionResult] = {}
530 for s in statuses:
531 cid = s.get("criterion_id")
532 if cid is None:
533 continue
534 # If the caller passes duplicates (older then newer), prefer the
535 # last entry — that's the most-recent-wins convention the engine
536 # uses everywhere else.
537 by_id[cid] = s
539 out: list[dict[str, Any]] = []
540 for c in criteria:
541 cid = c.get("criterion_id")
542 # Strip private cached AST and surface only the prompt-relevant fields.
543 public = {k: v for k, v in c.items() if not k.startswith("_")}
544 match = by_id.get(cid) if cid is not None else None
545 if match is None:
546 public["status"] = {
547 "status": "inconclusive",
548 "evidence": None,
549 }
550 else:
551 public["status"] = {
552 "status": match.get("status"),
553 "evidence": match.get("evidence"),
554 "evaluated_at": match.get("evaluated_at"),
555 }
556 out.append(public)
557 return out
560# ---------------------------------------------------------------------------
561# Tool allowlist rendering
562# ---------------------------------------------------------------------------
565def _render_tool_allowlist(
566 allowlist: Sequence[str],
567 docstrings: Mapping[str, str],
568 schemas: Mapping[str, Any] | None = None,
569) -> list[dict[str, Any]]:
570 """Pair every allowlisted tool name with its docstring and input schema.
572 Tools without a registered docstring get an empty string — the
573 prompt remains valid; the model just sees a tool name with no
574 inline description. Tools without a schema get ``null`` so the
575 model knows no args are required. Names are emitted in the
576 caller's allowlist order so the prompt is identical for two
577 callers that pass the same list.
578 """
579 rendered: list[dict[str, Any]] = []
580 for name in allowlist:
581 entry: dict[str, Any] = {
582 "tool_name": name,
583 "docstring": str(docstrings.get(name, "")),
584 }
585 if schemas:
586 schema = schemas.get(name)
587 if schema is not None:
588 entry["input_schema"] = schema
589 rendered.append(entry)
590 return rendered
593# ---------------------------------------------------------------------------
594# Budget context rendering
595# ---------------------------------------------------------------------------
598def _render_budget_context(
599 *,
600 remaining_iterations: int,
601 remaining_wall_clock_secs: float | None,
602 allow_scripts: bool,
603) -> dict[str, Any]:
604 """Render the budget context block. Stable shape regardless of inputs.
606 ``None`` for the wall-clock cap is rendered verbatim as JSON
607 ``null`` so the model can disambiguate "unbounded" from "0".
608 """
609 return {
610 "remaining_iterations": int(remaining_iterations),
611 "remaining_wall_clock_seconds": (
612 float(remaining_wall_clock_secs) if remaining_wall_clock_secs is not None else None
613 ),
614 "allow_scripted_strategies": bool(allow_scripts),
615 }
618# ---------------------------------------------------------------------------
619# SamplingPrompt — the public class
620# ---------------------------------------------------------------------------
623@dataclass(frozen=True)
624class SamplingPrompt:
625 """Bundle the bare data needed to assemble a sampling prompt string.
627 The dataclass is ``frozen=True`` so callers cannot mutate the inputs
628 between an :meth:`assemble` call and an :meth:`assemble_final_lessons`
629 call — both methods produce deterministic outputs from the same
630 bound state, which is the property the determinism tests pin down.
632 All inputs are required positionally or by keyword; defaults are
633 only provided where the design spec defines a default.
634 """
636 directive: str
637 success_criteria: Sequence[Criterion]
638 criteria_status: Sequence[CriterionResult]
639 recent_iterations: Sequence[IterationRecord]
640 tool_allowlist: Sequence[str]
641 tool_docstrings: Mapping[str, str]
642 remaining_iterations: int
643 remaining_wall_clock_secs: float | None
644 allow_scripts: bool = field(default=False)
645 #: Per-tool JSON Schema for the input parameters. Keyed by tool
646 #: name; values are the JSON-serialisable schema dict (or ``None``
647 #: for tools that take no args). Included in the prompt so the
648 #: Strategy_Revision model can propose valid ``args`` dicts.
649 tool_schemas: Mapping[str, Any] = field(default_factory=dict)
650 #: Optional snapshot of slow-moving live signals (per-region queue
651 #: depth, GPU utilisation, deployed-region list, reservation
652 #: counts, etc.) gathered once at session start and reused on
653 #: every iteration's prompt. ``None`` (the default) suppresses the
654 #: ``=== Environment context ===`` section entirely so the prompt
655 #: stays byte-identical to the pre-environment-context shape —
656 #: that's what every existing determinism test pins down.
657 environment_context: Mapping[str, Any] | None = field(default=None)
658 #: Optional list of similar past missions retrieved from the
659 #: mission-memory vector index (most-similar-first), gathered once
660 #: per engine wiring and reused on every iteration's prompt.
661 #: ``None`` (the default) suppresses the ``=== Prior similar
662 #: missions ===`` section entirely — the same byte-identical
663 #: contract as :attr:`environment_context`, and what keeps every
664 #: pre-memory prompt (and the determinism suite) unchanged.
665 prior_missions: Sequence[Mapping[str, Any]] | None = field(default=None)
667 # ---- Strategy_Revision rendering --------------------------------------
669 def assemble(self) -> str:
670 """Return the Strategy_Revision prompt string.
672 The output is capped at :data:`PROMPT_BYTE_BUDGET` UTF-8 bytes.
673 When the freshly-assembled prompt exceeds the cap, the oldest
674 Iteration summary is dropped and the prompt re-rendered. The
675 loop terminates because each drop monotonically shrinks the
676 prompt and there is a non-iteration baseline that fits well
677 under the cap on its own (the directive, criteria, allowlist,
678 budget block, and schema instruction together are ~6-10 KB
679 for any reasonable session shape).
680 """
681 # Defensive slice — the caller is asked to pass at most five,
682 # but if they pass more, take the most recent five.
683 iterations: list[IterationRecord] = list(self.recent_iterations[-RECENT_ITERATIONS_LIMIT:])
685 while True:
686 text = self._render(
687 schema=STRATEGY_REVISION_SCHEMA,
688 iterations=[_summarise_iteration(it) for it in iterations],
689 schema_purpose="strategy_revision",
690 )
691 if _utf8_len(text) <= PROMPT_BYTE_BUDGET or not iterations:
692 return text
693 # Drop the oldest iteration and try again.
694 iterations = iterations[1:]
696 # ---- Final_Report rendering -------------------------------------------
698 def assemble_final_lessons(self) -> str:
699 """Return the Final_Report ``lessons`` prompt string.
701 The shape parallels :meth:`assemble` but emits
702 :data:`FINAL_LESSONS_SCHEMA` and uses iteration **verdict
703 summaries only** instead of full Observation summaries. The same
704 :data:`PROMPT_BYTE_BUDGET` byte cap applies; the same
705 oldest-first drop policy kicks in if the cap is exceeded.
706 """
707 iterations: list[IterationRecord] = list(self.recent_iterations)
709 while True:
710 text = self._render(
711 schema=FINAL_LESSONS_SCHEMA,
712 iterations=[_summarise_iteration_for_lessons(it) for it in iterations],
713 schema_purpose="final_lessons",
714 )
715 if _utf8_len(text) <= PROMPT_BYTE_BUDGET or not iterations:
716 return text
717 iterations = iterations[1:]
719 # ---- Internal renderer ------------------------------------------------
721 def _render(
722 self,
723 *,
724 schema: dict[str, Any],
725 iterations: Sequence[Mapping[str, Any]],
726 schema_purpose: str,
727 ) -> str:
728 """Format the full prompt from the section blocks.
730 The text layout is fixed — every section is delimited by a
731 ``=== <name> ===`` header so the model can latch onto a
732 predictable structure. Section bodies are JSON wherever the
733 content is structured; the directive itself is rendered as
734 free text because that is how the operator wrote it.
735 """
736 criteria_block = _pair_criteria_with_status(self.success_criteria, self.criteria_status)
737 tool_block = _render_tool_allowlist(
738 self.tool_allowlist, self.tool_docstrings, self.tool_schemas
739 )
740 budget_block = _render_budget_context(
741 remaining_iterations=self.remaining_iterations,
742 remaining_wall_clock_secs=self.remaining_wall_clock_secs,
743 allow_scripts=self.allow_scripts,
744 )
746 if schema_purpose == "strategy_revision":
747 preamble = (
748 "You are advising a Mission goal-directed iteration loop. "
749 "Propose the next Strategy that moves the Mission toward "
750 "satisfying its Success_Criteria. The Verdict label, "
751 "budget enforcement, and Criteria evaluation are all "
752 "computed server-side and are unaffected by your output. "
753 "Your role is advisory: the rationale and next_strategy "
754 "you produce are validated against the Tool_Allowlist and "
755 "the remaining budget before being adopted.\n\n"
756 "IMPORTANT: You may propose MULTIPLE tool calls in a "
757 "single iteration by including multiple entries in the "
758 "tool_calls array. This is especially useful when the "
759 "unmet criteria require results from different tools — "
760 "calling them all in one iteration lets the evaluator "
761 "see all results together. Use the input_schema in the "
762 "Tool allowlist section to construct valid args for each "
763 "tool call."
764 )
765 recent_header = "Recent iterations (oldest first)"
766 else:
767 preamble = (
768 "You are advising a Mission goal-directed iteration loop "
769 "that has just reached a terminal Verdict. Produce the "
770 "lessons learned and the recommended follow-ups for the "
771 "operator. Your output is merged into the Final_Report; "
772 "the Verdict label, budget bookkeeping, and Criteria "
773 "evaluation that produced the terminal state are "
774 "deterministic server-side outputs and are not under "
775 "review."
776 )
777 recent_header = "Iteration verdict summary (oldest first)"
779 sections: list[str] = []
780 sections.append(preamble)
781 sections.append("")
782 sections.append("=== Mission directive ===")
783 sections.append(self.directive)
784 sections.append("")
785 sections.append("=== Success criteria with current status ===")
786 sections.append(_dumps(criteria_block, indent=2))
787 sections.append("")
788 sections.append("=== Tool allowlist ===")
789 sections.append(_dumps(tool_block, indent=2))
790 sections.append("")
791 sections.append("=== Budget context ===")
792 sections.append(_dumps(budget_block, indent=2))
793 sections.append("")
794 if self.environment_context is not None:
795 # Truncated + key-sorted — see :func:`_summarise_environment_context`.
796 # Emitting a header even for an empty dict means a session
797 # that opted in but had a probe failure still surfaces
798 # "we tried" so the operator can act on the gap.
799 env_summary = _summarise_environment_context(self.environment_context)
800 sections.append("=== Environment context (slow-moving live signals) ===")
801 sections.append(_dumps(env_summary, indent=2))
802 sections.append("")
803 if self.prior_missions is not None:
804 # Institutional memory: the closest past missions by directive
805 # similarity, with their lessons and verdicts. Advisory only —
806 # summarised and byte-capped in its own truncation domain.
807 sections.append("=== Prior similar missions (institutional memory) ===")
808 sections.append(
809 "Lessons and outcomes from the most similar past missions, "
810 "most similar first. Treat them as advisory context: they "
811 "may suggest which tools or query shapes worked before, or "
812 "what to avoid repeating."
813 )
814 sections.append(_dumps(_summarise_prior_missions(self.prior_missions), indent=2))
815 sections.append("")
816 sections.append(f"=== {recent_header} ===")
817 sections.append(_dumps(list(iterations), indent=2))
818 sections.append("")
819 sections.append("=== Output schema ===")
820 sections.append(
821 "Respond with a single JSON object that validates against "
822 "the JSON Schema below. Do not include any prose outside "
823 "the JSON object."
824 )
825 sections.append(_dumps(schema, indent=2))
827 return "\n".join(sections)
830# ---------------------------------------------------------------------------
831# Backend protocol and transport-error type
832# ---------------------------------------------------------------------------
835@runtime_checkable
836class SamplingBackend(Protocol):
837 """Transport-agnostic surface for the advisory LLM call.
839 Implementations bind a concrete transport (e.g., the MCP
840 ``Context.sample`` capability or ``bedrock-runtime:Converse``) and
841 expose a single async ``sample`` entry point. The protocol is
842 ``runtime_checkable`` so call sites — and tests — can use
843 ``isinstance(backend, SamplingBackend)`` to gate dispatch on a
844 duck-typed backend instance.
846 Attributes:
847 backend_name: Stable identifier the audit pipeline emits in the
848 ``sampling_backend`` field. Bedrock is the only transport
849 the system supports: MCP client sampling (``ctx.sample``)
850 left the protocol with FastMCP 4's sessionless era, so
851 missions sample server-side regardless of how they were
852 started.
853 model_id: The concrete model identifier the backend will route
854 the prompt to. Echoed in audit events so replay can
855 reproduce the exact request.
856 """
858 backend_name: Literal["bedrock"]
859 model_id: str
861 async def sample(self, prompt: SamplingPrompt) -> str:
862 """Render ``prompt`` through the bound transport and return the
863 raw model output text. Implementations raise
864 :class:`SamplingTransportError` (with a transport-tagged
865 ``code``) on any transport-layer failure so the engine's
866 fallback policy can branch on a single, well-typed exception.
867 """
868 ...
871class SamplingTransportError(Exception):
872 """Transport-layer failure raised by a :class:`SamplingBackend`.
874 The mandatory ``code`` attribute tags the failure class so the
875 engine's deterministic-fallback path can branch on a stable string
876 without parsing the message. The convention is
877 ``"<backend>_<error_class>"`` for backend-specific failures and a
878 short, snake-cased label for backend-agnostic failures.
880 Documented codes (used elsewhere in the Mission stack):
882 * ``"bedrock_AccessDeniedException"`` — IAM denied
883 ``bedrock:InvokeModel`` for the resolved model.
884 * ``"bedrock_malformed_response"`` — Converse returned a payload
885 that did not have the expected ``output.message.content[0].text``
886 shape.
887 * ``"bedrock_truncated_response"`` — the answer was cut off by an
888 output-token limit (``stopReason == "max_tokens"``), so its text
889 cannot be trusted to be complete.
890 * ``"bedrock_no_credentials"`` — the local ``boto3`` session could
891 not resolve credentials.
893 Args:
894 code: Mandatory failure tag (see examples above).
895 message: Optional human-readable detail. When present, it is
896 joined to ``code`` with ``": "`` for the string
897 representation; when absent, ``str(self)`` is just the
898 ``code``.
899 """
901 def __init__(self, code: str, message: str | None = None) -> None:
902 self.code: str = code
903 self.message: str | None = message
904 # Forward the most useful single-line representation to
905 # ``Exception.__init__`` so ``logging`` / ``traceback`` modules
906 # show the same string ``str(self)`` produces below.
907 if message is None:
908 super().__init__(code)
909 else:
910 super().__init__(f"{code}: {message}")
912 def __str__(self) -> str:
913 if self.message is None:
914 return self.code
915 return f"{self.code}: {self.message}"
918# ---------------------------------------------------------------------------
919# BedrockSamplingBackend — routes the prompt through bedrock-runtime:Converse
920# ---------------------------------------------------------------------------
923class BedrockSamplingBackend:
924 """Sampling backend that calls ``bedrock-runtime:Converse``.
926 The backend resolves its model id and region at construction time
927 from (in order of precedence) the explicit constructor argument,
928 the matching environment variable
929 (:data:`ENV_BEDROCK_MODEL_ID` / :data:`ENV_BEDROCK_REGION`), and
930 finally the ``cdk.json`` Mission default
931 (:func:`gco.bedrock.get_default_mission_model_id` /
932 :data:`DEFAULT_BEDROCK_REGION`). The ``boto3`` client itself is
933 constructed lazily on the first :meth:`sample` call so that
934 ``import mission.sampling`` does not pull ``boto3`` into the
935 import graph and so that test code can swap the import in via
936 ``unittest.mock.patch`` without paying for a real session at
937 construction time.
939 Failure modes:
941 * Missing or partial AWS credentials at client-construction time
942 surface as :class:`SamplingTransportError` with code
943 ``"bedrock_no_credentials"``; the original exception is chained
944 via ``__cause__``.
945 * A ``botocore.exceptions.ClientError`` from the ``Converse`` call
946 surfaces as :class:`SamplingTransportError` with code
947 ``"bedrock_<ErrorCode>"`` where ``<ErrorCode>`` is read from the
948 error envelope (defaulting to ``"Unknown"`` when the envelope is
949 malformed). The one exception is the Anthropic first-time-use
950 gate, which raises
951 :class:`gco.bedrock.BedrockFTUFormNotAcceptedError` instead of a
952 transport error so it is never absorbed by a deterministic
953 fallback. See ``docs/CUSTOMIZATION.md`` (Bedrock Model Selection).
954 * A response without a non-empty ``text`` block under
955 ``output.message.content`` — including reasoning-only and empty
956 ``content`` lists — surfaces as :class:`SamplingTransportError`
957 with code ``"bedrock_malformed_response"``.
958 * A response cut off by an output-token limit
959 (``stopReason == "max_tokens"``) surfaces as
960 :class:`SamplingTransportError` with code
961 ``"bedrock_truncated_response"``.
962 """
964 backend_name: Literal["bedrock"] = "bedrock"
966 def __init__(
967 self,
968 model_id: str | None = None,
969 region: str | None = None,
970 ) -> None:
971 """Resolve the model id and region; defer client construction.
973 Args:
974 model_id: Optional explicit model id. When ``None``, falls
975 back to the :data:`ENV_BEDROCK_MODEL_ID` environment
976 variable, then to the ``cdk.json`` Mission default from
977 :func:`gco.bedrock.get_default_mission_model_id`.
978 region: Optional explicit region. When ``None``, falls back
979 to :data:`ENV_BEDROCK_REGION`, then to
980 :data:`DEFAULT_BEDROCK_REGION`.
981 """
982 if model_id is not None:
983 self.model_id = model_id
984 self._uses_default_model = False
985 elif ENV_BEDROCK_MODEL_ID in os.environ:
986 # Preserve the existing explicit-environment semantics, including
987 # an intentionally empty value, without evaluating the fallback.
988 self.model_id = os.environ[ENV_BEDROCK_MODEL_ID]
989 self._uses_default_model = False
990 else:
991 self.model_id = get_default_mission_model_id()
992 self._uses_default_model = True
993 self._region: str = (
994 region
995 if region is not None
996 else os.environ.get(ENV_BEDROCK_REGION, DEFAULT_BEDROCK_REGION)
997 )
998 # The boto3 client is built on first ``sample`` call. ``None``
999 # here is the sentinel for "not yet constructed".
1000 self._client: Any = None
1002 @classmethod
1003 def from_canonical_default(
1004 cls,
1005 region: str | None = None,
1006 ) -> BedrockSamplingBackend:
1007 """Build a backend that deliberately applies canonical reasoning.
1009 Unlike ``cls(model_id=None)``, this bypasses the model environment
1010 override. Fixture capture uses it to reproduce the checked-in default
1011 exactly, while ordinary explicit model IDs retain override semantics.
1012 """
1013 backend = cls(model_id=get_default_mission_model_id(), region=region)
1014 backend._uses_default_model = True
1015 return backend
1017 def _get_client(self) -> Any:
1018 """Return the cached ``bedrock-runtime`` client, building it on first use.
1020 ``boto3`` and ``botocore.exceptions`` are imported here rather
1021 than at module top-level so that pure-Python consumers of this
1022 module (the prompt builder, the protocol, the error type) do
1023 not pay for the ``boto3`` import. This also lets tests patch
1024 ``mission.sampling.boto3`` after import.
1025 """
1026 if self._client is not None:
1027 return self._client
1028 # Local import — keeps the module's import surface boto3-free.
1029 import boto3
1030 from botocore.config import Config
1031 from botocore.exceptions import (
1032 NoCredentialsError,
1033 PartialCredentialsError,
1034 )
1036 try:
1037 self._client = boto3.Session().client(
1038 "bedrock-runtime",
1039 region_name=self._region,
1040 config=Config(read_timeout=BEDROCK_READ_TIMEOUT_SECONDS),
1041 )
1042 except (NoCredentialsError, PartialCredentialsError) as err:
1043 raise SamplingTransportError("bedrock_no_credentials") from err
1044 return self._client
1046 async def sample(self, prompt: SamplingPrompt) -> str:
1047 """Render ``prompt`` through ``Converse`` and return the response text.
1049 Raises:
1050 SamplingTransportError: On any transport-level failure.
1051 * ``bedrock_no_credentials`` — credentials could not be
1052 resolved by ``boto3`` at client-construction time.
1053 * ``bedrock_<ErrorCode>`` — the ``Converse`` call raised
1054 a ``ClientError``; ``<ErrorCode>`` is the AWS error
1055 code from the envelope.
1056 * ``bedrock_malformed_response`` — the response did not
1057 contain a non-empty final text content block.
1058 * ``bedrock_truncated_response`` — the response was cut
1059 off by an output-token limit and cannot be trusted to
1060 be complete.
1061 gco.bedrock.BedrockFTUFormNotAcceptedError: The account has
1062 not submitted Anthropic's one-time first-time-use case
1063 form. Raised instead of a transport error so callers
1064 cannot silently fall back past a permanent, one-line-fix
1065 misconfiguration.
1066 """
1067 # Local import — see ``_get_client`` for the rationale.
1068 from botocore.exceptions import ClientError
1070 client = self._get_client()
1072 text = prompt.assemble()
1073 converse_options = build_bedrock_converse_options(
1074 self.model_id,
1075 # Deliberately no maxTokens: the Converse default is the model's
1076 # own maximum output length, so a rationale can never be cut off
1077 # by a GCO-imposed cap. A cap is opt-in — pass maxTokens here to
1078 # restore one.
1079 inference_config={"temperature": BEDROCK_TEMPERATURE},
1080 apply_default_reasoning=self._uses_default_model,
1081 )
1082 try:
1083 response = await asyncio.to_thread(
1084 client.converse,
1085 modelId=self.model_id,
1086 messages=[{"role": "user", "content": [{"text": text}]}],
1087 **converse_options,
1088 )
1089 except ClientError as err:
1090 # A missing Anthropic FTU form is a permanent account-scoped
1091 # misconfiguration, not a transport fault: escalate it instead of
1092 # letting the deterministic-fallback path absorb it silently.
1093 raise_if_bedrock_ftu_form_error(err)
1094 # ``e.response`` is documented to be present on ClientError
1095 # but the envelope shape can vary; defend against missing
1096 # keys so the audit pipeline always sees a tagged code.
1097 envelope = getattr(err, "response", None) or {}
1098 error_block = envelope.get("Error", {}) if isinstance(envelope, dict) else {}
1099 code = (
1100 error_block.get("Code", "Unknown") if isinstance(error_block, dict) else "Unknown"
1101 )
1102 raise SamplingTransportError(f"bedrock_{code}") from err
1104 # Capture token usage from the Converse response for the audit
1105 # trail. The ``usage`` block is present on every successful
1106 # Converse response and carries ``inputTokens`` and
1107 # ``outputTokens``. Store on the instance so callers can read
1108 # it after each sample() call without changing the protocol.
1109 usage = response.get("usage") or {}
1110 self.last_input_tokens: int | None = usage.get("inputTokens")
1111 self.last_output_tokens: int | None = usage.get("outputTokens")
1113 try:
1114 return extract_bedrock_converse_text(response)
1115 except BedrockResponseTruncatedError as err:
1116 # A cut-off rationale is unusable; let the deterministic-fallback
1117 # path absorb it like any other transport-shaped fault.
1118 raise SamplingTransportError("bedrock_truncated_response") from err
1119 except (KeyError, IndexError, TypeError) as err:
1120 raise SamplingTransportError("bedrock_malformed_response") from err
1123# ---------------------------------------------------------------------------
1124# Backend resolver
1125# ---------------------------------------------------------------------------
1128def select_sampling_backend(model_id: str | None) -> SamplingBackend:
1129 """Construct the sampling backend for a session that opted into sampling.
1131 Bedrock is the only sampling transport. MCP client sampling
1132 (``ctx.sample``) left the protocol with FastMCP 4's sessionless era —
1133 per the v4 migration guidance, generation belongs server-side — so
1134 missions sample through ``bedrock-runtime:Converse`` with the server's
1135 own credentials regardless of whether they were started from the CLI
1136 or over MCP. Credential resolution is deferred to the first ``sample``
1137 call; a missing-credentials failure surfaces as
1138 :class:`SamplingTransportError` and the engine's deterministic
1139 fallback absorbs it.
1141 Args:
1142 model_id: Optional concrete model identifier. Forwarded to the
1143 backend constructor verbatim; ``None`` resolves through the
1144 environment and ``cdk.json`` Mission default.
1146 Returns:
1147 A :class:`BedrockSamplingBackend` bound to ``model_id``.
1148 """
1149 return BedrockSamplingBackend(model_id)
1152# ---------------------------------------------------------------------------
1153# Strategy-against-catalog validator
1154# ---------------------------------------------------------------------------
1157def _resolve_input_schema(tool: Any) -> Any:
1158 """Return the registered Pydantic input model for a Tool, or None.
1160 FastMCP exposes the model under ``input_schema`` in newer releases
1161 and ``inputSchema`` in older ones. Tolerate both. Tools that genuinely
1162 take no args (or test catalog mocks that omit the attribute) yield
1163 ``None``, in which case the caller skips per-call args validation.
1164 """
1165 schema = getattr(tool, "input_schema", None)
1166 if schema is None:
1167 schema = getattr(tool, "inputSchema", None)
1168 return schema
1171def _extract_tool_json_schemas(
1172 allowlist: Sequence[str],
1173 registered_tools: Mapping[str, Any],
1174) -> dict[str, Any]:
1175 """Extract JSON Schema dicts for each allowlisted tool's input parameters.
1177 Calls ``.model_json_schema()`` on the Pydantic model exposed by
1178 ``_resolve_input_schema``. Falls back gracefully: tools without a
1179 schema, tools whose schema isn't a Pydantic model, and any
1180 exception during schema extraction all yield ``None`` for that
1181 tool (omitted from the output dict). The caller renders the
1182 result into the Strategy_Revision prompt so the model can propose
1183 valid ``args`` dicts.
1184 """
1185 schemas: dict[str, Any] = {}
1186 for name in allowlist:
1187 tool = registered_tools.get(name)
1188 if tool is None:
1189 continue
1190 model = _resolve_input_schema(tool)
1191 if model is None:
1192 continue
1193 try:
1194 # Pydantic v2 models expose model_json_schema() as a classmethod.
1195 json_schema = model.model_json_schema()
1196 schemas[name] = json_schema
1197 except Exception:
1198 # Non-Pydantic schema, or a mock that doesn't support it.
1199 continue
1200 return schemas
1203def validate_strategy_against_catalog(
1204 strategy: Strategy,
1205 allowlist: list[str],
1206 registered_tools: dict[str, Any],
1207 allow_scripts: bool,
1208) -> None:
1209 """Validate a Strategy against the live tool catalog.
1211 Returns ``None`` on accept; raises :class:`MissionValidationError`
1212 with a structured ``details.reason`` enum on reject. The function
1213 layers catalog-aware checks on top of the structural validation in
1214 :func:`mission.validation.validate_strategy`:
1216 1. Mutual-exclusivity (exactly one of ``tool_calls`` / ``script``).
1217 2. Per-call ``tool_name`` is in ``allowlist``.
1218 3. Per-call ``args`` validates against the registered Pydantic model
1219 exposed under ``Tool.input_schema`` (or the older
1220 ``Tool.inputSchema``); calls whose tool has neither attribute or
1221 a ``None`` schema skip args validation.
1222 4. For scripted strategies, ``allow_scripts`` is True and the
1223 script's AST passes :func:`mission.sandbox.validate_script_ast`.
1225 Args:
1226 strategy: The Strategy dict to validate.
1227 allowlist: The session's resolved Tool_Allowlist.
1228 registered_tools: Mapping from tool name to a registered tool
1229 object (typed ``Any`` so the module imports without
1230 FastMCP). Read-only — only ``input_schema`` /
1231 ``inputSchema`` is consulted.
1232 allow_scripts: Session-level flag gating scripted strategies.
1233 """
1234 # 1. Structural validation: mutual exclusivity, script-allow gating,
1235 # and AST validation for scripts. Reuses the existing validator
1236 # so error shapes for those rejection classes stay aligned with
1237 # the rest of the input pipeline.
1238 _validation.validate_strategy(cast("dict[str, Any]", strategy), allowlist, allow_scripts)
1240 # The structural validator has already accepted exactly one of the
1241 # two shapes. Branch on which one is present.
1242 if "tool_calls" in strategy:
1243 tool_calls = strategy["tool_calls"]
1244 # Empty list is rejected by validate_strategy; this is a defence
1245 # in depth for callers that might bypass that path.
1246 if not tool_calls:
1247 raise MissionValidationError(
1248 "validation_error",
1249 details={
1250 "field": "strategy",
1251 "subfield": "tool_calls",
1252 "reason": "tool_calls_empty",
1253 },
1254 )
1256 # 2. Per-call name-in-allowlist check.
1257 for call in tool_calls:
1258 name = call.get("tool_name")
1259 if name not in allowlist:
1260 raise MissionValidationError(
1261 "validation_error",
1262 details={
1263 "field": "strategy",
1264 "subfield": "tool_calls",
1265 "tool_name": name,
1266 "reason": "tool_not_allowlisted",
1267 "allowlist": list(allowlist),
1268 },
1269 )
1271 # 3. Per-call args validation against the tool's Pydantic model.
1272 for call in tool_calls:
1273 name = call["tool_name"]
1274 tool = registered_tools.get(name)
1275 if tool is None:
1276 # Catalog could have a name in the allowlist that is not
1277 # currently registered (gating, dynamic load). Mirror the
1278 # unknown-tool shape used elsewhere.
1279 raise MissionValidationError(
1280 "validation_error",
1281 details={
1282 "field": "strategy",
1283 "subfield": "tool_calls",
1284 "tool_name": name,
1285 "reason": "tool_not_registered",
1286 },
1287 )
1288 schema = _resolve_input_schema(tool)
1289 if schema is None:
1290 # Either the tool genuinely takes no args, or the test
1291 # catalog omitted a model. Skip args validation rather
1292 # than reject — the design treats missing schema as
1293 # "trust the dispatcher".
1294 continue
1295 args = call.get("args", {})
1296 if not isinstance(args, dict):
1297 raise MissionValidationError(
1298 "validation_error",
1299 details={
1300 "field": "strategy",
1301 "subfield": "tool_calls",
1302 "tool_name": name,
1303 "reason": "tool_args_invalid",
1304 "errors": [
1305 {
1306 "type": "args_not_a_dict",
1307 "actual_type": type(args).__name__,
1308 }
1309 ],
1310 },
1311 )
1312 try:
1313 schema.model_validate(args)
1314 except Exception as exc: # noqa: BLE001 - pydantic ValidationError + similar
1315 # Pydantic v2 ValidationError exposes ``.errors()`` as a
1316 # list of structured dicts. Tolerate any other exception
1317 # type (e.g. older Pydantic, custom validators) by
1318 # falling back to ``str(exc)``.
1319 errors_method = getattr(exc, "errors", None)
1320 if callable(errors_method):
1321 try:
1322 errors_payload: Any = errors_method()
1323 except Exception: # noqa: BLE001 - defensive
1324 errors_payload = [{"type": "unknown", "msg": str(exc)}]
1325 else:
1326 errors_payload = [{"type": "unknown", "msg": str(exc)}]
1327 raise MissionValidationError(
1328 "validation_error",
1329 details={
1330 "field": "strategy",
1331 "subfield": "tool_calls",
1332 "tool_name": name,
1333 "reason": "tool_args_invalid",
1334 "errors": errors_payload,
1335 },
1336 ) from exc
1338 # 4. Cost estimation against remaining budget. Removed —
1339 # cost guardrails live out-of-band via AWS Budgets / Cost
1340 # Anomaly Detection rather than in the Mission cascade.
1341 # Scripted strategies: validate_strategy already ran allow_scripts
1342 # gating and the AST validator. No catalog-aware checks are layered
1343 # on top here — the script-side enforcement happens at execute time
1344 # via the in-script tool callable wrappers.
1347# ---------------------------------------------------------------------------
1348# Orchestration helpers — bind a backend to a SessionState and return either
1349# a used result or a deterministic fallback.
1350# ---------------------------------------------------------------------------
1352# Local imports kept inside this section so the prompt-builder /
1353# backend half above stays free of audit / decide dependencies.
1354from . import audit as _mission_audit # noqa: E402
1355from . import decide as _decide # noqa: E402
1357# Type alias used by the helpers below. ``SessionState`` is a TypedDict
1358# whose runtime value is just ``dict``; the alias keeps the signatures
1359# expressive without forcing the import to leak through ``__all__``.
1360from .types import SessionState as _SessionState # noqa: E402
1363@dataclass(frozen=True)
1364class SamplingUsed:
1365 """A successful sampling call's accepted output."""
1367 output_text: str
1368 """Raw model output (the text returned by the bound backend)."""
1370 parsed: dict[str, Any]
1371 """Parsed JSON payload that has cleared the schema and catalog checks."""
1373 backend_name: Literal["bedrock"]
1374 """Stable backend identifier — echoes the bound backend's tag."""
1376 model_id: str
1377 """The concrete model id the backend routed the prompt to."""
1380@dataclass(frozen=True)
1381class SamplingFallback:
1382 """A rejected or unavailable sampling call's deterministic substitute.
1384 Returned when the bound backend was ``None``, the transport raised,
1385 the model output failed to parse / validate, or any catalog or
1386 budget check rejected the proposed strategy. The ``rationale`` is
1387 a pure function of the bound :class:`SessionState` and the most
1388 recent :class:`IterationRecord`, so the engine can replay or
1389 reproduce a fallback exactly from persisted state.
1390 """
1392 rationale: str
1393 """Deterministic fallback text. Empty string for ``final_lessons``
1394 — the final-report writer fills in its own deterministic text in
1395 that case."""
1397 reason: str
1398 """Stable token tagging *why* the fallback fired. Examples:
1399 ``"transport_error"``, ``"json_parse"``, ``"schema_mismatch"``,
1400 ``"tool_not_allowlisted"``, ``"tool_args_invalid"``,
1401 ``"over_budget"``, ``"script_rejected"``,
1402 ``"no_backend_resolved"``, ``"disabled"``."""
1404 backend_name: Literal["bedrock", "none"]
1405 """The bound backend's tag, or ``"none"`` when no backend was
1406 resolved at the call site."""
1408 model_id: str | None
1409 """The bound backend's model id, or ``None`` when no backend was
1410 resolved."""
1413# ---------------------------------------------------------------------------
1414# JSON / schema helpers (private to the orchestration layer)
1415# ---------------------------------------------------------------------------
1418def _extract_json_object(text: str) -> dict[str, Any]:
1419 """Parse the first JSON object embedded in ``text``.
1421 Models routinely wrap JSON in prose. The implementation slices from
1422 the first ``{`` to the last ``}`` and feeds the result to
1423 :func:`json.loads`. When no balanced braces are present, or the
1424 sliced substring is not valid JSON, the function raises
1425 :class:`json.JSONDecodeError` so the caller can branch on a single
1426 well-typed exception.
1427 """
1428 start = text.find("{")
1429 end = text.rfind("}")
1430 if start == -1 or end == -1 or end < start:
1431 # No braces at all → treat as a parse error so the calling
1432 # branch surfaces ``reason="json_parse"``.
1433 raise json.JSONDecodeError("no JSON object found", text, 0)
1434 candidate = text[start : end + 1]
1435 parsed = json.loads(candidate)
1436 if not isinstance(parsed, dict):
1437 # The sliced substring parsed but is not an object — surface as
1438 # a parse error too, since downstream code requires a dict.
1439 raise json.JSONDecodeError("top-level JSON value is not an object", candidate, 0)
1440 return parsed
1443def _validate_revision_schema(parsed: dict[str, Any]) -> None:
1444 """Reject a parsed payload that is not a valid Strategy_Revision.
1446 Required keys: ``revision_rationale`` (non-empty str),
1447 ``next_strategy`` (dict), ``confidence`` (number in [0, 1]).
1448 """
1449 rationale = parsed.get("revision_rationale")
1450 if not isinstance(rationale, str) or not rationale:
1451 raise ValueError("schema_mismatch: revision_rationale must be non-empty str")
1452 next_strategy = parsed.get("next_strategy")
1453 if not isinstance(next_strategy, dict):
1454 raise ValueError("schema_mismatch: next_strategy must be a dict")
1455 confidence = parsed.get("confidence")
1456 # ``bool`` is excluded explicitly — it is a subclass of ``int`` in
1457 # Python and would otherwise sneak past the numeric check.
1458 if isinstance(confidence, bool) or not isinstance(confidence, (int, float)):
1459 raise ValueError("schema_mismatch: confidence must be a number")
1460 if not (0.0 <= float(confidence) <= 1.0):
1461 raise ValueError("schema_mismatch: confidence must be in [0, 1]")
1464def _validate_lessons_schema(parsed: dict[str, Any]) -> None:
1465 """Reject a parsed payload that is not a valid final-lessons dict.
1467 Required keys: ``lessons`` (non-empty list of non-empty str),
1468 ``recommended_followups`` (list of str — may be empty).
1469 """
1470 lessons = parsed.get("lessons")
1471 if not isinstance(lessons, list) or not lessons:
1472 raise ValueError("schema_mismatch: lessons must be a non-empty list")
1473 for item in lessons:
1474 if not isinstance(item, str) or not item:
1475 raise ValueError("schema_mismatch: each lesson must be a non-empty str")
1476 followups = parsed.get("recommended_followups")
1477 if not isinstance(followups, list):
1478 raise ValueError("schema_mismatch: recommended_followups must be a list")
1479 for item in followups:
1480 if not isinstance(item, str):
1481 raise ValueError("schema_mismatch: each follow-up must be a str")
1484# ---------------------------------------------------------------------------
1485# maybe_sample_strategy_revision
1486# ---------------------------------------------------------------------------
1489async def maybe_sample_strategy_revision(
1490 *,
1491 backend: SamplingBackend | None,
1492 session: _SessionState,
1493 iteration: IterationRecord,
1494 allowlist: list[str],
1495 registered_tools: dict[str, Any],
1496 tool_docstrings: dict[str, str],
1497 remaining_iterations: int,
1498 remaining_wall_clock_secs: float | None,
1499 allow_scripts: bool,
1500 environment_context: Mapping[str, Any] | None = None,
1501 prior_missions: Sequence[Mapping[str, Any]] | None = None,
1502) -> SamplingUsed | SamplingFallback:
1503 """Consult the advisory LLM for a Strategy_Revision, or fall back.
1505 Returns a :class:`SamplingUsed` when the bound backend produces a
1506 JSON object that clears schema validation and the catalog checks.
1507 Returns a :class:`SamplingFallback` carrying the deterministic
1508 rationale from
1509 :func:`mission.decide.build_revision_rationale_template` on every
1510 rejection class. Emits exactly one
1511 :func:`mission.audit.emit_sampling_event` per call.
1512 """
1513 session_id = session["session_id"]
1514 iteration_index = iteration["iteration_index"]
1515 template = _decide.build_revision_rationale_template(session, iteration)
1517 # ---- No backend resolved: short-circuit. ------------------------------
1518 if backend is None:
1519 _mission_audit.emit_sampling_event(
1520 session_id,
1521 iteration_index,
1522 sampling_purpose="strategy_revision",
1523 sampling_status="disabled",
1524 sampling_backend="none",
1525 )
1526 return SamplingFallback(
1527 rationale=template,
1528 reason="no_backend_resolved",
1529 backend_name="none",
1530 model_id=None,
1531 )
1533 backend_name = backend.backend_name
1534 model_id = backend.model_id
1536 # ---- Build the prompt. ------------------------------------------------
1537 # The in-progress iteration that triggered ``adjust`` is already in
1538 # ``session["iterations"][-1]``, so the most-recent-five window is a
1539 # plain slice; ``RECENT_ITERATIONS_LIMIT`` is enforced inside the
1540 # prompt builder as a defensive cap.
1541 recent_iterations = list(session["iterations"][-RECENT_ITERATIONS_LIMIT:])
1542 tool_schemas = _extract_tool_json_schemas(allowlist, registered_tools)
1543 prompt = SamplingPrompt(
1544 directive=session["directive_text"],
1545 success_criteria=session["criteria"],
1546 criteria_status=iteration["criteria_evaluation"],
1547 recent_iterations=recent_iterations,
1548 tool_allowlist=allowlist,
1549 tool_docstrings=tool_docstrings,
1550 remaining_iterations=remaining_iterations,
1551 remaining_wall_clock_secs=remaining_wall_clock_secs,
1552 allow_scripts=allow_scripts,
1553 tool_schemas=tool_schemas,
1554 environment_context=environment_context,
1555 prior_missions=prior_missions,
1556 )
1558 # ---- Transport: backend.sample. --------------------------------------
1559 try:
1560 output_text = await backend.sample(prompt)
1561 except SamplingTransportError as err:
1562 _mission_audit.emit_sampling_event(
1563 session_id,
1564 iteration_index,
1565 sampling_purpose="strategy_revision",
1566 sampling_status="rejected",
1567 sampling_backend=backend_name,
1568 sampling_model_id=model_id or None,
1569 validation_error=err.code,
1570 )
1571 return SamplingFallback(
1572 rationale=template,
1573 reason="transport_error",
1574 backend_name=backend_name,
1575 model_id=model_id,
1576 )
1578 # ---- Parse the output as JSON. ---------------------------------------
1579 try:
1580 parsed = _extract_json_object(output_text)
1581 except json.JSONDecodeError:
1582 _mission_audit.emit_sampling_event(
1583 session_id,
1584 iteration_index,
1585 sampling_purpose="strategy_revision",
1586 sampling_status="rejected",
1587 sampling_backend=backend_name,
1588 sampling_model_id=model_id or None,
1589 validation_error="json_parse",
1590 )
1591 return SamplingFallback(
1592 rationale=template,
1593 reason="json_parse",
1594 backend_name=backend_name,
1595 model_id=model_id,
1596 )
1598 # ---- Schema validation. ----------------------------------------------
1599 try:
1600 _validate_revision_schema(parsed)
1601 except ValueError:
1602 _mission_audit.emit_sampling_event(
1603 session_id,
1604 iteration_index,
1605 sampling_purpose="strategy_revision",
1606 sampling_status="rejected",
1607 sampling_backend=backend_name,
1608 sampling_model_id=model_id or None,
1609 validation_error="schema_mismatch",
1610 )
1611 return SamplingFallback(
1612 rationale=template,
1613 reason="schema_mismatch",
1614 backend_name=backend_name,
1615 model_id=model_id,
1616 )
1618 # ---- Catalog validation on the proposed next_strategy. ---------------
1619 try:
1620 validate_strategy_against_catalog(
1621 parsed["next_strategy"],
1622 allowlist,
1623 registered_tools,
1624 allow_scripts,
1625 )
1626 except MissionValidationError as err:
1627 # ``err.details["reason"]`` carries the structured rejection
1628 # token (e.g. ``"tool_not_allowlisted"``,
1629 # ``"tool_args_invalid"``). Fall back to a generic label when
1630 # the validator emits a rejection without a ``reason`` key.
1631 details = err.details or {}
1632 reason = details.get("reason", "validation_error")
1633 _mission_audit.emit_sampling_event(
1634 session_id,
1635 iteration_index,
1636 sampling_purpose="strategy_revision",
1637 sampling_status="rejected",
1638 sampling_backend=backend_name,
1639 sampling_model_id=model_id or None,
1640 validation_error=str(reason),
1641 )
1642 return SamplingFallback(
1643 rationale=template,
1644 reason=str(reason),
1645 backend_name=backend_name,
1646 model_id=model_id,
1647 )
1649 # ---- Success path. ---------------------------------------------------
1650 # Extract token usage from the backend if available (Bedrock backend
1651 # stores it as a side-channel after each sample() call).
1652 _input_tokens = getattr(backend, "last_input_tokens", None)
1653 _output_tokens = getattr(backend, "last_output_tokens", None)
1654 _mission_audit.emit_sampling_event(
1655 session_id,
1656 iteration_index,
1657 sampling_purpose="strategy_revision",
1658 sampling_status="used",
1659 sampling_backend=backend_name,
1660 sampling_model_id=model_id or None,
1661 model_output_bytes=len(output_text.encode("utf-8")),
1662 input_tokens=_input_tokens,
1663 output_tokens=_output_tokens,
1664 )
1665 return SamplingUsed(
1666 output_text=output_text,
1667 parsed=parsed,
1668 backend_name=backend_name,
1669 model_id=model_id,
1670 )
1673# ---------------------------------------------------------------------------
1674# maybe_sample_final_lessons
1675# ---------------------------------------------------------------------------
1678async def maybe_sample_final_lessons(
1679 *,
1680 backend: SamplingBackend | None,
1681 session: _SessionState,
1682 remaining_iterations: int = 0,
1683 remaining_wall_clock_secs: float | None = None,
1684 allow_scripts: bool = False,
1685 tool_docstrings: dict[str, str] | None = None,
1686 environment_context: Mapping[str, Any] | None = None,
1687) -> SamplingUsed | SamplingFallback:
1688 """Consult the advisory LLM for final lessons, or fall back.
1690 Returns a :class:`SamplingUsed` when the bound backend produces a
1691 JSON object that clears the lessons schema. Returns a
1692 :class:`SamplingFallback` with an *empty* rationale on every
1693 rejection class — the final-report writer is responsible for the
1694 deterministic-text path when sampling does not produce usable
1695 output. Emits exactly one
1696 :func:`mission.audit.emit_sampling_event` per call, with
1697 ``iteration_index_or_purpose=None`` since the call is out-of-loop.
1698 """
1699 session_id = session["session_id"]
1701 # ---- No backend resolved: short-circuit. ------------------------------
1702 if backend is None:
1703 _mission_audit.emit_sampling_event(
1704 session_id,
1705 None,
1706 sampling_purpose="final_lessons",
1707 sampling_status="disabled",
1708 sampling_backend="none",
1709 )
1710 return SamplingFallback(
1711 rationale="",
1712 reason="no_backend_resolved",
1713 backend_name="none",
1714 model_id=None,
1715 )
1717 backend_name = backend.backend_name
1718 model_id = backend.model_id
1720 # ---- Build the prompt. ------------------------------------------------
1721 # Pass *all* iterations; the prompt builder trims / drops as needed
1722 # to fit the byte budget.
1723 prompt = SamplingPrompt(
1724 directive=session["directive_text"],
1725 success_criteria=session["criteria"],
1726 # The lessons prompt has no per-iteration criteria status; the
1727 # builder still expects the field, so reuse the most recent
1728 # iteration's evaluation when available, else an empty list.
1729 criteria_status=(
1730 list(session["iterations"][-1]["criteria_evaluation"]) if session["iterations"] else []
1731 ),
1732 recent_iterations=list(session["iterations"]),
1733 tool_allowlist=session.get("tool_allowlist", []),
1734 tool_docstrings=tool_docstrings or {},
1735 remaining_iterations=remaining_iterations,
1736 remaining_wall_clock_secs=remaining_wall_clock_secs,
1737 allow_scripts=allow_scripts,
1738 environment_context=environment_context,
1739 )
1741 # ---- Transport: backend.sample (uses lessons assembler). -------------
1742 try:
1743 # We render the lessons-specific prompt here so the byte-cap
1744 # bookkeeping uses the right schema header. The backend's own
1745 # ``sample`` calls ``prompt.assemble()`` under the hood for the
1746 # Strategy_Revision flow, but for lessons we assemble here and
1747 # invoke a thin shim through the backend.
1748 rendered = prompt.assemble_final_lessons()
1749 output_text = await _sample_with_assembled_text(backend, rendered)
1750 except SamplingTransportError as err:
1751 _mission_audit.emit_sampling_event(
1752 session_id,
1753 None,
1754 sampling_purpose="final_lessons",
1755 sampling_status="rejected",
1756 sampling_backend=backend_name,
1757 sampling_model_id=model_id or None,
1758 validation_error=err.code,
1759 )
1760 return SamplingFallback(
1761 rationale="",
1762 reason="transport_error",
1763 backend_name=backend_name,
1764 model_id=model_id,
1765 )
1767 # ---- Parse the output as JSON. ---------------------------------------
1768 try:
1769 parsed = _extract_json_object(output_text)
1770 except json.JSONDecodeError:
1771 _mission_audit.emit_sampling_event(
1772 session_id,
1773 None,
1774 sampling_purpose="final_lessons",
1775 sampling_status="rejected",
1776 sampling_backend=backend_name,
1777 sampling_model_id=model_id or None,
1778 validation_error="json_parse",
1779 )
1780 return SamplingFallback(
1781 rationale="",
1782 reason="json_parse",
1783 backend_name=backend_name,
1784 model_id=model_id,
1785 )
1787 # ---- Schema validation. ----------------------------------------------
1788 try:
1789 _validate_lessons_schema(parsed)
1790 except ValueError:
1791 _mission_audit.emit_sampling_event(
1792 session_id,
1793 None,
1794 sampling_purpose="final_lessons",
1795 sampling_status="rejected",
1796 sampling_backend=backend_name,
1797 sampling_model_id=model_id or None,
1798 validation_error="schema_mismatch",
1799 )
1800 return SamplingFallback(
1801 rationale="",
1802 reason="schema_mismatch",
1803 backend_name=backend_name,
1804 model_id=model_id,
1805 )
1807 # ---- Success path. ---------------------------------------------------
1808 _input_tokens = getattr(backend, "last_input_tokens", None)
1809 _output_tokens = getattr(backend, "last_output_tokens", None)
1810 _mission_audit.emit_sampling_event(
1811 session_id,
1812 None,
1813 sampling_purpose="final_lessons",
1814 sampling_status="used",
1815 sampling_backend=backend_name,
1816 sampling_model_id=model_id or None,
1817 model_output_bytes=len(output_text.encode("utf-8")),
1818 input_tokens=_input_tokens,
1819 output_tokens=_output_tokens,
1820 )
1821 return SamplingUsed(
1822 output_text=output_text,
1823 parsed=parsed,
1824 backend_name=backend_name,
1825 model_id=model_id,
1826 )
1829async def _sample_with_assembled_text(backend: SamplingBackend, rendered: str) -> str:
1830 """Route a pre-assembled prompt string through a backend.
1832 Both shipped backends accept a :class:`SamplingPrompt` and call
1833 ``assemble`` themselves to render the strategy-revision shape. For
1834 the final-lessons path we render the lessons-shaped prompt here
1835 and need to deliver that exact text to the transport. The shim
1836 builds a tiny prompt-shaped wrapper whose :meth:`assemble` returns
1837 the pre-rendered text and forwards it to the backend.
1838 """
1839 pre_rendered = rendered
1841 class _PreRendered:
1842 """Thin :class:`SamplingPrompt` look-alike with a fixed assemble()."""
1844 def assemble(self) -> str:
1845 return pre_rendered
1847 # The two shipped backends only call ``prompt.assemble()`` so the
1848 # duck-typed wrapper above is enough to drive either of them.
1849 return await backend.sample(_PreRendered()) # type: ignore[arg-type]
1852# ---------------------------------------------------------------------------
1853# Session-start sampling-state resolver
1854# ---------------------------------------------------------------------------
1857def _bedrock_credentials_available() -> bool:
1858 """Lightweight probe: do local AWS credentials resolve?
1860 Instantiates a ``boto3.Session()`` and asks for ``get_credentials()``
1861 without making any network call. ``boto3`` is imported inside the
1862 function so the module's top-level import surface stays free of
1863 SDK dependencies — and so a host that has no ``boto3`` installed
1864 (or any other unexpected import-time failure) cleanly degrades to
1865 "no credentials available" rather than crashing the helper.
1866 """
1867 try:
1868 import boto3
1870 session = boto3.Session()
1871 creds = session.get_credentials()
1872 return creds is not None
1873 except Exception:
1874 return False
1877def resolve_sampling_state(
1878 use_sampling_param: bool | None,
1879) -> tuple[bool, Literal["bedrock", "none"]]:
1880 """Decide whether sampling is enabled for a session and which backend resolves.
1882 Bedrock is the only sampling transport (MCP client sampling left the
1883 protocol with FastMCP 4), so resolution no longer depends on how the
1884 session was started — CLI and MCP callers probe the same server-side
1885 credentials.
1887 Resolution precedence (first match wins):
1889 1. ``use_sampling_param is False`` — caller explicitly disabled
1890 sampling, so the result is ``(False, "none")`` regardless of
1891 any capability the environment advertises.
1892 2. Local AWS credentials resolve — ``(True, "bedrock")``.
1893 3. No credentials — ``(True, "none")`` if the caller opted in
1894 explicitly with ``use_sampling_param is True`` (so the caller can
1895 decide whether to error or proceed deterministic-only), and
1896 ``(False, "none")`` otherwise.
1898 Args:
1899 use_sampling_param: Three-state opt-in flag. ``None`` means the
1900 caller did not specify and the helper should auto-detect.
1901 ``False`` short-circuits to a disabled state. ``True`` means
1902 the caller explicitly opted in; the backend is auto-detected
1903 and ``"none"`` is allowed when no concrete backend resolves.
1905 Returns:
1906 A ``(use_sampling, backend)`` tuple. The caller persists both
1907 values on its ``SessionState`` so the audit pipeline can stamp
1908 every later sampling event with the resolved backend.
1909 """
1910 # 1. Explicit opt-out wins outright.
1911 if use_sampling_param is False:
1912 return (False, "none")
1914 # 2. Probe server-side AWS credentials.
1915 if _bedrock_credentials_available():
1916 return (True, "bedrock")
1918 # 3. No credentials — only honour an explicit True.
1919 if use_sampling_param is True:
1920 return (True, "none")
1921 return (False, "none")