358 lines
16 KiB
Python
358 lines
16 KiB
Python
"""argparse and the commands; each launch is record -> one `claude -p` call -> record."""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import string
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from grm import call, launch, store, stream
|
|
|
|
SKILL_DIR = Path(__file__).resolve().parent.parent
|
|
NAME_RE = re.compile(r"^[A-Za-z0-9._-]+$")
|
|
EXCERPT_CHARS = 500
|
|
LAUNCHES = ("train", "spawn", "rebase")
|
|
|
|
|
|
class Refusal(Exception):
|
|
"""Refused before or instead of a launch: exit 2, payload printed as the JSON object."""
|
|
|
|
def __init__(self, error: str, **fields):
|
|
super().__init__(error)
|
|
self.payload = {"error": error, **fields}
|
|
|
|
|
|
class _Parser(argparse.ArgumentParser):
|
|
def error(self, message):
|
|
raise Refusal("usage", detail=message)
|
|
|
|
|
|
def _emit(obj: dict) -> None:
|
|
sys.stdout.write(json.dumps(obj) + "\n")
|
|
sys.stdout.flush()
|
|
|
|
|
|
def render_template(name: str, values: dict) -> str:
|
|
with open(SKILL_DIR / "templates" / name, "r", encoding="utf-8") as f:
|
|
return string.Template(f.read()).safe_substitute(values)
|
|
|
|
|
|
def git(args, cwd=None, timeout=60) -> subprocess.CompletedProcess:
|
|
"""Every helper git call: `git --no-optional-locks …`, output captured."""
|
|
return subprocess.run(["git", "--no-optional-locks", *args], cwd=cwd, stdin=subprocess.DEVNULL,
|
|
capture_output=True, timeout=timeout)
|
|
|
|
|
|
def _git_out(args, cwd) -> str | None:
|
|
try:
|
|
r = git(args, cwd=cwd)
|
|
except (OSError, subprocess.SubprocessError):
|
|
return None
|
|
out = r.stdout.decode("utf-8", "replace").strip()
|
|
return out if r.returncode == 0 and out else None
|
|
|
|
|
|
def git_head(cwd) -> str | None:
|
|
return _git_out(["rev-parse", "HEAD"], cwd)
|
|
|
|
|
|
def delta_header(cwd, base_commit: str | None) -> str:
|
|
"""Commits since the base's commit plus the working tree's shortstat against it."""
|
|
if not base_commit:
|
|
return ("No base commit is recorded (the base's cwd was not a git repository, or git failed "
|
|
"at training). Treat everything on disk as possibly changed since training.")
|
|
parts = []
|
|
try:
|
|
lg = git(["log", "--oneline", "%s..HEAD" % base_commit], cwd=cwd)
|
|
st = git(["diff", "--shortstat", base_commit], cwd=cwd)
|
|
log_txt = lg.stdout.decode("utf-8", "replace").strip() if lg.returncode == 0 else \
|
|
"(git log failed: %s)" % lg.stderr.decode("utf-8", "replace").strip()
|
|
st_txt = st.stdout.decode("utf-8", "replace").strip() if st.returncode == 0 else \
|
|
"(git diff failed: %s)" % st.stderr.decode("utf-8", "replace").strip()
|
|
except (OSError, subprocess.SubprocessError) as e:
|
|
log_txt = st_txt = "(git failed: %s)" % e
|
|
parts.append("Base commit: `%s`." % base_commit)
|
|
parts.append("Commits since then (`git log --oneline %s..HEAD`):\n\n```\n%s\n```"
|
|
% (base_commit[:12], log_txt or "(none)"))
|
|
parts.append("Working tree against the base commit (`git diff --shortstat %s`):\n\n```\n%s\n```"
|
|
% (base_commit[:12], st_txt or "(no changes)"))
|
|
return "\n\n".join(parts)
|
|
|
|
|
|
def _name(kind: str, value: str) -> str:
|
|
if not NAME_RE.match(value or ""):
|
|
raise Refusal("bad %s name" % kind, **{kind: value})
|
|
return value
|
|
|
|
|
|
def _read_brief(path: str) -> str:
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return f.read()
|
|
except OSError as e:
|
|
raise Refusal("cannot read brief", brief=path, detail=str(e))
|
|
|
|
|
|
def _load_run(run: str) -> tuple[Path, dict]:
|
|
"""State dir and run.json of a launch; the stored binary must still exist (no job file yet)."""
|
|
state = store.state_dir(run)
|
|
r = store.read_json(state / "run.json")
|
|
if r is None:
|
|
raise Refusal("no run.json for this run; call init first", state=str(state))
|
|
if not os.path.isfile(r["argv"][0]):
|
|
raise Refusal("claude binary missing", binary=r["argv"][0])
|
|
return state, r
|
|
|
|
|
|
def _load_base(state: Path, b: str) -> dict:
|
|
base = store.read_json(state / "bases" / (b + ".json"))
|
|
if base is None:
|
|
raise Refusal("unknown base", base=b, bases=sorted(p.stem for p in (state / "bases").glob("*.json")))
|
|
return base
|
|
|
|
|
|
def _warmth(lines: list, base: dict, now: float) -> dict:
|
|
"""B6: from the ledger line of this base whose call read the base prefix last."""
|
|
hits = [ln for ln in lines if ln.get("base") == base["base"] and stream.parse_iso(ln.get("hit_at"))]
|
|
if not hits:
|
|
return {**stream.warmth(None, None, now), "ttl_s": base.get("ttl_s")}
|
|
last = max(hits, key=lambda ln: stream.parse_iso(ln["hit_at"]))
|
|
ttl = stream.infer_ttl(last.get("tokens"), base.get("ttl_s")) or stream.TTL_1H
|
|
return {**stream.warmth(stream.parse_iso(last["hit_at"]), ttl, now), "ttl_s": ttl}
|
|
|
|
|
|
def _execute(run: str, state: Path, r: dict, *, job, kind, base, sid, argv, prompt,
|
|
before=(None, 0.0), base_ctx=None):
|
|
"""One call: job file `running` -> run_claude -> job file final + one ledger line.
|
|
Returns (job record, StreamState|None, usage|None, ledger line|None)."""
|
|
fp = r["fingerprint"]
|
|
jpath = state / "jobs" / (job + ".json")
|
|
raw = state / "raw" / (job + ".jsonl")
|
|
store.atomic_write_bytes(state / "briefs" / (job + ".md"), prompt.encode("utf-8"))
|
|
t0 = time.time()
|
|
rec = {"job": job, "kind": kind, "base": base, "session_id": sid, "status": "running",
|
|
"reason": None, "started": stream.iso(t0), "ended": None, "child_pid": None,
|
|
"raw": str(raw), "note": None, "cache": None, "result_excerpt": None}
|
|
store.write_json(jpath, rec)
|
|
|
|
def on_start(pid):
|
|
rec["child_pid"] = pid
|
|
store.write_json(jpath, rec)
|
|
try:
|
|
ss, code, stopped = call.run_claude(argv, prompt, launch.scrubbed_env(fp["binary"]),
|
|
fp["cwd"], raw, on_start)
|
|
except OSError as e:
|
|
rec.update(status="error", reason="launch_error", ended=stream.iso(),
|
|
result_excerpt=str(e)[:EXCERPT_CHARS])
|
|
store.write_json(jpath, rec)
|
|
return rec, None, None, None
|
|
wall = time.time() - t0
|
|
status, reason = stream.classify(ss, code, stopped)
|
|
u = stream.usage_from_result(ss.result, *before) if ss.result else None
|
|
cache = stream.cache_verdict(ss.first_usage, base_ctx) if kind != "train" else None
|
|
hit_t = None
|
|
if kind in ("train", "rebase") and status == "ok":
|
|
hit_t = ss.last_assistant_at
|
|
elif kind == "spawn" and cache in ("hit", "partial"):
|
|
hit_t = ss.first_at
|
|
note = None
|
|
if kind == "spawn":
|
|
p = state / "notes" / (job + ".md")
|
|
if p.is_file():
|
|
note = str(p)
|
|
text = (ss.result or {}).get("result")
|
|
rec.update(session_id=ss.session_id or sid, status=status, reason=reason, ended=stream.iso(),
|
|
cache=cache, note=note,
|
|
result_excerpt=text[:EXCERPT_CHARS] if isinstance(text, str) else None)
|
|
store.write_json(jpath, rec)
|
|
models = u["models_used"] if u else []
|
|
line = {"ts": rec["ended"], "run": r["run"], "job": job, "kind": kind, "base": base,
|
|
"session_id": rec["session_id"], "model": models[0] if models else None,
|
|
"models_used": models,
|
|
"tokens": u["tokens"] if u else {k: None for k in stream.TOKEN_KEYS},
|
|
"cost_delta_usd": u["cost_delta"] if u else None,
|
|
"cost_cumulative_usd": u["total"] if u else None, "cache": cache, "status": status,
|
|
"wall_s": round(wall, 1), "hit_at": stream.iso(hit_t) if hit_t else None}
|
|
store.ledger_append(run, line)
|
|
return rec, ss, u, line
|
|
|
|
|
|
def _cost_basis(ss) -> float | None:
|
|
"""The base's cost basis stays unrounded: deltas against a rounded basis drift by 1e-6."""
|
|
total = (ss.result or {}).get("total_cost_usd")
|
|
return None if total is None else float(total)
|
|
|
|
|
|
def _out(rec: dict, line: dict | None) -> dict:
|
|
out = dict(rec)
|
|
if line:
|
|
out.update(model=line["model"], models_used=line["models_used"], tokens=line["tokens"],
|
|
cost=line["cost_delta_usd"], cost_cumulative_usd=line["cost_cumulative_usd"])
|
|
return out
|
|
|
|
|
|
def cmd_init(args) -> dict:
|
|
"""B1: validate, then write the frozen preamble copy, the empty MCP config and run.json once."""
|
|
state = store.state_dir(args.run)
|
|
if (state / "run.json").exists():
|
|
raise Refusal("run.json exists; init never overwrites a run", state=str(state))
|
|
binary = launch.resolve_binary(args.binary)
|
|
if not binary or not os.path.isfile(binary):
|
|
raise Refusal("no claude binary (no --binary, no $CLAUDE_CODE_EXECPATH, none on PATH)",
|
|
binary=binary)
|
|
cwd = os.path.abspath(args.cwd) if args.cwd else _git_out(["rev-parse", "--show-toplevel"], os.getcwd())
|
|
if not cwd or not os.path.isdir(cwd):
|
|
raise Refusal("no usable cwd (pass --cwd, or run inside a git repository)", cwd=cwd)
|
|
if not launch.MODEL_RE.match(args.model): # before anything is written
|
|
raise Refusal("model must be a full id (aliases are refused)", model=args.model)
|
|
state = store.init_state(args.run)
|
|
pre = state / "preamble.md"
|
|
store.atomic_write_bytes(pre, (SKILL_DIR / "preamble.md").read_bytes())
|
|
mcp = state / "mcp-empty.json"
|
|
store.write_json(mcp, {"mcpServers": {}})
|
|
fpd = launch.build_fingerprint(binary=binary, model=args.model, effort=args.effort,
|
|
cwd=os.path.abspath(cwd), agent=args.agent,
|
|
preamble_path=str(pre), mcp_config=str(mcp))
|
|
run = {"run": os.path.abspath(args.run), "created": stream.iso(), **fpd}
|
|
store.write_json(state / "run.json", run)
|
|
return {"state": str(state), **run}
|
|
|
|
|
|
def cmd_train(args) -> dict:
|
|
"""B1: one blocking call with `--session-id`; the base file only on `ok`."""
|
|
b = _name("base", args.base)
|
|
state, r = _load_run(args.run)
|
|
bpath = state / "bases" / (b + ".json")
|
|
if bpath.exists():
|
|
raise Refusal("base exists", base=b)
|
|
brief = _read_brief(args.brief)
|
|
sid, job = str(uuid.uuid4()), b + "-train"
|
|
prompt = render_template("train.md", {
|
|
"base": b, "run": r["run"], "state_dir": str(state), "notes_dir": str(state / "notes"),
|
|
"logs_dir": str(state / "logs"), "brief": brief})
|
|
rec, ss, u, line = _execute(args.run, state, r, job=job, kind="train", base=b, sid=sid,
|
|
argv=launch.train_argv(r, sid), prompt=prompt)
|
|
if rec["status"] == "ok":
|
|
store.write_json(bpath, {
|
|
"base": b, "session_id": rec["session_id"], "train_job": job,
|
|
"commit": git_head(r["fingerprint"]["cwd"]), "trained_at": rec["ended"],
|
|
"context_tokens": stream.context_tokens(ss.last_usage), "model_usage": u["model_usage"],
|
|
"cost_cumulative_usd": _cost_basis(ss), "ttl_s": stream.infer_ttl(u["tokens"], None),
|
|
"rebases": []})
|
|
return _out(rec, line)
|
|
|
|
|
|
def cmd_spawn(args) -> dict:
|
|
"""B2, B3, B5: one blocking fork of the base; never writes the base file."""
|
|
b, j = _name("base", args.base), _name("job", args.job)
|
|
state, r = _load_run(args.run)
|
|
base = _load_base(state, b)
|
|
if (state / "jobs" / (j + ".json")).exists():
|
|
raise Refusal("job exists; pick a new --job", job=j)
|
|
brief = _read_brief(args.brief)
|
|
sid = str(uuid.uuid4())
|
|
prompt = render_template("spawn.md", {
|
|
"block": j, "base": b, "run": r["run"],
|
|
"delta": delta_header(r["fingerprint"]["cwd"], base.get("commit")),
|
|
"note_path": str(state / "notes" / (j + ".md")), "logs_dir": str(state / "logs"),
|
|
"brief": brief})
|
|
rec, _, _, line = _execute(args.run, state, r, job=j, kind="spawn", base=b, sid=sid,
|
|
argv=launch.fork_argv(r, base["session_id"], sid), prompt=prompt,
|
|
before=(base.get("model_usage"), base.get("cost_cumulative_usd")),
|
|
base_ctx=base.get("context_tokens"))
|
|
return _out(rec, line)
|
|
|
|
|
|
def cmd_rebase(args) -> dict:
|
|
"""B7: resume the base without a fork; on `ok` the base takes the result's session id."""
|
|
b = _name("base", args.base)
|
|
state, r = _load_run(args.run)
|
|
base = _load_base(state, b)
|
|
brief = _read_brief(args.brief)
|
|
running = sorted(rec["job"] for rec in map(store.read_json, (state / "jobs").glob("*.json"))
|
|
if rec and rec.get("base") == b and rec.get("status") == "running")
|
|
warnings = ["jobs of %s still marked running: %s" % (b, ", ".join(running))] if running else []
|
|
job = "%s-rebase%d" % (b, len(base.get("rebases") or []) + 1)
|
|
cwd = r["fingerprint"]["cwd"]
|
|
prompt = render_template("rebase.md", {
|
|
"base": b, "run": r["run"], "notes_dir": str(state / "notes"),
|
|
"delta": delta_header(cwd, base.get("commit")), "brief": brief})
|
|
rec, ss, u, line = _execute(args.run, state, r, job=job, kind="rebase", base=b, sid=None,
|
|
argv=launch.resume_argv(r, base["session_id"]), prompt=prompt,
|
|
before=(base.get("model_usage"), base.get("cost_cumulative_usd")),
|
|
base_ctx=base.get("context_tokens"))
|
|
if rec["status"] == "ok":
|
|
base["rebases"] = list(base.get("rebases") or []) + [{
|
|
"job": job, "at": rec["ended"], "from_commit": base.get("commit"),
|
|
"from_session": base["session_id"]}]
|
|
base.update(session_id=rec["session_id"] or base["session_id"], commit=git_head(cwd),
|
|
context_tokens=stream.context_tokens(ss.last_usage), model_usage=u["model_usage"],
|
|
cost_cumulative_usd=_cost_basis(ss), ttl_s=stream.infer_ttl(u["tokens"], base.get("ttl_s")))
|
|
store.write_json(state / "bases" / (b + ".json"), base)
|
|
return {**_out(rec, line), "warnings": warnings}
|
|
|
|
|
|
def cmd_status(args) -> dict:
|
|
"""B6 (warmth per base) and B3 (re-reading a job's result)."""
|
|
state = store.state_dir(args.run)
|
|
if not (state / "run.json").exists():
|
|
raise Refusal("no run.json for this run; call init first", state=str(state))
|
|
if args.job:
|
|
rec = store.read_json(state / "jobs" / (_name("job", args.job) + ".json"))
|
|
if rec is None:
|
|
raise Refusal("unknown job", job=args.job)
|
|
return rec
|
|
lines, now = store.ledger_lines(args.run), time.time()
|
|
bases = {}
|
|
for p in sorted((state / "bases").glob("*.json")):
|
|
base = store.read_json(p)
|
|
bases[base["base"]] = {**{k: base.get(k) for k in ("session_id", "commit", "context_tokens")},
|
|
**_warmth(lines, base, now)}
|
|
jobs = [store.read_json(p) for p in sorted((state / "jobs").glob("*.json"))]
|
|
return {"state": str(state), "bases": bases, "jobs": jobs}
|
|
|
|
|
|
COMMANDS = {"init": cmd_init, "train": cmd_train, "spawn": cmd_spawn, "rebase": cmd_rebase,
|
|
"status": cmd_status}
|
|
|
|
|
|
def _parser() -> argparse.ArgumentParser:
|
|
p = _Parser(prog="gremlins.py", description="spire:gremlins helper; prints one JSON object.")
|
|
sub = p.add_subparsers(dest="cmd", required=True)
|
|
s = sub.add_parser("init")
|
|
s.add_argument("--model", required=True)
|
|
s.add_argument("--effort", default="high", choices=launch.EFFORTS)
|
|
s.add_argument("--cwd")
|
|
s.add_argument("--agent")
|
|
s.add_argument("--binary")
|
|
for name in LAUNCHES:
|
|
s2 = sub.add_parser(name)
|
|
s2.add_argument("--base", required=True)
|
|
if name == "spawn":
|
|
s2.add_argument("--job", required=True)
|
|
s2.add_argument("--brief", required=True)
|
|
sst = sub.add_parser("status")
|
|
sst.add_argument("--job")
|
|
for sp in sub.choices.values():
|
|
sp.add_argument("--run", required=True)
|
|
return p
|
|
|
|
|
|
def main(argv=None) -> int:
|
|
"""Exit 0 ok, 1 a call that ended `error`/`stopped`, 2 refused or a usage error."""
|
|
try:
|
|
args = _parser().parse_args(argv)
|
|
out = COMMANDS[args.cmd](args)
|
|
except Refusal as e:
|
|
_emit(e.payload)
|
|
return 2
|
|
except (ValueError, OSError) as e:
|
|
_emit({"error": str(e), "type": type(e).__name__})
|
|
return 2
|
|
_emit(out)
|
|
return 1 if args.cmd in LAUNCHES and out.get("status") != "ok" else 0
|