session-manager.py

#!/usr/bin/env python3
"""AgentBox session manager — manages tmux sessions + per-session ttyd instances.

Stdlib-only HTTP server on port 8081. No pip dependencies.
Endpoints:
  GET  /api/sessions           — list sessions
  POST /api/sessions           — create session {"name": "..."}
  DELETE /api/sessions/{name}  — destroy session
"""

import json
import os
import re
import signal
import subprocess
import time
from http.server import HTTPServer, BaseHTTPRequestHandler
from pathlib import Path
from threading import Lock

PORT_MIN, PORT_MAX = 9001, 9050
NGINX_SESSIONS_DIR = Path("/etc/nginx/conf.d/sessions")
SESSIONS_META_DIR = Path("/home/agentbox/.sessions")

_lock = Lock()
# port -> session_name mapping
_port_map: dict[int, str] = {}
# session_name -> {"port": int, "pid": int}
_sessions: dict[str, dict] = {}


def _run(cmd: str, user: str | None = None) -> tuple[int, str]:
    """Run a shell command, optionally as a specific user."""
    if user:
        cmd = f"su - {user} -c {_shell_quote(cmd)}"
    r = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    return r.returncode, r.stdout.strip()


def _shell_quote(s: str) -> str:
    return "'" + s.replace("'", "'\\''") + "'"


def _allocate_port() -> int | None:
    for p in range(PORT_MIN, PORT_MAX + 1):
        if p not in _port_map:
            return p
    return None


def _valid_name(name: str) -> bool:
    return bool(re.match(r'^[a-zA-Z0-9][a-zA-Z0-9_-]{0,29}$', name))


def _list_tmux_sessions() -> list[dict]:
    fmt = "#{session_name}|#{session_windows}|#{session_created}|#{session_attached}"
    rc, out = _run(f"tmux list-sessions -F '{fmt}' 2>/dev/null", user="agentbox")
    if rc != 0 or not out:
        return []
    sessions = []
    for line in out.splitlines():
        parts = line.split("|")
        if len(parts) < 4:
            continue
        name = parts[0]
        # Check for copilot running in any pane
        _, panes = _run(f"tmux list-panes -t {_shell_quote(name)} -F '#{{pane_current_command}}' 2>/dev/null", user="agentbox")
        copilot_running = "copilot" in panes.lower() if panes else False

        sessions.append({
            "name": name,
            "windows": int(parts[1]) if parts[1].isdigit() else 0,
            "createdAt": int(parts[2]) if parts[2].isdigit() else 0,
            "attached": int(parts[3]) > 0 if parts[3].isdigit() else False,
            "copilotRunning": copilot_running,
            "ttydPort": _sessions.get(name, {}).get("port"),
        })
    return sessions


def _write_nginx_conf(name: str, port: int):
    """Write nginx location fragment for a session-specific ttyd."""
    NGINX_SESSIONS_DIR.mkdir(parents=True, exist_ok=True)
    conf = f"""location /terminal/sessions/{name}/ {{
    if ($agentbox_auth_ok = 0) {{
        return 403 "Forbidden: missing or invalid proxy token";
    }}
    proxy_pass http://127.0.0.1:{port}/;
    proxy_set_header Host $host;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection upgrade;
    proxy_http_version 1.1;
    sub_filter_types text/html;
    sub_filter_once off;
    sub_filter "</head>" "<meta name=\\"viewport\\" content=\\"width=device-width,initial-scale=1.0,maximum-scale=1.0,interactive-widget=resizes-content\\"></head>";
    sub_filter "</body>" "<script src=\\"../../../mobile-keys.js\\"></script></body>";
    proxy_set_header Accept-Encoding "";
}}
"""
    (NGINX_SESSIONS_DIR / f"{name}.conf").write_text(conf)
    subprocess.run(["nginx", "-s", "reload"], capture_output=True)


def _remove_nginx_conf(name: str):
    conf = NGINX_SESSIONS_DIR / f"{name}.conf"
    if conf.exists():
        conf.unlink()
    subprocess.run(["nginx", "-s", "reload"], capture_output=True)


