ariadne/agent-ops/skills/common/orchestrate-agent-task-loop/scripts/execution_target_policy.py

379 lines
14 KiB
Python

#!/usr/bin/env python3
"""Runtime-injected execution-target catalog and route policy.
This common module intentionally owns no agent or model catalog. A caller
supplies a JSON catalog at runtime; this module validates it and resolves one
ordered route without interpreting provider-specific identities.
"""
from __future__ import annotations
import hashlib
import json
from dataclasses import dataclass
from datetime import datetime, time
from pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
CATALOG_SCHEMA_VERSION = "1.0"
VALID_STAGES = {"worker", "review"}
VALID_LANES = {"local", "cloud"}
VALID_EXECUTION_CLASSES = {"local_model", "cloud_model"}
VALID_OUTPUT_FORMATS = {"jsonl", "text"}
ALLOWED_TEMPLATE_FIELDS = {
"agent",
"attempt_dir",
"model",
"prompt",
"resume_session",
"session_id",
"target_id",
"workspace",
}
class CatalogError(ValueError):
"""The injected execution catalog is missing or malformed."""
@dataclass(frozen=True)
class RouteTarget:
catalog_id: str
agent: str
model: str
execution_class: str
selfcheck_required: bool
runtime: dict[str, Any]
@dataclass(frozen=True)
class ExecutionTargetCatalog:
source: Path
revision: str
targets: dict[str, RouteTarget]
routes: dict[str, dict[str, dict[str, Any]]]
@dataclass(frozen=True)
class PolicyDecision:
route_id: str
rule_id: str
policy_priority: int
reason_codes: tuple[str, ...]
time_window: str
catalog_revision: str
candidates: tuple[RouteTarget, ...]
def _require_string(value: object, label: str) -> str:
if not isinstance(value, str) or not value:
raise CatalogError(f"{label} must be a non-empty string")
return value
def _validate_template(parts: object, label: str) -> tuple[str, ...]:
if not isinstance(parts, list) or not parts:
raise CatalogError(f"{label} must be a non-empty string list")
if not all(isinstance(part, str) and part for part in parts):
raise CatalogError(f"{label} must contain only non-empty strings")
for part in parts:
offset = 0
while True:
start = part.find("{", offset)
if start < 0:
break
end = part.find("}", start + 1)
if end < 0:
raise CatalogError(f"{label} contains an unmatched '{{': {part!r}")
field = part[start + 1 : end]
if field not in ALLOWED_TEMPLATE_FIELDS:
raise CatalogError(
f"{label} uses unsupported template field {field!r}"
)
offset = end + 1
return tuple(parts)
def _validate_runtime(value: object, label: str) -> dict[str, Any]:
if not isinstance(value, dict):
raise CatalogError(f"{label} must be an object")
unknown = set(value) - {
"command",
"resume_command",
"preflight_command",
"environment",
"output_format",
"session_path",
"native_session_monitor",
"auxiliary_logs",
}
if unknown:
raise CatalogError(f"{label} has unsupported keys: {sorted(unknown)}")
command = list(_validate_template(value.get("command"), f"{label}.command"))
if "{" in command[0] or "}" in command[0]:
raise CatalogError(f"{label}.command executable must be a literal path or name")
runtime: dict[str, Any] = {
"command": command,
"output_format": value.get("output_format", "text"),
}
if runtime["output_format"] not in VALID_OUTPUT_FORMATS:
raise CatalogError(
f"{label}.output_format must be one of {sorted(VALID_OUTPUT_FORMATS)}"
)
for field in ("resume_command", "preflight_command"):
if field in value:
template = list(
_validate_template(value[field], f"{label}.{field}")
)
if "{" in template[0] or "}" in template[0]:
raise CatalogError(
f"{label}.{field} executable must be a literal path or name"
)
runtime[field] = template
environment = value.get("environment", {})
if not isinstance(environment, dict) or not all(
isinstance(key, str)
and key
and isinstance(item, str)
for key, item in environment.items()
):
raise CatalogError(f"{label}.environment must be a string map")
runtime["environment"] = dict(environment)
session_path = value.get("session_path")
if session_path is not None:
runtime["session_path"] = _require_string(
session_path, f"{label}.session_path"
)
_validate_template([session_path], f"{label}.session_path")
monitor = value.get("native_session_monitor", False)
if not isinstance(monitor, bool):
raise CatalogError(f"{label}.native_session_monitor must be a boolean")
runtime["native_session_monitor"] = monitor
auxiliary_logs = value.get("auxiliary_logs", [])
if not isinstance(auxiliary_logs, list) or not all(
isinstance(item, str) and item for item in auxiliary_logs
):
raise CatalogError(f"{label}.auxiliary_logs must be a string list")
for index, item in enumerate(auxiliary_logs):
_validate_template([item], f"{label}.auxiliary_logs[{index}]")
runtime["auxiliary_logs"] = list(auxiliary_logs)
return runtime
def _validate_target(target_id: str, value: object) -> RouteTarget:
label = f"targets.{target_id}"
if not isinstance(value, dict):
raise CatalogError(f"{label} must be an object")
unknown = set(value) - {
"agent",
"model",
"execution_class",
"selfcheck_required",
"runtime",
}
if unknown:
raise CatalogError(f"{label} has unsupported keys: {sorted(unknown)}")
execution_class = value.get("execution_class")
if execution_class not in VALID_EXECUTION_CLASSES:
raise CatalogError(
f"{label}.execution_class must be one of "
f"{sorted(VALID_EXECUTION_CLASSES)}"
)
selfcheck_required = value.get("selfcheck_required", False)
if not isinstance(selfcheck_required, bool):
raise CatalogError(f"{label}.selfcheck_required must be a boolean")
return RouteTarget(
catalog_id=target_id,
agent=_require_string(value.get("agent"), f"{label}.agent"),
model=_require_string(value.get("model"), f"{label}.model"),
execution_class=execution_class,
selfcheck_required=selfcheck_required,
runtime=_validate_runtime(value.get("runtime"), f"{label}.runtime"),
)
def _validate_window(value: object, label: str) -> dict[str, Any]:
if not isinstance(value, dict):
raise CatalogError(f"{label} must be an object")
required = {"timezone", "start", "end", "candidates"}
missing = required - set(value)
if missing:
raise CatalogError(f"{label} missing keys: {sorted(missing)}")
timezone_name = _require_string(value["timezone"], f"{label}.timezone")
try:
ZoneInfo(timezone_name)
except ZoneInfoNotFoundError as exc:
raise CatalogError(f"{label}.timezone is unknown: {timezone_name}") from exc
for field in ("start", "end"):
raw = _require_string(value[field], f"{label}.{field}")
try:
time.fromisoformat(raw)
except ValueError as exc:
raise CatalogError(f"{label}.{field} must be HH:MM[:SS]") from exc
return dict(value)
def _validate_route(
value: object,
label: str,
target_ids: set[str],
) -> dict[str, Any]:
if not isinstance(value, dict):
raise CatalogError(f"{label} must be an object")
unknown = set(value) - {
"candidates",
"rule_id",
"policy_priority",
"reason_codes",
"windows",
}
if unknown:
raise CatalogError(f"{label} has unsupported keys: {sorted(unknown)}")
candidates = value.get("candidates")
windows = value.get("windows")
if (candidates is None) == (windows is None):
raise CatalogError(
f"{label} must define exactly one of candidates or windows"
)
normalized = dict(value)
if windows is not None:
if not isinstance(windows, list) or not windows:
raise CatalogError(f"{label}.windows must be a non-empty list")
normalized["windows"] = [
_validate_window(item, f"{label}.windows[{index}]")
for index, item in enumerate(windows)
]
candidate_lists = [item["candidates"] for item in normalized["windows"]]
else:
candidate_lists = [candidates]
for index, candidate_list in enumerate(candidate_lists):
item_label = f"{label}.candidates[{index}]"
if not isinstance(candidate_list, list) or not candidate_list:
raise CatalogError(f"{item_label} must be a non-empty list")
if len(candidate_list) != len(set(candidate_list)):
raise CatalogError(f"{item_label} must not contain duplicates")
unknown_targets = [item for item in candidate_list if item not in target_ids]
if unknown_targets:
raise CatalogError(
f"{item_label} references unknown targets: {unknown_targets}"
)
priority = value.get("policy_priority", 0)
if isinstance(priority, bool) or not isinstance(priority, int):
raise CatalogError(f"{label}.policy_priority must be an integer")
reasons = value.get("reason_codes", [])
if not isinstance(reasons, list) or not all(
isinstance(item, str) and item for item in reasons
):
raise CatalogError(f"{label}.reason_codes must be a string list")
return normalized
def load_catalog(path: str | Path) -> ExecutionTargetCatalog:
source = Path(path).expanduser().resolve()
try:
raw = source.read_bytes()
except OSError as exc:
raise CatalogError(f"execution catalog is unreadable: {source}: {exc}") from exc
try:
value = json.loads(raw)
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
raise CatalogError(f"execution catalog is not valid UTF-8 JSON: {source}") from exc
if not isinstance(value, dict):
raise CatalogError("execution catalog root must be an object")
if set(value) != {"schema_version", "targets", "routes"}:
raise CatalogError(
"execution catalog root must contain exactly schema_version, targets, routes"
)
if value["schema_version"] != CATALOG_SCHEMA_VERSION:
raise CatalogError(
f"execution catalog schema_version must be {CATALOG_SCHEMA_VERSION!r}"
)
raw_targets = value["targets"]
if not isinstance(raw_targets, dict) or not raw_targets:
raise CatalogError("execution catalog targets must be a non-empty object")
targets = {
_require_string(target_id, "target id"): _validate_target(target_id, item)
for target_id, item in raw_targets.items()
}
raw_routes = value["routes"]
if not isinstance(raw_routes, dict) or set(raw_routes) != VALID_STAGES:
raise CatalogError(
f"execution catalog routes must contain exactly {sorted(VALID_STAGES)}"
)
routes: dict[str, dict[str, dict[str, Any]]] = {}
required_route_ids = {
f"{lane}-G{grade:02d}"
for lane in VALID_LANES
for grade in range(1, 11)
}
for stage in sorted(VALID_STAGES):
stage_routes = raw_routes[stage]
if not isinstance(stage_routes, dict) or set(stage_routes) != required_route_ids:
missing = sorted(required_route_ids - set(stage_routes or {}))
extra = sorted(set(stage_routes or {}) - required_route_ids)
raise CatalogError(
f"routes.{stage} must cover local/cloud G01..G10 exactly; "
f"missing={missing}, extra={extra}"
)
routes[stage] = {
route_id: _validate_route(
route, f"routes.{stage}.{route_id}", set(targets)
)
for route_id, route in stage_routes.items()
}
revision = hashlib.sha256(raw).hexdigest()
return ExecutionTargetCatalog(source, revision, targets, routes)
def canonical_target(catalog: ExecutionTargetCatalog, target_id: str) -> RouteTarget | None:
return catalog.targets.get(target_id)
def _window_matches(window: dict[str, Any], evaluated_at: datetime) -> bool:
local_time = evaluated_at.astimezone(ZoneInfo(window["timezone"])).time()
start = time.fromisoformat(window["start"])
end = time.fromisoformat(window["end"])
return start <= local_time < end if start < end else local_time >= start or local_time < end
def select_policy(
*,
catalog: ExecutionTargetCatalog,
stage: str,
lane: str,
grade: int,
evaluated_at: datetime,
) -> PolicyDecision:
if stage not in VALID_STAGES:
raise ValueError(f"unsupported stage: {stage}")
if lane not in VALID_LANES:
raise ValueError(f"unsupported lane: {lane}")
if not 1 <= grade <= 10:
raise ValueError(f"grade must be in G01..G10: {grade}")
if evaluated_at.tzinfo is None or evaluated_at.utcoffset() is None:
raise ValueError("evaluated_at must be timezone-aware")
route_id = f"{lane}-G{grade:02d}"
route = catalog.routes[stage][route_id]
selected_route = route
time_window = "not_applicable"
if "windows" in route:
matches = [item for item in route["windows"] if _window_matches(item, evaluated_at)]
if len(matches) != 1:
raise CatalogError(
f"routes.{stage}.{route_id}.windows must match exactly once; matches={len(matches)}"
)
selected_route = {**route, **matches[0]}
time_window = (
f"{matches[0]['timezone']}:{matches[0]['start']}-{matches[0]['end']}"
)
candidate_ids = selected_route["candidates"]
return PolicyDecision(
route_id=route_id,
rule_id=str(selected_route.get("rule_id") or f"{stage}-{route_id}"),
policy_priority=int(selected_route.get("policy_priority", 0)),
reason_codes=tuple(selected_route.get("reason_codes", [])),
time_window=time_window,
catalog_revision=catalog.revision,
candidates=tuple(catalog.targets[target_id] for target_id in candidate_ids),
)