314 lines
15 KiB
Python
314 lines
15 KiB
Python
|
"""Initial tamq command-line surface."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import asyncio
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import signal
|
||
|
|
import fcntl
|
||
|
|
from uuid import uuid4
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
from datetime import datetime, timedelta, timezone
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from . import __version__
|
||
|
|
from .config import db_path, pid_path, lock_path, config_path, socket_path, state_dir, setting
|
||
|
|
from .service import Service, ping, request
|
||
|
|
from .store import Store
|
||
|
|
from .tmux import TmuxError, TmuxManager
|
||
|
|
from .broker import BrokerIdentity, InputBroker
|
||
|
|
from .ptytap import PtyTap
|
||
|
|
from .control import ControlModeClient
|
||
|
|
from .diagnostics import configure
|
||
|
|
from .policy import load_profile
|
||
|
|
|
||
|
|
|
||
|
|
def parse_size(value: str) -> int:
|
||
|
|
units = {"B": 1, "KB": 1000, "MB": 1000**2, "GB": 1000**3}
|
||
|
|
text = value.upper()
|
||
|
|
for unit, multiplier in sorted(units.items(), key=lambda item: -len(item[0])):
|
||
|
|
if text.endswith(unit):
|
||
|
|
return int(float(text[:-len(unit)]) * multiplier)
|
||
|
|
return int(text)
|
||
|
|
|
||
|
|
|
||
|
|
def history_advisory(store: Store) -> str | None:
|
||
|
|
size, oldest = store.history_stats()
|
||
|
|
limit_text = setting("history_max_size", "100MB")
|
||
|
|
limit = parse_size(limit_text)
|
||
|
|
if size <= limit:
|
||
|
|
return None
|
||
|
|
since = time.strftime("%y%m%d", time.localtime(oldest)) if oldest else "unknown"
|
||
|
|
return f"Size of history is {size / 1000**2:.0f} MB since {since}, exceeding {limit_text}. Consider using 'tamq purge'."
|
||
|
|
|
||
|
|
|
||
|
|
def ensure_service() -> bool:
|
||
|
|
if asyncio.run(ping()):
|
||
|
|
return True
|
||
|
|
path = lock_path()
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
with path.open("a+") as lock:
|
||
|
|
fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
|
||
|
|
if asyncio.run(ping()):
|
||
|
|
return True
|
||
|
|
subprocess.Popen(
|
||
|
|
[sys.executable, "-m", "tamq.cli", "serve"],
|
||
|
|
stdin=subprocess.DEVNULL,
|
||
|
|
stdout=subprocess.DEVNULL,
|
||
|
|
stderr=subprocess.DEVNULL,
|
||
|
|
start_new_session=True,
|
||
|
|
env=os.environ.copy(),
|
||
|
|
)
|
||
|
|
for _ in range(20):
|
||
|
|
if asyncio.run(ping()):
|
||
|
|
return True
|
||
|
|
time.sleep(0.05)
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def build_parser() -> argparse.ArgumentParser:
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
prog="tamq",
|
||
|
|
description="Tmux Agentic Message Queueing for local repository workers.",
|
||
|
|
)
|
||
|
|
parser.add_argument("--version", "-V", action="version", version=__version__)
|
||
|
|
parser.add_argument("--orwell", action="store_true", help="enable unsafe local diagnostics")
|
||
|
|
parser.add_argument("--verbose", action="store_true", help="enable debug diagnostics")
|
||
|
|
parser.add_argument("--policy-profile", default=None, help="select configured policy profile")
|
||
|
|
subparsers = parser.add_subparsers(dest="command")
|
||
|
|
|
||
|
|
start = subparsers.add_parser("start", help="start or reuse a tamq endpoint")
|
||
|
|
start.add_argument("repos", nargs="*", help="gita-registered repository slugs")
|
||
|
|
start.add_argument("--cmd", default="codex", help="agent command for new windows")
|
||
|
|
start.add_argument("--detach", action="store_true", help="detach after startup")
|
||
|
|
start.add_argument("--no-service", action="store_true", help="skip service startup")
|
||
|
|
attach = subparsers.add_parser("attach", help="attach to the managed tmux session")
|
||
|
|
attach.add_argument("--session", default="tamq", help="tmux session name")
|
||
|
|
|
||
|
|
subparsers.add_parser("serve", help="run the local service in the foreground")
|
||
|
|
subparsers.add_parser("stop", help="stop the local service")
|
||
|
|
subparsers.add_parser("status", help="show endpoint health")
|
||
|
|
subparsers.add_parser("ping", help="check local service liveness")
|
||
|
|
history = subparsers.add_parser("history", help="inspect local message history")
|
||
|
|
history.add_argument("--repo", dest="target_repo")
|
||
|
|
history.add_argument("--state")
|
||
|
|
inspect = subparsers.add_parser("inspect", help="inspect one message")
|
||
|
|
inspect.add_argument("message_id")
|
||
|
|
ack = subparsers.add_parser("ack", help="acknowledge one injected message")
|
||
|
|
ack.add_argument("message_id")
|
||
|
|
send = subparsers.add_parser("send", help="queue a direct message")
|
||
|
|
send.add_argument("address", help="@repo: message")
|
||
|
|
send.add_argument("body", nargs="*", help="message body when address is a repo slug")
|
||
|
|
send.add_argument("--endpoint-id")
|
||
|
|
export = subparsers.add_parser("export", help="export history as JSONL")
|
||
|
|
export.add_argument("--output", required=True)
|
||
|
|
export.add_argument("--repo", dest="target_repo")
|
||
|
|
export.add_argument("--state")
|
||
|
|
replay = subparsers.add_parser("replay", help="replay a JSONL message file")
|
||
|
|
replay.add_argument("file")
|
||
|
|
replay.add_argument("--endpoint-id", help="target a registered tmux-amq endpoint")
|
||
|
|
purge = subparsers.add_parser("purge", help="dry-run or remove history")
|
||
|
|
purge.add_argument("--before", default=None)
|
||
|
|
purge.add_argument("--max-size", default=None)
|
||
|
|
purge.add_argument("--yes", action="store_true")
|
||
|
|
purge.add_argument("--orwell", action="store_true")
|
||
|
|
completion = subparsers.add_parser("completion", help="emit shell completion script")
|
||
|
|
completion.add_argument("shell", choices=("bash", "zsh", "fish"))
|
||
|
|
subparsers.add_parser("db-version", help="show SQLite schema version")
|
||
|
|
subparsers.add_parser("config", help="show effective configuration paths and policy")
|
||
|
|
tap = subparsers.add_parser("tap", help="run an agent behind the full-duplex PTY tap")
|
||
|
|
tap.add_argument("--repo", required=True, help="source gita repository slug")
|
||
|
|
tap.add_argument("--endpoint", required=True, help="tmux-amq endpoint identity")
|
||
|
|
tap.add_argument("agent_command", nargs=argparse.REMAINDER, help="agent command after --")
|
||
|
|
return parser
|
||
|
|
|
||
|
|
|
||
|
|
def main(argv: list[str] | None = None) -> int:
|
||
|
|
parser = build_parser()
|
||
|
|
args = parser.parse_args(argv)
|
||
|
|
configure(args.orwell, args.verbose)
|
||
|
|
try:
|
||
|
|
profile = load_profile(selected=args.policy_profile or os.environ.get("TAMQ_POLICY_PROFILE", "default"))
|
||
|
|
except ValueError as exc:
|
||
|
|
print(f"tamq: {exc}", file=sys.stderr)
|
||
|
|
return 2
|
||
|
|
if args.command is None:
|
||
|
|
parser.print_help()
|
||
|
|
return 0
|
||
|
|
if args.command == "db-version":
|
||
|
|
store = Store(db_path()); print(store.db.execute("SELECT value FROM metadata WHERE key='schema_version'").fetchone()[0]); store.close(); return 0
|
||
|
|
if args.command == "ping":
|
||
|
|
return 0 if asyncio.run(ping()) else 1
|
||
|
|
if args.command == "stop":
|
||
|
|
try:
|
||
|
|
pid = int(pid_path().read_text(encoding="utf-8"))
|
||
|
|
os.kill(pid, signal.SIGTERM)
|
||
|
|
deadline = time.monotonic() + 2.0
|
||
|
|
while time.monotonic() < deadline:
|
||
|
|
try:
|
||
|
|
os.kill(pid, 0)
|
||
|
|
except ProcessLookupError:
|
||
|
|
break
|
||
|
|
time.sleep(0.05)
|
||
|
|
print(f"stopped tamq service {pid}")
|
||
|
|
return 0
|
||
|
|
except ProcessLookupError:
|
||
|
|
pid_path().unlink(missing_ok=True)
|
||
|
|
print("tamq service is not running", file=sys.stderr)
|
||
|
|
return 1
|
||
|
|
except (FileNotFoundError, ValueError):
|
||
|
|
print("tamq service is not running", file=sys.stderr)
|
||
|
|
return 1
|
||
|
|
if args.command == "serve":
|
||
|
|
try: asyncio.run(Service().run())
|
||
|
|
except KeyboardInterrupt: return 0
|
||
|
|
return 0
|
||
|
|
if args.command == "attach":
|
||
|
|
result = subprocess.run(["tmux", "attach-session", "-t", args.session], check=False)
|
||
|
|
return result.returncode
|
||
|
|
if args.command == "completion":
|
||
|
|
print(completion_script(args.shell), end="")
|
||
|
|
return 0
|
||
|
|
if args.command == "config":
|
||
|
|
print(json.dumps({"config": str(config_path()), "state_dir": str(state_dir()), "database": str(db_path()), "socket": str(socket_path()), "pidfile": str(pid_path()), "lockfile": str(lock_path()), "policy_profile": profile.name, "delivery_ack_mode": profile.delivery_ack_mode}, sort_keys=True))
|
||
|
|
return 0
|
||
|
|
if args.command == "tap":
|
||
|
|
command = list(args.agent_command)
|
||
|
|
if command and command[0] == "--":
|
||
|
|
command = command[1:]
|
||
|
|
if not command:
|
||
|
|
print("tamq tap requires an agent command after --", file=sys.stderr)
|
||
|
|
return 2
|
||
|
|
store = Store(db_path())
|
||
|
|
try:
|
||
|
|
return PtyTap(command, InputBroker(store, BrokerIdentity(args.endpoint, args.repo))).run()
|
||
|
|
finally:
|
||
|
|
store.close()
|
||
|
|
if args.command == "start":
|
||
|
|
if not args.no_service and not ensure_service():
|
||
|
|
print("tamq service failed to start; use --no-service to open repos without messaging", file=sys.stderr)
|
||
|
|
return 1
|
||
|
|
try:
|
||
|
|
endpoint = TmuxManager().ensure(args.repos, args.cmd)
|
||
|
|
except (TmuxError, OSError) as exc:
|
||
|
|
print(f"tamq: {exc}", file=sys.stderr)
|
||
|
|
return 2
|
||
|
|
registration = asyncio.run(request({"op": "register", "endpoint_id": endpoint.endpoint_id, "instance_id": endpoint.instance_key, "pid": endpoint.pid, "session": endpoint.session, "repos": endpoint.repos}))
|
||
|
|
if not registration.get("ok"):
|
||
|
|
print(f"tamq: endpoint registration failed: {registration.get('error', 'unknown error')}", file=sys.stderr)
|
||
|
|
return 1
|
||
|
|
control = ControlModeClient(endpoint.session)
|
||
|
|
queue_store = Store(db_path())
|
||
|
|
try:
|
||
|
|
control.start()
|
||
|
|
delivered = InputBroker(queue_store, BrokerIdentity(endpoint.instance_key, endpoint.repos[0])).deliver_pending(
|
||
|
|
control, window_for_repo=lambda repo: f"{endpoint.session}:{repo}"
|
||
|
|
)
|
||
|
|
finally:
|
||
|
|
control.close()
|
||
|
|
queue_store.close()
|
||
|
|
print(json.dumps({"endpoint_id": endpoint.endpoint_id, "instance_id": endpoint.instance_key, "session": endpoint.session, "repos": endpoint.repos, "delivered": delivered}))
|
||
|
|
return 0
|
||
|
|
store = Store(db_path())
|
||
|
|
try:
|
||
|
|
advisory = history_advisory(store)
|
||
|
|
if advisory and args.command in {"start", "serve", "status", "history", "send"}:
|
||
|
|
print(advisory, file=sys.stderr)
|
||
|
|
if args.command == "send":
|
||
|
|
text = " ".join([args.address, *args.body]).strip()
|
||
|
|
if not text.startswith("@") or ":" not in text:
|
||
|
|
print("tamq send expects @repo: message", file=sys.stderr); return 2
|
||
|
|
target, body = text[1:].split(":", 1); body = body.strip()
|
||
|
|
if not target or not body: print("tamq send requires a target and body", file=sys.stderr); return 2
|
||
|
|
if args.endpoint_id:
|
||
|
|
if not asyncio.run(ping()):
|
||
|
|
print("tamq: service is not running", file=sys.stderr); return 1
|
||
|
|
response = asyncio.run(request({"op": "send", "endpoint_id": args.endpoint_id, "sender_repo": "local", "target_repo": target, "body": body}))
|
||
|
|
if not response.get("ok"):
|
||
|
|
print(f"tamq: {response.get('error', 'send failed')}", file=sys.stderr); return 1
|
||
|
|
print(response["message_id"]); return 0
|
||
|
|
try:
|
||
|
|
print(store.add("local", target, body))
|
||
|
|
except ValueError as exc:
|
||
|
|
print(f"tamq: {exc}", file=sys.stderr); return 2
|
||
|
|
return 0
|
||
|
|
if args.command == "history":
|
||
|
|
for row in store.list(args.target_repo, args.state): print(json.dumps(dict(row), sort_keys=True))
|
||
|
|
return 0
|
||
|
|
if args.command == "inspect":
|
||
|
|
rows = [row for row in store.list() if row["message_id"] == args.message_id]
|
||
|
|
if not rows:
|
||
|
|
print("tamq: message not found", file=sys.stderr)
|
||
|
|
return 1
|
||
|
|
print(json.dumps(dict(rows[0]), sort_keys=True)); return 0
|
||
|
|
if args.command == "ack":
|
||
|
|
if not store.acknowledge(args.message_id):
|
||
|
|
print("tamq: message not found", file=sys.stderr); return 1
|
||
|
|
print(args.message_id); return 0
|
||
|
|
if args.command == "export":
|
||
|
|
store.export(store.list(args.target_repo, args.state), Path(args.output)); return 0
|
||
|
|
if args.command == "replay":
|
||
|
|
batch_id = f"replay-{uuid4()}"
|
||
|
|
count = 0
|
||
|
|
if args.endpoint_id and store.endpoint(args.endpoint_id) is None:
|
||
|
|
print(f"tamq: endpoint is not registered: {args.endpoint_id}", file=sys.stderr)
|
||
|
|
return 1
|
||
|
|
try:
|
||
|
|
lines = Path(args.file).read_text(encoding="utf-8").splitlines()
|
||
|
|
for line in lines:
|
||
|
|
if not line.strip():
|
||
|
|
continue
|
||
|
|
item = json.loads(line)
|
||
|
|
provenance = json.dumps({"batch_id": batch_id, "original_message_id": item.get("message_id")}, sort_keys=True)
|
||
|
|
endpoint = args.endpoint_id or item.get("endpoint_id")
|
||
|
|
store.add(item.get("sender_repo", "replay"), item["target_repo"], item["body"], endpoint=endpoint, provenance=provenance)
|
||
|
|
count += 1
|
||
|
|
except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
||
|
|
print(f"tamq: cannot replay {args.file}: {exc}", file=sys.stderr)
|
||
|
|
return 2
|
||
|
|
print(json.dumps({"batch_id": batch_id, "count": count})); return 0
|
||
|
|
if args.command == "purge":
|
||
|
|
before_text = args.before or setting("purge_before", "365d")
|
||
|
|
max_size_text = args.max_size or setting("purge_max_size", "100MB")
|
||
|
|
before = time.time() - int(before_text[:-1]) * 86400 if before_text.endswith("d") else None
|
||
|
|
if not args.yes:
|
||
|
|
matching = [row for row in store.list() if before is None or row["created_at"] < before]
|
||
|
|
print(f"Dry run: {len(matching)} message(s) match; re-run with --yes to purge.")
|
||
|
|
return 0
|
||
|
|
count = store.purge(before=before, max_bytes=parse_size(max_size_text))
|
||
|
|
print(f"Purged {count} message(s)."); return 0
|
||
|
|
if args.command == "status":
|
||
|
|
live = asyncio.run(ping())
|
||
|
|
endpoints = asyncio.run(request({"op": "endpoints"})) if live else {"endpoints": []}
|
||
|
|
print(json.dumps({"version": __version__, "db": str(db_path()), "messages": len(store.list()), "pending": len(store.list(state="pending")), "leases": store.db.execute("SELECT COUNT(*) FROM leases").fetchone()[0], "service": live, "policy_profile": profile.name, "delivery_ack_mode": profile.delivery_ack_mode, "endpoints": endpoints.get("endpoints", [])}, sort_keys=True)); return 0
|
||
|
|
print(f"tamq {__version__}: command '{args.command}' is not implemented yet", file=sys.stderr)
|
||
|
|
return 2
|
||
|
|
finally:
|
||
|
|
store.close()
|
||
|
|
|
||
|
|
|
||
|
|
def completion_script(shell: str) -> str:
|
||
|
|
commands = "start attach serve stop status ping history inspect ack send export replay purge completion db-version config tap"
|
||
|
|
if shell == "bash":
|
||
|
|
return f"""_tamq_complete() {{
|
||
|
|
local commands=\"{commands}\"
|
||
|
|
if [[ ${{COMP_WORDS[1]}} == start || ${{COMP_WORDS[1]}} == send ]]; then
|
||
|
|
COMPREPLY=( $(compgen -W \"$(gita ls 2>/dev/null)\" -- \"${{COMP_WORDS[COMP_CWORD]}}\") )
|
||
|
|
else
|
||
|
|
COMPREPLY=( $(compgen -W \"$commands\" -- \"${{COMP_WORDS[COMP_CWORD]}}\") )
|
||
|
|
fi
|
||
|
|
}}
|
||
|
|
complete -F _tamq_complete tamq
|
||
|
|
"""
|
||
|
|
if shell == "zsh":
|
||
|
|
return f"#compdef tamq\n_arguments '1:command:(({commands}))' '*:repo:->repos'\nif [[ $state == repos ]]; then _values 'gita repo' $(gita ls 2>/dev/null); fi\n"
|
||
|
|
return f"complete -c tamq -f -n '__fish_seen_subcommand_from start send' -a '(gita ls 2>/dev/null)'\ncomplete -c tamq -f -a '{commands}'\n"
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
raise SystemExit(main())
|