tq.py

#!/usr/bin/env python3
"""
tq - Agent Task Queue CLI

CLI to inspect and run commands through the Agent Task Queue.
"""

import argparse
import json
import os
import shlex
import signal
import sqlite3
import subprocess
import sys
import time
import uuid
from datetime import datetime
from pathlib import Path

# Import shared queue infrastructure
from queue_core import (
    QueuePaths,
    TaskOrigin,
    collect_task_origin,
    get_db,
    init_db,
    ensure_db,
    cleanup_queue as _cleanup_queue,
    cleanup_targets_for_queue,
    log_metric as _log_metric,
    release_lock,
    is_process_alive,
    kill_process_tree,
    insert_waiting_task,
    normalize_queue_name,
    parse_queue_capacities,
    attempt_task_start,
    task_origin_kwargs,
    POLL_INTERVAL_WAITING,
    DEFAULT_MAX_LOCK_AGE_MINUTES,
    DEFAULT_MAX_METRICS_SIZE_MB,
)

# Unique identifier for this CLI instance - used to detect orphaned tasks
# from previous CLI instances even if the PID is reused
CLI_INSTANCE_ID = str(uuid.uuid4())[:8]


def get_paths(args) -> QueuePaths:
    """Get queue paths from args or environment."""
    if args.data_dir:
        data_dir = Path(args.data_dir)
    else:
        data_dir = Path(os.environ.get("TASK_QUEUE_DATA_DIR", "/tmp/agent-task-queue"))
    return QueuePaths.from_data_dir(data_dir)


def get_queue_capacities(args) -> dict[str, int]:
    """Parse queue capacity overrides from CLI args."""
    return parse_queue_capacities(getattr(args, "queue_capacity", []))


def cmd_list(args):
    """List all tasks in the queue."""
    paths = get_paths(args)
    json_output = getattr(args, "json", False)

    if not paths.db_path.exists():
        if json_output:
            print(json.dumps({"tasks": [], "summary": {"total": 0, "running": 0, "waiting": 0}}))
        else:
            print(f"No queue database found at {paths.db_path}")
            print("Queue is empty (no tasks have been run yet)")
        return

    conn = sqlite3.connect(paths.db_path, timeout=5.0)
    conn.row_factory = sqlite3.Row

    try:
        rows = conn.execute(
            "SELECT * FROM queue ORDER BY queue_name, id"
        ).fetchall()

        if json_output:
            tasks = []
            running_count = 0
            waiting_count = 0
            for row in rows:
                task = {
                    "id": row["id"],
                    "queue_name": row["queue_name"],
                    "status": row["status"],
                    "command": row["command"] if "command" in row.keys() else None,
                    "pid": row["pid"],
                    "child_pid": row["child_pid"],
                    "created_at": row["created_at"],
                    "updated_at": row["updated_at"],
                }
                tasks.append(task)
                if row["status"] == "running":
                    running_count += 1
                elif row["status"] == "waiting":
                    waiting_count += 1

            output = {
                "tasks": tasks,
                "summary": {
                    "total": len(tasks),
                    "running": running_count,
                    "waiting": waiting_count,
                },
            }
            print(json.dumps(output))
            return

        if not rows:
            print("Queue is empty")
            return

        # Group by queue name
        queues = {}
        for row in rows:
            qname = row["queue_name"]
            if qname not in queues:
                queues[qname] = []
            queues[qname].append(row)

        for qname, tasks in queues.items():
            print(f"\n[{qname}] ({len(tasks)} task(s))")
            print("-" * 50)

            for task in tasks:
                status = task["status"].upper()
                task_id = task["id"]
                pid = task["pid"] or "-"
                child_pid = task["child_pid"] or "-"
                created = task["created_at"]

                # Format timestamp
                if created:
                    try:
                        dt = datetime.fromisoformat(created)
                        created = dt.strftime("%H:%M:%S")
                    except ValueError:
                        pass

                status_icon = "🔄" if status == "RUNNING" else "⏳"
                print(f"  {status_icon} #{task_id} {status} (pid={pid}, child={child_pid}) @ {created}")

    finally:
        conn.close()


