claude-hr

#!/usr/bin/env python3
"""
claude-hr — Launch Claude Code through the Headroom compression proxy.

Cross-platform (macOS / Linux / Windows) rewrite of the original zsh function.

What it does, in order:
  1. Resolve a proxy PORT:
       --port N / -p N        explicit
       else walk up dirs for  .claude-hr-port
       else pick a free port, ask (or auto-pick), save it to the repo root
  2. Make sure a Headroom proxy is listening on that port (start one if not).
  3. Run:  headroom wrap claude --port N --no-proxy -- --dangerously-skip-permissions <args>

With no resolvable port it falls back to Headroom's own managed proxy:
       headroom wrap claude --memory -- --dangerously-skip-permissions <args>

Usage:
  claude-hr [--port N] [claude args...]
"""

import os
import random
import shutil
import signal
import socket
import subprocess
import sys
import time

PORT_FILE = ".claude-hr-port"


def which(name: str) -> str:
    return shutil.which(name) or name


def port_in_use(port: int) -> bool:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.settimeout(0.3)
        return s.connect_ex(("127.0.0.1", port)) == 0


def free_port() -> int:
    while True:
        p = random.randint(10000, 65535)
        if not port_in_use(p):
            return p


def find_port_file(start: str) -> str | None:
    """Walk up from `start` looking for a .claude-hr-port file."""
    d = os.path.abspath(start)
    while True:
        f = os.path.join(d, PORT_FILE)
        if os.path.isfile(f):
            return f
        parent = os.path.dirname(d)
        if parent == d:
            return None
        d = parent


def git_root() -> str:
    try:
        out = subprocess.run(["git", "rev-parse", "--show-toplevel"],
                             capture_output=True, text=True)
        if out.returncode == 0 and out.stdout.strip():
            return out.stdout.strip()
    except Exception:
        pass
    return os.getcwd()


def start_proxy(port: int) -> None:
    """Spawn `headroom proxy --port N` detached, log to temp dir."""
    import tempfile
    log = open(os.path.join(tempfile.gettempdir(),
                            f"headroom-proxy-{port}.log"), "ab")
    kwargs: dict = {"stdout": log, "stderr": log, "stdin": subprocess.DEVNULL}
    if os.name == "nt":
        kwargs["creationflags"] = (subprocess.CREATE_NEW_PROCESS_GROUP
                                   | subprocess.DETACHED_PROCESS)
    else:
        kwargs["start_new_session"] = True
    print(f"[hr] starting proxy on port {port}...")
    subprocess.Popen([which("headroom"), "proxy", "--port", str(port)], **kwargs)
    # give it a moment to bind
    import time
    for _ in range(20):
        if port_in_use(port):
            break
        time.sleep(0.1)


def _listener_pid(port: int) -> str | None:
    """PID listening on 127.0.0.1:port (POSIX only)."""
    if os.name == "nt":
        return None
    try:
        out = subprocess.run(["lsof", f"-tiTCP:{port}", "-sTCP:LISTEN"],
                             capture_output=True, text=True)
    except FileNotFoundError:
        return None
    pids = out.stdout.split()
    return pids[0] if pids else None


def _other_client_pids(port: int, listener: str | None) -> list[str]:
    """PIDs with an ESTABLISHED conn to the port, excluding the proxy's own listener.
    On loopback each connection shows two rows; the server-side row is owned by the
    listener, so excluding it leaves only real clients (claude, VS Code, ...)."""
    if os.name == "nt":
        return []
    try:
        out = subprocess.run(["lsof", "-nP", f"-iTCP:{port}", "-sTCP:ESTABLISHED"],
                             capture_output=True, text=True)
    except FileNotFoundError:
        return []
    pids = set()
    for line in out.stdout.splitlines()[1:]:
        parts = line.split()
        if len(parts) > 1 and parts[1] != listener:
            pids.add(parts[1])
    return list(pids)


def stop_proxy_if_idle(port: int) -> None:
    """When Claude exits, stop this port's proxy unless another client is still attached.
    ponytail: POSIX-only (needs lsof); on Windows this no-ops — stop proxies with `hr-proxy stop`."""
    if os.name == "nt":
        return
    time.sleep(0.3)  # let the closed TCP socket clear
    listener = _listener_pid(port)
    if not listener:
        return
    if _other_client_pids(port, listener):
        return  # someone else (another session, editor) still using it
    print(f"[hr] no other clients on port {port} — stopping proxy")
    r = subprocess.run([which("hr-proxy"), "stop", str(port)],
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if r.returncode != 0:
        try:
            os.kill(int(listener), signal.SIGTERM)
        except (ValueError, ProcessLookupError, PermissionError):
            pass


def resolve_port(args: list[str]) -> tuple[int | None, list[str]]:
    """Return (port, remaining_args). port is None → use Headroom-managed proxy."""
    port: int | None = None
    rest: list[str] = []
    i = 0
    while i < len(args):
        a = args[i]
        if a.startswith("--port="):
            port = int(a.split("=", 1)[1]); i += 1
        elif a in ("--port", "-p") and i + 1 < len(args):
            port = int(args[i + 1]); i += 2
        else:
            rest.append(a); i += 1

    if port is not None:
        return port, rest

    # 2. .claude-hr-port up the tree
    pf = find_port_file(os.getcwd())
    if pf:
        try:
            with open(pf) as f:
                return int(f.read().strip()), rest
        except (ValueError, OSError):
            pass

    # 3. no port anywhere → pick one, persist to repo root
    root = git_root()
    suggested = free_port()
    if sys.stdin.isatty():
        entered = input(
            f"[hr] no {PORT_FILE} in {root} (or parents). "
            f"Enter port [{suggested}]: ").strip()
        port = int(entered) if entered.isdigit() else suggested
    else:
        port = suggested
        print(f"[hr] no {PORT_FILE} found; using random port {port}")

    if 1 <= port <= 65535:
        try:
            with open(os.path.join(root, PORT_FILE), "w") as f:
                f.write(str(port))
            print(f"[hr] saved port {port} to {os.path.join(root, PORT_FILE)}")
        except OSError as e:
            print(f"[hr] could not save port file: {e}")
        return port, rest

    print(f"[hr] invalid port {port!r}")
    return None, rest


def main() -> int:
    port, rest = resolve_port(sys.argv[1:])
    headroom = which("headroom")

    if port is not None:
        if not port_in_use(port):
            start_proxy(port)
        cmd = [headroom, "wrap", "claude", "--port", str(port), "--no-proxy",
               "--no-serena", "--", "--dangerously-skip-permissions", *rest]
    else:
        cmd = [headroom, "wrap", "claude", "--memory", "--no-serena",
               "--", "--dangerously-skip-permissions", *rest]

    rc = subprocess.run(cmd).returncode
    if port is not None:
        stop_proxy_if_idle(port)
    return rc


if __name__ == "__main__":
    sys.exit(main())