tracer-cloud--opensre/tests/core/agent/scenario_loader.py

"""Load turn scenario directories into typed fixtures for pytest."""

from __future__ import annotations

import math
import os
import random
from dataclasses import dataclass
from pathlib import Path
from typing import Any, cast

import yaml

from surfaces.interactive_shell.command_registry import SLASH_COMMANDS
from tests.core.agent._planned_action import (
    default_target_surface,
)
from tools.interactive_shell.actions.synthetic import (
    list_rds_postgres_scenarios,
)

TESTS_DIR = Path(__file__).resolve().parent
SCENARIOS_DIR = TESTS_DIR / "scenarios"

INTENT_CLASSES = frozenset(
    {
        "chat_handoff",
        "local_execution",
        "investigation",
        "complex_shell_prompts",
        "compound",
        "remote",
        "follow_up",
        "non_actionable",
    }
)
VALID_TOOL_ACTION_SURFACES = frozenset({"dispatch", "gather"})
VALID_GATHER_EXPECTS = frozenset(
    {
        "not_called",
        "called",
        "call_any",
        "valid_data",
        "valid_data_any",
    }
)
VALID_ACTION_KINDS = frozenset(
    {
        "llm_provider",
        "slash",
        "shell",
        "sample_alert",
        "investigation",
        "synthetic_test",
        "task_cancel",
        "cli_command",
        "implementation",
        "assistant_handoff",
    }
)
VALID_ACTION_SOURCES = frozenset({"deterministic", "llm"})
VALID_TARGET_SURFACES = frozenset({"slash", "terminal", "investigation", "implementation"})

INTENT_TO_BEHAVIOR_CLASS: dict[str, str] = {
    "chat_handoff": "chat_handoff",
    "local_execution": "local_execution",
    "investigation": "investigations",
    "complex_shell_prompts": "complex_shell_prompts",
    "compound": "compound",
    "remote": "remote",
    "follow_up": "follow_up",
    "non_actionable": "non_actionable",
}


@dataclass(frozen=True)
class ScenarioInput:
    prompt: str


@dataclass(frozen=True)
class ScenarioSession:
    has_prior_state: bool
    configured_integrations: tuple[str, ...]
    resolved_integrations: dict[str, Any] | None = None


@dataclass(frozen=True)
class ScenarioCapabilities:
    """Per-scenario planner capability constraints (three-state).

    Each field carries one of three states that map directly onto the runtime
    capability gate (``capability_not_explicitly_disabled``):

    * ``None`` — the capability key is absent; the tool stays available, which
      matches the production default (``Session()`` has no capability
      constraints).
    * ``()`` — an explicit empty list; the tool is explicitly disabled (hidden
      from the planner specs and blocked at dispatch).
    * a non-empty tuple — an allowlist; the tool is available and the action
      normalizer drops proposed actions outside the list.
    """

    slash_commands: tuple[str, ...] | None
    cli_commands: tuple[str, ...] | None
    synthetic_suites: tuple[str, ...] | None
    llm_provider: tuple[str, ...] | None


@dataclass(frozen=True)
class Scenario:
    id: str
    title: str
    intent_class: str
    input: ScenarioInput
    session: ScenarioSession
    available_capabilities: ScenarioCapabilities
    notes: tuple[str, ...]
    behavior_class: str
    scenario_dir: Path


@dataclass(frozen=True)
class AnswerTurn:
    expected_kind: str


@dataclass(frozen=True)
class AnswerPolicy:
    """Execution expectation for the action-agent tool path only.

    ``executes_terminal_action`` is true when the turn is expected to run at
    least one shell action AgentTool -- a slash command, shell command, sample
    alert, investigation start, synthetic run, etc. It is false for
    conversational turns that answer in chat without executing a terminal
    action.

    This flag does NOT describe the conversational data-gathering path
    (``gather_tool_evidence``), where the assistant may query configured
    integrations (Sentry, GitHub, PostHog, ...) while composing a chat answer.
    That path is not modeled as planned/executed actions; it is asserted via
    ``response_contract`` text and by execution-layer tests. See the ``Answer``
    docstring for the full two-path model.
    """

    executes_terminal_action: bool


