plane-agentd/agentd/cli.py

"""agentd command line. Tokens come from the environment (agent-access.py exec --secret-env ...)."""
import argparse
import os
import re
import sys
import time

from . import config as config_mod
from .core import Agentd
from .plane import Plane, brief_to_html, html_to_text, text_to_html
from .store import Busy


def identity(cfg, wanted=None):
    """Terminal whose token the CLI uses: --as, then agentd.cli_terminal, then the first with a token."""
    for term in [wanted, cfg.agentd.get("cli_terminal")] + list(cfg.terminals):
        if term and term in cfg.terminals and cfg.token(term):
            return term
    sys.exit("no Plane token in the environment; run under agent-access.py exec %s" % " ".join(cfg.secret_env_args()))


def client(cfg, term):
    return Plane(cfg.plane["base_url"], cfg.plane["workspace"], cfg.token(term), cfg.plane.get("timeout", 60))


def seq_of(ref):
    m = re.search(r"(\d+)$", ref)
    if not m:
        sys.exit("bad work item reference %r (use 12 or GW-12)" % ref)
    return int(m.group(1))


def cmd_tick(cfg, a):
    d = Agentd(cfg)
    if a.loop:
        d.loop(a.interval or cfg.agentd.get("interval", 120))
    try:
        d.tick(project=a.project)
    except Busy as exc:
        print("skip: %s" % exc)
        sys.exit(75)   # same code as the cron script's flock conflict: the webhook receiver retries later


def cmd_webhook_serve(cfg, a):
    if not cfg.webhook:
        sys.exit("no webhook section in %s" % cfg.path)
    from .webhook import serve
    serve(cfg, a.bind, a.port)


def cmd_new(cfg, a, pl=None, quiet=False):
    p = cfg.project(a.project)
    pl = pl or client(cfg, identity(cfg, a.as_))
    text = sys.stdin.read() if a.file == "-" else open(a.file, encoding="utf-8").read()
    body = {"name": a.title, "description_html": brief_to_html(text), "state": pl.states(p["id"])[a.state]}
    if a.labels:
        labels = pl.labels(p["id"])
        body["labels"] = [labels[x] for x in a.labels.split(",")]
    if a.assign:
        body["assignees"] = [pl.members()[a.assign]]
    issue = pl.create_issue(p["id"], body)
    if a.module:
        mods = pl.modules(p["id"])
        mid = mods.get(a.module) or pl.create_module(p["id"], a.module)["id"]
        pl.add_to_module(p["id"], mid, issue["id"])
    if not quiet:
        print("created %s-%s %s" % (p.get("key_prefix", a.project), issue["sequence_id"], issue["id"]))
    return issue


def cmd_show(cfg, a):
    p = cfg.project(a.project)
    pl = client(cfg, identity(cfg, a.as_))
    i = pl.find_issue(p["id"], seq_of(a.ref))
    sn = {v: k for k, v in pl.states(p["id"]).items()}
    ln = {v: k for k, v in pl.labels(p["id"]).items()}
    un = {v: k for k, v in pl.members().items()}
    print("%s-%s %s | state=%s | assignees=%s | labels=%s" % (
        p.get("key_prefix", a.project), i["sequence_id"], i["name"], sn.get(i.get("state"), i.get("state")),
        [un.get(x, x) for x in i.get("assignees") or []], [ln.get(x, x) for x in i.get("labels") or []]))
    cs = pl.list_comments(p["id"], i["id"])
    for c in cs[-a.n:] if a.n else []:
        who = c.get("actor") or c.get("created_by")
        print("--- %s (%s)" % (c.get("created_at"), un.get(who, who)))
        print(html_to_text(c.get("comment_html"))[:2500])


def cmd_comment(cfg, a):
    p = cfg.project(a.project)
    pl = client(cfg, identity(cfg, a.as_))
    i = pl.find_issue(p["id"], seq_of(a.ref))
    text = sys.stdin.read() if a.file == "-" else open(a.file, encoding="utf-8").read()
    pl.post_comment(p["id"], i["id"], text_to_html(text))
    print("commented %s" % a.ref)


def cmd_state(cfg, a):
    p = cfg.project(a.project)
    pl = client(cfg, identity(cfg, a.as_))
    i = pl.find_issue(p["id"], seq_of(a.ref))
    name = p["states"].get(a.state, a.state)  # accepts keys (todo) or Plane names (Todo)
    pl.set_state(p["id"], i["id"], pl.states(p["id"])[name])
    print("%s -> %s" % (a.ref, name))


def cmd_assign(cfg, a):
    p = cfg.project(a.project)
    pl = client(cfg, identity(cfg, a.as_))
    i = pl.find_issue(p["id"], seq_of(a.ref))
    pl.set_assignees(p["id"], i["id"], [pl.members()[a.user]])
    print("%s assigned %s" % (a.ref, a.user))


SMOKE_BRIEF = ("<p>agentd 冒烟测试。不要读写任何文件，不要运行任何命令。</p>"
               "<p>你的最后一条消息必须恰好是下面这一行（不要加任何别的内容）：</p><pre>[review] ok</pre>")