def cmd_clear(args):
    """Clear all tasks from the queue."""
    paths = get_paths(args)
    json_output = getattr(args, "json", False)

    if not paths.db_path.exists():
        if json_output:
            print(json.dumps({"cleared": 0, "success": True}))
        else:
            print("No queue database found")
        return

    conn = sqlite3.connect(paths.db_path, timeout=5.0)
    try:
        # Check how many tasks exist
        count = conn.execute("SELECT COUNT(*) FROM queue").fetchone()[0]
        if count == 0:
            if json_output:
                print(json.dumps({"cleared": 0, "success": True}))
            else:
                print("Queue is already empty")
            return

        # JSON mode skips confirmation (implies --force)
        if not json_output:
            response = input(f"Clear {count} task(s) from queue? [y/N] ")
            if response.lower() != 'y':
                print("Cancelled")
                return

        cursor = conn.execute("DELETE FROM queue")
        conn.commit()

        if json_output:
            print(json.dumps({"cleared": cursor.rowcount, "success": True}))
        else:
            print(f"Cleared {cursor.rowcount} task(s) from queue")
    finally:
        conn.close()


def cmd_logs(args):
    """Show recent log entries."""
    paths = get_paths(args)
    json_output = getattr(args, "json", False)

    if not paths.metrics_path.exists():
        if json_output:
            print(json.dumps({"entries": []}))
        else:
            print(f"No log file found at {paths.metrics_path}")
        return

    lines = paths.metrics_path.read_text().strip().split("\n")
    recent = lines[-args.n:] if len(lines) > args.n else lines

    if json_output:
        entries = []
        for line in recent:
            try:
                entry = json.loads(line)
                entries.append(entry)
            except json.JSONDecodeError:
                # Skip malformed lines in JSON mode
                pass
        print(json.dumps({"entries": entries}))
        return

    for line in recent:
        try:
            entry = json.loads(line)
            ts = entry.get("timestamp", "")[:19].replace("T", " ")
            event = entry.get("event", "unknown")
            task_id = entry.get("task_id", "")
            queue = entry.get("queue_name", "")

            # Format based on event type
            if event == "task_completed":
                exit_code = entry.get("exit_code", "?")
                duration = entry.get("duration_seconds", "?")
                print(f"{ts} [{queue}] #{task_id} completed exit={exit_code} {duration}s")
            elif event == "task_started":
                wait = entry.get("wait_time_seconds", 0)
                print(f"{ts} [{queue}] #{task_id} started (waited {wait}s)")
            elif event == "task_queued":
                print(f"{ts} [{queue}] #{task_id} queued")
            elif event == "task_timeout":
                print(f"{ts} [{queue}] #{task_id} TIMEOUT")
            elif event == "task_error":
                error = entry.get("error", "?")
                print(f"{ts} [{queue}] #{task_id} ERROR: {error}")
            elif event == "zombie_cleared":
                reason = entry.get("reason", "?")
                print(f"{ts} [{queue}] #{task_id} zombie cleared ({reason})")
            elif event == "orphan_cleared":
                reason = entry.get("reason", "?")
                print(f"{ts} [{queue}] #{task_id} orphan cleared ({reason})")
            else:
                print(f"{ts} {event}")
        except json.JSONDecodeError:
            print(line)


# --- Run Command Implementation ---

def log_metric(paths: QueuePaths, event: str, **kwargs):
    """Log metric using paths (wrapper for CLI)."""
    _log_metric(paths.metrics_path, event, DEFAULT_MAX_METRICS_SIZE_MB, **kwargs)


