335 lines
11 KiB
Python
335 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Run one fresh target from a runtime-injected execution catalog."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from datetime import datetime, timezone
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from typing import Any, Iterable
|
|
import uuid
|
|
|
|
|
|
CATALOG_ENV = "AGENT_TASK_EXECUTION_CATALOG"
|
|
LABEL_PATTERN = re.compile(r"^[A-Za-z0-9._-]+$")
|
|
PROBE_EXPECTED = "MILESTONE_AGENT_READY"
|
|
|
|
|
|
class AgentRunError(RuntimeError):
|
|
"""One-shot runner contract error."""
|
|
|
|
|
|
def load_policy_module():
|
|
path = (
|
|
Path(__file__).resolve().parents[2]
|
|
/ "orchestrate-agent-task-loop"
|
|
/ "scripts"
|
|
/ "execution_target_policy.py"
|
|
)
|
|
spec = importlib.util.spec_from_file_location("epic_execution_target_policy", path)
|
|
if spec is None or spec.loader is None:
|
|
raise AgentRunError(f"execution catalog policy not found: {path}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
def emit(event: str, **payload: Any) -> None:
|
|
print(json.dumps({"event": event, **payload}, ensure_ascii=False, sort_keys=True), flush=True)
|
|
|
|
|
|
def atomic_json(path: Path, value: dict[str, Any]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.with_suffix(path.suffix + f".tmp.{os.getpid()}")
|
|
temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
os.replace(temporary, path)
|
|
|
|
|
|
def process_start_token(pid: int) -> str | None:
|
|
stat = Path(f"/proc/{pid}/stat")
|
|
try:
|
|
remainder = stat.read_text(encoding="utf-8").rsplit(")", 1)[1].split()
|
|
return f"proc:{remainder[19]}"
|
|
except (OSError, IndexError):
|
|
return None
|
|
|
|
|
|
def workspace_root(raw: str) -> Path:
|
|
workspace = Path(raw).expanduser().resolve()
|
|
if not workspace.is_dir():
|
|
raise AgentRunError(f"workspace directory not found: {workspace}")
|
|
result = subprocess.run(
|
|
["git", "rev-parse", "--show-toplevel"],
|
|
cwd=workspace,
|
|
text=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
check=False,
|
|
)
|
|
if result.returncode != 0 or Path(result.stdout.strip()).resolve() != workspace:
|
|
raise AgentRunError(f"workspace must be a git repository root: {workspace}")
|
|
return workspace
|
|
|
|
|
|
def state_root(workspace: Path) -> Path:
|
|
result = subprocess.run(
|
|
["git", "rev-parse", "--git-common-dir"],
|
|
cwd=workspace,
|
|
text=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
check=False,
|
|
)
|
|
if result.returncode == 0:
|
|
raw = Path(result.stdout.strip())
|
|
common = (workspace / raw).resolve() if not raw.is_absolute() else raw.resolve()
|
|
if os.access(common, os.W_OK):
|
|
return common / "epic-work-preparation"
|
|
fallback = Path(os.environ.get("XDG_STATE_HOME", str(Path.home() / ".local" / "state")))
|
|
identity = hashlib.sha256(str(workspace).encode()).hexdigest()[:16]
|
|
return fallback / "epic-work-preparation" / identity
|
|
|
|
|
|
def result_file(workspace: Path, raw: str | None) -> Path | None:
|
|
if raw is None:
|
|
return None
|
|
path = Path(raw).expanduser().resolve()
|
|
root = state_root(workspace).resolve()
|
|
try:
|
|
path.relative_to(root)
|
|
except ValueError as exc:
|
|
raise AgentRunError(f"--result-file must be inside runner state root: {root}") from exc
|
|
return path
|
|
|
|
|
|
def persist(locator: Path, result: Path | None, record: dict[str, Any]) -> None:
|
|
atomic_json(locator, record)
|
|
if result is not None:
|
|
atomic_json(result, record)
|
|
|
|
|
|
def prompt_text(args: argparse.Namespace) -> str:
|
|
if args.probe:
|
|
return (
|
|
"Reply only with the result of joining MILESTONE, _AGENT, and _READY. "
|
|
"Do not inspect files, call tools, or modify the workspace."
|
|
)
|
|
if args.prompt is not None:
|
|
return args.prompt
|
|
if args.prompt_file is None:
|
|
raise AgentRunError("--prompt or --prompt-file is required")
|
|
path = Path(args.prompt_file).expanduser().resolve()
|
|
if not path.is_file():
|
|
raise AgentRunError(f"prompt file not found: {path}")
|
|
return path.read_text(encoding="utf-8")
|
|
|
|
|
|
def resolve_target(catalog_path: str, target_id: str):
|
|
policy = load_policy_module()
|
|
try:
|
|
catalog = policy.load_catalog(catalog_path)
|
|
except (OSError, ValueError) as exc:
|
|
raise AgentRunError(f"invalid execution catalog: {exc}") from exc
|
|
target = policy.canonical_target(catalog, target_id)
|
|
if target is None:
|
|
raise AgentRunError(f"execution catalog target not found: {target_id}")
|
|
return catalog, target
|
|
|
|
|
|
def template_values(*, target, prompt: str, workspace: Path, session_id: str, attempt_dir: Path) -> dict[str, str]:
|
|
return {
|
|
"agent": target.agent,
|
|
"attempt_dir": str(attempt_dir),
|
|
"model": target.model,
|
|
"prompt": prompt,
|
|
"resume_session": "",
|
|
"session_id": session_id,
|
|
"target_id": target.catalog_id,
|
|
"workspace": str(workspace),
|
|
}
|
|
|
|
|
|
def build_command(*, target, prompt: str, workspace: Path, session_id: str, attempt_dir: Path) -> list[str]:
|
|
values = template_values(
|
|
target=target,
|
|
prompt=prompt,
|
|
workspace=workspace,
|
|
session_id=session_id,
|
|
attempt_dir=attempt_dir,
|
|
)
|
|
return [str(item).format_map(values) for item in target.runtime["command"]]
|
|
|
|
|
|
def sanitized_command(command: list[str], prompt: str) -> list[str]:
|
|
return ["<prompt>" if value == prompt else value for value in command]
|
|
|
|
|
|
def parser() -> argparse.ArgumentParser:
|
|
value = argparse.ArgumentParser(description=__doc__)
|
|
value.add_argument("--execution-catalog", default=os.environ.get(CATALOG_ENV))
|
|
value.add_argument("--target-id", required=True)
|
|
value.add_argument("--workspace", required=True)
|
|
prompt_group = value.add_mutually_exclusive_group()
|
|
prompt_group.add_argument("--prompt")
|
|
prompt_group.add_argument("--prompt-file")
|
|
value.add_argument("--label", default="one-shot")
|
|
value.add_argument("--probe", action="store_true")
|
|
value.add_argument("--result-file")
|
|
return value
|
|
|
|
|
|
def execute(args: argparse.Namespace) -> int:
|
|
workspace = workspace_root(args.workspace)
|
|
if not args.execution_catalog:
|
|
raise AgentRunError(
|
|
f"--execution-catalog or {CATALOG_ENV} is required"
|
|
)
|
|
catalog, target = resolve_target(args.execution_catalog, args.target_id)
|
|
if not LABEL_PATTERN.fullmatch(args.label):
|
|
raise AgentRunError("--label may contain only letters, digits, dot, underscore, and hyphen")
|
|
prompt = prompt_text(args)
|
|
result = result_file(workspace, args.result_file)
|
|
executable = target.runtime["command"][0]
|
|
if shutil.which(executable) is None:
|
|
raise AgentRunError(
|
|
f"target command not found: target_id={target.catalog_id} command={executable}"
|
|
)
|
|
|
|
execution_id = f"{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')}-{uuid.uuid4().hex[:12]}"
|
|
root = state_root(workspace)
|
|
attempt_dir = root / "runs" / f"{args.label}-{execution_id}"
|
|
attempt_dir.mkdir(parents=True, exist_ok=False)
|
|
stream = attempt_dir / "stream.log"
|
|
locator = attempt_dir / "locator.json"
|
|
session_id = str(uuid.uuid4())
|
|
values = template_values(
|
|
target=target,
|
|
prompt=prompt,
|
|
workspace=workspace,
|
|
session_id=session_id,
|
|
attempt_dir=attempt_dir,
|
|
)
|
|
command = build_command(
|
|
target=target,
|
|
prompt=prompt,
|
|
workspace=workspace,
|
|
session_id=session_id,
|
|
attempt_dir=attempt_dir,
|
|
)
|
|
environment = {
|
|
str(key): str(item).format_map(values)
|
|
for key, item in target.runtime.get("environment", {}).items()
|
|
}
|
|
record: dict[str, Any] = {
|
|
"execution_id": execution_id,
|
|
"label": args.label,
|
|
"workspace": str(workspace),
|
|
"catalog": {
|
|
"source": str(catalog.source),
|
|
"revision": catalog.revision,
|
|
"schema_version": "1.0",
|
|
},
|
|
"target_id": target.catalog_id,
|
|
"agent": target.agent,
|
|
"model": target.model,
|
|
"command": sanitized_command(command, prompt),
|
|
"prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(),
|
|
"session_id": session_id,
|
|
"stream_log": str(stream),
|
|
"locator": str(locator),
|
|
"started_at": now(),
|
|
"status": "starting",
|
|
}
|
|
persist(locator, result, record)
|
|
emit(
|
|
"AGENT_STARTED",
|
|
agent=target.agent,
|
|
execution_id=execution_id,
|
|
label=args.label,
|
|
locator=str(locator),
|
|
model=target.model,
|
|
target_id=target.catalog_id,
|
|
)
|
|
with stream.open("wb") as output:
|
|
try:
|
|
process = subprocess.Popen(
|
|
command,
|
|
cwd=workspace,
|
|
env={
|
|
**os.environ,
|
|
"MILESTONE_PREPARATION_EXECUTION_ID": execution_id,
|
|
**environment,
|
|
},
|
|
stdout=output,
|
|
stderr=subprocess.STDOUT,
|
|
start_new_session=True,
|
|
)
|
|
except OSError as exc:
|
|
record.update(status="failed", finished_at=now(), exit_code=127, error=str(exc))
|
|
persist(locator, result, record)
|
|
emit("AGENT_FINISHED", execution_id=execution_id, label=args.label, result="failed", exit_code=127)
|
|
return 127
|
|
record.update(
|
|
status="running",
|
|
agent_pid=process.pid,
|
|
agent_process_start_token=process_start_token(process.pid),
|
|
)
|
|
persist(locator, result, record)
|
|
try:
|
|
exit_code = process.wait()
|
|
except KeyboardInterrupt:
|
|
record.update(status="tracking", interrupted_at=now(), agent_pid=process.pid)
|
|
persist(locator, result, record)
|
|
emit(
|
|
"AGENT_TRACKING",
|
|
execution_id=execution_id,
|
|
label=args.label,
|
|
locator=str(locator),
|
|
pid=process.pid,
|
|
)
|
|
return 3
|
|
|
|
size = stream.stat().st_size
|
|
status = "succeeded" if exit_code == 0 and size > 0 else "failed"
|
|
if args.probe and status == "succeeded":
|
|
content = stream.read_text(encoding="utf-8", errors="replace")
|
|
if PROBE_EXPECTED not in content.upper():
|
|
status = "failed"
|
|
exit_code = 2
|
|
record.update(status=status, finished_at=now(), exit_code=exit_code, output_bytes=size)
|
|
persist(locator, result, record)
|
|
emit(
|
|
"AGENT_FINISHED",
|
|
execution_id=execution_id,
|
|
exit_code=exit_code,
|
|
label=args.label,
|
|
locator=str(locator),
|
|
result=status,
|
|
)
|
|
return exit_code if status == "succeeded" else (exit_code or 2)
|
|
|
|
|
|
def main(argv: Iterable[str] | None = None) -> int:
|
|
args = parser().parse_args(argv)
|
|
try:
|
|
return execute(args)
|
|
except (AgentRunError, OSError, ValueError) as exc:
|
|
emit("AGENT_FINISHED", label=getattr(args, "label", "one-shot"), result="failed", reason=str(exc))
|
|
return 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|