@dataclass(frozen=True)
class GatheredToolsContract:
    """Assertions on which registered tools fire during the conversational
    ``gather_tool_evidence`` loop for a turn.

    A turn's conversational data-gathering pass runs the same registered tools
    the investigation uses. This contract lets a scenario assert that the right
    tools were (or were not) invoked when grounding a chat answer:

    * ``must_call_any`` — at least one of these tool names must be invoked.
    * ``must_call_all`` — every one of these tool names must be invoked.
    * ``must_not_call`` — none of these tool names may be invoked.
    * ``must_return_valid_data`` — every one of these tool names must be invoked
      AND return a successful result (a real integration response, not an error
      or an ``available: false`` placeholder). This is strictly stronger than
      ``must_call_all``: it fails on a credential 401, a malformed-param 400, or
      any other errored call, so it can only pass when the tool actually reached
      the live integration and got valid data back.
    * ``must_return_valid_data_any`` — at least one of these tool names must be
      invoked AND return valid data (same success criteria as
      ``must_return_valid_data``).

    For ``must_call_any``, ``must_call_all``, and ``must_not_call`` a tool counts
    as "called" when it appears in ``AgentRunResult.executed`` regardless of
    whether the call succeeded. ``must_return_valid_data`` additionally inspects
    the tool's output and only counts a call that returned valid data.
    """

    must_call_any: tuple[str, ...]
    must_call_all: tuple[str, ...]
    must_not_call: tuple[str, ...]
    must_return_valid_data: tuple[str, ...]
    must_return_valid_data_any: tuple[str, ...]


@dataclass(frozen=True)
class Answer:
    """Expected behavior for one turn scenario.

    A turn can resolve down one of two independent execution paths, and these
    fields only describe the first:

    1. Action agent -> AgentTool execution (the "execution" path). Covered by
       ``policy.executes_terminal_action``, ``planned_actions``, and dispatch
       entries in ``tool_actions`` (``surface: dispatch``). An empty
       ``planned_actions`` means the action agent is expected to hand the turn to
       the conversational assistant (an ``assistant_handoff``), i.e. no terminal
       action runs.

    2. Conversational answer + ``gather_tool_evidence`` tool loop (the "chat"
       path). Assert gather behaviour via ``tool_actions`` entries with
       ``surface: gather`` and an ``expect`` mode (``not_called``, ``called``,
       ``valid_data``, etc.). ``response_contract`` still covers reply text.
    """

    turn: AnswerTurn
    policy: AnswerPolicy
    planned_actions: tuple[dict[str, Any], ...]
    executed_actions: tuple[dict[str, Any], ...]
    response_contract: dict[str, list[str]]
    history_expected: tuple[dict[str, Any], ...]
    runs: int
    gathered_tools_contract: GatheredToolsContract | None = None


@dataclass(frozen=True)
class ScenarioCase:
    scenario: Scenario
    answer: Answer


def _require_mapping(raw: object, *, label: str) -> dict[str, Any]:
    if not isinstance(raw, dict):
        msg = f"{label} must be a mapping, got {type(raw).__name__}."
        raise ValueError(msg)
    return cast(dict[str, Any], raw)


def _optional_mapping(raw: object, *, label: str) -> dict[str, Any] | None:
    """Parse an optional mapping field.

    Returns ``None`` when the key is absent or explicitly null (preserving the
    "use the real resolved store" default), and the mapping itself when present
    (including an explicit empty ``{}`` that forces an isolated, empty store).
    """
    if raw is None:
        return None
    if not isinstance(raw, dict):
        msg = f"{label} must be a mapping, got {type(raw).__name__}."
        raise ValueError(msg)
    return cast(dict[str, Any], raw)


def _gather_tool_names(entry: dict[str, Any], *, label: str) -> tuple[str, ...]:
    tool = entry.get("tool")
    tools = entry.get("tools")
    if tool is not None and tools is not None:
        msg = f"{label}: set either 'tool' or 'tools', not both."
        raise ValueError(msg)
    if isinstance(tool, str) and tool.strip():
        return (tool.strip(),)
    if tools is not None:
        return _string_list(tools, label=f"{label}.tools")
    msg = f"{label}: gather action requires 'tool' or 'tools'."
    raise ValueError(msg)