def cleanup_queue(
    conn,
    queue_name: str,
    paths: QueuePaths,
    queue_capacities: dict[str, int] | None = None,
):
    """Clean up queue (wrapper for CLI)."""
    for target_queue in cleanup_targets_for_queue(conn, queue_name, queue_capacities):
        _cleanup_queue(conn, target_queue, paths.metrics_path, DEFAULT_MAX_LOCK_AGE_MINUTES)

        # Additional cleanup: Tasks with our PID but DIFFERENT instance_id (from old CLI instance)
        # This handles the edge case where PID is reused after CLI crash
        my_pid = os.getpid()
        stale_tasks = conn.execute(
            "SELECT id, status, child_pid, server_id FROM queue WHERE queue_name = ? AND pid = ? AND server_id IS NOT NULL AND server_id != ?",
            (target_queue, my_pid, CLI_INSTANCE_ID),
        ).fetchall()

        for task in stale_tasks:
            if task["child_pid"] and is_process_alive(task["child_pid"]):
                print(f"[tq] WARNING: Killing orphaned subprocess {task['child_pid']} from old CLI instance")
                kill_process_tree(task["child_pid"])

            conn.execute("DELETE FROM queue WHERE id = ?", (task["id"],))
            log_metric(
                paths,
                "orphan_cleared",
                task_id=task["id"],
                queue_name=target_queue,
                status=task["status"],
                old_instance_id=task["server_id"],
                reason="stale_cli_instance",
            )
            print(f"[tq] WARNING: Cleared task from old CLI instance (ID: {task['id']}, old_instance: {task['server_id']})")

    if conn.in_transaction:
        conn.commit()


def register_task(
    conn,
    queue_name: str,
    paths: QueuePaths,
    command: str = None,
    task_origin: TaskOrigin | None = None,
) -> int:
    """Register a task in the queue. Returns task_id immediately."""
    my_pid = os.getpid()

    task_id = insert_waiting_task(
        conn,
        queue_name,
        my_pid,
        CLI_INSTANCE_ID,
        command=command,
        task_origin=task_origin,
    )
    conn.commit()

    log_metric(
        paths,
        "task_queued",
        task_id=task_id,
        queue_name=queue_name,
        pid=my_pid,
        **task_origin_kwargs(task_origin),
    )
    print(f"[tq] Task #{task_id} queued in '{queue_name}'")
    return task_id


def wait_for_turn(
    conn,
    queue_name: str,
    task_id: int,
    paths: QueuePaths,
    queue_capacities: dict[str, int],
    task_origin: TaskOrigin | None = None,
) -> None:
    """Wait for the task's turn to run. Task must already be registered."""
    my_pid = os.getpid()
    queued_at = time.time()

    last_pos = -1

    while True:
        try:
            cleanup_queue(conn, queue_name, paths, queue_capacities)

            started, pos = attempt_task_start(
                conn,
                task_id,
                queue_name,
                queue_capacities,
                my_pid,
            )

            if not started:

                if pos != last_pos:
                    print(f"[tq] Position #{pos} in queue. Waiting...")
                    last_pos = pos

                time.sleep(POLL_INTERVAL_WAITING)
                continue

            wait_time = time.time() - queued_at
            log_metric(
                paths,
                "task_started",
                task_id=task_id,
                queue_name=queue_name,
                pid=my_pid,
                wait_time_seconds=round(wait_time, 2),
                **task_origin_kwargs(task_origin),
            )
            if wait_time > 1:
                print(f"[tq] Lock acquired after {wait_time:.1f}s wait")
            else:
                print("[tq] Lock acquired")
            return  # Lock acquired, task_id was passed in
        except sqlite3.OperationalError as exc:
            if "database is locked" not in str(exc).lower():
                raise

        time.sleep(POLL_INTERVAL_WAITING)