def cmd_smoke(cfg, a):
    p = cfg.project(a.project)
    if a.terminal not in cfg.terminals:
        sys.exit("unknown terminal %s" % a.terminal)
    pl = client(cfg, identity(cfg, a.as_))
    user = cfg.terminals[a.terminal]["plane_user"]
    body = {"name": "agentd smoke %s %s" % (a.terminal, time.strftime("%Y-%m-%d %H:%M")),
            "description_html": SMOKE_BRIEF, "state": pl.states(p["id"])[p["states"]["todo"]],
            "assignees": [pl.members()[user]]}
    if a.labels:
        labels = pl.labels(p["id"])
        body["labels"] = [labels[x] for x in a.labels.split(",")]
    issue = pl.create_issue(p["id"], body)
    key = "%s-%s" % (p.get("key_prefix", a.project), issue["sequence_id"])
    print("created %s; %s every %ss (timeout %s min)" % (
        key, "watching (the scheduled dispatcher runs it)" if a.watch else "ticking", a.interval, a.timeout), flush=True)
    d = None if a.watch else Agentd(cfg)
    states = pl.states(p["id"])
    done_states = {states[p["states"][k]] for k in ("review", "waiting")}
    deadline = time.time() + a.timeout * 60
    state = None
    while time.time() < deadline:
        try:
            if d:
                d.tick(project=a.project, only_issue=issue["id"])
        except Busy as exc:
            print("skip: %s" % exc)
        state = pl.get_issue(p["id"], issue["id"]).get("state")
        if state in done_states:
            break
        time.sleep(a.interval)
    last = ""
    for c in reversed(pl.list_comments(p["id"], issue["id"])):
        text = html_to_text(c.get("comment_html"))
        if not text.startswith("[worker]"):
            last = text
            break
    ok = state == states[p["states"]["review"]] and re.search(r"(?m)^\[review\] ok\s*$", last)
    print("%s %s: state=%s\n%s" % ("PASS" if ok else "FAIL", key, {v: k for k, v in states.items()}.get(state, state), last[-800:]))
    if ok:
        pl.set_state(p["id"], issue["id"], states[p["states"]["done"]])
    sys.exit(0 if ok else 1)


def cmd_check_config(cfg, a):
    print("config %s is valid" % cfg.path)
    print("projects: %s" % ", ".join(cfg.projects))
    for name, t in cfg.terminals.items():
        env = cfg.token_env(name)
        print("terminal %-10s %-13s %-8s executors=%s token %s: %s" % (
            name, t["transport"], t["os"], ",".join(t["executors"]), env, "set" if os.environ.get(env) else "MISSING"))
    print("cron/exec secret args: %s" % " ".join(cfg.secret_env_args()))
    if cfg.webhook:
        k = cfg.webhook["secret_key"]
        print("webhook-serve secret arg: --secret-env %s=%s (%s)" % (k, k, "set" if os.environ.get(k) else "not in this env"))


def parser():
    ap = argparse.ArgumentParser(prog="agentd", description=__doc__)
    ap.add_argument("--config", help="YAML config (default: $AGENTD_CONFIG, config/agentd.local.yaml, config/agentd.yaml)")
    sub = ap.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("tick", help="run the dispatcher once or in a loop")
    g = s.add_mutually_exclusive_group()
    g.add_argument("--once", action="store_true")
    g.add_argument("--loop", action="store_true")
    s.add_argument("--interval", type=int)
    s.add_argument("--project")
    s.set_defaults(fn=cmd_tick)

    def item(name, fn, help_):
        s = sub.add_parser(name, help=help_)
        s.add_argument("project")
        s.add_argument("ref", help="sequence number or KEY-n")
        s.add_argument("--as", dest="as_", help="terminal identity to act as")
        s.set_defaults(fn=fn)
        return s

    s = sub.add_parser("new", help="create a work item from a brief file")
    s.add_argument("project")
    s.add_argument("title")
    s.add_argument("--file", required=True, help="brief file (HTML passes through, text goes in <pre>); - = stdin")
    s.add_argument("--assign")
    s.add_argument("--labels", default="")
    s.add_argument("--module")
    s.add_argument("--state", default="Backlog")
    s.add_argument("--as", dest="as_")
    s.set_defaults(fn=cmd_new)
    item("show", cmd_show, "show a work item and its last n comments").add_argument("n", nargs="?", type=int, default=1)
    item("comment", cmd_comment, "post a comment from a file (- = stdin)").add_argument("file")
    item("state", cmd_state, "set state (Plane name or config key)").add_argument("state")
    item("assign", cmd_assign, "set the assignee").add_argument("user")
    s = sub.add_parser("smoke", help="end-to-end round on a terminal: expects '[review] ok'")
    s.add_argument("project")
    s.add_argument("terminal")
    s.add_argument("--interval", type=int, default=30)
    s.add_argument("--timeout", type=int, default=30, help="minutes")
    s.add_argument("--labels", default="", help="comma-separated Plane labels, e.g. model:sol")
    s.add_argument("--watch", action="store_true",
                   help="do not tick here; only wait for the scheduled dispatcher (use when it runs on another host)")
    s.add_argument("--as", dest="as_")
    s.set_defaults(fn=cmd_smoke)
    s = sub.add_parser("webhook-serve", help="Plane webhook receiver: each delivery triggers a tick (agentd/webhook.py)")
    s.add_argument("--bind", help="address (default: webhook.bind)")
    s.add_argument("--port", type=int, help="port (default: webhook.port)")
    s.set_defaults(fn=cmd_webhook_serve)
    sub.add_parser("check-config", help="validate the config and list token env vars").set_defaults(fn=cmd_check_config)
    return ap


def main(argv=None):
    a = parser().parse_args(argv)
    try:
        cfg = config_mod.load(a.config)
    except config_mod.ConfigError as exc:
        sys.exit(str(exc))
    a.fn(cfg, a)