def _parse_tool_actions(
    raw: object,
    *,
    label: str,
    scenario_id: str,
    executes_terminal_action: bool,
) -> tuple[tuple[dict[str, Any], ...], GatheredToolsContract | None]:
    """Parse unified ``tool_actions`` into dispatch + gather contract views.

    ``surface: dispatch`` entries become ``executed_actions`` shapes. ``surface:
    gather`` entries aggregate into a :class:`GatheredToolsContract` by
    ``expect`` mode.
    """
    if raw is None:
        return (), None
    if not isinstance(raw, list):
        msg = f"{label} must be a list, got {type(raw).__name__}."
        raise ValueError(msg)

    executed: list[dict[str, Any]] = []
    must_call_any: list[str] = []
    must_call_all: list[str] = []
    must_not_call: list[str] = []
    must_return_valid_data: list[str] = []
    must_return_valid_data_any: list[str] = []

    for index, item in enumerate(raw):
        entry_label = f"{label}[{index}]"
        if not isinstance(item, dict):
            msg = f"{entry_label} must be a mapping."
            raise ValueError(msg)
        entry = cast(dict[str, Any], item)
        surface = str(entry.get("surface", "")).strip()
        if surface not in VALID_TOOL_ACTION_SURFACES:
            msg = (
                f"{entry_label}: surface must be one of "
                f"{sorted(VALID_TOOL_ACTION_SURFACES)!r}, got {surface!r}."
            )
            raise ValueError(msg)

        if surface == "dispatch":
            if "expect" in entry:
                msg = f"{entry_label}: dispatch actions must not set 'expect'."
                raise ValueError(msg)
            dispatch_action = {key: value for key, value in entry.items() if key != "surface"}
            validate_action_shape(
                dispatch_action,
                prefix=f"{scenario_id} tool_actions[{index}]",
                require_source=False,
            )
            executed.append(dispatch_action)
            continue

        expect = str(entry.get("expect", "")).strip()
        if expect not in VALID_GATHER_EXPECTS:
            msg = (
                f"{entry_label}: expect must be one of {sorted(VALID_GATHER_EXPECTS)!r}, "
                f"got {expect!r}."
            )
            raise ValueError(msg)
        tool_names = _gather_tool_names(entry, label=entry_label)
        if expect == "not_called":
            must_not_call.extend(tool_names)
        elif expect == "called":
            must_call_all.extend(tool_names)
        elif expect == "call_any":
            must_call_any.extend(tool_names)
        elif expect == "valid_data":
            must_return_valid_data.extend(tool_names)
        else:
            must_return_valid_data_any.extend(tool_names)

    if not executes_terminal_action and executed:
        msg = f"{label}: executes_terminal_action=false requires no dispatch tool_actions."
        raise ValueError(msg)

    contract = GatheredToolsContract(
        must_call_any=tuple(must_call_any),
        must_call_all=tuple(must_call_all),
        must_not_call=tuple(must_not_call),
        must_return_valid_data=tuple(must_return_valid_data),
        must_return_valid_data_any=tuple(must_return_valid_data_any),
    )
    if not (
        contract.must_call_any
        or contract.must_call_all
        or contract.must_not_call
        or contract.must_return_valid_data
        or contract.must_return_valid_data_any
    ):
        return tuple(executed), None
    return tuple(executed), contract


def _string_list(raw: object, *, label: str) -> tuple[str, ...]:
    if raw is None:
        return ()
    if not isinstance(raw, list):
        msg = f"{label} must be a list, got {type(raw).__name__}."
        raise ValueError(msg)
    values: list[str] = []
    for index, item in enumerate(raw):
        if not isinstance(item, str) or not item.strip():
            msg = f"{label}[{index}] must be a non-empty string."
            raise ValueError(msg)
        values.append(item.strip())
    return tuple(values)


def _optional_string_list(raw: object, *, label: str) -> tuple[str, ...] | None:
    """Parse a capability allowlist while preserving the absent-vs-empty split.

    Returns ``None`` when the key is absent or explicitly null (no constraint;
    the tool stays available, matching the production default), ``()`` for an
    explicit empty list (the capability is explicitly disabled), and a tuple of
    non-empty strings for an allowlist.
    """
    if raw is None:
        return None
    return _string_list(raw, label=label)


def _action_list(raw: object, *, label: str) -> tuple[dict[str, Any], ...]:
    if raw is None:
        return ()
    if not isinstance(raw, list):
        msg = f"{label} must be a list, got {type(raw).__name__}."
        raise ValueError(msg)
    actions: list[dict[str, Any]] = []
    for index, item in enumerate(raw):
        if not isinstance(item, dict):
            msg = f"{label}[{index}] must be a mapping."
            raise ValueError(msg)
        actions.append(cast(dict[str, Any], item))
    return tuple(actions)


def _slash_content(command: str, args: list[str]) -> str:
    return " ".join([command, *args]) if args else command


def _normalize_planned_action(action: dict[str, Any]) -> dict[str, Any]:
    """Backfill derived fields so YAMLs can omit redundant data."""
    kind = str(action.get("kind", "")).strip()
    if kind == "slash":
        command = str(action.get("command", "")).strip()
        raw_args = action.get("args") or []
        args = [str(arg).strip() for arg in raw_args] if isinstance(raw_args, list) else []
        if "content" not in action and command:
            action["content"] = _slash_content(command, args)
    elif kind == "synthetic_test":
        suite = str(action.get("suite", "")).strip()
        scenario = str(action.get("scenario", "")).strip()
        if "content" not in action and suite and scenario:
            action["content"] = f"{suite}:{scenario}"
    elif kind == "cli_command":
        payload = str(action.get("payload", "")).strip()
        if "content" not in action and payload:
            action["content"] = payload
    elif kind == "sample_alert":
        if "content" not in action and "template" in action:
            action["content"] = str(action["template"]).strip()
    return action


