Coverage for gco_mcp / mission / sampling.py: 100.00%
448 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +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 Iterable, 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-23T11:36:55Z
64# Generated from Git commit: c28a947089a90cba72adeefc0fa4debf25662e9d
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 "ENVIRONMENT_CONTEXT_BYTE_CAP",
77 "ENV_BEDROCK_MODEL_ID",
78 "ENV_BEDROCK_REGION",
79 "FINAL_LESSONS_SCHEMA",
80 "OBSERVATION_FIELD_BYTE_CAP",
81 "OBSERVATION_FIELD_TRUNCATE_TO",
82 "PRIOR_MISSIONS_BYTE_CAP",
83 "PROMPT_BYTE_BUDGET",
84 "RECENT_ITERATIONS_LIMIT",
85 "STRATEGY_REVISION_SCHEMA",
86 "STRATEGY_SHAPE_SCHEMA",
87 "TRUNCATION_MARKER",
88 "BedrockSamplingBackend",
89 "MissionValidationError",
90 "SamplingBackend",
91 "SamplingFallback",
92 "SamplingPrompt",
93 "SamplingTransportError",
94 "SamplingUsed",
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.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: Iterable[Mapping[str, Any]],
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 ``statuses`` is normally the persisted ``criteria_evaluation`` list
526 of :class:`CriterionResult` rows, but it is read back from disk, so
527 the helper takes any mappings and skips an entry that lacks a
528 ``criterion_id`` rather than trusting the row shape.
530 The ``_parsed_ast`` private key on a ``predicate`` criterion is
531 stripped — it is a Python ``ast.Expression`` object that is not
532 JSON-serialisable and that the model has no use for.
533 """
534 by_id: dict[str, Mapping[str, Any]] = {}
535 for s in statuses:
536 cid = s.get("criterion_id")
537 if cid is None:
538 continue
539 # If the caller passes duplicates (older then newer), prefer the
540 # last entry — that's the most-recent-wins convention the engine
541 # uses everywhere else.
542 by_id[cid] = s
544 out: list[dict[str, Any]] = []
545 for c in criteria:
546 cid = c.get("criterion_id")
547 # Strip private cached AST and surface only the prompt-relevant fields.
548 public = {k: v for k, v in c.items() if not k.startswith("_")}
549 match = by_id.get(cid) if cid is not None else None
550 if match is None:
551 public["status"] = {
552 "status": "inconclusive",
553 "evidence": None,
554 }
555 else:
556 public["status"] = {
557 "status": match.get("status"),
558 "evidence": match.get("evidence"),
559 "evaluated_at": match.get("evaluated_at"),
560 }
561 out.append(public)
562 return out
565# ---------------------------------------------------------------------------
566# Tool allowlist rendering
567# ---------------------------------------------------------------------------
570def _render_tool_allowlist(
571 allowlist: Sequence[str],
572 docstrings: Mapping[str, str],
573 schemas: Mapping[str, Any] | None = None,
574) -> list[dict[str, Any]]:
575 """Pair every allowlisted tool name with its docstring and input schema.
577 Tools without a registered docstring get an empty string — the
578 prompt remains valid; the model just sees a tool name with no
579 inline description. Tools without a schema get ``null`` so the
580 model knows no args are required. Names are emitted in the
581 caller's allowlist order so the prompt is identical for two
582 callers that pass the same list.
583 """
584 rendered: list[dict[str, Any]] = []
585 for name in allowlist:
586 entry: dict[str, Any] = {
587 "tool_name": name,
588 "docstring": str(docstrings.get(name, "")),
589 }
590 if schemas:
591 schema = schemas.get(name)
592 if schema is not None:
593 entry["input_schema"] = schema
594 rendered.append(entry)
595 return rendered
598# ---------------------------------------------------------------------------
599# Budget context rendering
600# ---------------------------------------------------------------------------
603def _render_budget_context(
604 *,
605 remaining_iterations: int,
606 remaining_wall_clock_secs: float | None,
607 allow_scripts: bool,
608) -> dict[str, Any]:
609 """Render the budget context block. Stable shape regardless of inputs.
611 ``None`` for the wall-clock cap is rendered verbatim as JSON
612 ``null`` so the model can disambiguate "unbounded" from "0".
613 """
614 return {
615 "remaining_iterations": int(remaining_iterations),
616 "remaining_wall_clock_seconds": (
617 float(remaining_wall_clock_secs) if remaining_wall_clock_secs is not None else None
618 ),
619 "allow_scripted_strategies": bool(allow_scripts),
620 }
623# ---------------------------------------------------------------------------
624# SamplingPrompt — the public class
625# ---------------------------------------------------------------------------
628@dataclass(frozen=True)
629class SamplingPrompt:
630 """Bundle the bare data needed to assemble a sampling prompt string.
632 The dataclass is ``frozen=True`` so callers cannot mutate the inputs
633 between an :meth:`assemble` call and an :meth:`assemble_final_lessons`
634 call — both methods produce deterministic outputs from the same
635 bound state, which is the property the determinism tests pin down.
637 All inputs are required positionally or by keyword; defaults exist
638 only for the fields whose contract documents a default value.
639 """
641 directive: str
642 success_criteria: Sequence[Criterion]
643 criteria_status: Sequence[CriterionResult]
644 recent_iterations: Sequence[IterationRecord]
645 tool_allowlist: Sequence[str]
646 tool_docstrings: Mapping[str, str]
647 remaining_iterations: int
648 remaining_wall_clock_secs: float | None
649 allow_scripts: bool = field(default=False)
650 #: Per-tool JSON Schema for the input parameters. Keyed by tool
651 #: name; values are the JSON-serialisable schema dict (or ``None``
652 #: for tools that take no args). Included in the prompt so the
653 #: Strategy_Revision model can propose valid ``args`` dicts.
654 tool_schemas: Mapping[str, Any] = field(default_factory=dict)
655 #: Optional snapshot of slow-moving live signals (per-region queue
656 #: depth, GPU utilisation, deployed-region list, reservation
657 #: counts, etc.) gathered once at session start and reused on
658 #: every iteration's prompt. ``None`` (the default) suppresses the
659 #: ``=== Environment context ===`` section entirely so the prompt
660 #: stays byte-identical to the pre-environment-context shape —
661 #: that's what every existing determinism test pins down.
662 environment_context: Mapping[str, Any] | None = field(default=None)
663 #: Optional list of similar past missions retrieved from the
664 #: mission-memory vector index (most-similar-first), gathered once
665 #: per engine wiring and reused on every iteration's prompt.
666 #: ``None`` (the default) suppresses the ``=== Prior similar
667 #: missions ===`` section entirely — the same byte-identical
668 #: contract as :attr:`environment_context`, and what keeps every
669 #: pre-memory prompt (and the determinism suite) unchanged.
670 prior_missions: Sequence[Mapping[str, Any]] | None = field(default=None)
672 # ---- Strategy_Revision rendering --------------------------------------
674 def assemble(self) -> str:
675 """Return the Strategy_Revision prompt string.
677 The output is capped at :data:`PROMPT_BYTE_BUDGET` UTF-8 bytes.
678 When the freshly-assembled prompt exceeds the cap, the oldest
679 Iteration summary is dropped and the prompt re-rendered. The
680 loop terminates because each drop monotonically shrinks the
681 prompt and there is a non-iteration baseline that fits well
682 under the cap on its own (the directive, criteria, allowlist,
683 budget block, and schema instruction together are ~6-10 KB
684 for any reasonable session shape).
685 """
686 # Defensive slice — the caller is asked to pass at most five,
687 # but if they pass more, take the most recent five.
688 iterations: list[IterationRecord] = list(self.recent_iterations[-RECENT_ITERATIONS_LIMIT:])
690 while True:
691 text = self._render(
692 schema=STRATEGY_REVISION_SCHEMA,
693 iterations=[_summarise_iteration(it) for it in iterations],
694 schema_purpose="strategy_revision",
695 )
696 if _utf8_len(text) <= PROMPT_BYTE_BUDGET or not iterations:
697 return text
698 # Drop the oldest iteration and try again.
699 iterations = iterations[1:]
701 # ---- Final_Report rendering -------------------------------------------
703 def assemble_final_lessons(self) -> str:
704 """Return the Final_Report ``lessons`` prompt string.
706 The shape parallels :meth:`assemble` but emits
707 :data:`FINAL_LESSONS_SCHEMA` and uses iteration **verdict
708 summaries only** instead of full Observation summaries. The same
709 :data:`PROMPT_BYTE_BUDGET` byte cap applies; the same
710 oldest-first drop policy kicks in if the cap is exceeded.
711 """
712 iterations: list[IterationRecord] = list(self.recent_iterations)
714 while True:
715 text = self._render(
716 schema=FINAL_LESSONS_SCHEMA,
717 iterations=[_summarise_iteration_for_lessons(it) for it in iterations],
718 schema_purpose="final_lessons",
719 )
720 if _utf8_len(text) <= PROMPT_BYTE_BUDGET or not iterations:
721 return text
722 iterations = iterations[1:]
724 # ---- Internal renderer ------------------------------------------------
726 def _render(
727 self,
728 *,
729 schema: dict[str, Any],
730 iterations: Sequence[Mapping[str, Any]],
731 schema_purpose: str,
732 ) -> str:
733 """Format the full prompt from the section blocks.
735 The text layout is fixed — every section is delimited by a
736 ``=== <name> ===`` header so the model can latch onto a
737 predictable structure. Section bodies are JSON wherever the
738 content is structured; the directive itself is rendered as
739 free text because that is how the operator wrote it.
740 """
741 criteria_block = _pair_criteria_with_status(self.success_criteria, self.criteria_status)
742 tool_block = _render_tool_allowlist(
743 self.tool_allowlist, self.tool_docstrings, self.tool_schemas
744 )
745 budget_block = _render_budget_context(
746 remaining_iterations=self.remaining_iterations,
747 remaining_wall_clock_secs=self.remaining_wall_clock_secs,
748 allow_scripts=self.allow_scripts,
749 )
751 if schema_purpose == "strategy_revision":
752 preamble = (
753 "You are advising a Mission goal-directed iteration loop. "
754 "Propose the next Strategy that moves the Mission toward "
755 "satisfying its Success_Criteria. The Verdict label, "
756 "budget enforcement, and Criteria evaluation are all "
757 "computed server-side and are unaffected by your output. "
758 "Your role is advisory: the rationale and next_strategy "
759 "you produce are validated against the Tool_Allowlist and "
760 "the remaining budget before being adopted.\n\n"
761 "IMPORTANT: You may propose MULTIPLE tool calls in a "
762 "single iteration by including multiple entries in the "
763 "tool_calls array. This is especially useful when the "
764 "unmet criteria require results from different tools — "
765 "calling them all in one iteration lets the evaluator "
766 "see all results together. Use the input_schema in the "
767 "Tool allowlist section to construct valid args for each "
768 "tool call."
769 )
770 recent_header = "Recent iterations (oldest first)"
771 else:
772 preamble = (
773 "You are advising a Mission goal-directed iteration loop "
774 "that has just reached a terminal Verdict. Produce the "
775 "lessons learned and the recommended follow-ups for the "
776 "operator. Your output is merged into the Final_Report; "
777 "the Verdict label, budget bookkeeping, and Criteria "
778 "evaluation that produced the terminal state are "
779 "deterministic server-side outputs and are not under "
780 "review."
781 )
782 recent_header = "Iteration verdict summary (oldest first)"
784 sections: list[str] = []
785 sections.append(preamble)
786 sections.append("")
787 sections.append("=== Mission directive ===")
788 sections.append(self.directive)
789 sections.append("")
790 sections.append("=== Success criteria with current status ===")
791 sections.append(_dumps(criteria_block, indent=2))
792 sections.append("")
793 sections.append("=== Tool allowlist ===")
794 sections.append(_dumps(tool_block, indent=2))
795 sections.append("")
796 sections.append("=== Budget context ===")
797 sections.append(_dumps(budget_block, indent=2))
798 sections.append("")
799 if self.environment_context is not None:
800 # Truncated + key-sorted — see :func:`_summarise_environment_context`.
801 # Emitting a header even for an empty dict means a session
802 # that opted in but had a probe failure still surfaces
803 # "we tried" so the operator can act on the gap.
804 env_summary = _summarise_environment_context(self.environment_context)
805 sections.append("=== Environment context (slow-moving live signals) ===")
806 sections.append(_dumps(env_summary, indent=2))
807 sections.append("")
808 if self.prior_missions is not None:
809 # Institutional memory: the closest past missions by directive
810 # similarity, with their lessons and verdicts. Advisory only —
811 # summarised and byte-capped in its own truncation domain.
812 sections.append("=== Prior similar missions (institutional memory) ===")
813 sections.append(
814 "Lessons and outcomes from the most similar past missions, "
815 "most similar first. Treat them as advisory context: they "
816 "may suggest which tools or query shapes worked before, or "
817 "what to avoid repeating."
818 )
819 sections.append(_dumps(_summarise_prior_missions(self.prior_missions), indent=2))
820 sections.append("")
821 sections.append(f"=== {recent_header} ===")
822 sections.append(_dumps(list(iterations), indent=2))
823 sections.append("")
824 sections.append("=== Output schema ===")
825 sections.append(
826 "Respond with a single JSON object that validates against "
827 "the JSON Schema below. Do not include any prose outside "
828 "the JSON object."
829 )
830 sections.append(_dumps(schema, indent=2))
832 return "\n".join(sections)
835# ---------------------------------------------------------------------------
836# Backend protocol and transport-error type
837# ---------------------------------------------------------------------------
840@runtime_checkable
841class SamplingBackend(Protocol):
842 """Transport-agnostic surface for the advisory LLM call.
844 Implementations bind a concrete transport (e.g., the MCP
845 ``Context.sample`` capability or ``bedrock-runtime:Converse``) and
846 expose a single async ``sample`` entry point. The protocol is
847 ``runtime_checkable`` so call sites — and tests — can use
848 ``isinstance(backend, SamplingBackend)`` to gate dispatch on a
849 duck-typed backend instance.
851 Attributes:
852 backend_name: Stable identifier the audit pipeline emits in the
853 ``sampling_backend`` field. Bedrock is the only transport
854 the system supports: MCP client sampling (``ctx.sample``)
855 left the protocol with FastMCP 4's sessionless era, so
856 missions sample server-side regardless of how they were
857 started.
858 model_id: The concrete model identifier the backend will route
859 the prompt to. Echoed in audit events so replay can
860 reproduce the exact request.
861 """
863 backend_name: Literal["bedrock"]
864 model_id: str
866 async def sample(self, prompt: SamplingPrompt) -> str:
867 """Render ``prompt`` through the bound transport and return the
868 raw model output text. Implementations raise
869 :class:`SamplingTransportError` (with a transport-tagged
870 ``code``) on any transport-layer failure so the engine's
871 fallback policy can branch on a single, well-typed exception.
872 """
873 ...
876class SamplingTransportError(Exception):
877 """Transport-layer failure raised by a :class:`SamplingBackend`.
879 The mandatory ``code`` attribute tags the failure class so the
880 engine's deterministic-fallback path can branch on a stable string
881 without parsing the message. The convention is
882 ``"<backend>_<error_class>"`` for backend-specific failures and a
883 short, snake-cased label for backend-agnostic failures.
885 Documented codes (used elsewhere in the Mission stack):
887 * ``"bedrock_AccessDeniedException"`` — IAM denied
888 ``bedrock:InvokeModel`` for the resolved model.
889 * ``"bedrock_malformed_response"`` — Converse returned a payload
890 that did not have the expected ``output.message.content[0].text``
891 shape.
892 * ``"bedrock_truncated_response"`` — the answer was cut off by an
893 output-token limit (``stopReason == "max_tokens"``), so its text
894 cannot be trusted to be complete.
895 * ``"bedrock_no_credentials"`` — the local ``boto3`` session could
896 not resolve credentials.
898 Args:
899 code: Mandatory failure tag (see examples above).
900 message: Optional human-readable detail. When present, it is
901 joined to ``code`` with ``": "`` for the string
902 representation; when absent, ``str(self)`` is just the
903 ``code``.
904 """
906 def __init__(self, code: str, message: str | None = None) -> None:
907 self.code: str = code
908 self.message: str | None = message
909 # Forward the most useful single-line representation to
910 # ``Exception.__init__`` so ``logging`` / ``traceback`` modules
911 # show the same string ``str(self)`` produces below.
912 if message is None:
913 super().__init__(code)
914 else:
915 super().__init__(f"{code}: {message}")
917 def __str__(self) -> str:
918 if self.message is None:
919 return self.code
920 return f"{self.code}: {self.message}"
923# ---------------------------------------------------------------------------
924# BedrockSamplingBackend — routes the prompt through bedrock-runtime:Converse
925# ---------------------------------------------------------------------------
928class BedrockSamplingBackend:
929 """Sampling backend that calls ``bedrock-runtime:Converse``.
931 The backend resolves its model id and region at construction time
932 from (in order of precedence) the explicit constructor argument,
933 the matching environment variable
934 (:data:`ENV_BEDROCK_MODEL_ID` / :data:`ENV_BEDROCK_REGION`), and
935 finally the ``cdk.json`` Mission default
936 (:func:`gco.bedrock.get_default_mission_model_id` /
937 :data:`DEFAULT_BEDROCK_REGION`). The ``boto3`` client itself is
938 constructed lazily on the first :meth:`sample` call so that
939 ``import mission.sampling`` does not pull ``boto3`` into the
940 import graph and so that test code can swap the import in via
941 ``unittest.mock.patch`` without paying for a real session at
942 construction time.
944 Failure modes:
946 * Missing or partial AWS credentials at client-construction time
947 surface as :class:`SamplingTransportError` with code
948 ``"bedrock_no_credentials"``; the original exception is chained
949 via ``__cause__``.
950 * A ``botocore.exceptions.ClientError`` from the ``Converse`` call
951 surfaces as :class:`SamplingTransportError` with code
952 ``"bedrock_<ErrorCode>"`` where ``<ErrorCode>`` is read from the
953 error envelope (defaulting to ``"Unknown"`` when the envelope is
954 malformed). The one exception is the Anthropic first-time-use
955 gate, which raises
956 :class:`gco.bedrock.BedrockFTUFormNotAcceptedError` instead of a
957 transport error so it is never absorbed by a deterministic
958 fallback. See ``docs/CUSTOMIZATION.md`` (Bedrock Model Selection).
959 * A response without a non-empty ``text`` block under
960 ``output.message.content`` — including reasoning-only and empty
961 ``content`` lists — surfaces as :class:`SamplingTransportError`
962 with code ``"bedrock_malformed_response"``.
963 * A response cut off by an output-token limit
964 (``stopReason == "max_tokens"``) surfaces as
965 :class:`SamplingTransportError` with code
966 ``"bedrock_truncated_response"``.
967 """
969 backend_name: Literal["bedrock"] = "bedrock"
971 def __init__(
972 self,
973 model_id: str | None = None,
974 region: str | None = None,
975 ) -> None:
976 """Resolve the model id and region; defer client construction.
978 Args:
979 model_id: Optional explicit model id. When ``None``, falls
980 back to the :data:`ENV_BEDROCK_MODEL_ID` environment
981 variable, then to the ``cdk.json`` Mission default from
982 :func:`gco.bedrock.get_default_mission_model_id`.
983 region: Optional explicit region. When ``None``, falls back
984 to :data:`ENV_BEDROCK_REGION`, then to
985 :data:`DEFAULT_BEDROCK_REGION`.
986 """
987 if model_id is not None:
988 self.model_id = model_id
989 self._uses_default_model = False
990 elif ENV_BEDROCK_MODEL_ID in os.environ:
991 # Preserve the existing explicit-environment semantics, including
992 # an intentionally empty value, without evaluating the fallback.
993 self.model_id = os.environ[ENV_BEDROCK_MODEL_ID]
994 self._uses_default_model = False
995 else:
996 self.model_id = get_default_mission_model_id()
997 self._uses_default_model = True
998 self._region: str = (
999 region
1000 if region is not None
1001 else os.environ.get(ENV_BEDROCK_REGION, DEFAULT_BEDROCK_REGION)
1002 )
1003 # The boto3 client is built on first ``sample`` call. ``None``
1004 # here is the sentinel for "not yet constructed".
1005 self._client: Any = None
1007 @classmethod
1008 def from_canonical_default(
1009 cls,
1010 region: str | None = None,
1011 ) -> BedrockSamplingBackend:
1012 """Build a backend that deliberately applies canonical reasoning.
1014 Unlike ``cls(model_id=None)``, this bypasses the model environment
1015 override. Fixture capture uses it to reproduce the checked-in default
1016 exactly, while ordinary explicit model IDs retain override semantics.
1017 """
1018 backend = cls(model_id=get_default_mission_model_id(), region=region)
1019 backend._uses_default_model = True
1020 return backend
1022 def _get_client(self) -> Any:
1023 """Return the cached ``bedrock-runtime`` client, building it on first use.
1025 ``boto3`` and ``botocore.exceptions`` are imported here rather
1026 than at module top-level so that pure-Python consumers of this
1027 module (the prompt builder, the protocol, the error type) do
1028 not pay for the ``boto3`` import. This also lets tests patch
1029 ``mission.sampling.boto3`` after import.
1030 """
1031 if self._client is not None:
1032 return self._client
1033 # Local import — keeps the module's import surface boto3-free.
1034 import boto3
1035 from botocore.config import Config
1036 from botocore.exceptions import (
1037 NoCredentialsError,
1038 PartialCredentialsError,
1039 )
1041 try:
1042 self._client = boto3.Session().client(
1043 "bedrock-runtime",
1044 region_name=self._region,
1045 config=Config(read_timeout=BEDROCK_READ_TIMEOUT_SECONDS),
1046 )
1047 except (NoCredentialsError, PartialCredentialsError) as err:
1048 raise SamplingTransportError("bedrock_no_credentials") from err
1049 return self._client
1051 async def sample(self, prompt: SamplingPrompt) -> str:
1052 """Render ``prompt`` through ``Converse`` and return the response text.
1054 Raises:
1055 SamplingTransportError: On any transport-level failure.
1056 * ``bedrock_no_credentials`` — credentials could not be
1057 resolved by ``boto3`` at client-construction time.
1058 * ``bedrock_<ErrorCode>`` — the ``Converse`` call raised
1059 a ``ClientError``; ``<ErrorCode>`` is the AWS error
1060 code from the envelope.
1061 * ``bedrock_malformed_response`` — the response did not
1062 contain a non-empty final text content block.
1063 * ``bedrock_truncated_response`` — the response was cut
1064 off by an output-token limit and cannot be trusted to
1065 be complete.
1066 gco.bedrock.BedrockFTUFormNotAcceptedError: The account has
1067 not submitted Anthropic's one-time first-time-use case
1068 form. Raised instead of a transport error so callers
1069 cannot silently fall back past a permanent, one-line-fix
1070 misconfiguration.
1071 """
1072 # Local import — see ``_get_client`` for the rationale.
1073 from botocore.exceptions import ClientError
1075 client = self._get_client()
1077 text = prompt.assemble()
1078 converse_options = build_bedrock_converse_options(
1079 self.model_id,
1080 # Deliberately no maxTokens: the Converse default is the model's
1081 # own maximum output length, so a rationale can never be cut off
1082 # by a GCO-imposed cap. A cap is opt-in — pass maxTokens here to
1083 # restore one.
1084 inference_config={"temperature": BEDROCK_TEMPERATURE},
1085 apply_default_reasoning=self._uses_default_model,
1086 )
1087 try:
1088 response = await asyncio.to_thread(
1089 client.converse,
1090 modelId=self.model_id,
1091 messages=[{"role": "user", "content": [{"text": text}]}],
1092 **converse_options,
1093 )
1094 except ClientError as err:
1095 # A missing Anthropic FTU form is a permanent account-scoped
1096 # misconfiguration, not a transport fault: escalate it instead of
1097 # letting the deterministic-fallback path absorb it silently.
1098 raise_if_bedrock_ftu_form_error(err)
1099 # ``e.response`` is documented to be present on ClientError
1100 # but the envelope shape can vary; defend against missing
1101 # keys so the audit pipeline always sees a tagged code.
1102 envelope = getattr(err, "response", None) or {}
1103 error_block = envelope.get("Error", {}) if isinstance(envelope, dict) else {}
1104 code = (
1105 error_block.get("Code", "Unknown") if isinstance(error_block, dict) else "Unknown"
1106 )
1107 raise SamplingTransportError(f"bedrock_{code}") from err
1109 # Capture token usage from the Converse response for the audit
1110 # trail. The ``usage`` block is present on every successful
1111 # Converse response and carries ``inputTokens`` and
1112 # ``outputTokens``. Store on the instance so callers can read
1113 # it after each sample() call without changing the protocol.
1114 usage = response.get("usage") or {}
1115 self.last_input_tokens: int | None = usage.get("inputTokens")
1116 self.last_output_tokens: int | None = usage.get("outputTokens")
1118 try:
1119 return extract_bedrock_converse_text(response)
1120 except BedrockResponseTruncatedError as err:
1121 # A cut-off rationale is unusable; let the deterministic-fallback
1122 # path absorb it like any other transport-shaped fault.
1123 raise SamplingTransportError("bedrock_truncated_response") from err
1124 except (KeyError, IndexError, TypeError) as err:
1125 raise SamplingTransportError("bedrock_malformed_response") from err
1128# ---------------------------------------------------------------------------
1129# Backend resolver
1130# ---------------------------------------------------------------------------
1133def select_sampling_backend(model_id: str | None) -> SamplingBackend:
1134 """Construct the sampling backend for a session that opted into sampling.
1136 Bedrock is the only sampling transport. MCP client sampling
1137 (``ctx.sample``) left the protocol with FastMCP 4's sessionless era —
1138 per the v4 migration guidance, generation belongs server-side — so
1139 missions sample through ``bedrock-runtime:Converse`` with the server's
1140 own credentials regardless of whether they were started from the CLI
1141 or over MCP. Credential resolution is deferred to the first ``sample``
1142 call; a missing-credentials failure surfaces as
1143 :class:`SamplingTransportError` and the engine's deterministic
1144 fallback absorbs it.
1146 Args:
1147 model_id: Optional concrete model identifier. Forwarded to the
1148 backend constructor verbatim; ``None`` resolves through the
1149 environment and ``cdk.json`` Mission default.
1151 Returns:
1152 A :class:`BedrockSamplingBackend` bound to ``model_id``.
1153 """
1154 return BedrockSamplingBackend(model_id)
1157# ---------------------------------------------------------------------------
1158# Strategy-against-catalog validator
1159# ---------------------------------------------------------------------------
1162def _resolve_input_schema(tool: Any) -> Any:
1163 """Return the registered Pydantic input model for a Tool, or None.
1165 FastMCP exposes the model under ``input_schema`` in newer releases
1166 and ``inputSchema`` in older ones. Tolerate both. Tools that genuinely
1167 take no args (or test catalog mocks that omit the attribute) yield
1168 ``None``, in which case the caller skips per-call args validation.
1169 """
1170 schema = getattr(tool, "input_schema", None)
1171 if schema is None:
1172 schema = getattr(tool, "inputSchema", None)
1173 return schema
1176def _extract_tool_json_schemas(
1177 allowlist: Sequence[str],
1178 registered_tools: Mapping[str, Any],
1179) -> dict[str, Any]:
1180 """Extract JSON Schema dicts for each allowlisted tool's input parameters.
1182 Calls ``.model_json_schema()`` on the Pydantic model exposed by
1183 ``_resolve_input_schema``. Falls back gracefully: tools without a
1184 schema, tools whose schema isn't a Pydantic model, and any
1185 exception during schema extraction all yield ``None`` for that
1186 tool (omitted from the output dict). The caller renders the
1187 result into the Strategy_Revision prompt so the model can propose
1188 valid ``args`` dicts.
1189 """
1190 schemas: dict[str, Any] = {}
1191 for name in allowlist:
1192 tool = registered_tools.get(name)
1193 if tool is None:
1194 continue
1195 model = _resolve_input_schema(tool)
1196 if model is None:
1197 continue
1198 try:
1199 # Pydantic v2 models expose model_json_schema() as a classmethod.
1200 json_schema = model.model_json_schema()
1201 schemas[name] = json_schema
1202 except Exception: # nosec B112 # a tool without a schema has none to render
1203 # Non-Pydantic schema, or a mock that doesn't support it.
1204 continue
1205 return schemas
1208def validate_strategy_against_catalog(
1209 strategy: Strategy,
1210 allowlist: list[str],
1211 registered_tools: dict[str, Any],
1212 allow_scripts: bool,
1213) -> None:
1214 """Validate a Strategy against the live tool catalog.
1216 Returns ``None`` on accept; raises :class:`MissionValidationError`
1217 with a structured ``details.reason`` enum on reject. The function
1218 layers catalog-aware checks on top of the structural validation in
1219 :func:`mission.validation.validate_strategy`:
1221 1. Mutual-exclusivity (exactly one of ``tool_calls`` / ``script``).
1222 2. Per-call ``tool_name`` is in ``allowlist``.
1223 3. Per-call ``args`` validates against the registered Pydantic model
1224 exposed under ``Tool.input_schema`` (or the older
1225 ``Tool.inputSchema``); calls whose tool has neither attribute or
1226 a ``None`` schema skip args validation.
1227 4. For scripted strategies, ``allow_scripts`` is True and the
1228 script's AST passes :func:`mission.sandbox.validate_script_ast`.
1230 Args:
1231 strategy: The Strategy dict to validate.
1232 allowlist: The session's resolved Tool_Allowlist.
1233 registered_tools: Mapping from tool name to a registered tool
1234 object (typed ``Any`` so the module imports without
1235 FastMCP). Read-only — only ``input_schema`` /
1236 ``inputSchema`` is consulted.
1237 allow_scripts: Session-level flag gating scripted strategies.
1238 """
1239 # 1. Structural validation: mutual exclusivity, script-allow gating,
1240 # and AST validation for scripts. Reuses the existing validator
1241 # so error shapes for those rejection classes stay aligned with
1242 # the rest of the input pipeline.
1243 _validation.validate_strategy(cast("dict[str, Any]", strategy), allowlist, allow_scripts)
1245 # The structural validator has already accepted exactly one of the
1246 # two shapes. Branch on which one is present.
1247 if "tool_calls" in strategy:
1248 tool_calls = strategy["tool_calls"]
1249 # Empty list is rejected by validate_strategy; this is a defence
1250 # in depth for callers that might bypass that path.
1251 if not tool_calls:
1252 raise MissionValidationError(
1253 "validation_error",
1254 details={
1255 "field": "strategy",
1256 "subfield": "tool_calls",
1257 "reason": "tool_calls_empty",
1258 },
1259 )
1261 # 2. Per-call name-in-allowlist check.
1262 for call in tool_calls:
1263 name = call.get("tool_name")
1264 if name not in allowlist:
1265 raise MissionValidationError(
1266 "validation_error",
1267 details={
1268 "field": "strategy",
1269 "subfield": "tool_calls",
1270 "tool_name": name,
1271 "reason": "tool_not_allowlisted",
1272 "allowlist": list(allowlist),
1273 },
1274 )
1276 # 3. Per-call args validation against the tool's Pydantic model.
1277 for call in tool_calls:
1278 name = call["tool_name"]
1279 tool = registered_tools.get(name)
1280 if tool is None:
1281 # Catalog could have a name in the allowlist that is not
1282 # currently registered (gating, dynamic load). Mirror the
1283 # unknown-tool shape used elsewhere.
1284 raise MissionValidationError(
1285 "validation_error",
1286 details={
1287 "field": "strategy",
1288 "subfield": "tool_calls",
1289 "tool_name": name,
1290 "reason": "tool_not_registered",
1291 },
1292 )
1293 schema = _resolve_input_schema(tool)
1294 if schema is None:
1295 # Either the tool genuinely takes no args, or the test
1296 # catalog omitted a model. Skip args validation rather
1297 # than reject — the design treats missing schema as
1298 # "trust the dispatcher".
1299 continue
1300 args = call.get("args", {})
1301 if not isinstance(args, dict):
1302 raise MissionValidationError(
1303 "validation_error",
1304 details={
1305 "field": "strategy",
1306 "subfield": "tool_calls",
1307 "tool_name": name,
1308 "reason": "tool_args_invalid",
1309 "errors": [
1310 {
1311 "type": "args_not_a_dict",
1312 "actual_type": type(args).__name__,
1313 }
1314 ],
1315 },
1316 )
1317 try:
1318 schema.model_validate(args)
1319 except Exception as exc: # pydantic ValidationError + similar
1320 # Pydantic v2 ValidationError exposes ``.errors()`` as a
1321 # list of structured dicts. Tolerate any other exception
1322 # type (e.g. older Pydantic, custom validators) by
1323 # falling back to ``str(exc)``.
1324 errors_method = getattr(exc, "errors", None)
1325 if callable(errors_method):
1326 try:
1327 errors_payload: Any = errors_method()
1328 except Exception: # defensive
1329 errors_payload = [{"type": "unknown", "msg": str(exc)}]
1330 else:
1331 errors_payload = [{"type": "unknown", "msg": str(exc)}]
1332 raise MissionValidationError(
1333 "validation_error",
1334 details={
1335 "field": "strategy",
1336 "subfield": "tool_calls",
1337 "tool_name": name,
1338 "reason": "tool_args_invalid",
1339 "errors": errors_payload,
1340 },
1341 ) from exc
1343 # 4. Cost estimation against remaining budget. Removed —
1344 # cost guardrails live out-of-band via AWS Budgets / Cost
1345 # Anomaly Detection rather than in the Mission cascade.
1346 # Scripted strategies: validate_strategy already ran allow_scripts
1347 # gating and the AST validator. No catalog-aware checks are layered
1348 # on top here — the script-side enforcement happens at execute time
1349 # via the in-script tool callable wrappers.
1352# ---------------------------------------------------------------------------
1353# Orchestration helpers — bind a backend to a SessionState and return either
1354# a used result or a deterministic fallback.
1355# ---------------------------------------------------------------------------
1357# Local imports kept inside this section so the prompt-builder /
1358# backend half above stays free of audit / decide dependencies.
1359from . import audit as _mission_audit # noqa: E402
1360from . import decide as _decide # noqa: E402
1362# Type alias used by the helpers below. ``SessionState`` is a TypedDict
1363# whose runtime value is just ``dict``; the alias keeps the signatures
1364# expressive without forcing the import to leak through ``__all__``.
1365from .types import SessionState as _SessionState # noqa: E402
1368@dataclass(frozen=True)
1369class SamplingUsed:
1370 """A successful sampling call's accepted output."""
1372 output_text: str
1373 """Raw model output (the text returned by the bound backend)."""
1375 parsed: dict[str, Any]
1376 """Parsed JSON payload that has cleared the schema and catalog checks."""
1378 backend_name: Literal["bedrock"]
1379 """Stable backend identifier — echoes the bound backend's tag."""
1381 model_id: str
1382 """The concrete model id the backend routed the prompt to."""
1385@dataclass(frozen=True)
1386class SamplingFallback:
1387 """A rejected or unavailable sampling call's deterministic substitute.
1389 Returned when the bound backend was ``None``, the transport raised,
1390 the model output failed to parse / validate, or any catalog or
1391 budget check rejected the proposed strategy. The ``rationale`` is
1392 a pure function of the bound :class:`SessionState` and the most
1393 recent :class:`IterationRecord`, so the engine can replay or
1394 reproduce a fallback exactly from persisted state.
1395 """
1397 rationale: str
1398 """Deterministic fallback text. Empty string for ``final_lessons``
1399 — the final-report writer fills in its own deterministic text in
1400 that case."""
1402 reason: str
1403 """Stable token tagging *why* the fallback fired. Examples:
1404 ``"transport_error"``, ``"json_parse"``, ``"schema_mismatch"``,
1405 ``"tool_not_allowlisted"``, ``"tool_args_invalid"``,
1406 ``"over_budget"``, ``"script_rejected"``,
1407 ``"no_backend_resolved"``, ``"disabled"``."""
1409 backend_name: Literal["bedrock", "none"]
1410 """The bound backend's tag, or ``"none"`` when no backend was
1411 resolved at the call site."""
1413 model_id: str | None
1414 """The bound backend's model id, or ``None`` when no backend was
1415 resolved."""
1418# ---------------------------------------------------------------------------
1419# JSON / schema helpers (private to the orchestration layer)
1420# ---------------------------------------------------------------------------
1423def _extract_json_object(text: str) -> dict[str, Any]:
1424 """Parse the first JSON object embedded in ``text``.
1426 Models routinely wrap JSON in prose. The implementation slices from
1427 the first ``{`` to the last ``}`` and feeds the result to
1428 :func:`json.loads`. When no balanced braces are present, or the
1429 sliced substring is not valid JSON, the function raises
1430 :class:`json.JSONDecodeError` so the caller can branch on a single
1431 well-typed exception.
1432 """
1433 start = text.find("{")
1434 end = text.rfind("}")
1435 if start == -1 or end == -1 or end < start:
1436 # No braces at all → treat as a parse error so the calling
1437 # branch surfaces ``reason="json_parse"``.
1438 raise json.JSONDecodeError("no JSON object found", text, 0)
1439 candidate = text[start : end + 1]
1440 parsed = json.loads(candidate)
1441 if not isinstance(parsed, dict):
1442 # The sliced substring parsed but is not an object — surface as
1443 # a parse error too, since downstream code requires a dict.
1444 raise json.JSONDecodeError("top-level JSON value is not an object", candidate, 0)
1445 return parsed
1448def _validate_revision_schema(parsed: dict[str, Any]) -> None:
1449 """Reject a parsed payload that is not a valid Strategy_Revision.
1451 Required keys: ``revision_rationale`` (non-empty str),
1452 ``next_strategy`` (dict), ``confidence`` (number in [0, 1]).
1453 """
1454 rationale = parsed.get("revision_rationale")
1455 if not isinstance(rationale, str) or not rationale:
1456 raise ValueError("schema_mismatch: revision_rationale must be non-empty str")
1457 next_strategy = parsed.get("next_strategy")
1458 if not isinstance(next_strategy, dict):
1459 raise ValueError("schema_mismatch: next_strategy must be a dict")
1460 confidence = parsed.get("confidence")
1461 # ``bool`` is excluded explicitly — it is a subclass of ``int`` in
1462 # Python and would otherwise sneak past the numeric check.
1463 if isinstance(confidence, bool) or not isinstance(confidence, (int, float)):
1464 raise ValueError("schema_mismatch: confidence must be a number")
1465 if not (0.0 <= float(confidence) <= 1.0):
1466 raise ValueError("schema_mismatch: confidence must be in [0, 1]")
1469def _validate_lessons_schema(parsed: dict[str, Any]) -> None:
1470 """Reject a parsed payload that is not a valid final-lessons dict.
1472 Required keys: ``lessons`` (non-empty list of non-empty str),
1473 ``recommended_followups`` (list of str — may be empty).
1474 """
1475 lessons = parsed.get("lessons")
1476 if not isinstance(lessons, list) or not lessons:
1477 raise ValueError("schema_mismatch: lessons must be a non-empty list")
1478 for item in lessons:
1479 if not isinstance(item, str) or not item:
1480 raise ValueError("schema_mismatch: each lesson must be a non-empty str")
1481 followups = parsed.get("recommended_followups")
1482 if not isinstance(followups, list):
1483 raise ValueError("schema_mismatch: recommended_followups must be a list")
1484 for item in followups:
1485 if not isinstance(item, str):
1486 raise ValueError("schema_mismatch: each follow-up must be a str")
1489# ---------------------------------------------------------------------------
1490# maybe_sample_strategy_revision
1491# ---------------------------------------------------------------------------
1494async def maybe_sample_strategy_revision(
1495 *,
1496 backend: SamplingBackend | None,
1497 session: _SessionState,
1498 iteration: IterationRecord,
1499 allowlist: list[str],
1500 registered_tools: dict[str, Any],
1501 tool_docstrings: dict[str, str],
1502 remaining_iterations: int,
1503 remaining_wall_clock_secs: float | None,
1504 allow_scripts: bool,
1505 environment_context: Mapping[str, Any] | None = None,
1506 prior_missions: Sequence[Mapping[str, Any]] | None = None,
1507) -> SamplingUsed | SamplingFallback:
1508 """Consult the advisory LLM for a Strategy_Revision, or fall back.
1510 Returns a :class:`SamplingUsed` when the bound backend produces a
1511 JSON object that clears schema validation and the catalog checks.
1512 Returns a :class:`SamplingFallback` carrying the deterministic
1513 rationale from
1514 :func:`mission.decide.build_revision_rationale_template` on every
1515 rejection class. Emits exactly one
1516 :func:`mission.audit.emit_sampling_event` per call.
1517 """
1518 session_id = session["session_id"]
1519 iteration_index = iteration["iteration_index"]
1520 template = _decide.build_revision_rationale_template(session, iteration)
1522 # ---- No backend resolved: short-circuit. ------------------------------
1523 if backend is None:
1524 _mission_audit.emit_sampling_event(
1525 session_id,
1526 iteration_index,
1527 sampling_purpose="strategy_revision",
1528 sampling_status="disabled",
1529 sampling_backend="none",
1530 )
1531 return SamplingFallback(
1532 rationale=template,
1533 reason="no_backend_resolved",
1534 backend_name="none",
1535 model_id=None,
1536 )
1538 backend_name = backend.backend_name
1539 model_id = backend.model_id
1541 # ---- Build the prompt. ------------------------------------------------
1542 # The in-progress iteration that triggered ``adjust`` is already in
1543 # ``session["iterations"][-1]``, so the most-recent-five window is a
1544 # plain slice; ``RECENT_ITERATIONS_LIMIT`` is enforced inside the
1545 # prompt builder as a defensive cap.
1546 recent_iterations = list(session["iterations"][-RECENT_ITERATIONS_LIMIT:])
1547 tool_schemas = _extract_tool_json_schemas(allowlist, registered_tools)
1548 prompt = SamplingPrompt(
1549 directive=session["directive_text"],
1550 success_criteria=session["criteria"],
1551 criteria_status=iteration["criteria_evaluation"],
1552 recent_iterations=recent_iterations,
1553 tool_allowlist=allowlist,
1554 tool_docstrings=tool_docstrings,
1555 remaining_iterations=remaining_iterations,
1556 remaining_wall_clock_secs=remaining_wall_clock_secs,
1557 allow_scripts=allow_scripts,
1558 tool_schemas=tool_schemas,
1559 environment_context=environment_context,
1560 prior_missions=prior_missions,
1561 )
1563 # ---- Transport: backend.sample. --------------------------------------
1564 try:
1565 output_text = await backend.sample(prompt)
1566 except SamplingTransportError as err:
1567 _mission_audit.emit_sampling_event(
1568 session_id,
1569 iteration_index,
1570 sampling_purpose="strategy_revision",
1571 sampling_status="rejected",
1572 sampling_backend=backend_name,
1573 sampling_model_id=model_id or None,
1574 validation_error=err.code,
1575 )
1576 return SamplingFallback(
1577 rationale=template,
1578 reason="transport_error",
1579 backend_name=backend_name,
1580 model_id=model_id,
1581 )
1583 # ---- Parse the output as JSON. ---------------------------------------
1584 try:
1585 parsed = _extract_json_object(output_text)
1586 except json.JSONDecodeError:
1587 _mission_audit.emit_sampling_event(
1588 session_id,
1589 iteration_index,
1590 sampling_purpose="strategy_revision",
1591 sampling_status="rejected",
1592 sampling_backend=backend_name,
1593 sampling_model_id=model_id or None,
1594 validation_error="json_parse",
1595 )
1596 return SamplingFallback(
1597 rationale=template,
1598 reason="json_parse",
1599 backend_name=backend_name,
1600 model_id=model_id,
1601 )
1603 # ---- Schema validation. ----------------------------------------------
1604 try:
1605 _validate_revision_schema(parsed)
1606 except ValueError:
1607 _mission_audit.emit_sampling_event(
1608 session_id,
1609 iteration_index,
1610 sampling_purpose="strategy_revision",
1611 sampling_status="rejected",
1612 sampling_backend=backend_name,
1613 sampling_model_id=model_id or None,
1614 validation_error="schema_mismatch",
1615 )
1616 return SamplingFallback(
1617 rationale=template,
1618 reason="schema_mismatch",
1619 backend_name=backend_name,
1620 model_id=model_id,
1621 )
1623 # ---- Catalog validation on the proposed next_strategy. ---------------
1624 try:
1625 validate_strategy_against_catalog(
1626 parsed["next_strategy"],
1627 allowlist,
1628 registered_tools,
1629 allow_scripts,
1630 )
1631 except MissionValidationError as err:
1632 # ``err.details["reason"]`` carries the structured rejection
1633 # token (e.g. ``"tool_not_allowlisted"``,
1634 # ``"tool_args_invalid"``). Fall back to a generic label when
1635 # the validator emits a rejection without a ``reason`` key.
1636 details = err.details or {}
1637 reason = details.get("reason", "validation_error")
1638 _mission_audit.emit_sampling_event(
1639 session_id,
1640 iteration_index,
1641 sampling_purpose="strategy_revision",
1642 sampling_status="rejected",
1643 sampling_backend=backend_name,
1644 sampling_model_id=model_id or None,
1645 validation_error=str(reason),
1646 )
1647 return SamplingFallback(
1648 rationale=template,
1649 reason=str(reason),
1650 backend_name=backend_name,
1651 model_id=model_id,
1652 )
1654 # ---- Success path. ---------------------------------------------------
1655 # Extract token usage from the backend if available (Bedrock backend
1656 # stores it as a side-channel after each sample() call).
1657 _input_tokens = getattr(backend, "last_input_tokens", None)
1658 _output_tokens = getattr(backend, "last_output_tokens", None)
1659 _mission_audit.emit_sampling_event(
1660 session_id,
1661 iteration_index,
1662 sampling_purpose="strategy_revision",
1663 sampling_status="used",
1664 sampling_backend=backend_name,
1665 sampling_model_id=model_id or None,
1666 model_output_bytes=len(output_text.encode("utf-8")),
1667 input_tokens=_input_tokens,
1668 output_tokens=_output_tokens,
1669 )
1670 return SamplingUsed(
1671 output_text=output_text,
1672 parsed=parsed,
1673 backend_name=backend_name,
1674 model_id=model_id,
1675 )
1678# ---------------------------------------------------------------------------
1679# maybe_sample_final_lessons
1680# ---------------------------------------------------------------------------
1683async def maybe_sample_final_lessons(
1684 *,
1685 backend: SamplingBackend | None,
1686 session: _SessionState,
1687 remaining_iterations: int = 0,
1688 remaining_wall_clock_secs: float | None = None,
1689 allow_scripts: bool = False,
1690 tool_docstrings: dict[str, str] | None = None,
1691 environment_context: Mapping[str, Any] | None = None,
1692) -> SamplingUsed | SamplingFallback:
1693 """Consult the advisory LLM for final lessons, or fall back.
1695 Returns a :class:`SamplingUsed` when the bound backend produces a
1696 JSON object that clears the lessons schema. Returns a
1697 :class:`SamplingFallback` with an *empty* rationale on every
1698 rejection class — the final-report writer is responsible for the
1699 deterministic-text path when sampling does not produce usable
1700 output. Emits exactly one
1701 :func:`mission.audit.emit_sampling_event` per call, with
1702 ``iteration_index_or_purpose=None`` since the call is out-of-loop.
1703 """
1704 session_id = session["session_id"]
1706 # ---- No backend resolved: short-circuit. ------------------------------
1707 if backend is None:
1708 _mission_audit.emit_sampling_event(
1709 session_id,
1710 None,
1711 sampling_purpose="final_lessons",
1712 sampling_status="disabled",
1713 sampling_backend="none",
1714 )
1715 return SamplingFallback(
1716 rationale="",
1717 reason="no_backend_resolved",
1718 backend_name="none",
1719 model_id=None,
1720 )
1722 backend_name = backend.backend_name
1723 model_id = backend.model_id
1725 # ---- Build the prompt. ------------------------------------------------
1726 # Pass *all* iterations; the prompt builder trims / drops as needed
1727 # to fit the byte budget.
1728 prompt = SamplingPrompt(
1729 directive=session["directive_text"],
1730 success_criteria=session["criteria"],
1731 # The lessons prompt has no per-iteration criteria status; the
1732 # builder still expects the field, so reuse the most recent
1733 # iteration's evaluation when available, else an empty list.
1734 criteria_status=(
1735 list(session["iterations"][-1]["criteria_evaluation"]) if session["iterations"] else []
1736 ),
1737 recent_iterations=list(session["iterations"]),
1738 tool_allowlist=session.get("tool_allowlist", []),
1739 tool_docstrings=tool_docstrings or {},
1740 remaining_iterations=remaining_iterations,
1741 remaining_wall_clock_secs=remaining_wall_clock_secs,
1742 allow_scripts=allow_scripts,
1743 environment_context=environment_context,
1744 )
1746 # ---- Transport: backend.sample (uses lessons assembler). -------------
1747 try:
1748 # We render the lessons-specific prompt here so the byte-cap
1749 # bookkeeping uses the right schema header. The backend's own
1750 # ``sample`` calls ``prompt.assemble()`` under the hood for the
1751 # Strategy_Revision flow, but for lessons we assemble here and
1752 # invoke a thin shim through the backend.
1753 rendered = prompt.assemble_final_lessons()
1754 output_text = await _sample_with_assembled_text(backend, rendered)
1755 except SamplingTransportError as err:
1756 _mission_audit.emit_sampling_event(
1757 session_id,
1758 None,
1759 sampling_purpose="final_lessons",
1760 sampling_status="rejected",
1761 sampling_backend=backend_name,
1762 sampling_model_id=model_id or None,
1763 validation_error=err.code,
1764 )
1765 return SamplingFallback(
1766 rationale="",
1767 reason="transport_error",
1768 backend_name=backend_name,
1769 model_id=model_id,
1770 )
1772 # ---- Parse the output as JSON. ---------------------------------------
1773 try:
1774 parsed = _extract_json_object(output_text)
1775 except json.JSONDecodeError:
1776 _mission_audit.emit_sampling_event(
1777 session_id,
1778 None,
1779 sampling_purpose="final_lessons",
1780 sampling_status="rejected",
1781 sampling_backend=backend_name,
1782 sampling_model_id=model_id or None,
1783 validation_error="json_parse",
1784 )
1785 return SamplingFallback(
1786 rationale="",
1787 reason="json_parse",
1788 backend_name=backend_name,
1789 model_id=model_id,
1790 )
1792 # ---- Schema validation. ----------------------------------------------
1793 try:
1794 _validate_lessons_schema(parsed)
1795 except ValueError:
1796 _mission_audit.emit_sampling_event(
1797 session_id,
1798 None,
1799 sampling_purpose="final_lessons",
1800 sampling_status="rejected",
1801 sampling_backend=backend_name,
1802 sampling_model_id=model_id or None,
1803 validation_error="schema_mismatch",
1804 )
1805 return SamplingFallback(
1806 rationale="",
1807 reason="schema_mismatch",
1808 backend_name=backend_name,
1809 model_id=model_id,
1810 )
1812 # ---- Success path. ---------------------------------------------------
1813 _input_tokens = getattr(backend, "last_input_tokens", None)
1814 _output_tokens = getattr(backend, "last_output_tokens", None)
1815 _mission_audit.emit_sampling_event(
1816 session_id,
1817 None,
1818 sampling_purpose="final_lessons",
1819 sampling_status="used",
1820 sampling_backend=backend_name,
1821 sampling_model_id=model_id or None,
1822 model_output_bytes=len(output_text.encode("utf-8")),
1823 input_tokens=_input_tokens,
1824 output_tokens=_output_tokens,
1825 )
1826 return SamplingUsed(
1827 output_text=output_text,
1828 parsed=parsed,
1829 backend_name=backend_name,
1830 model_id=model_id,
1831 )
1834async def _sample_with_assembled_text(backend: SamplingBackend, rendered: str) -> str:
1835 """Route a pre-assembled prompt string through a backend.
1837 Both shipped backends accept a :class:`SamplingPrompt` and call
1838 ``assemble`` themselves to render the strategy-revision shape. For
1839 the final-lessons path we render the lessons-shaped prompt here
1840 and need to deliver that exact text to the transport. The shim
1841 builds a tiny prompt-shaped wrapper whose :meth:`assemble` returns
1842 the pre-rendered text and forwards it to the backend.
1843 """
1844 pre_rendered = rendered
1846 class _PreRendered:
1847 """Thin :class:`SamplingPrompt` look-alike with a fixed assemble()."""
1849 def assemble(self) -> str:
1850 return pre_rendered
1852 # The two shipped backends only call ``prompt.assemble()`` so the
1853 # duck-typed wrapper above is enough to drive either of them.
1854 return await backend.sample(_PreRendered()) # type: ignore[arg-type]
1857# ---------------------------------------------------------------------------
1858# Session-start sampling-state resolver
1859# ---------------------------------------------------------------------------
1862def _bedrock_credentials_available() -> bool:
1863 """Lightweight probe: do local AWS credentials resolve?
1865 Instantiates a ``boto3.Session()`` and asks for ``get_credentials()``
1866 without making any network call. ``boto3`` is imported inside the
1867 function so the module's top-level import surface stays free of
1868 SDK dependencies — and so a host that has no ``boto3`` installed
1869 (or any other unexpected import-time failure) cleanly degrades to
1870 "no credentials available" rather than crashing the helper.
1871 """
1872 try:
1873 import boto3
1875 session = boto3.Session()
1876 creds = session.get_credentials()
1877 return creds is not None
1878 except Exception:
1879 return False
1882def resolve_sampling_state(
1883 use_sampling_param: bool | None,
1884) -> tuple[bool, Literal["bedrock", "none"]]:
1885 """Decide whether sampling is enabled for a session and which backend resolves.
1887 Bedrock is the only sampling transport (MCP client sampling left the
1888 protocol with FastMCP 4), so resolution no longer depends on how the
1889 session was started — CLI and MCP callers probe the same server-side
1890 credentials.
1892 Resolution precedence (first match wins):
1894 1. ``use_sampling_param is False`` — caller explicitly disabled
1895 sampling, so the result is ``(False, "none")`` regardless of
1896 any capability the environment advertises.
1897 2. Local AWS credentials resolve — ``(True, "bedrock")``.
1898 3. No credentials — ``(True, "none")`` if the caller opted in
1899 explicitly with ``use_sampling_param is True`` (so the caller can
1900 decide whether to error or proceed deterministic-only), and
1901 ``(False, "none")`` otherwise.
1903 Args:
1904 use_sampling_param: Three-state opt-in flag. ``None`` means the
1905 caller did not specify and the helper should auto-detect.
1906 ``False`` short-circuits to a disabled state. ``True`` means
1907 the caller explicitly opted in; the backend is auto-detected
1908 and ``"none"`` is allowed when no concrete backend resolves.
1910 Returns:
1911 A ``(use_sampling, backend)`` tuple. The caller persists both
1912 values on its ``SessionState`` so the audit pipeline can stamp
1913 every later sampling event with the resolved backend.
1914 """
1915 # 1. Explicit opt-out wins outright.
1916 if use_sampling_param is False:
1917 return (False, "none")
1919 # 2. Probe server-side AWS credentials.
1920 if _bedrock_credentials_available():
1921 return (True, "bedrock")
1923 # 3. No credentials — only honour an explicit True.
1924 if use_sampling_param is True:
1925 return (True, "none")
1926 return (False, "none")