local-rig/local_rig/cli.py

"""rig CLI — 一条命令管理单机 Agent 团队。

用法:
  rig start -f rig.yaml [--port 27900]
  rig status -f rig.yaml
  rig assign -f rig.yaml --to dev.lead "修复登录 bug"
  rig tasks -f rig.yaml [--status pending]
  rig stop -f rig.yaml
"""
from __future__ import annotations

import argparse
import json
import sys
import time
import urllib.request
import ssl

from .rig_schema import parse_rig, RigSchemaError
from .daemon import RigDaemon


def _api(port: int, path: str, method: str = "GET", data: dict | None = None) -> dict:
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    url = f"http://127.0.0.1:{port}{path}"
    body = json.dumps(data).encode() if data else None
    req = urllib.request.Request(url, data=body, method=method,
                                 headers={"Content-Type": "application/json"} if body else {})
    try:
        with urllib.request.urlopen(req, context=ctx, timeout=10) as resp:
            return json.loads(resp.read())
    except Exception as e:
        print(f"❌ daemon 未响应 ({url}): {e}", file=sys.stderr)
        print("   先运行: rig start -f rig.yaml", file=sys.stderr)
        sys.exit(1)


def cmd_start(args):
    rig = parse_rig(args.file)
    print(f"🚀 启动 rig '{rig.name}' ({len(rig.all_members())} members)")
    d = RigDaemon(args.file, host="127.0.0.1", port=args.port)
    d.start(background=False)  # 前台运行
    print(f"   http://127.0.0.1:{args.port}/health")
    print(f"   members: {', '.join(rig.all_members())}")
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n⏹ 停止...")
        d.stop()


def cmd_stop(args):
    _api(args.port, "/stop", "POST")
    print("⏹ rig 已停止")


def cmd_status(args):
    health = _api(args.port, "/health")
    print(f"📡 {health['rig']} — {health['status']}")
    status = _api(args.port, "/status")
    for mid, st in status.items():
        icon = "🟢" if st == "running" else "⚪"
        print(f"   {icon} {mid:20s} {st}")
    tasks = _api(args.port, "/tasks?limit=10")
    if tasks:
        print(f"\n📋 最近任务:")
        for t in tasks[:10]:
            print(f"   [{t['status']:8s}] {t['id'][:20]} → {t['assignee']}: {t['prompt'][:50]}")


def cmd_assign(args):
    result = _api(args.port, "/tasks", "POST", {
        "assignee": args.to,
        "prompt": args.prompt,
        "deadline_s": args.deadline,
    })
    print(f"✅ 任务已入队: {result['id']} → {args.to}")


def cmd_tasks(args):
    path = "/tasks"
    params = []
    if args.status:
        params.append(f"status={args.status}")
    if params:
        path += "?" + "&".join(params)
    tasks = _api(args.port, path)
    if not tasks:
        print("(无任务)")
        return
    for t in tasks:
        print(f"[{t['status']:8s}] {t['id']}")
        print(f"   → {t['assignee']}")
        print(f"   {t['prompt']}")
        if t.get("result"):
            print(f"   ✅ {t['result'][:80]}")
        if t.get("error"):
            print(f"   ❌ {t['error'][:80]}")
        print()


def main():
    p = argparse.ArgumentParser(prog="rig", description="单机轻量 Agent 团队")
    p.add_argument("-f", "--file", default="rig.yaml", help="rig.yaml 路径")
    p.add_argument("--port", type=int, default=27900, help="daemon 端口")
    sub = p.add_subparsers(dest="cmd", required=True)

    sub.add_parser("start", help="启动 daemon + 所有 member")
    sub.add_parser("stop", help="停止 daemon + 所有 member")
    sub.add_parser("status", help="查看 rig 状态")
    sub.add_parser("tasks", help="查看任务")

    pa = sub.add_parser("assign", help="派发任务")
    pa.add_argument("--to", required=True, help="目标 member (pod.member)")
    pa.add_argument("prompt", help="任务描述")
    pa.add_argument("--deadline", type=float, default=600, help="超时秒数")

    pa = sub.add_parser("tasks")
    pa.add_argument("--status", help="过滤: pending/running/done/failed/timeout")

    args = p.parse_args()
    try:
        parse_rig(args.file)  # 预校验
    except RigSchemaError as e:
        print(f"❌ rig.yaml 错误: {e}", file=sys.stderr)
        sys.exit(1)

    {"start": cmd_start, "stop": cmd_stop, "status": cmd_status,
     "assign": cmd_assign, "tasks": cmd_tasks}[args.cmd](args)


if __name__ == "__main__":
    main()