def validate_action_shape(
    action: dict[str, Any],
    *,
    prefix: str,
    require_source: bool,
) -> None:
    kind = str(action.get("kind", "")).strip()
    if kind not in VALID_ACTION_KINDS:
        msg = f"{prefix} has invalid kind {kind!r}."
        raise ValueError(msg)

    if require_source and kind != "assistant_handoff":
        source = str(action.get("source", "")).strip()
        if source not in VALID_ACTION_SOURCES:
            msg = f"{prefix} has invalid source {source!r}."
            raise ValueError(msg)
        target_surface = str(action.get("target_surface", "")).strip()
        if target_surface not in VALID_TARGET_SURFACES:
            msg = f"{prefix} has invalid target_surface {target_surface!r}."
            raise ValueError(msg)
        canonical = default_target_surface(kind)  # type: ignore[arg-type]
        if target_surface != canonical:
            msg = (
                f"{prefix} target_surface {target_surface!r} "
                f"must be {canonical!r} for kind {kind!r}."
            )
            raise ValueError(msg)

    if kind == "slash":
        command = str(action.get("command", "")).strip()
        raw_args = action.get("args")
        if not command.startswith("/"):
            msg = f"{prefix} slash command must start with '/'."
            raise ValueError(msg)
        source = str(action.get("source", "")).strip()
        if require_source and source == "llm" and command not in SLASH_COMMANDS:
            msg = f"{prefix} references unknown slash command {command!r}."
            raise ValueError(msg)
        if not isinstance(raw_args, list):
            msg = f"{prefix} slash action must define args list."
            raise ValueError(msg)
        args = [str(arg).strip() for arg in raw_args]
        content = str(action.get("content", "")).strip()
        if content and content != _slash_content(command, args):
            msg = f"{prefix} content must match command+args when set."
            raise ValueError(msg)
    elif kind == "synthetic_test":
        suite = str(action.get("suite", "")).strip()
        scenario = str(action.get("scenario", "")).strip()
        if not suite or not scenario:
            msg = f"{prefix} synthetic_test requires suite and scenario."
            raise ValueError(msg)
        available = set(list_rds_postgres_scenarios())
        if scenario not in available:
            msg = f"{prefix} unknown synthetic scenario {scenario!r}."
            raise ValueError(msg)
        content = str(action.get("content", "")).strip()
        if content and content != f"{suite}:{scenario}":
            msg = f"{prefix} content must match suite:scenario when set."
            raise ValueError(msg)
    elif kind == "cli_command":
        payload = str(action.get("payload", "")).strip()
        if not payload:
            msg = f"{prefix} cli_command requires payload."
            raise ValueError(msg)
        if payload.lower().startswith("opensre "):
            msg = f"{prefix} cli_command payload must not include opensre prefix."
            raise ValueError(msg)