def _spawn_ttyd(name: str, port: int) -> int | None:
    """Spawn a ttyd process for a specific tmux session. Returns PID."""
    cmd = [
        "su", "-", "agentbox", "-c",
        f"ttyd --writable -p {port} -i 127.0.0.1 tmux attach-session -t {name}"
    ]
    proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    return proc.pid


def _stop_ttyd(name: str):
    info = _sessions.get(name)
    if info and info.get("pid"):
        try:
            os.kill(info["pid"], signal.SIGTERM)
        except ProcessLookupError:
            pass


def _save_meta(name: str, data: dict):
    SESSIONS_META_DIR.mkdir(parents=True, exist_ok=True)
    (SESSIONS_META_DIR / f"{name}.json").write_text(json.dumps(data))


def _remove_meta(name: str):
    meta = SESSIONS_META_DIR / f"{name}.json"
    if meta.exists():
        meta.unlink()


def create_session(name: str) -> tuple[int, dict]:
    with _lock:
        if name in _sessions:
            return 409, {"error": f"Session '{name}' already exists"}
        if not _valid_name(name):
            return 400, {"error": "Invalid name (alphanumeric/hyphens, 1-30 chars)"}

        port = _allocate_port()
        if port is None:
            return 503, {"error": "No ports available"}

        # Create tmux session
        rc, _ = _run(f"tmux new-session -d -s {_shell_quote(name)} -c /home/agentbox/workspace", user="agentbox")
        if rc != 0:
            return 500, {"error": "Failed to create tmux session"}

        # Spawn ttyd
        pid = _spawn_ttyd(name, port)
        time.sleep(0.3)

        _port_map[port] = name
        _sessions[name] = {"port": port, "pid": pid}
        _save_meta(name, {"port": port, "createdAt": int(time.time())})

        # Write nginx config and reload
        _write_nginx_conf(name, port)

        return 201, {"name": name, "port": port, "status": "active"}


def delete_session(name: str) -> tuple[int, dict]:
    with _lock:
        _stop_ttyd(name)

        # Kill tmux session
        _run(f"tmux kill-session -t {_shell_quote(name)} 2>/dev/null", user="agentbox")

        info = _sessions.pop(name, None)
        if info:
            _port_map.pop(info["port"], None)

        _remove_nginx_conf(name)
        _remove_meta(name)

        return 200, {"deleted": name}


def _recover_existing_sessions():
    """On startup, detect existing tmux sessions and spawn ttyd for each."""
    sessions = _list_tmux_sessions()
    for s in sessions:
        name = s["name"]
        if name in _sessions:
            continue
        port = _allocate_port()
        if port is None:
            break
        pid = _spawn_ttyd(name, port)
        _port_map[port] = name
        _sessions[name] = {"port": port, "pid": pid}
        _write_nginx_conf(name, port)
        time.sleep(0.2)


class Handler(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        pass  # Suppress request logging

    def _send_json(self, status: int, data):
        body = json.dumps(data).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _read_body(self) -> dict:
        length = int(self.headers.get("Content-Length", 0))
        if length == 0:
            return {}
        return json.loads(self.rfile.read(length))

    def do_GET(self):
        if self.path.rstrip("/") == "/api/sessions":
            sessions = _list_tmux_sessions()
            self._send_json(200, sessions)
        else:
            self._send_json(404, {"error": "Not found"})

    def do_POST(self):
        if self.path.rstrip("/") == "/api/sessions":
            body = self._read_body()
            name = body.get("name", "").strip()
            if not name:
                self._send_json(400, {"error": "name is required"})
                return
            status, data = create_session(name)
            self._send_json(status, data)
        else:
            self._send_json(404, {"error": "Not found"})

    def do_DELETE(self):
        # DELETE /api/sessions/{name}
        parts = self.path.rstrip("/").split("/")
        if len(parts) == 4 and parts[1] == "api" and parts[2] == "sessions":
            name = parts[3]
            status, data = delete_session(name)
            self._send_json(status, data)
        else:
            self._send_json(404, {"error": "Not found"})


if __name__ == "__main__":
    NGINX_SESSIONS_DIR.mkdir(parents=True, exist_ok=True)
    SESSIONS_META_DIR.mkdir(parents=True, exist_ok=True)

    _recover_existing_sessions()

    server = HTTPServer(("127.0.0.1", 8081), Handler)
    print(f"Session manager listening on port 8081")
    server.serve_forever()