jan-cli 0.27.0

#!/usr/bin/env python3
"""Jan warm shell worker: spawn isolated shell children per job."""

from __future__ import annotations

import base64
import json
import os
import signal
import socket
import subprocess
import sys
import traceback


def _apply_env(env: dict) -> dict:
    out = {}
    for k, v in (env or {}).items():
        if v is None:
            continue
        out[str(k)] = str(v)
    return out


def _run_job(job: dict) -> tuple[int, bytes, bytes]:
    cwd = job.get("cwd") or os.getcwd()
    env = _apply_env(job.get("env") or {})
    source = job.get("source") or {}
    kind = source.get("kind")
    value = source.get("value") or ""
    shell = (job.get("shell") or "bash").lower()
    argv = job.get("argv") or []

    if shell == "zsh":
        prog = "zsh"
        # Skip user/system zshrcs for speed + isolation from interactive config.
        base = [prog, "-f"]
    elif shell == "sh":
        prog = "sh"
        base = [prog]
    else:
        prog = "bash"
        base = [prog, "--norc", "--noprofile"]

    if kind == "path":
        cmd = base + [value] + list(argv)
    elif kind == "inline":
        # Match jan's insert of $0 before user args for -c forms.
        argv0 = job.get("argv0") or "jan"
        c_flag = "-lc" if shell == "bash" else "-c"
        # zsh -f already set; for -c use plain -c
        if shell == "bash":
            cmd = ["bash", "--norc", "--noprofile", c_flag, value, argv0] + list(argv)
        elif shell == "zsh":
            cmd = ["zsh", "-f", "-c", value, argv0] + list(argv)
        else:
            cmd = ["sh", "-c", value, argv0] + list(argv)
    else:
        return 2, b"", f"unknown source kind {kind!r}\n".encode()

    # Full child process — crash/exit cannot take down the daemon.
    proc = subprocess.run(
        cmd,
        cwd=cwd,
        env=env if env else None,
        capture_output=True,
        check=False,
    )
    return proc.returncode, proc.stdout or b"", proc.stderr or b""


def _handle_client(conn: socket.socket) -> None:
    buf = b""
    while b"\n" not in buf:
        chunk = conn.recv(65536)
        if not chunk:
            return
        buf += chunk
        if len(buf) > 16 * 1024 * 1024:
            conn.sendall(b'{"ok":false,"error":"request too large"}\n')
            return
    line, _ = buf.split(b"\n", 1)
    try:
        job = json.loads(line.decode("utf-8"))
    except Exception as e:
        conn.sendall(
            json.dumps({"ok": False, "error": f"invalid json: {e}"}).encode() + b"\n"
        )
        return

    try:
        code, out_b, err_b = _run_job(job)
        resp = {
            "ok": True,
            "exit_code": int(code) & 0xFF if code is not None else 255,
            "stdout_b64": base64.b64encode(out_b).decode("ascii"),
            "stderr_b64": base64.b64encode(err_b).decode("ascii"),
        }
    except Exception:
        resp = {
            "ok": True,
            "exit_code": 1,
            "stdout_b64": "",
            "stderr_b64": base64.b64encode(traceback.format_exc().encode()).decode(
                "ascii"
            ),
        }
    conn.sendall(json.dumps(resp).encode("utf-8") + b"\n")


def main() -> int:
    if len(sys.argv) < 2:
        print("usage: shell_worker.py <sock-path>", file=sys.stderr)
        return 2
    sock_path = sys.argv[1]
    try:
        os.unlink(sock_path)
    except FileNotFoundError:
        pass

    signal.signal(signal.SIGPIPE, signal.SIG_IGN)
    srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    srv.bind(sock_path)
    os.chmod(sock_path, 0o600)
    srv.listen(64)
    sys.stdout.write(f"ready {sock_path}\n")
    sys.stdout.flush()

    while True:
        conn, _ = srv.accept()
        try:
            _handle_client(conn)
        except Exception:
            traceback.print_exc()
        finally:
            try:
                conn.close()
            except Exception:
                pass


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