def _parse_scenario_yaml(
    scenario_path: Path,
    *,
    behavior_class: str,
) -> Scenario:
    raw = yaml.safe_load(scenario_path.read_text(encoding="utf-8"))
    data = _require_mapping(raw, label=str(scenario_path))

    scenario_id = str(data.get("id", "")).strip()
    if not scenario_id:
        msg = f"{scenario_path}: missing id."
        raise ValueError(msg)

    title = str(data.get("title", "")).strip()
    if not title:
        msg = f"{scenario_path}: missing title."
        raise ValueError(msg)

    intent_class = str(data.get("intent_class", "")).strip()
    if intent_class not in INTENT_CLASSES:
        msg = f"{scenario_path}: invalid intent_class {intent_class!r}."
        raise ValueError(msg)

    expected_behavior = INTENT_TO_BEHAVIOR_CLASS.get(intent_class)
    if expected_behavior != behavior_class:
        msg = (
            f"{scenario_path}: intent_class {intent_class!r} "
            f"does not match directory behavior class {behavior_class!r}."
        )
        raise ValueError(msg)

    input_raw = _require_mapping(data.get("input"), label=f"{scenario_path} input")
    prompt = str(input_raw.get("prompt", "")).strip()
    if not prompt:
        msg = f"{scenario_path}: input.prompt must be non-empty."
        raise ValueError(msg)

    session_raw = _require_mapping(data.get("session"), label=f"{scenario_path} session")
    capabilities_raw = _require_mapping(
        data.get("available_capabilities", {}),
        label=f"{scenario_path} available_capabilities",
    )

    return Scenario(
        id=scenario_id,
        title=title,
        intent_class=intent_class,
        input=ScenarioInput(prompt=prompt),
        session=ScenarioSession(
            has_prior_state=bool(session_raw.get("has_prior_state", False)),
            configured_integrations=_string_list(
                session_raw.get("configured_integrations"),
                label=f"{scenario_path} session.configured_integrations",
            ),
            resolved_integrations=_optional_mapping(
                session_raw.get("resolved_integrations"),
                label=f"{scenario_path} session.resolved_integrations",
            ),
        ),
        available_capabilities=ScenarioCapabilities(
            slash_commands=_optional_string_list(
                capabilities_raw.get("slash_commands"),
                label=f"{scenario_path} slash_commands",
            ),
            cli_commands=_optional_string_list(
                capabilities_raw.get("cli_commands"),
                label=f"{scenario_path} cli_commands",
            ),
            synthetic_suites=_optional_string_list(
                capabilities_raw.get("synthetic_suites"),
                label=f"{scenario_path} synthetic_suites",
            ),
            llm_provider=_optional_string_list(
                capabilities_raw.get("llm_provider"),
                label=f"{scenario_path} llm_provider",
            ),
        ),
        notes=_string_list(data.get("notes"), label=f"{scenario_path} notes"),
        behavior_class=behavior_class,
        scenario_dir=scenario_path,
    )


def _parse_answer_yaml(answer_path: Path, *, scenario_id: str) -> Answer:
    raw = yaml.safe_load(answer_path.read_text(encoding="utf-8"))
    data = _require_mapping(raw, label=str(answer_path))

    turn_raw = _require_mapping(data.get("turn"), label=f"{answer_path} turn")
    policy_raw = _require_mapping(data.get("policy"), label=f"{answer_path} policy")
    response_raw = _require_mapping(
        data.get("response_contract", {}),
        label=f"{answer_path} response_contract",
    )
    history_raw = _require_mapping(data.get("history", {}), label=f"{answer_path} history")

    expected_kind = str(turn_raw.get("expected_kind", "")).strip()
    if expected_kind != "agent":
        msg = f"{answer_path}: invalid turn.expected_kind {expected_kind!r}."
        raise ValueError(msg)
    if "expected_signals" in turn_raw:
        msg = f"{answer_path}: turn.expected_signals was removed; drop it from the fixture."
        raise ValueError(msg)
    if "expected_command_text" in turn_raw:
        msg = (
            f"{answer_path}: turn.expected_command_text was removed along with the "
            "deterministic command-detection layer; drop it from the fixture."
        )
        raise ValueError(msg)

    for removed_key in ("should_execute", "has_unhandled_clause", "fail_closed"):
        if removed_key in policy_raw:
            msg = (
                f"{answer_path}: policy.{removed_key!r} was removed; "
                "use policy.executes_terminal_action instead."
            )
            raise ValueError(msg)
    executes_terminal_action = bool(policy_raw.get("executes_terminal_action", False))

    if "executed_actions" in data or "gathered_tools_contract" in data:
        msg = (
            f"{answer_path}: executed_actions and gathered_tools_contract were removed; "
            "use tool_actions with surface dispatch|gather instead."
        )
        raise ValueError(msg)

    planned_actions = tuple(
        _normalize_planned_action(dict(item))
        for item in _action_list(
            data.get("planned_actions"), label=f"{answer_path} planned_actions"
        )
    )

    for index, action in enumerate(planned_actions):
        validate_action_shape(
            action,
            prefix=f"{scenario_id} planned_actions[{index}]",
            require_source=True,
        )

    executed_actions, gathered_tools_contract = _parse_tool_actions(
        data.get("tool_actions"),
        label=f"{answer_path} tool_actions",
        scenario_id=scenario_id,
        executes_terminal_action=executes_terminal_action,
    )

    must_contain_any = list(
        _string_list(
            response_raw.get("must_contain_any", response_raw.get("any_of_contains")),
            label=f"{answer_path} response_contract.must_contain_any",
        )
    )
    must_contain_all = list(
        _string_list(
            response_raw.get("must_contain_all"),
            label=f"{answer_path} response_contract.must_contain_all",
        )
    )
    must_not_contain = list(
        _string_list(
            response_raw.get("must_not_contain"),
            label=f"{answer_path} response_contract.must_not_contain",
        )
    )
    forbidden_actions = list(
        _string_list(
            response_raw.get("forbidden_actions"),
            label=f"{answer_path} response_contract.forbidden_actions",
        )
    )
    # Validate that forbidden_actions entries reference known action kinds.
    for entry in forbidden_actions:
        if entry not in VALID_ACTION_KINDS:
            msg = f"{answer_path}: forbidden_actions entry {entry!r} is not a valid action kind."
            raise ValueError(msg)

    if not executes_terminal_action and "$ /" not in must_not_contain:
        must_not_contain.append("$ /")

    runs_raw = data.get("runs", 1)
    runs = int(runs_raw) if isinstance(runs_raw, int | str) else 1
    if runs < 1:
        msg = f"{answer_path}: runs must be >= 1."
        raise ValueError(msg)

    history_expected = _action_list(
        history_raw.get("expected"),
        label=f"{answer_path} history.expected",
    )

    return Answer(
        turn=AnswerTurn(
            expected_kind=expected_kind,
        ),
        policy=AnswerPolicy(
            executes_terminal_action=executes_terminal_action,
        ),
        planned_actions=planned_actions,
        executed_actions=executed_actions,
        response_contract={
            "must_contain_any": must_contain_any,
            "must_contain_all": must_contain_all,
            "must_not_contain": must_not_contain,
            "forbidden_actions": forbidden_actions,
        },
        history_expected=history_expected,
        runs=runs,
        gathered_tools_contract=gathered_tools_contract,
    )


