tui: add stacked RAM-breakdown bar for local providers
ober
7765f99b96aa7eac67023ac011effb2f2786fe93
--- a/build-binary.ss +++ b/build-binary.ss @@ -128,6 +128,7 @@ "lib/jcode/ui/tui-message" "lib/jcode/ui/tui-status" "lib/jcode/ui/tui-input" + "lib/jcode/ui/tui-memstats" "lib/jcode/ui/tui-sidebar" "lib/jcode/ui/tui-dialog" "lib/jcode/ui/tui-toast" --- a/src/jcode/core/agent.ss +++ b/src/jcode/core/agent.ss @@ -7,7 +7,8 @@ current-tool-cb current-usage-cb current-provider-override - current-model-override) + current-model-override + get-current-provider) (import :std/text/json :std/misc/thread --- a/src/jcode/mcp/client.ss +++ b/src/jcode/mcp/client.ss @@ -5,13 +5,15 @@ (export init-mcp-tools mcp-conn? mcp-conn-name + mcp-conn-pid mcp-start mcp-initialize mcp-list-tools mcp-call-tool mcp-stop! mcp-stop-all! - mcp-active-servers) + mcp-active-servers + mcp-server-pids) (import :std/text/json :std/misc/string @@ -75,6 +77,16 @@ (set! *mcp-servers* '()) (set! *mcp-tool-counts* (make-hash-table))) +(def (mcp-server-pids) + "Return a list of PIDs for all active MCP server subprocesses." + (let lp ((cs *mcp-servers*) (acc '())) + (cond + ((null? cs) acc) + (else + (let ((pid (mcp-conn-pid (car cs)))) + (lp (cdr cs) + (if (and pid (number? pid)) (cons pid acc) acc))))))) + (def (mcp-active-servers) "Return list of (name . tool-count) for active MCP servers. Reads the cached count populated at init time — never does a --- a/src/jcode/provider/provider.ss +++ b/src/jcode/provider/provider.ss @@ -7,7 +7,8 @@ provider-list-models provider-list-pricing provider-name - provider-model) + provider-model + provider-base-url) (import :std/text/json :std/net/request new file mode 100644 --- /dev/null +++ b/src/jcode/ui/tui-memstats.ss @@ -0,0 +1,314 @@ +;;; jcode TUI memstats — RAM breakdown for local model providers +;;; Tracks model weights, KV cache (estimated), inference-server "other" +;;; (activations + framework), jcode RSS, MCP children RSS, system free. +;;; Active only when the configured provider is local (mlx/ollama at a +;;; loopback host). On remote providers, memstats-active? returns #f and +;;; the sidebar suppresses the bar. + +(export + make-fresh-memstats memstats? + memstats-update! memstats-active? + memstats-segments memstats-total-bytes memstats-used-bytes + memstats-set-prompt-tokens! + memstats-server-kind memstats-prompt-tokens + memstats-weight-bytes memstats-kv-bytes + memstats-server-rss memstats-own-rss memstats-mcps-rss) + +(import :jerboa/core + :jerboa/runtime + :std/os/path + :std/os/shell + :std/text/json + :std/misc/string + :jcode/core/log + :jcode/core/agent + :jcode/provider/provider + :jcode/mcp/client) + +;; ---- Struct ---- + +(defstruct memstats + (own-pid ;; jcode pid (cached) + server-pid ;; mlx/ollama PID, refreshed on death + server-host ;; "127.0.0.1" or #f + server-port ;; integer or #f + server-kind ;; 'mlx | 'ollama | #f + model-dir ;; "/path/.../Qwen3-Coder-..." or #f + weight-bytes ;; sum of *.safetensors (or 0) + kv-bytes-per-token ;; from config.json (or 0) + prompt-tokens ;; latest from usage callback + own-rss ;; bytes + server-rss ;; bytes + mcps-rss ;; bytes (sum) + total-bytes ;; system total + last-error) ;; last failure cause for debugging + transparent: #t) + +(def (make-fresh-memstats) + (let ((pid (try (get-process-id) (catch (e) 0)))) + (make-memstats pid #f #f #f #f #f 0 0 0 0 0 0 0 #f))) + +(def (memstats-active? m) + (and m + (memstats-server-pid m) + (> (memstats-weight-bytes m) 0))) + +(def (memstats-set-prompt-tokens! m n) + (when (and m (number? n) (>= n 0)) + (memstats-prompt-tokens-set! m n))) + +(def (memstats-kv-bytes m) + (* (memstats-kv-bytes-per-token m) + (memstats-prompt-tokens m))) + +(def (memstats-used-bytes m) + (+ (memstats-server-rss m) + (memstats-own-rss m) + (memstats-mcps-rss m))) + +;; ---- Public segments ---- +;; Returns list of (label face-name bytes) — ordered for stacked render. +;; Faces are looked up by tui-theme; tui-sidebar provides defaults if a +;; theme doesn't define them. + +(def (memstats-segments m) + ;; Clamp the breakdown so the three server-side slices sum to server-RSS. + ;; Computed weights and KV are estimates; for an idle server with no + ;; activations, RSS is mostly weights and `other` rounds to 0. During + ;; generation, activations push RSS up, eating into `other` first. + (let* ((srv-rss (memstats-server-rss m)) + (raw-weights (memstats-weight-bytes m)) + (raw-kv (memstats-kv-bytes m)) + (weights (min srv-rss raw-weights)) + (after-w (- srv-rss weights)) + (kv (min after-w raw-kv)) + (other (- after-w kv)) + (jc (memstats-own-rss m)) + (mcps (memstats-mcps-rss m))) + ;; Face names are existing theme faces, picked for color contrast: + ;; syntax-keyword blue — model weights (dominant, static) + ;; syntax-type cyan — KV cache (grows with context) + ;; code-block green — server activations / framework + ;; heading yellow — jcode itself + ;; syntax-string orange — MCP children + (list + (list "weights" 'syntax-keyword weights) + (list "kv" 'syntax-type kv) + (list "other" 'code-block other) + (list "jcode" 'heading jc) + (list "mcp" 'syntax-string mcps)))) + +;; ---- Update ---- + +(def (memstats-update! m) + (try + (memstats-total-bytes-set! m (host-total-bytes)) + (memstats-own-rss-set! m (rss-of-pid (memstats-own-pid m))) + (memstats-mcps-rss-set! m (sum-mcp-rss)) + (resolve-server! m) + (when (memstats-server-pid m) + (memstats-server-rss-set! m (rss-of-pid (memstats-server-pid m)))) + (when (and (memstats-server-pid m) + (not (memstats-model-dir m))) + (resolve-model! m)) + (catch (e) + (memstats-last-error-set! m (try (err->string e) (catch (_) "?")))))) + +(def (resolve-server! m) + (let ((pid (memstats-server-pid m))) + (when (and pid (not (pid-alive? pid))) + (memstats-server-pid-set! m #f) + (memstats-model-dir-set! m #f) + (memstats-weight-bytes-set! m 0) + (memstats-kv-bytes-per-token-set! m 0) + (memstats-server-rss-set! m 0))) + (unless (memstats-server-pid m) + (let-values (((host port kind) (current-local-endpoint))) + (when (and host port (> port 0)) + (memstats-server-host-set! m host) + (memstats-server-port-set! m port) + (memstats-server-kind-set! m kind) + (let ((p (pid-listening-on port))) + (when p (memstats-server-pid-set! m p))))))) + +(def (resolve-model! m) + (let* ((pid (memstats-server-pid m)) + (cmd (and pid (cmdline-of-pid pid))) + (dir (and cmd (extract-model-arg cmd)))) + (when (and dir (file-exists? dir)) + (memstats-model-dir-set! m dir) + (memstats-weight-bytes-set! m (sum-safetensors dir)) + (memstats-kv-bytes-per-token-set! m + (or (kv-bytes-per-token-from dir) 0))))) + +;; ---- Provider URL parsing ---- + +(def (current-local-endpoint) + ;; Returns (values host port kind-symbol) or (values #f #f #f). + (try + (let* ((p (get-current-provider)) + (name (provider-name p)) + (url (provider-base-url p))) + (cond + ((not (or (equal? name "mlx") (equal? name "ollama"))) + (values #f #f #f)) + (else + (let-values (((host port) (parse-host-port url))) + (cond + ((not (loopback? host)) (values #f #f #f)) + (else (values host port (string->symbol name)))))))) + (catch (e) (values #f #f #f)))) + +(def (parse-host-port url) + ;; "http://127.0.0.1:8080/v1" -> (values "127.0.0.1" 8080) + (let* ((after-scheme + (cond + ((string-prefix? "http://" url) (substring url 7 (string-length url))) + ((string-prefix? "https://" url) (substring url 8 (string-length url))) + (else url))) + (slash (string-index after-scheme #\/)) + (hostport (if slash (substring after-scheme 0 slash) after-scheme)) + (colon (last-index-of hostport #\:))) + (if colon + (let ((h (substring hostport 0 colon)) + (p (string->number (substring hostport (+ colon 1) (string-length hostport))))) + (values h (or p 0))) + (values hostport 0)))) + +(def (last-index-of s ch) + (let lp ((i (- (string-length s) 1))) + (cond + ((< i 0) #f) + ((char=? (string-ref s i) ch) i) + (else (lp (- i 1)))))) + +(def (loopback? host) + (and host + (or (equal? host "127.0.0.1") + (equal? host "localhost") + (equal? host "::1") + (equal? host "0.0.0.0")))) + +;; ---- shell-out helpers ---- + +(def (run-cmd cmd) + (try + (let-values (((out err code) (shell/status cmd))) + (if (= code 0) (string-trim out) "")) + (catch (e) ""))) + +(def (pid-listening-on port) + (let* ((cmd (format "lsof -nP -iTCP:~a -sTCP:LISTEN -t 2>/dev/null | head -1" port)) + (out (run-cmd cmd))) + (and (not (string-empty? out)) + (string->number out)))) + +(def (pid-alive? pid) + (try + (let-values (((out err code) (shell/status (format "kill -0 ~a 2>/dev/null" pid)))) + (= code 0)) + (catch (e) #f))) + +(def (rss-of-pid pid) + ;; ps -o rss returns kibibytes + (let ((out (run-cmd (format "ps -o rss= -p ~a 2>/dev/null" pid)))) + (cond + ((string-empty? out) 0) + (else (let ((kb (string->number (string-trim out)))) + (if kb (* kb 1024) 0)))))) + +(def (cmdline-of-pid pid) + ;; macOS ps truncates 'command' by default; pass ww (or use -ww in `ps -ww -o command=`) + ;; -ww = unlimited width on macOS and Linux. + (let ((out (run-cmd (format "ps -ww -o command= -p ~a 2>/dev/null" pid)))) + (and (not (string-empty? out)) out))) + +(def (extract-model-arg cmd) + (let lp ((toks (tokenize-cmd cmd))) + (cond + ((null? toks) #f) + ((or (equal? (car toks) "--model") + (equal? (car toks) "--model-path")) + (and (pair? (cdr toks)) (cadr toks))) + ((string-prefix? "--model=" (car toks)) + (substring (car toks) 8 (string-length (car toks)))) + ((string-prefix? "--model-path=" (car toks)) + (substring (car toks) 13 (string-length (car toks)))) + (else (lp (cdr toks)))))) + +(def (tokenize-cmd s) + (filter (lambda (x) (not (string-empty? x))) + (string-split s #\space))) + +;; ---- Host total memory ---- + +(def (host-total-bytes) + (let ((mac (run-cmd "sysctl -n hw.memsize 2>/dev/null"))) + (cond + ((and (not (string-empty? mac)) (string->number mac)) + (string->number mac)) + (else + (let ((linux (run-cmd "awk '/MemTotal:/ {print $2*1024}' /proc/meminfo 2>/dev/null"))) + (or (and (not (string-empty? linux)) (string->number linux)) + 0)))))) + +;; ---- MCP RSS ---- + +(def (sum-mcp-rss) + (try + (let lp ((pids (mcp-server-pids)) (acc 0)) + (cond + ((null? pids) acc) + (else (lp (cdr pids) (+ acc (rss-of-pid (car pids))))))) + (catch (e) 0))) + +;; ---- Model dir scanning ---- + +(def (sum-safetensors dir) + (try + (let lp ((files (directory-list dir)) (acc 0)) + (cond + ((null? files) acc) + ((string-suffix? ".safetensors" (car files)) + (lp (cdr files) (+ acc (file-size (path-join dir (car files)))))) + (else (lp (cdr files) acc)))) + (catch (e) 0))) + +(def (file-size path) + (try + (let* ((p (open-file-input-port path)) + (sz (port-length p))) + (close-port p) + (or sz 0)) + (catch (e) 0))) + +;; ---- KV per-token from config.json ---- +;; Formula: 2 (K and V) * num_layers * num_kv_heads * head_dim * dtype_bytes. +;; mlx_lm keeps the KV cache in fp16 even for 4-bit quantized weights, so we +;; assume 2 bytes/element. This is the dominant approximation; if a future +;; mlx_lm release exposes a /v1/stats endpoint we should prefer that. + +(def (kv-bytes-per-token-from dir) + (try + (let* ((cfg (path-join dir "config.json")) + (h (and (file-exists? cfg) (read-json-file cfg))) + (layers (and h (hash-get h "num_hidden_layers"))) + (kv-heads (and h (or (hash-get h "num_key_value_heads") + (hash-get h "num_attention_heads")))) + (head-dim (and h (or (hash-get h "head_dim") + (and (hash-get h "hidden_size") + (hash-get h "num_attention_heads") + (> (hash-get h "num_attention_heads") 0) + (quotient (hash-get h "hidden_size") + (hash-get h "num_attention_heads")))))) + (dtype-bytes 2)) + (and layers kv-heads head-dim + (number? layers) (number? kv-heads) (number? head-dim) + (* 2 layers kv-heads head-dim dtype-bytes))) + (catch (e) #f))) + +(def (read-json-file path) + (let* ((p (open-input-file path)) + (obj (read-json p))) + (close-port p) + obj)) --- a/src/jcode/ui/tui-sidebar.ss +++ b/src/jcode/ui/tui-sidebar.ss @@ -14,6 +14,7 @@ :jerboa/core :jerboa/runtime :jcode/ui/tui-theme + :jcode/ui/tui-memstats :std/os/sysmon) ;; ---- Sidebar state ---- @@ -32,7 +33,7 @@ ;; ---- Rendering ---- -(def (render-sidebar! sidebar sysmon x y width height) +(def (render-sidebar! sidebar sysmon memstats x y width height) "Render the sidebar at (x, y) with given width and height." (let ((bg (face-bg-attr 'sidebar-bg)) (fg (face-fg-attr 'sidebar-item)) @@ -76,7 +77,7 @@ ;; System section (let ((row (+ row 1))) (let ((row (render-section! "System" cx row cw max-row tfg tbg))) - (render-system! sysmon cx row cw max-row fg bg)))))))))))))))))) + (render-system! sysmon memstats cx row cw max-row fg bg)))))))))))))))))) (def (render-section! title x row width max-row tfg tbg) (when (< row max-row) @@ -153,7 +154,7 @@ ;; ---- System utilization (CPU/MEM/GPU bars) ---- -(def (render-system! sysmon x row width max-row fg bg) +(def (render-system! sysmon memstats x row width max-row fg bg) (cond ((not sysmon) row) ((>= row max-row) row) @@ -161,13 +162,74 @@ (render-bar! x row width "CPU" (sysmon-cpu-percent sysmon) fg bg) (set! row (+ row 1)) (when (< row max-row) - (render-bar! x row width "MEM" (sysmon-mem-percent sysmon) fg bg) + (cond + ;; Local provider with weights+KV resolved — show stacked breakdown. + ((and memstats (memstats-active? memstats)) + (render-memstats-bar! x row width memstats fg bg)) + (else + (render-bar! x row width "MEM" (sysmon-mem-percent sysmon) fg bg))) (set! row (+ row 1))) (when (< row max-row) (render-bar! x row width "GPU" (sysmon-gpu-percent sysmon) fg bg) (set! row (+ row 1))) row))) +;; ---- Stacked memory-breakdown bar (local provider only) ---- +;; Layout matches render-bar!: +;; "MEM " + colored-blocks-per-segment + dim-blocks-for-free + " 21G " +;; Segments come from (memstats-segments mon) in fixed order. +(def (render-memstats-bar! x y total-width memstats fg bg) + (let* ((label-w 4) + (suffix-w 5) + (bar-w (max 1 (- total-width label-w suffix-w))) + (total (max 1 (memstats-total-bytes memstats))) + (used (max 0 (min total (memstats-used-bytes memstats)))) + (used-cells (max 0 (min bar-w + (exact (round (* bar-w (/ used total))))))) + (free-cells (- bar-w used-cells)) + (segs (memstats-segments memstats)) + (seg-total (max 1 (sum-seg-bytes segs)))) + (tb-print! x y fg bg "MEM ") + ;; Paint segments proportional to their share of `used-cells`. + (let lp ((segs segs) (col (+ x label-w)) (remaining used-cells)) + (cond + ((or (null? segs) (= remaining 0)) (void)) + (else + (let* ((seg (car segs)) + (face (cadr seg)) + (sb (caddr seg)) + (cells (cond + ((null? (cdr segs)) remaining) ;; last gets remainder + (else (max 0 (min remaining + (exact (round (* used-cells (/ sb seg-total)))))))) )) + (when (> cells 0) + (tb-print! col y (face-fg-attr face) bg (make-string cells #\█))) + (lp (cdr segs) (+ col cells) (- remaining cells)))))) + (when (> free-cells 0) + (tb-print! (+ x label-w used-cells) y + (face-fg-attr 'horizontal-rule) bg + (make-string free-cells #\░))) + (tb-print! (+ x label-w bar-w) y fg bg (format-bytes-short total)))) + +(def (sum-seg-bytes segs) + (let lp ((s segs) (acc 0)) + (cond ((null? s) acc) + (else (lp (cdr s) (+ acc (caddr (car s)))))))) + +(def (format-bytes-short bytes) + ;; 5-char field: " N/A" " 1.4G" " 64G" " 128G". + (cond + ((or (not bytes) (<= bytes 0)) " N/A") + (else + (let* ((gb (/ bytes 1073741824.0)) + (gb-i (exact (round gb)))) + (cond + ((>= gb 100) (format " ~aG" gb-i)) + ((>= gb 10) (format " ~aG" gb-i)) + (else (format " ~a.~aG" + (exact (truncate gb)) + (modulo (exact (round (* gb 10))) 10)))))))) + (def (render-bar! x y total-width label pct fg bg) ;; Format: "CPU " + filled-blocks + empty-blocks + " 100%" ;; 4 varies varies 5 --- a/src/jcode/ui/tui.ss +++ b/src/jcode/ui/tui.ss @@ -15,6 +15,7 @@ :jcode/ui/tui-status :jcode/ui/tui-input :jcode/ui/tui-sidebar + :jcode/ui/tui-memstats :jcode/ui/tui-dialog :jcode/ui/tui-toast :jcode/core/config @@ -126,12 +127,15 @@ sidebar ;; sidebar-state dialog ;; #f or active dialog tool-counts ;; hash-table: name → count - sysmon) ;; system utilization sampler + sysmon ;; system utilization sampler + memstats) ;; local-model RAM breakdown sampler transparent: #t) (def (make-fresh-state w h) - (let ((mon (make-sysmon))) + (let ((mon (make-sysmon)) + (mem (make-fresh-memstats))) (sysmon-update! mon) ;; seed counters; first read is meaningful + (try (memstats-update! mem) (catch (e) (void))) (make-app-state w h (>= w 100) ;; sidebar visible if wide enough @@ -151,7 +155,8 @@ (make-fresh-sidebar) #f ;; dialog (make-hash-table) - mon))) + mon + mem))) ;; ---- Layout calculations ---- @@ -274,7 +279,10 @@ (let ((mon (app-state-sysmon state))) (when mon (try (sysmon-update! mon) (catch (e) (void))) - (app-state-dirty?-set! state #t)))) + (app-state-dirty?-set! state #t))) + (let ((mem (app-state-memstats state))) + (when mem + (try (memstats-update! mem) (catch (e) (void)))))) ;; Drain pending events from agent worker thread and apply them. ;; All state mutation happens here, on the main thread, not from @@ -946,7 +954,11 @@ (lambda (pair) (case (car pair) ((tokens-in) (app-state-tokens-in-set! state - (+ (app-state-tokens-in state) (cdr pair)))) + (+ (app-state-tokens-in state) (cdr pair))) + ;; Also feed live prompt size into memstats so the KV + ;; cache slice tracks the current request's context. + (let ((mem (app-state-memstats state))) + (when mem (memstats-set-prompt-tokens! mem (cdr pair))))) ((tokens-out) (app-state-tokens-out-set! state (+ (app-state-tokens-out state) (cdr pair)))) ((cost) (app-state-cost-set! state @@ -1048,6 +1060,7 @@ (when (app-state-sidebar-visible? state) (render-sidebar! (app-state-sidebar state) (app-state-sysmon state) + (app-state-memstats state) (sidebar-x state) 0 (app-state-sidebar-width state) (- (app-state-height state) 1))) ;; don't overlap status bar