jan-cli 0.24.0

#!/usr/bin/env python3
"""Jan warm Python worker: fork-per-job isolation over a Unix socket."""

from __future__ import annotations

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


def _apply_env(env: dict) -> None:
    os.environ.clear()
    for k, v in env.items():
        if v is None:
            continue
        os.environ[str(k)] = str(v)


def _run_job(job: dict) -> int:
    cwd = job.get("cwd") or os.getcwd()
    os.chdir(cwd)
    env = job.get("env") or {}
    _apply_env(env)

    argv = job.get("argv") or []
    source = job.get("source") or {}
    kind = source.get("kind")
    value = source.get("value") or ""

    # Match `python3 script.py a b` / `python3 -c code a b` argv layout.
    if kind == "path":
        sys.argv = [value] + list(argv)
        runpy.run_path(value, run_name="__main__")
    elif kind == "inline":
        sys.argv = ["-c"] + list(argv)
        # Fresh globals per job (also fresh because we are in a forked child).
        g = {"__name__": "__main__", "__builtins__": __builtins__}
        exec(compile(value, "<jan-inline>", "exec"), g, g)
    else:
        print(f"jan python worker: unknown source kind {kind!r}", file=sys.stderr)
        return 2
    return 0


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, _rest = 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

    r_out, w_out = os.pipe()
    r_err, w_err = os.pipe()
    pid = os.fork()
    if pid == 0:
        # Child: run the job with captured stdio.
        try:
            os.close(r_out)
            os.close(r_err)
            os.dup2(w_out, 1)
            os.dup2(w_err, 2)
            os.close(w_out)
            os.close(w_err)
            conn.close()
            # Pipes are block-buffered; replace stdio with line-buffered wrappers.
            sys.stdout = os.fdopen(1, "w", buffering=1)
            sys.stderr = os.fdopen(2, "w", buffering=1)
            code = _run_job(job)
            try:
                sys.stdout.flush()
                sys.stderr.flush()
            except Exception:
                pass
            os._exit(code & 0xFF)
        except SystemExit as e:
            code = e.code if isinstance(e.code, int) else (0 if e.code is None else 1)
            os._exit(code & 0xFF)
        except BaseException:
            traceback.print_exc()
            os._exit(1)

    # Parent
    os.close(w_out)
    os.close(w_err)

    def _read_all(fd: int) -> bytes:
        chunks = []
        while True:
            data = os.read(fd, 65536)
            if not data:
                break
            chunks.append(data)
        return b"".join(chunks)

    stdout_b = _read_all(r_out)
    stderr_b = _read_all(r_err)
    os.close(r_out)
    os.close(r_err)
    _pid, status = os.waitpid(pid, 0)
    if os.WIFEXITED(status):
        exit_code = os.WEXITSTATUS(status)
    elif os.WIFSIGNALED(status):
        exit_code = 128 + os.WTERMSIG(status)
    else:
        exit_code = 255

    resp = {
        "ok": True,
        "exit_code": exit_code,
        "stdout_b64": base64.b64encode(stdout_b).decode("ascii"),
        "stderr_b64": base64.b64encode(stderr_b).decode("ascii"),
    }
    conn.sendall(json.dumps(resp).encode("utf-8") + b"\n")


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

    # Ignore SIGCHLD leftovers from waitpid path; we wait explicitly.
    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)
    # Ready marker for the supervisor.
    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)