def load_scenario_case(scenario_file: Path, *, behavior_class: str) -> ScenarioCase:
    """Load one scenario file into a ScenarioCase."""
    if not scenario_file.is_file():
        msg = f"Missing scenario file: {scenario_file}"
        raise FileNotFoundError(msg)

    scenario = _parse_scenario_yaml(scenario_file, behavior_class=behavior_class)
    if scenario.scenario_dir.stem != scenario.id:
        msg = (
            f"{scenario_file}: file stem {scenario.scenario_dir.stem!r} "
            f"does not match scenario id {scenario.id!r}."
        )
        raise ValueError(msg)

    answer = _parse_answer_yaml(scenario_file, scenario_id=scenario.id)
    return ScenarioCase(scenario=scenario, answer=answer)


def load_all_scenarios() -> list[ScenarioCase]:
    """Discover and load every scenario under scenarios/<behavior_class>/*.yml."""
    if not SCENARIOS_DIR.is_dir():
        return []

    cases: list[ScenarioCase] = []
    seen_ids: set[str] = set()

    for behavior_dir in sorted(SCENARIOS_DIR.iterdir()):
        if not behavior_dir.is_dir():
            continue
        behavior_class = behavior_dir.name
        for scenario_file in sorted(behavior_dir.iterdir()):
            if not scenario_file.is_file() or scenario_file.suffix != ".yml":
                continue
            case = load_scenario_case(scenario_file, behavior_class=behavior_class)
            if case.scenario.id in seen_ids:
                msg = f"Duplicate scenario id {case.scenario.id!r}."
                raise ValueError(msg)
            seen_ids.add(case.scenario.id)
            cases.append(case)

    return cases


def load_scenarios_for_class(behavior_class: str) -> list[ScenarioCase]:
    """Load scenarios for one behavior-class directory."""
    return [case for case in load_all_scenarios() if case.scenario.behavior_class == behavior_class]


def read_shard_config() -> tuple[int, int]:
    """Read TURN_SHARD_TOTAL and TURN_SHARD_INDEX from the environment."""
    total = int(os.getenv("TURN_SHARD_TOTAL", "1"))
    index = int(os.getenv("TURN_SHARD_INDEX", "0"))
    if total < 1:
        msg = "TURN_SHARD_TOTAL must be >= 1"
        raise ValueError(msg)
    if index < 0 or index >= total:
        msg = "TURN_SHARD_INDEX must satisfy 0 <= index < TURN_SHARD_TOTAL"
        raise ValueError(msg)
    return total, index


def iter_scenarios_for_shard(
    cases: list[ScenarioCase],
    *,
    total: int | None = None,
    index: int | None = None,
) -> list[ScenarioCase]:
    """Return the shard subset of cases using stable offset modulo sharding."""
    shard_total, shard_index = (
        (total, index) if total is not None and index is not None else read_shard_config()
    )
    return [case for offset, case in enumerate(cases) if offset % shard_total == shard_index]


_INTENT_COMPLEXITY_WEIGHT: dict[str, float] = {
    "compound": 5.0,
    "complex_shell_prompts": 4.0,
    "remote": 4.0,
    "investigation": 3.0,
    "follow_up": 2.0,
    "local_execution": 1.5,
    "chat_handoff": 1.0,
    "non_actionable": 0.5,
}
_LIVE_INTEGRATION_SENTINEL = "@live"
_SELECT_MODES = frozenset({"sample", "complex"})
_DEFAULT_SELECT_FRACTION = 0.05

