#!/usr/bin/env python3

from __future__ import annotations

import argparse

import json

import os

import re

from dataclasses import asdict, dataclass

from pathlib import Path

from typing import Iterable

DEFAULT_SOCKET = f"/run/user/{os.getuid()}/jcode.sock"

@dataclass

class ProcMem:

pid: int

role: str

cmd: str

session_id: str | None

socket_path: str | None

rss_mb: float

pss_mb: float

anon_mb: float

shared_clean_mb: float

private_clean_mb: float

private_dirty_mb: float

swap_mb: float

@dataclass

class Totals:

count: int

rss_mb: float

pss_mb: float

anon_mb: float

private_dirty_mb: float

shared_clean_mb: float

swap_mb: float

SMAPS_KEYS = {

"Rss": "rss_mb",

"Pss": "pss_mb",

"Anonymous": "anon_mb",

"Shared_Clean": "shared_clean_mb",

"Private_Clean": "private_clean_mb",

"Private_Dirty": "private_dirty_mb",

"Swap": "swap_mb",

}

def parse_args() -> argparse.Namespace:

p = argparse.ArgumentParser(description="Summarize jcode server/client process memory using smaps_rollup")

p.add_argument("--json", action="store_true", help="Print JSON instead of a human summary")

p.add_argument(

"--include-aux",

action="store_true",

help="Include non-default-socket helper/test jcode processes in the output",

)

p.add_argument(

"--socket",

default=DEFAULT_SOCKET,

help=f"Main jcode socket to treat as the primary instance (default: {DEFAULT_SOCKET})",

)

return p.parse_args()

def read_text(path: Path, binary: bool = False) -> str | None:

try:

if binary:

return path.read_bytes().replace(b"\x00", b" ").decode("utf-8", "ignore").strip()

return path.read_text().strip()

except Exception:

return None

def read_argv(path: Path) -> list[str] | None:

try:

raw = path.read_bytes()

except Exception:

return None

if not raw:

return None

return [part.decode("utf-8", "ignore") for part in raw.split(b"\x00") if part]

def parse_socket_path(cmd: str) -> str | None:

m = re.search(r"(?:^| )--socket\s+(\S+)", cmd)

return m.group(1) if m else None

def parse_session_id(cmd: str) -> str | None:

m = re.search(r"--resume\s+(session_[^\s]+)", cmd)

return m.group(1) if m else None

def first_non_option(argv: list[str]) -> str | None:

skip_next = False

for arg in argv[1:]:

if skip_next:

skip_next = False

continue

if arg in {"--socket", "--resume", "-C", "--chdir", "--model", "--provider"}:

skip_next = True

continue

if arg.startswith("-"):

continue

return arg

return None

def classify_process(argv: list[str], cmd: str, main_socket: str) -> tuple[str | None, bool]:

argv0 = Path(argv[0]).name if argv else ""

if not argv0.startswith("jcode"):

return None, False

if "jcode serve" in cmd:

socket_path = parse_socket_path(cmd) or main_socket

if socket_path == main_socket:

return "server", True

return "server_aux", False

socket_path = parse_socket_path(cmd) or main_socket

is_main = socket_path == main_socket

if "cargo build" in cmd or "rustc --crate-name jcode" in cmd or "handle_resume_session" in cmd:

return None, False

subcommand = first_non_option(argv)

if subcommand in {"auth", "login", "logout", "serve", "self-dev"} and "--resume session_" not in cmd:

if subcommand == "self-dev" and "--resume session_" in cmd:

pass

else:

return None, False

if "--resume session_" in cmd or " --fresh-spawn " in cmd:

return ("client_session" if is_main else "client_aux"), is_main

return ("client_interactive" if is_main else "client_aux"), is_main

def parse_smaps_rollup(pid: int) -> dict[str, float] | None:

path = Path(f"/proc/{pid}/smaps_rollup")

txt = read_text(path)

if not txt:

return None

out = {value: 0.0 for value in SMAPS_KEYS.values()}

for line in txt.splitlines():

for key, out_key in SMAPS_KEYS.items():

if line.startswith(f"{key}:"):

