501 lines
19 KiB
Python
501 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
"""Select an execution target from a runtime-injected catalog.
|
|
|
|
The selector performs no quota lookup. Every initial candidate is eligible;
|
|
runtime failures such as ``provider-quota`` advance to the next catalog entry.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
|
|
SCHEMA_VERSION = "2.0"
|
|
CATALOG_ENV = "AGENT_TASK_EXECUTION_CATALOG"
|
|
TIMEZONE_NAME = "UTC"
|
|
_FILENAME_RE = re.compile(r"^(PLAN|CODE_REVIEW)-(local|cloud)-G(\d{2})\.md$")
|
|
_MILESTONE_TASK_ID_PATTERN = r"[A-Za-z0-9]+(?:[-_+=][A-Za-z0-9]+){0,3}"
|
|
_MILESTONE_TASK_ID_RE = re.compile(rf"\A{_MILESTONE_TASK_ID_PATTERN}\Z")
|
|
_HEADER_RE = re.compile(
|
|
r"\A<!--\s*task=(?P<task>\S+)\s+plan=(?P<plan>\d+)\s+tag=(?P<tag>\S+)"
|
|
r"(?:\s+milestone-task=(?P<milestone_task>[^,\s]+(?:,[^,\s]+)*))?"
|
|
r"\s*-->[ \t]*(?:\r?\n|\Z)"
|
|
)
|
|
_STAGE_BY_KIND = {"PLAN": "worker", "CODE_REVIEW": "review"}
|
|
_VALID_TRANSITIONS = {"initial", "resume", "failover"}
|
|
_QUALIFIED_FAILOVER_FAILURES = {
|
|
"provider-quota",
|
|
"context-limit",
|
|
"model-unavailable",
|
|
"provider-stream-disconnect",
|
|
"provider-connection",
|
|
}
|
|
|
|
|
|
def _load_policy():
|
|
path = Path(__file__).resolve().parent / "execution_target_policy.py"
|
|
spec = importlib.util.spec_from_file_location("execution_target_policy", path)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError(f"failed to load execution target policy: {path}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
policy = _load_policy()
|
|
|
|
|
|
class SelectorInputError(Exception):
|
|
"""Input contract violation returned as stderr JSON with a non-zero exit."""
|
|
|
|
def __init__(self, code: str, message: str) -> None:
|
|
super().__init__(message)
|
|
self.code = code
|
|
|
|
|
|
def resolve_catalog_path(value: str | Path | None = None) -> Path:
|
|
raw = str(value) if value is not None else os.environ.get(CATALOG_ENV, "")
|
|
if not raw:
|
|
raise SelectorInputError(
|
|
"missing_execution_catalog",
|
|
f"inject the execution catalog with --catalog or {CATALOG_ENV}",
|
|
)
|
|
return Path(raw).expanduser().resolve()
|
|
|
|
|
|
def load_runtime_catalog(value: str | Path | None = None):
|
|
try:
|
|
return policy.load_catalog(resolve_catalog_path(value))
|
|
except SelectorInputError:
|
|
raise
|
|
except (OSError, ValueError) as exc:
|
|
raise SelectorInputError("invalid_execution_catalog", str(exc)) from exc
|
|
|
|
|
|
def _parse_filename(task_file: Path) -> tuple[str, str, int]:
|
|
name = Path(task_file).name
|
|
match = _FILENAME_RE.match(name)
|
|
if match is None:
|
|
raise SelectorInputError(
|
|
"invalid_task_filename",
|
|
f"task file must match (PLAN|CODE_REVIEW)-(local|cloud)-GNN.md: {name!r}",
|
|
)
|
|
kind, lane, grade_str = match.group(1), match.group(2), match.group(3)
|
|
grade = int(grade_str)
|
|
if not 1 <= grade <= 10:
|
|
raise SelectorInputError("invalid_grade", f"grade must be G01..G10: G{grade_str}")
|
|
return kind, lane, grade
|
|
|
|
|
|
def _parse_header(task_file: Path) -> tuple[str, int, str, str | None]:
|
|
try:
|
|
with Path(task_file).open("rb") as handle:
|
|
text = handle.read(1024).decode("utf-8", errors="replace")
|
|
except OSError as exc:
|
|
raise SelectorInputError("task_file_unreadable", str(exc)) from exc
|
|
match = _HEADER_RE.search(text)
|
|
if match is None:
|
|
raise SelectorInputError(
|
|
"malformed_header",
|
|
"first line must contain <!-- task=... plan=N tag=... "
|
|
"[milestone-task=id[,id...]] -->",
|
|
)
|
|
task = match.group("task")
|
|
milestone_task = match.group("milestone_task")
|
|
task_ids = tuple(milestone_task.split(",")) if milestone_task else ()
|
|
invalid_ids = [item for item in task_ids if _MILESTONE_TASK_ID_RE.fullmatch(item) is None]
|
|
if invalid_ids:
|
|
raise SelectorInputError(
|
|
"invalid_milestone_task",
|
|
"milestone-task ids must follow the Milestone item-id grammar: "
|
|
+ ", ".join(invalid_ids),
|
|
)
|
|
if len(task_ids) != len(set(task_ids)):
|
|
raise SelectorInputError(
|
|
"duplicate_milestone_task",
|
|
"milestone-task must contain unique comma-separated Task ids",
|
|
)
|
|
milestone_group = task.split("/", 1)[0].startswith("m-")
|
|
if milestone_group and not milestone_task:
|
|
raise SelectorInputError(
|
|
"missing_milestone_task",
|
|
"m-* task headers require milestone-task=id[,id...]",
|
|
)
|
|
if not milestone_group and milestone_task:
|
|
raise SelectorInputError(
|
|
"unexpected_milestone_task",
|
|
"non-milestone task headers must omit milestone-task",
|
|
)
|
|
return task, int(match.group("plan")), match.group("tag"), milestone_task
|
|
|
|
|
|
def _work_unit_id(header: tuple[str, int, str, str | None]) -> str:
|
|
task, plan, tag, milestone_task = header
|
|
result = f"{task}::plan-{plan}::tag-{tag}"
|
|
if milestone_task:
|
|
result += f"::milestone-task-{milestone_task}"
|
|
return result
|
|
|
|
|
|
def _target_snapshot(target) -> dict:
|
|
return {
|
|
"target_id": target.catalog_id,
|
|
"agent": target.agent,
|
|
"model": target.model,
|
|
"execution_class": target.execution_class,
|
|
"selfcheck_required": target.selfcheck_required,
|
|
}
|
|
|
|
|
|
def _candidate_snapshot(target, rank: int) -> dict:
|
|
return {"candidate_rank": rank, **_target_snapshot(target)}
|
|
|
|
|
|
def _validate_target_snapshot(value: object, prefix: str) -> dict:
|
|
code = "malformed_prior_decision"
|
|
if not isinstance(value, dict):
|
|
raise SelectorInputError(code, f"{prefix} must be an object")
|
|
required = {
|
|
"target_id",
|
|
"agent",
|
|
"model",
|
|
"execution_class",
|
|
"selfcheck_required",
|
|
}
|
|
missing = required - set(value)
|
|
if missing:
|
|
raise SelectorInputError(code, f"{prefix} missing keys: {sorted(missing)}")
|
|
for field in ("target_id", "agent", "model"):
|
|
if not isinstance(value[field], str) or not value[field]:
|
|
raise SelectorInputError(code, f"{prefix}.{field} must be a non-empty string")
|
|
if value["execution_class"] not in policy.VALID_EXECUTION_CLASSES:
|
|
raise SelectorInputError(
|
|
code,
|
|
f"{prefix}.execution_class must be one of {sorted(policy.VALID_EXECUTION_CLASSES)}",
|
|
)
|
|
if not isinstance(value["selfcheck_required"], bool):
|
|
raise SelectorInputError(code, f"{prefix}.selfcheck_required must be a boolean")
|
|
return value
|
|
|
|
|
|
def _validate_prior_decision(value: object) -> dict:
|
|
code = "malformed_prior_decision"
|
|
if not isinstance(value, dict):
|
|
raise SelectorInputError(code, "prior_decision must be an object")
|
|
required = {
|
|
"schema_version",
|
|
"work_unit_id",
|
|
"stage",
|
|
"lane",
|
|
"grade",
|
|
"catalog",
|
|
"selected",
|
|
"candidates",
|
|
"decision",
|
|
"transition",
|
|
}
|
|
missing = required - set(value)
|
|
if missing:
|
|
raise SelectorInputError(code, f"prior_decision missing keys: {sorted(missing)}")
|
|
if value["schema_version"] != SCHEMA_VERSION:
|
|
raise SelectorInputError(code, f"prior_decision.schema_version must be {SCHEMA_VERSION!r}")
|
|
if value["stage"] not in policy.VALID_STAGES or value["lane"] not in policy.VALID_LANES:
|
|
raise SelectorInputError(code, "prior_decision stage/lane is invalid")
|
|
grade = value["grade"]
|
|
if isinstance(grade, bool) or not isinstance(grade, int) or not 1 <= grade <= 10:
|
|
raise SelectorInputError(code, "prior_decision.grade must be G01..G10")
|
|
if not isinstance(value["work_unit_id"], str) or not value["work_unit_id"]:
|
|
raise SelectorInputError(code, "prior_decision.work_unit_id must be a non-empty string")
|
|
catalog = value["catalog"]
|
|
if not isinstance(catalog, dict):
|
|
raise SelectorInputError(code, "prior_decision.catalog must be an object")
|
|
for field in ("schema_version", "revision", "source", "route_id"):
|
|
if not isinstance(catalog.get(field), str) or not catalog[field]:
|
|
raise SelectorInputError(code, f"prior_decision.catalog.{field} must be a non-empty string")
|
|
_validate_target_snapshot(value["selected"], "prior_decision.selected")
|
|
candidates = value["candidates"]
|
|
if not isinstance(candidates, list) or not candidates:
|
|
raise SelectorInputError(code, "prior_decision.candidates must be a non-empty list")
|
|
for index, candidate in enumerate(candidates, 1):
|
|
_validate_target_snapshot(candidate, f"prior_decision.candidates[{index - 1}]")
|
|
if candidate.get("candidate_rank") != index:
|
|
raise SelectorInputError(
|
|
code,
|
|
f"prior_decision.candidates[{index - 1}].candidate_rank must be {index}",
|
|
)
|
|
decision = value["decision"]
|
|
if not isinstance(decision, dict):
|
|
raise SelectorInputError(code, "prior_decision.decision must be an object")
|
|
for field in ("rule_id", "evaluated_at", "timezone", "time_window"):
|
|
if not isinstance(decision.get(field), str) or not decision[field]:
|
|
raise SelectorInputError(code, f"prior_decision.decision.{field} must be a non-empty string")
|
|
if not isinstance(decision.get("policy_priority"), int) or isinstance(
|
|
decision["policy_priority"], bool
|
|
):
|
|
raise SelectorInputError(code, "prior_decision.decision.policy_priority must be an integer")
|
|
if not isinstance(decision.get("reason_codes"), list) or not all(
|
|
isinstance(item, str) and item for item in decision["reason_codes"]
|
|
):
|
|
raise SelectorInputError(code, "prior_decision.decision.reason_codes must be a string list")
|
|
if not isinstance(decision.get("pinned"), bool):
|
|
raise SelectorInputError(code, "prior_decision.decision.pinned must be a boolean")
|
|
transition = value["transition"]
|
|
if not isinstance(transition, dict) or transition.get("trigger") not in _VALID_TRANSITIONS:
|
|
raise SelectorInputError(code, "prior_decision.transition is invalid")
|
|
return value
|
|
|
|
|
|
def _catalog_matches_prior(catalog, prior: dict, decision) -> None:
|
|
code = "catalog_revision_mismatch"
|
|
evidence = prior["catalog"]
|
|
if evidence["schema_version"] != policy.CATALOG_SCHEMA_VERSION:
|
|
raise SelectorInputError(code, "persisted catalog schema is unsupported")
|
|
if evidence["revision"] != catalog.revision:
|
|
raise SelectorInputError(
|
|
code,
|
|
"the injected execution catalog changed after this work unit was selected",
|
|
)
|
|
if evidence["route_id"] != decision.route_id:
|
|
raise SelectorInputError(code, "persisted catalog route does not match the task route")
|
|
|
|
|
|
def _validate_prior_candidate_identity(prior: dict, *, catalog, decision) -> None:
|
|
code = "malformed_prior_decision"
|
|
expected = [_candidate_snapshot(item, rank) for rank, item in enumerate(decision.candidates, 1)]
|
|
if prior["candidates"] != expected:
|
|
raise SelectorInputError(code, "prior_decision candidates do not match the injected catalog route")
|
|
selected = prior["selected"]
|
|
if selected not in [{key: value for key, value in item.items() if key != "candidate_rank"} for item in expected]:
|
|
raise SelectorInputError(code, "prior_decision selected target is not in the injected route")
|
|
|
|
|
|
def _base_decision(
|
|
*,
|
|
catalog,
|
|
route,
|
|
work_unit_id: str,
|
|
stage: str,
|
|
lane: str,
|
|
grade: int,
|
|
evaluated_at: datetime,
|
|
selected,
|
|
pinned: bool,
|
|
previous_target: dict | None,
|
|
trigger: str,
|
|
) -> dict:
|
|
selected_snapshot = _target_snapshot(selected)
|
|
return {
|
|
"schema_version": SCHEMA_VERSION,
|
|
"work_unit_id": work_unit_id,
|
|
"stage": stage,
|
|
"lane": lane,
|
|
"grade": grade,
|
|
"catalog": {
|
|
"schema_version": policy.CATALOG_SCHEMA_VERSION,
|
|
"revision": catalog.revision,
|
|
"source": str(catalog.source),
|
|
"route_id": route.route_id,
|
|
},
|
|
"selected": selected_snapshot,
|
|
"candidates": [
|
|
_candidate_snapshot(item, rank)
|
|
for rank, item in enumerate(route.candidates, 1)
|
|
],
|
|
"decision": {
|
|
"rule_id": route.rule_id,
|
|
"policy_priority": route.policy_priority,
|
|
"reason_codes": list(route.reason_codes),
|
|
"evaluated_at": evaluated_at.astimezone(timezone.utc).isoformat(),
|
|
"timezone": TIMEZONE_NAME,
|
|
"time_window": route.time_window,
|
|
"pinned": pinned,
|
|
},
|
|
"transition": {
|
|
"previous_target": previous_target,
|
|
"next_target": selected_snapshot,
|
|
"trigger": trigger,
|
|
"context_transfer": "logical" if trigger == "failover" else "none",
|
|
},
|
|
}
|
|
|
|
|
|
def select_execution_target_for_route(
|
|
*,
|
|
work_unit_id: str,
|
|
stage: str,
|
|
lane: str,
|
|
grade: int,
|
|
evaluated_at: datetime,
|
|
catalog_path: str | Path | None = None,
|
|
transition: str = "initial",
|
|
prior_decision: dict | None = None,
|
|
failure_class: str | None = None,
|
|
) -> dict:
|
|
if transition not in _VALID_TRANSITIONS:
|
|
raise SelectorInputError("invalid_transition", f"unsupported transition: {transition}")
|
|
if evaluated_at.tzinfo is None or evaluated_at.utcoffset() is None:
|
|
raise SelectorInputError("naive_evaluated_at", "evaluated_at must be timezone-aware")
|
|
catalog = load_runtime_catalog(catalog_path)
|
|
try:
|
|
route = policy.select_policy(
|
|
catalog=catalog,
|
|
stage=stage,
|
|
lane=lane,
|
|
grade=grade,
|
|
evaluated_at=evaluated_at,
|
|
)
|
|
except ValueError as exc:
|
|
raise SelectorInputError("invalid_route", str(exc)) from exc
|
|
if transition == "initial":
|
|
if prior_decision is not None:
|
|
raise SelectorInputError("unexpected_prior_decision", "initial transition must not include prior_decision")
|
|
return _base_decision(
|
|
catalog=catalog,
|
|
route=route,
|
|
work_unit_id=work_unit_id,
|
|
stage=stage,
|
|
lane=lane,
|
|
grade=grade,
|
|
evaluated_at=evaluated_at,
|
|
selected=route.candidates[0],
|
|
pinned=False,
|
|
previous_target=None,
|
|
trigger="initial",
|
|
)
|
|
prior = _validate_prior_decision(prior_decision)
|
|
expected_identity = (work_unit_id, stage, lane, grade)
|
|
actual_identity = (
|
|
prior["work_unit_id"],
|
|
prior["stage"],
|
|
prior["lane"],
|
|
prior["grade"],
|
|
)
|
|
if actual_identity != expected_identity:
|
|
raise SelectorInputError("prior_decision_mismatch", "prior_decision belongs to a different work unit or route")
|
|
_catalog_matches_prior(catalog, prior, route)
|
|
_validate_prior_candidate_identity(prior, catalog=catalog, decision=route)
|
|
selected_id = prior["selected"]["target_id"]
|
|
selected_index = [item.catalog_id for item in route.candidates].index(selected_id)
|
|
if transition == "resume":
|
|
selected = route.candidates[selected_index]
|
|
return _base_decision(
|
|
catalog=catalog,
|
|
route=route,
|
|
work_unit_id=work_unit_id,
|
|
stage=stage,
|
|
lane=lane,
|
|
grade=grade,
|
|
evaluated_at=evaluated_at,
|
|
selected=selected,
|
|
pinned=True,
|
|
previous_target=_target_snapshot(selected),
|
|
trigger="resume",
|
|
)
|
|
if failure_class not in _QUALIFIED_FAILOVER_FAILURES:
|
|
raise SelectorInputError(
|
|
"unqualified_failover",
|
|
f"failure_class does not qualify for target failover: {failure_class!r}",
|
|
)
|
|
next_index = selected_index + 1
|
|
if next_index >= len(route.candidates):
|
|
raise SelectorInputError("no_failover_candidate", "the injected route has no unused next target")
|
|
return _base_decision(
|
|
catalog=catalog,
|
|
route=route,
|
|
work_unit_id=work_unit_id,
|
|
stage=stage,
|
|
lane=lane,
|
|
grade=grade,
|
|
evaluated_at=evaluated_at,
|
|
selected=route.candidates[next_index],
|
|
pinned=False,
|
|
previous_target=dict(prior["selected"]),
|
|
trigger="failover",
|
|
)
|
|
|
|
|
|
def select_execution_target(
|
|
task_file: Path,
|
|
*,
|
|
stage: str | None = None,
|
|
evaluated_at: datetime | None = None,
|
|
catalog_path: str | Path | None = None,
|
|
transition: str = "initial",
|
|
prior_decision: dict | None = None,
|
|
failure_class: str | None = None,
|
|
) -> dict:
|
|
kind, lane, grade = _parse_filename(Path(task_file))
|
|
inferred_stage = _STAGE_BY_KIND[kind]
|
|
if stage is not None and stage != inferred_stage:
|
|
raise SelectorInputError(
|
|
"stage_mismatch",
|
|
f"stage {stage!r} does not match task filename stage {inferred_stage!r}",
|
|
)
|
|
return select_execution_target_for_route(
|
|
work_unit_id=_work_unit_id(_parse_header(Path(task_file))),
|
|
stage=inferred_stage,
|
|
lane=lane,
|
|
grade=grade,
|
|
evaluated_at=evaluated_at or datetime.now(timezone.utc),
|
|
catalog_path=catalog_path,
|
|
transition=transition,
|
|
prior_decision=prior_decision,
|
|
failure_class=failure_class,
|
|
)
|
|
|
|
|
|
def to_json(payload: dict) -> str:
|
|
return json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
|
|
|
|
|
def _load_json_arg(value: str | None):
|
|
if value is None:
|
|
return None
|
|
path = Path(value)
|
|
try:
|
|
return json.loads(path.read_text(encoding="utf-8")) if path.is_file() else json.loads(value)
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
raise SelectorInputError("invalid_json_argument", str(exc)) from exc
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("task_file", type=Path)
|
|
parser.add_argument("--stage", choices=sorted(policy.VALID_STAGES))
|
|
parser.add_argument("--catalog")
|
|
parser.add_argument("--evaluated-at")
|
|
parser.add_argument("--transition", choices=sorted(_VALID_TRANSITIONS), default="initial")
|
|
parser.add_argument("--prior-decision")
|
|
parser.add_argument("--failure-class")
|
|
args = parser.parse_args(argv)
|
|
try:
|
|
evaluated_at = datetime.fromisoformat(args.evaluated_at) if args.evaluated_at else None
|
|
payload = select_execution_target(
|
|
args.task_file,
|
|
stage=args.stage,
|
|
evaluated_at=evaluated_at,
|
|
catalog_path=args.catalog,
|
|
transition=args.transition,
|
|
prior_decision=_load_json_arg(args.prior_decision),
|
|
failure_class=args.failure_class,
|
|
)
|
|
except (SelectorInputError, ValueError) as exc:
|
|
print(
|
|
to_json({"error": {"code": getattr(exc, "code", "invalid_input"), "message": str(exc)}}),
|
|
file=sys.stderr,
|
|
)
|
|
return 2
|
|
print(to_json(payload))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|