# Spec values that explicitly request the FULL suite (opt out of the default
# representative downsample). Accepted by ``--turn-select`` / ``TURN_SELECT``.
FULL_SELECT_SENTINELS = frozenset({"all", "full", "everything", "*"})

# Default gate: the live suite is downsampled everywhere (local AND CI) to a
# small, stratified, representative subset so a run stays fast and cheap. Pick
# this many of the most complex scenarios per behaviour class; classes with
# fewer scenarios contribute all of theirs. Override with ``TURN_SELECT=all``
# to run the complete suite, or ``--turn-select`` for a different subset.
DEFAULT_GATE_PER_CLASS = 2

# Env var capping the majority-vote ``runs`` of each scenario. Defaults to 1 so
# a downsampled run does a single LLM call per test; set ``TURN_MAX_RUNS=0`` (or
# ``all``/``off``) to honour each fixture's ``runs`` (CI keeps full majority
# voting this way).
_TURN_MAX_RUNS_ENV = "TURN_MAX_RUNS"
_DEFAULT_MAX_RUNS_CAP = 1
_UNCAPPED_RUNS_TOKENS = frozenset({"", "0", "all", "none", "off", "uncapped"})


def is_full_selection(spec: str | None) -> bool:
    """True when ``spec`` explicitly requests the full (non-downsampled) suite."""
    return spec is not None and spec.strip().lower() in FULL_SELECT_SENTINELS


def select_representative(
    cases: list[ScenarioCase],
    *,
    per_class: int = DEFAULT_GATE_PER_CLASS,
) -> list[ScenarioCase]:
    """Return a small, deterministic, behaviour-class-stratified subset.

    For each behaviour class, keep the ``per_class`` most complex scenarios
    (ties broken by id), so every intent class stays represented while the total
    stays tiny. Selected cases keep their original ordering for stable test ids.
    This is the default gate applied when no explicit selection is requested.
    """
    if per_class < 1:
        msg = "per_class must be >= 1"
        raise ValueError(msg)
    if not cases:
        return []
    by_class: dict[str, list[ScenarioCase]] = {}
    for case in cases:
        by_class.setdefault(case.scenario.behavior_class, []).append(case)
    chosen_ids: set[str] = set()
    for behavior_class in sorted(by_class):
        ranked = sorted(
            by_class[behavior_class],
            key=lambda case: (scenario_complexity(case), case.scenario.id),
        )
        for case in ranked[max(0, len(ranked) - per_class) :]:
            chosen_ids.add(case.scenario.id)
    return [case for case in cases if case.scenario.id in chosen_ids]


def max_runs_cap(value: str | None = None) -> int | None:
    """Resolve the majority-vote ``runs`` cap from ``TURN_MAX_RUNS``.

    Returns ``None`` when uncapped (honour each fixture's ``runs``) and a
    positive int otherwise. Defaults to ``_DEFAULT_MAX_RUNS_CAP`` (1) when the
    env var is unset, so downsampled runs do a single LLM call per test.
    """
    raw = value if value is not None else os.getenv(_TURN_MAX_RUNS_ENV)
    if raw is None:
        return _DEFAULT_MAX_RUNS_CAP
    text = raw.strip().lower()
    if text in _UNCAPPED_RUNS_TOKENS:
        return None
    parsed = int(text)
    if parsed < 1:
        return None
    return parsed


def effective_runs(answer_runs: int) -> int:
    """Apply the ``TURN_MAX_RUNS`` cap to a fixture's ``runs`` value."""
    base = max(1, answer_runs)
    cap = max_runs_cap()
    return base if cap is None else min(base, cap)


def scenario_complexity(case: ScenarioCase) -> float:
    """Heuristic difficulty/cost score for ranking live turn scenarios.

    Higher means "more worth running when you can only afford a few": multi-step
    plans, majority-vote ``runs`` (the dominant live cost), gather-loop tool
    contracts, real ``@live`` integration calls, and prior-state/long prompts all
    push the score up. Used by ``select_cases`` in ``complex`` mode.
    """
    scenario = case.scenario
    answer = case.answer
    score = _INTENT_COMPLEXITY_WEIGHT.get(scenario.intent_class, 1.0)
    score += 3.0 * len(answer.planned_actions)
    score += 2.0 * len(answer.executed_actions)
    # Majority-vote fixtures issue ``runs`` LLM calls per test in BOTH the
    # planning and oracle suites, so this is the single biggest time multiplier.
    score += 2.0 * max(0, answer.runs - 1)
    contract = answer.gathered_tools_contract
    if contract is not None:
        score += float(
            len(contract.must_call_any)
            + len(contract.must_call_all)
            + len(contract.must_not_call)
            + len(contract.must_return_valid_data)
            + len(contract.must_return_valid_data_any)
        )
    override = scenario.session.resolved_integrations or {}
    score += 2.0 * sum(
        1 for value in override.values() if str(value).strip() == _LIVE_INTEGRATION_SENTINEL
    )
    if scenario.session.has_prior_state:
        score += 1.0
    score += min(len(scenario.input.prompt) / 200.0, 2.0)
    return score