def cmd_run(args):
    """Run a command through the task queue."""
    if not args.run_command:
        print("Error: No command specified", file=sys.stderr)
        sys.exit(1)

    # Use shlex.join to properly quote arguments with spaces
    command = shlex.join(args.run_command)
    working_dir = os.path.abspath(args.dir) if args.dir else os.getcwd()
    try:
        queue_name = normalize_queue_name(args.queue)
        queue_capacities = get_queue_capacities(args)
    except ValueError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        sys.exit(1)

    timeout = args.timeout

    if not os.path.exists(working_dir):
        print(f"Error: Working directory does not exist: {working_dir}", file=sys.stderr)
        sys.exit(1)

    paths = get_paths(args)
    paths.data_dir.mkdir(parents=True, exist_ok=True)
    task_origin = collect_task_origin(working_dir)

    # Ensure database exists and is valid (recover if corrupted)
    ensure_db(paths)

    # Get database connection
    with get_db(paths.db_path) as conn:
        # Initialize schema if needed (idempotent via IF NOT EXISTS)
        init_db(paths)

    # Open connection for the duration of the run
    conn = sqlite3.connect(paths.db_path, timeout=60.0)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout=60000")
    conn.row_factory = sqlite3.Row

    task_id = None
    proc = None
    cleaned_up = False

    def cleanup_handler(signum, frame):
        """Handle Ctrl+C - clean up and exit."""
        nonlocal cleaned_up
        if cleaned_up:
            return
        cleaned_up = True

        print("\n[tq] Interrupted. Cleaning up...")
        if proc and proc.poll() is None:
            try:
                os.killpg(proc.pid, signal.SIGTERM)
                proc.wait(timeout=5)
            except Exception:
                try:
                    os.killpg(proc.pid, signal.SIGKILL)
                except Exception:
                    pass
        if task_id:
            try:
                release_lock(conn, task_id)
            except Exception:
                pass
        try:
            conn.close()
        except Exception:
            pass
        sys.exit(130)

    signal.signal(signal.SIGINT, cleanup_handler)
    signal.signal(signal.SIGTERM, cleanup_handler)

    try:
        # Run cleanup BEFORE inserting - this clears orphaned tasks that would otherwise
        # block the queue forever (since cleanup only runs during polling)
        cleanup_queue(conn, queue_name, paths, queue_capacities)

        # Register task first so task_id is available for cleanup if interrupted
        task_id = register_task(conn, queue_name, paths, command=command, task_origin=task_origin)
        wait_for_turn(conn, queue_name, task_id, paths, queue_capacities, task_origin=task_origin)

        print(f"[tq] Running: {command}")
        print(f"[tq] Directory: {working_dir}")
        print("-" * 60)

        start = time.time()

        # Run subprocess in passthrough mode - direct terminal connection
        # This preserves rich output (progress bars, colors, etc.)
        # nosec B602: shell=True is intentional - this CLI tool executes user-provided
        # commands, similar to bash -c or make. Users control their own CLI arguments.
        # Shell features (pipes, redirects, globs) are required for build commands.
        proc = subprocess.Popen(
            command,
            shell=True,  # nosec B602
            cwd=working_dir,
            start_new_session=True,  # For clean process group kill
        )

        # Record child PID for zombie protection
        conn.execute(
            "UPDATE queue SET child_pid = ? WHERE id = ?", (proc.pid, task_id)
        )
        conn.commit()

        # Wait for process (Ctrl+C will trigger cleanup_handler)
        try:
            proc.wait(timeout=timeout if timeout else None)
        except subprocess.TimeoutExpired:
            print(f"\n[tq] TIMEOUT after {timeout}s")
            try:
                os.killpg(proc.pid, signal.SIGKILL)
            except OSError:
                pass
            proc.wait()
            log_metric(
                paths,
                "task_timeout",
                task_id=task_id,
                queue_name=queue_name,
                pid=os.getpid(),
                command=command,
                timeout_seconds=timeout,
                **task_origin_kwargs(task_origin),
            )
            return 124  # Standard timeout exit code

        duration = time.time() - start
        exit_code = proc.returncode

        print("-" * 60)
        if exit_code == 0:
            print(f"[tq] SUCCESS in {duration:.1f}s")
        else:
            print(f"[tq] FAILED exit={exit_code} in {duration:.1f}s")

        log_metric(
            paths,
            "task_completed",
            task_id=task_id,
            queue_name=queue_name,
            pid=os.getpid(),
            command=command,
            exit_code=exit_code,
            duration_seconds=round(duration, 2),
            **task_origin_kwargs(task_origin),
        )

        return exit_code

    except Exception as e:
        print(f"[tq] Error: {e}", file=sys.stderr)
        if task_id:
            log_metric(
                paths,
                "task_error",
                task_id=task_id,
                queue_name=queue_name,
                pid=os.getpid(),
                error=str(e),
                **task_origin_kwargs(task_origin),
            )
        return 1

    finally:
        if not cleaned_up:
            if task_id:
                try:
                    release_lock(conn, task_id)
                except Exception:
                    pass
            try:
                conn.close()
            except Exception:
                pass