parts = line.split()

if len(parts) >= 2:

out[out_key] = round(int(parts[1]) / 1024.0, 1)

return out

def iter_jcode_processes(main_socket: str, include_aux: bool) -> Iterable[ProcMem]:

for pid_dir in Path("/proc").iterdir():

if not pid_dir.name.isdigit():

continue

pid = int(pid_dir.name)

argv = read_argv(pid_dir / "cmdline")

if not argv:

continue

cmd = " ".join(argv)

if "jcode" not in cmd:

continue

role, is_main = classify_process(argv, cmd, main_socket)

if role is None:

continue

if not include_aux and not is_main and role not in {"server"}:

continue

smaps = parse_smaps_rollup(pid)

if not smaps:

continue

yield ProcMem(

pid=pid,

role=role,

cmd=cmd,

session_id=parse_session_id(cmd),

socket_path=parse_socket_path(cmd),

**smaps,

)

def sum_totals(procs: list[ProcMem]) -> Totals:

return Totals(

count=len(procs),

rss_mb=round(sum(p.rss_mb for p in procs), 1),

pss_mb=round(sum(p.pss_mb for p in procs), 1),

anon_mb=round(sum(p.anon_mb for p in procs), 1),

private_dirty_mb=round(sum(p.private_dirty_mb for p in procs), 1),

shared_clean_mb=round(sum(p.shared_clean_mb for p in procs), 1),

swap_mb=round(sum(p.swap_mb for p in procs), 1),

)

def print_human(server: list[ProcMem], clients: list[ProcMem], aux: list[ProcMem]) -> None:

def print_group(name: str, procs: list[ProcMem]) -> None:

totals = sum_totals(procs)

print(f"\n{name} ({totals.count})")

print("-" * len(f"{name} ({totals.count})"))

print(

f"total: RSS {totals.rss_mb:.1f} MB | PSS {totals.pss_mb:.1f} MB | "

f"Anon {totals.anon_mb:.1f} MB | Private dirty {totals.private_dirty_mb:.1f} MB"

)

for p in sorted(procs, key=lambda x: (x.role, -x.pss_mb, x.pid)):

label = p.session_id or p.role

print(

f" pid {p.pid:<7} {label:<48} RSS {p.rss_mb:>6.1f} | PSS {p.pss_mb:>6.1f} | "

f"Anon {p.anon_mb:>6.1f} | PrivDirty {p.private_dirty_mb:>6.1f}"

)

print_group("Server", server)

print_group("Clients", clients)

if aux:

print_group("Auxiliary", aux)

grand = sum_totals(server + clients)

print("\nPrimary total (server + clients)")

print("--------------------------------")

print(

f"RSS {grand.rss_mb:.1f} MB | PSS {grand.pss_mb:.1f} MB | "

f"Anon {grand.anon_mb:.1f} MB | Private dirty {grand.private_dirty_mb:.1f} MB | "

f"Shared clean {grand.shared_clean_mb:.1f} MB"

)

print(

f"Overcount if you sum RSS instead of PSS: {round(grand.rss_mb - grand.pss_mb, 1):.1f} MB"

)

def main() -> int:

args = parse_args()

procs = list(iter_jcode_processes(args.socket, args.include_aux))

server = [p for p in procs if p.role == "server"]

clients = [p for p in procs if p.role.startswith("client_") and p.role != "client_aux"]

aux = [p for p in procs if p.role.endswith("aux")]

if args.json:

payload = {

"socket": args.socket,

"server": [asdict(p) for p in sorted(server, key=lambda x: x.pid)],

"clients": [asdict(p) for p in sorted(clients, key=lambda x: x.pid)],

"auxiliary": [asdict(p) for p in sorted(aux, key=lambda x: x.pid)],

"totals": {

"server": asdict(sum_totals(server)),

"clients": asdict(sum_totals(clients)),

"primary": asdict(sum_totals(server + clients)),

"auxiliary": asdict(sum_totals(aux)),

},

}

print(json.dumps(payload, indent=2))

else:

print_human(server, clients, aux)

return 0

if __name__ == "__main__":

raise SystemExit(main())