atlas-iac/services/hermes/scripts/herdr_dispatch.py

250 lines
8.1 KiB
Python

#!/usr/bin/env python3
"""Plan and launch difficulty-aware Codex or Claude Code workers through Herdr."""
from __future__ import annotations
import argparse
import json
import os
import re
import subprocess
from pathlib import Path
from typing import Any
ALLOWED_EFFORTS = ("low", "medium", "high", "xhigh")
ROUTING_PATH = Path("/opt/data/workspace/coordinator/model-routing.json")
HERDR_BIN = Path("/opt/data/tools/bin/herdr")
CODEX_AUTH = Path("/opt/data/home/.codex/auth.json")
PROMPT_READY_MARKERS = {
"codex": "OpenAI Codex",
"claude": "Claude Code",
}
def _load_routes(path: Path) -> dict[str, Any]:
"""Load the non-secret route status generated by the model steward."""
try:
value = json.loads(path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as error:
raise RuntimeError(f"routing status unavailable: {path}") from error
if not isinstance(value, dict) or not isinstance(value.get("routes"), dict):
raise RuntimeError(f"routing status is malformed: {path}")
return value
def _split_route(route: str) -> tuple[str, str]:
"""Split a provider/model route without corrupting model punctuation."""
provider, separator, model = route.partition("/")
if not separator or not provider or not model:
raise RuntimeError(f"invalid route: {route}")
return provider, model
def select_plan(
status: dict[str, Any], shape: str, effort: str, provider: str | None = None
) -> dict[str, Any]:
"""Select a capped route and retain the complete capacity fallback chain."""
if effort not in ALLOWED_EFFORTS:
raise ValueError(f"effort must be one of: {', '.join(ALLOWED_EFFORTS)}")
default_provider = "codex" if shape == "implementation" else "claude"
selected_provider = provider or default_provider
if selected_provider not in {"codex", "claude"}:
raise ValueError("provider must be codex or claude")
profile = f"{selected_provider}-{effort}"
chain = status["routes"].get(profile)
if not isinstance(chain, list) or not chain:
raise RuntimeError(f"route profile unavailable: {profile}")
primary_provider, model = _split_route(str(chain[0]))
return {
"shape": shape,
"effort": effort,
"profile": profile,
"worker": selected_provider,
"provider": primary_provider,
"model": model,
"fallback_chain": [str(route) for route in chain[1:]],
}
def _slug(value: str, limit: int = 32) -> str:
"""Return a Herdr-safe stable label."""
cleaned = re.sub(r"[^a-z0-9_-]+", "-", value.lower()).strip("-")
if not cleaned or not cleaned[0].isalpha():
cleaned = f"task-{cleaned}"
return cleaned[:limit].rstrip("-")
def _run(command: list[str], env: dict[str, str]) -> dict[str, Any]:
"""Run one Herdr JSON command and surface a concise error."""
completed = subprocess.run(
command,
env=env,
text=True,
capture_output=True,
timeout=320,
check=False,
)
if completed.returncode != 0:
detail = completed.stderr.strip() or completed.stdout.strip()
raise RuntimeError(detail or f"command failed with {completed.returncode}")
try:
value = json.loads(completed.stdout)
except json.JSONDecodeError as error:
raise RuntimeError("Herdr returned non-JSON output") from error
if not isinstance(value, dict):
raise RuntimeError("Herdr returned an unexpected response")
return value
def launch_worker(
plan: dict[str, Any], project: Path, task: str, prompt: str | None
) -> dict[str, Any]:
"""Create an isolated Herdr workspace and launch the selected worker."""
if not HERDR_BIN.is_file():
raise RuntimeError("Herdr is not installed in the agent tools volume")
if not project.is_dir():
raise RuntimeError(f"project workspace does not exist: {project}")
if plan["worker"] == "codex" and not CODEX_AUTH.is_file():
raise RuntimeError(
"Codex CLI needs its one-time device login. Ask Hermes to run "
"`codex login --device-auth`, complete the code, then retry."
)
env = os.environ.copy()
env.update(
{
"HOME": "/opt/data/home",
"CODEX_HOME": "/opt/data/home/.codex",
"CLAUDE_CONFIG_DIR": "/opt/data/home/.claude",
"HERDR_CONFIG_PATH": "/opt/data/home/.config/herdr/config.toml",
"HERDR_SOCKET_PATH": "/opt/data/herdr/herdr.sock",
"PATH": "/opt/data/tools/bin:" + env.get("PATH", ""),
}
)
project_label = _slug(project.name, 12)
task_label = _slug(task, 14)
label = _slug(f"{project_label}-{plan['worker']}-{task_label}")
created = _run(
[
str(HERDR_BIN),
"workspace",
"create",
"--cwd",
str(project),
"--label",
label,
"--no-focus",
],
env,
)
try:
pane = str(created["result"]["root_pane"]["pane_id"])
except (KeyError, TypeError) as error:
raise RuntimeError("Herdr workspace response omitted the root pane") from error
agent_name = label
command = [
str(HERDR_BIN),
"agent",
"start",
agent_name,
"--kind",
plan["worker"],
"--pane",
pane,
"--timeout",
"120000",
"--",
]
if plan["worker"] == "codex":
command.extend(
[
"-m",
plan["model"],
"-c",
f'model_reasoning_effort="{plan["effort"]}"',
"--approve-for-me",
"--sandbox",
"workspace-write",
]
)
else:
command.extend(
[
"--model",
plan["model"],
"--effort",
plan["effort"],
"--permission-mode",
"auto",
]
)
started = _run(command, env)
result = {**plan, "agent": agent_name, "pane": pane, "started": started}
if prompt:
# Herdr can detect the process before the full-screen prompt has finished
# drawing. Wait for a pinned CLI marker so the submitted Enter is not lost.
_run(
[
str(HERDR_BIN),
"pane",
"wait-output",
pane,
"--match",
PROMPT_READY_MARKERS[plan["worker"]],
"--source",
"recent",
"--lines",
"120",
"--timeout",
"120000",
],
env,
)
result["prompted"] = _run(
[
str(HERDR_BIN),
"agent",
"prompt",
agent_name,
prompt,
"--wait",
"--until",
"working",
"--until",
"done",
"--until",
"blocked",
"--timeout",
"15000",
],
env,
)
return result
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--shape", choices=("implementation", "architecture", "review"), required=True)
parser.add_argument("--effort", choices=ALLOWED_EFFORTS, required=True)
parser.add_argument("--provider", choices=("codex", "claude"))
parser.add_argument("--routes", type=Path, default=ROUTING_PATH)
parser.add_argument("--start", action="store_true")
parser.add_argument("--project", type=Path)
parser.add_argument("--task", default="objective")
parser.add_argument("--prompt")
args = parser.parse_args()
plan = select_plan(_load_routes(args.routes), args.shape, args.effort, args.provider)
if args.start:
if args.project is None:
parser.error("--project is required with --start")
plan = launch_worker(plan, args.project.resolve(), args.task, args.prompt)
print(json.dumps(plan, indent=2, sort_keys=True))
return 0
if __name__ == "__main__":
raise SystemExit(main())