def main():
    parser = argparse.ArgumentParser(
        prog="tq",
        description="Agent Task Queue CLI - inspect and manage the task queue",
    )
    parser.add_argument(
        "--data-dir",
        help="Data directory (default: $TASK_QUEUE_DATA_DIR or /tmp/agent-task-queue)",
    )
    parser.add_argument(
        "--queue-capacity",
        action="append",
        default=[],
        metavar="SCOPE=CAPACITY",
        help=(
            "Hierarchical queue capacity override. Repeatable. "
            "Example: --queue-capacity=gradle=2 --queue-capacity=gradle/emu-5557=1"
        ),
    )

    subparsers = parser.add_subparsers(dest="command", help="Commands")

    # run
    run_parser = subparsers.add_parser("run", help="Run a command through the queue")
    run_parser.add_argument("-q", "--queue", default="global", help="Queue name (default: global)")
    run_parser.add_argument("-t", "--timeout", type=int, default=1200, help="Timeout in seconds (default: 1200)")
    run_parser.add_argument("-C", "--dir", help="Working directory (default: current)")
    run_parser.add_argument("run_command", nargs=argparse.REMAINDER, metavar="COMMAND", help="Command to run")

    # list
    list_parser = subparsers.add_parser("list", help="List tasks in queue")
    list_parser.add_argument("--json", action="store_true", help="Output in JSON format")

    # clear
    clear_parser = subparsers.add_parser("clear", help="Clear all tasks from queue")
    clear_parser.add_argument("--json", action="store_true", help="Output in JSON format and skip confirmation")

    # logs
    logs_parser = subparsers.add_parser("logs", help="Show recent log entries")
    logs_parser.add_argument("-n", type=int, default=20, help="Number of entries (default: 20)")
    logs_parser.add_argument("--json", action="store_true", help="Output in JSON format")

    # Handle implicit run: tq ./gradlew build -> tq run ./gradlew build
    # Pre-process argv to insert 'run' if needed
    known_subcommands = {"run", "list", "clear", "logs"}
    args_list = sys.argv[1:]

    # Find the first non-option argument (skip --data-dir and its value)
    first_positional_idx = None
    i = 0
    while i < len(args_list):
        arg = args_list[i]
        if arg.startswith("--data-dir") or arg.startswith("--queue-capacity"):
            # Skip --data-dir=value, --queue-capacity=value, or the following value.
            if "=" not in arg:
                i += 1  # Skip the next arg (value)
            i += 1
            continue
        if arg in ("-h", "--help"):
            i += 1
            continue
        # Found first positional argument
        first_positional_idx = i
        break

    # If first positional is not a known subcommand, insert 'run'
    if first_positional_idx is not None and args_list[first_positional_idx] not in known_subcommands:
        args_list.insert(first_positional_idx, "run")

    args = parser.parse_args(args_list)

    if args.command == "run":
        exit_code = cmd_run(args)
        sys.exit(exit_code if exit_code else 0)
    elif args.command == "list":
        cmd_list(args)
    elif args.command == "clear":
        cmd_clear(args)
    elif args.command == "logs":
        cmd_logs(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()