@dataclass(frozen=True)
class SelectionSpec:
    """Parsed ``--turn-select`` / ``TURN_SELECT`` request.

    Exactly one of ``count`` (absolute) or ``fraction`` (0 < f <= 1) is set.
    """

    mode: str
    fraction: float | None = None
    count: int | None = None


def parse_selection_spec(spec: str | None) -> SelectionSpec | None:
    """Parse a selection spec like ``complex:5``, ``sample:0.1``, or ``complex``.

    Returns ``None`` for an empty/unset spec (meaning "run everything"). The
    count component may be an absolute integer (``6``), a percentage (``5%``), or
    a fraction (``0.05``); a bare ``complex``/``sample`` defaults to 5%.
    """
    if spec is None:
        return None
    text = spec.strip().lower()
    if not text:
        return None
    mode, sep, raw = text.partition(":")
    mode = mode.strip()
    if mode not in _SELECT_MODES:
        msg = f"Invalid turn selection mode {mode!r}; expected one of {sorted(_SELECT_MODES)}."
        raise ValueError(msg)
    raw = raw.strip()
    if not sep or not raw:
        return SelectionSpec(mode=mode, fraction=_DEFAULT_SELECT_FRACTION)
    if raw.endswith("%"):
        percent = float(raw[:-1])
        if not 0.0 < percent <= 100.0:
            msg = f"Turn selection percentage must be in (0, 100]; got {raw!r}."
            raise ValueError(msg)
        return SelectionSpec(mode=mode, fraction=percent / 100.0)
    value = float(raw)
    if value <= 0.0:
        msg = f"Turn selection count must be positive; got {raw!r}."
        raise ValueError(msg)
    if value < 1.0:
        return SelectionSpec(mode=mode, fraction=value)
    return SelectionSpec(mode=mode, count=int(value))


def select_cases(
    cases: list[ScenarioCase],
    *,
    spec: str | SelectionSpec | None,
    seed: int = 1337,
) -> list[ScenarioCase]:
    """Return a subset of ``cases`` for fast local iteration.

    ``spec`` selects either the most complex cases (``complex:N``) or a random
    sample (``sample:N``). ``None``/empty returns every case (the default, so CI
    and the full local suite are unchanged). Selected cases keep their original
    ordering for stable, readable test ids.
    """
    parsed = spec if isinstance(spec, SelectionSpec) else parse_selection_spec(spec)
    if parsed is None or not cases:
        return list(cases)
    if parsed.count is not None:
        count = parsed.count
    else:
        fraction = parsed.fraction if parsed.fraction is not None else _DEFAULT_SELECT_FRACTION
        count = math.ceil(len(cases) * fraction)
    count = max(1, min(count, len(cases)))
    if parsed.mode == "complex":
        ranked = sorted(cases, key=lambda case: (scenario_complexity(case), case.scenario.id))
        chosen = ranked[len(ranked) - count :]
    else:
        chosen = random.Random(seed).sample(cases, count)
    order = {case.scenario.id: index for index, case in enumerate(cases)}
    return sorted(chosen, key=lambda case: order[case.scenario.id])


__all__ = [
    "Answer",
    "AnswerPolicy",
    "AnswerTurn",
    "GatheredToolsContract",
    "SCENARIOS_DIR",
    "Scenario",
    "ScenarioCapabilities",
    "ScenarioCase",
    "ScenarioInput",
    "ScenarioSession",
    "SelectionSpec",
    "DEFAULT_GATE_PER_CLASS",
    "FULL_SELECT_SENTINELS",
    "effective_runs",
    "is_full_selection",
    "load_all_scenarios",
    "load_scenario_case",
    "load_scenarios_for_class",
    "iter_scenarios_for_shard",
    "max_runs_cap",
    "parse_selection_spec",
    "read_shard_config",
    "scenario_complexity",
    "select_cases",
    "select_representative",
    "validate_action_shape",
]