jan-cli 0.24.0

#!/usr/bin/env python3
"""Jan warm Kotlin/JVM worker: run compiled main classes in isolated subprocesses.

The worker itself stays warm (socket + classpath string ready). Each job is a
fresh `java`/`kotlin` child so a crash cannot kill the daemon. `.kts` script
jobs are rejected so the client can cold-fall back to `kotlinc -script`.
"""

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")
    argv = job.get("argv") or []

    if kind == "kts":
        return (
            2,
            b"",
            b"kotlin worker does not run .kts; use cold kotlinc -script\n",
        )

    # Job carries a full argv already prepared by jan (java/kotlin + -cp + main).
    if kind == "argv":
        cmd = list(argv)
        if not cmd:
            return 2, b"", b"empty kotlin argv\n"
    elif kind == "main":
        classpath = source.get("classpath") or ""
        main_class = source.get("main_class") or ""
        java = source.get("java") or "java"
        cmd = [java]
        if classpath:
            cmd += ["-cp", classpath]
        cmd.append(main_class)
        cmd.extend(argv)
    else:
        return 2, b"", f"unknown source kind {kind!r}\n".encode()

    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: kotlin_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)