colibri/k3: give the mux engines a remote service too

These two were the exception in the last commit: they are driven over a
stdin/stdout mux protocol, so there was no URL to redirect. This adds the
missing HTTP boundary.

tools/mux_service.py hosts a real MuxEngine (colibri or k3 — both use the
same class) and exposes /health, /run (NDJSON: text chunks then stats),
/pause and /resume. codai/api/mux_remote.RemoteMuxEngine is the client
and presents the exact surface the backends use — run/is_alive/pause/
resume/close/log_tail/family — so ensure_engine() can hand one back and
nothing downstream can tell local from remote. Cancellation travels as a
dropped connection: the service's write fails and it stops generating,
matching the local `cancelled` callback semantics.

`service_url` is now a real config field on all five engine backends
(ds4, colibri, k3, ktransformers, vllm), settable through the admin API
and persisted, rather than only an env var.

Tested end-to-end against a stub engine: streaming, stats, mid-stream
cancellation, pause/resume, and error propagation all round-trip.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY

codai/admin/routes.py

@@ -3522,6 +3522,7 @@ def build_settings_dict(c, gpu_cards):
         },
         "ds4": {
             "enabled": c.ds4.enabled,
+            "service_url": c.ds4.service_url,
             "repo_url": c.ds4.repo_url,
             "install_dir": c.ds4.install_dir,
             "build_target": c.ds4.build_target,
@@ -3543,6 +3544,7 @@ def build_settings_dict(c, gpu_cards):
         },
         "colibri": {
             "enabled": c.colibri.enabled,
+            "service_url": c.colibri.service_url,
             "repo_url": c.colibri.repo_url,
             "install_dir": c.colibri.install_dir,
             "build_target": c.colibri.build_target,
@@ -3558,6 +3560,7 @@ def build_settings_dict(c, gpu_cards):
         },
         "k3": {
             "enabled": c.k3.enabled,
+            "service_url": c.k3.service_url,
             "repo_url": c.k3.repo_url,
             "install_dir": c.k3.install_dir,
             "model_path": c.k3.model_path,
@@ -3574,6 +3577,7 @@ def build_settings_dict(c, gpu_cards):
         },
         "ktransformers": {
             "enabled": c.ktransformers.enabled,
+            "service_url": c.ktransformers.service_url,
             "repo_url": c.ktransformers.repo_url,
             "install_dir": c.ktransformers.install_dir,
             "model_path": c.ktransformers.model_path,
@@ -3588,6 +3592,7 @@ def build_settings_dict(c, gpu_cards):
         },
         "vllm": {
             "enabled": c.vllm.enabled,
+            "service_url": c.vllm.service_url,
             "venv": c.vllm.venv,
             "model_path": c.vllm.model_path,
             "model_id": c.vllm.model_id,
@@ -3954,6 +3959,8 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
 
     if "ds4" in data:
         d = data["ds4"]
+        if "service_url" in d:
+            c.ds4.service_url = (d.get("service_url") or "").strip()
         c.ds4.enabled = bool(d.get("enabled", c.ds4.enabled))
         if "repo_url" in d:
             c.ds4.repo_url = (d.get("repo_url") or c.ds4.repo_url or "").strip()
@@ -4003,6 +4010,8 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
 
     if "colibri" in data:
         d = data["colibri"]
+        if "service_url" in d:
+            c.colibri.service_url = (d.get("service_url") or "").strip()
         c.colibri.enabled = bool(d.get("enabled", c.colibri.enabled))
         if "repo_url" in d:
             c.colibri.repo_url = (d.get("repo_url") or c.colibri.repo_url or "").strip()
@@ -4037,6 +4046,8 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
 
     if "k3" in data:
         d = data["k3"]
+        if "service_url" in d:
+            c.k3.service_url = (d.get("service_url") or "").strip()
         c.k3.enabled = bool(d.get("enabled", c.k3.enabled))
         if "repo_url" in d:
             c.k3.repo_url = (d.get("repo_url") or c.k3.repo_url or "").strip()
@@ -4076,6 +4087,8 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
 
     if "ktransformers" in data:
         d = data["ktransformers"]
+        if "service_url" in d:
+            c.ktransformers.service_url = (d.get("service_url") or "").strip()
         c.ktransformers.enabled = bool(d.get("enabled", c.ktransformers.enabled))
         if "repo_url" in d:
             c.ktransformers.repo_url = (d.get("repo_url") or c.ktransformers.repo_url or "").strip()
@@ -4108,6 +4121,8 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
 
     if "vllm" in data:
         d = data["vllm"]
+        if "service_url" in d:
+            c.vllm.service_url = (d.get("service_url") or "").strip()
         v = c.vllm
         v.enabled = bool(d.get("enabled", v.enabled))
         if "venv" in d: v.venv = (d.get("venv") or "").strip()

codai/api/colibri_worker.py

@@ -822,6 +822,18 @@ def ensure_engine(cfg, model_dir: Optional[str] = None, ctx: Optional[int] = Non
     directory (or ``cfg.model_path`` is set) the engine loads THAT. ``ctx`` sizes the
     per-turn generation budget (NGEN). Returns a live :class:`MuxEngine`.
     """
+    # An explicit `service_url` points at a colibri engine hosted by
+    # tools/mux_service.py somewhere else — another host, a rented pod — so
+    # nothing is built or launched here. The client speaks the same surface the
+    # backends use, so callers cannot tell it from a local MuxEngine.
+    _remote = (str(getattr(cfg, "service_url", "") or "")
+               or os.environ.get("CODERAI_COLIBRI_SERVICE_URL") or "").strip()
+    if _remote:
+        from codai.api.mux_remote import remote_engine_for
+        eng = remote_engine_for("colibri", _remote)
+        print(f"[colibri] using the configured remote engine at {eng.url}", flush=True)
+        return eng
+
     resolved, svc_key = resolve_service_key(cfg, model_dir)
     # Family is inferred from the resolved container (config.json / dir name) and
     # decides the binary, the make target, and the serve protocol.

codai/api/k3_worker.py

@@ -158,6 +158,18 @@ def _build_env(cfg, ctx: int) -> dict:
 def ensure_engine(cfg, model_dir: Optional[str] = None, ctx: Optional[int] = None,
                   ready_timeout: float = 3600.0) -> MuxEngine:
     """Build (as needed), then start (or reuse) the k3 engine for a checkpoint."""
+    # An explicit `service_url` points at a k3 engine hosted by
+    # tools/mux_service.py somewhere else — another host, a rented pod — so
+    # nothing is built or launched here. The client speaks the same surface the
+    # backends use, so callers cannot tell it from a local MuxEngine.
+    _remote = (str(getattr(cfg, "service_url", "") or "")
+               or os.environ.get("CODERAI_K3_SERVICE_URL") or "").strip()
+    if _remote:
+        from codai.api.mux_remote import remote_engine_for
+        eng = remote_engine_for("k3", _remote, family="kimi_k3")
+        print(f"[k3] using the configured remote engine at {eng.url}", flush=True)
+        return eng
+
     ckpt, svc_key = resolve_service_key(cfg, model_dir)
     with _lock:
         eng = _services.get(svc_key)

codai/api/mux_remote.py

+# CoderAI - OpenAI-compatible API server
+# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+
+"""Client for a colibri/k3 engine hosted by tools/mux_service.py.
+
+colibri and k3 are driven over a stdin/stdout mux protocol, so unlike every
+other worker they had no URL to redirect. :class:`RemoteMuxEngine` supplies the
+missing half: it speaks HTTP to a remote ``mux_service`` while presenting the
+exact surface the backends use on a local ``MuxEngine`` — ``run``, ``is_alive``,
+``pause``, ``resume``, ``close``, ``log_tail`` and ``family`` — so
+``colibri_worker.ensure_engine()`` can hand one back and neither backend nor
+model manager can tell the difference.
+
+Cancellation travels as a dropped connection: when ``cancelled()`` goes true we
+stop reading and close the response, and the service sees its write fail and
+aborts generation. That is the same contract the local engine has.
+"""
+
+import json
+from typing import Callable, Optional
+
+
+class RemoteMuxEngine:
+    """A colibri/k3 engine running behind tools/mux_service.py."""
+
+    def __init__(self, url: str, kind: str = "colibri", family: str = "",
+                 timeout: float = 3600.0):
+        self.url = (url or "").rstrip("/")
+        self.kind = kind
+        self.family = family or "glm"
+        self.timeout = float(timeout)
+        self._last_error = ""
+
+    # ---- lifecycle -------------------------------------------------------
+    def is_alive(self) -> bool:
+        info = self.health()
+        return bool(info and info.get("alive"))
+
+    def health(self) -> Optional[dict]:
+        import requests
+        try:
+            r = requests.get(self.url + "/health", timeout=10)
+            return r.json() if r.ok else None
+        except Exception as exc:
+            self._last_error = f"health check failed: {exc}"
+            return None
+
+    def log_tail(self) -> str:
+        return self._last_error
+
+    def close(self) -> None:
+        """A remote engine outlives this process — never kill what we don't own."""
+        return None
+
+    def pause(self):
+        return self._post_simple("/pause")
+
+    def resume(self):
+        return self._post_simple("/resume")
+
+    def _post_simple(self, path: str):
+        import requests
+        try:
+            requests.post(self.url + path, timeout=120)
+        except Exception as exc:
+            self._last_error = f"{path} failed: {exc}"
+
+    # ---- generation ------------------------------------------------------
+    def run(self, prompt: str, max_tokens: int, temperature: float, top_p: float,
+            on_text: Callable[[str], None],
+            cancelled: Optional[Callable[[], bool]] = None) -> dict:
+        """Submit one rendered prompt; stream text to ``on_text``; return stats."""
+        import requests
+        payload = {"prompt": prompt, "max_tokens": int(max_tokens),
+                   "temperature": float(temperature), "top_p": float(top_p)}
+        stats: dict = {}
+        emitted = 0
+        try:
+            resp = requests.post(self.url + "/run", json=payload,
+                                 stream=True, timeout=(30, self.timeout))
+        except Exception as exc:
+            self._last_error = str(exc)
+            raise RuntimeError(f"remote {self.kind} engine at {self.url} is unreachable: {exc}")
+        if not resp.ok:
+            detail = ""
+            try:
+                detail = resp.json().get("error", "")
+            except Exception:
+                detail = (resp.text or "")[:400]
+            self._last_error = detail
+            raise RuntimeError(f"remote {self.kind} engine returned HTTP {resp.status_code}"
+                               + (f": {detail}" if detail else ""))
+        try:
+            for line in resp.iter_lines(decode_unicode=False):
+                if not line:
+                    continue
+                try:
+                    obj = json.loads(line)
+                except Exception:
+                    continue
+                if "t" in obj:
+                    on_text(obj["t"])
+                    emitted += 1
+                    # Closing the response is how the far side learns to stop.
+                    if cancelled is not None and cancelled():
+                        break
+                elif "stats" in obj:
+                    stats = obj["stats"] or {}
+                elif "error" in obj:
+                    self._last_error = str(obj["error"])
+                    raise RuntimeError(f"remote {self.kind} engine: {obj['error']}")
+        finally:
+            try:
+                resp.close()
+            except Exception:
+                pass
+        if not stats:
+            stats = {"completion_tokens": emitted, "prompt_tokens": 0,
+                     "total_tokens": emitted, "tokens_per_second": 0.0,
+                     "cache_hit_percent": 0.0, "rss_gb": 0.0, "length_limited": False}
+        return stats
+
+
+def remote_engine_for(kind: str, url: str, family: str = "") -> RemoteMuxEngine:
+    """Build a client and prove the far side is answering before returning it."""
+    eng = RemoteMuxEngine(url, kind=kind, family=family)
+    info = eng.health()
+    if not info:
+        raise RuntimeError(
+            f"configured service_url {eng.url} is not answering its health check")
+    if not info.get("alive"):
+        raise RuntimeError(
+            f"remote {kind} service at {eng.url} is up but its engine is not running")
+    eng.family = family or info.get("family") or eng.family
+    return eng

codai/config.py

@@ -368,6 +368,10 @@ class Ds4Config:
     contains ``deepseek-v4``) is routed to ds4 instead of the normal backends.
     """
     enabled: bool = False
+    # Point this at an engine already running somewhere else (another host, a
+    # container, a rented pod) and coderai proxies to it instead of building,
+    # downloading or launching anything locally. Blank = manage it here.
+    service_url: str = ""
     repo_url: str = "https://github.com/antirez/ds4"
     install_dir: Optional[str] = None      # None = ~/.coderai/ds4
     build_target: str = "auto"             # auto|cuda-generic|cuda-spark|metal|cpu
@@ -426,6 +430,10 @@ class ColibriConfig:
     ``colibri``) is routed to the colibri engine instead of the normal backends.
     """
     enabled: bool = False
+    # Point this at an engine already running somewhere else (another host, a
+    # container, a rented pod) and coderai proxies to it instead of building,
+    # downloading or launching anything locally. Blank = manage it here.
+    service_url: str = ""
     repo_url: str = "https://github.com/JustVugg/colibri"
     install_dir: Optional[str] = None      # None = ~/.coderai/colibri
     build_target: str = "auto"             # auto|cuda|hip|cpu (auto: CUDA if nvcc present)
@@ -468,6 +476,10 @@ class K3Config:
     fast local storage.
     """
     enabled: bool = False
+    # Point this at an engine already running somewhere else (another host, a
+    # container, a rented pod) and coderai proxies to it instead of building,
+    # downloading or launching anything locally. Blank = manage it here.
+    service_url: str = ""
     repo_url: str = "https://github.com/FareedKhan-dev/kimi-k3-in-c"
     install_dir: Optional[str] = None      # None = ~/.coderai/kimi-k3-in-c
     model_path: str = ""                   # the Kimi-K3 checkpoint directory (config + tokenizer + shards)
@@ -501,6 +513,10 @@ class KtransformersConfig:
     weights. Best throughput needs AMX/AVX-512 CPUs.
     """
     enabled: bool = False
+    # Point this at an engine already running somewhere else (another host, a
+    # container, a rented pod) and coderai proxies to it instead of building,
+    # downloading or launching anything locally. Blank = manage it here.
+    service_url: str = ""
     repo_url: str = "https://github.com/kvcache-ai/ktransformers"
     install_dir: Optional[str] = None      # None = ~/.coderai/ktransformers (kt-kernel build)
     model_path: str = ""                   # HF model directory (--model)
@@ -536,6 +552,10 @@ class VllmConfig:
     Also reused by the OCR subsystem to serve Surya2 (a VLM) with continuous batching.
     """
     enabled: bool = False
+    # Point this at an engine already running somewhere else (another host, a
+    # container, a rented pod) and coderai proxies to it instead of building,
+    # downloading or launching anything locally. Blank = manage it here.
+    service_url: str = ""
     venv: str = ""                       # isolated venv dir; blank = auto (baked/cache/home)
     model_path: str = ""                 # OPTIONAL single-model HF dir/repo id (--model); blank = use the model list
     model_id: str = ""                   # OPTIONAL single-model alias/served-name; blank = per-model from the list
@@ -1050,6 +1070,7 @@ class ConfigManager:
             },
             "ds4": {
                 "enabled": self.config.ds4.enabled,
+                "service_url": self.config.ds4.service_url,
                 "repo_url": self.config.ds4.repo_url,
                 "install_dir": self.config.ds4.install_dir,
                 "build_target": self.config.ds4.build_target,
@@ -1071,6 +1092,7 @@ class ConfigManager:
             },
             "colibri": {
                 "enabled": self.config.colibri.enabled,
+                "service_url": self.config.colibri.service_url,
                 "repo_url": self.config.colibri.repo_url,
                 "install_dir": self.config.colibri.install_dir,
                 "build_target": self.config.colibri.build_target,
@@ -1086,6 +1108,7 @@ class ConfigManager:
             },
             "k3": {
                 "enabled": self.config.k3.enabled,
+                "service_url": self.config.k3.service_url,
                 "repo_url": self.config.k3.repo_url,
                 "install_dir": self.config.k3.install_dir,
                 "model_path": self.config.k3.model_path,
@@ -1102,6 +1125,7 @@ class ConfigManager:
             },
             "ktransformers": {
                 "enabled": self.config.ktransformers.enabled,
+                "service_url": self.config.ktransformers.service_url,
                 "repo_url": self.config.ktransformers.repo_url,
                 "install_dir": self.config.ktransformers.install_dir,
                 "model_path": self.config.ktransformers.model_path,
@@ -1116,6 +1140,7 @@ class ConfigManager:
             },
             "vllm": {
                 "enabled": self.config.vllm.enabled,
+                "service_url": self.config.vllm.service_url,
                 "venv": self.config.vllm.venv,
                 "model_path": self.config.vllm.model_path,
                 "model_id": self.config.vllm.model_id,

tools/mux_service.py

+#!/usr/bin/env python3
+# CoderAI - OpenAI-compatible API server
+# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+
+"""Mux-engine service — hosts a colibri or kimi-k3-in-c engine over HTTP.
+
+colibri and k3 are the two engines coderai drives DIRECTLY over a stdin/stdout
+mux protocol rather than through an HTTP server, which is exactly what stopped
+them from being pointed at a remote box the way every other worker can be. This
+service is the missing HTTP boundary: it runs a real ``MuxEngine`` locally and
+exposes it, so a remote host (a rented pod, another machine with the container
+on fast storage) can serve the engine while coderai stays a thin client.
+
+The client is :class:`codai.api.mux_remote.RemoteMuxEngine`, which implements
+the same surface the backends use (``run``/``is_alive``/``pause``/``resume``/
+``close``/``family``), so nothing downstream can tell local from remote.
+
+  GET  /health   -> {"ok": true, "family": "glm", "alive": true}
+  POST /run      -> NDJSON stream: {"t": "…"} chunks, then {"stats": {…}}
+                    (body: {prompt, max_tokens, temperature, top_p})
+  POST /pause    -> {"ok": true}     release the engine's VRAM
+  POST /resume   -> {"ok": true}
+
+Cancellation is the dropped connection: when the client stops reading, the
+write fails and generation is cancelled — the same semantics a local engine
+gets from its ``cancelled`` callback.
+"""
+
+import argparse
+import json
+import os
+import sys
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from types import SimpleNamespace
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+ENGINE = None
+KIND = "colibri"
+_LOCK = threading.Lock()
+
+
+def _build_engine(kind: str, cfg: dict, model_dir: str, ctx: int):
+    """Start the real engine through the normal worker path."""
+    ns = SimpleNamespace(**(cfg or {}))
+    if kind == "k3":
+        from codai.api import k3_worker as worker
+    else:
+        from codai.api import colibri_worker as worker
+    return worker.ensure_engine(ns, model_dir=model_dir or None, ctx=ctx or None)
+
+
+class Handler(BaseHTTPRequestHandler):
+    protocol_version = "HTTP/1.1"
+
+    def log_message(self, fmt, *args):   # quieter than the stdlib default
+        pass
+
+    def _json(self, code: int, payload: dict):
+        body = json.dumps(payload).encode()
+        self.send_response(code)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_GET(self):
+        if self.path.rstrip("/") != "/health":
+            return self._json(404, {"error": "not found"})
+        alive = bool(ENGINE and ENGINE.is_alive())
+        self._json(200, {"ok": alive, "kind": KIND, "alive": alive,
+                         "family": getattr(ENGINE, "family", "")})
+
+    def do_POST(self):
+        path = self.path.rstrip("/")
+        try:
+            n = int(self.headers.get("Content-Length") or 0)
+            req = json.loads(self.rfile.read(n) or b"{}") if n else {}
+        except Exception as exc:
+            return self._json(400, {"error": f"bad request body: {exc}"})
+
+        if path in ("/pause", "/resume"):
+            try:
+                getattr(ENGINE, path[1:])()
+            except Exception as exc:
+                return self._json(500, {"error": str(exc)})
+            return self._json(200, {"ok": True})
+        if path != "/run":
+            return self._json(404, {"error": "not found"})
+        if ENGINE is None or not ENGINE.is_alive():
+            return self._json(503, {"error": "engine is not running: "
+                                             + (ENGINE.log_tail() if ENGINE else "not started")})
+
+        self.send_response(200)
+        self.send_header("Content-Type", "application/x-ndjson")
+        self.send_header("Transfer-Encoding", "chunked")
+        self.end_headers()
+
+        dead = threading.Event()
+
+        def _emit(obj):
+            line = (json.dumps(obj) + "\n").encode()
+            try:
+                self.wfile.write(b"%x\r\n" % len(line) + line + b"\r\n")
+                self.wfile.flush()
+            except Exception:
+                # Client hung up — that is the cancel signal.
+                dead.set()
+
+        # One generation at a time: a mux engine's slots are its own business, but
+        # this service exists to front a single engine for a single coderai.
+        with _LOCK:
+            try:
+                stats = ENGINE.run(
+                    req.get("prompt") or "",
+                    int(req.get("max_tokens") or 1024),
+                    float(req.get("temperature") or 0.0),
+                    float(req.get("top_p") or 1.0),
+                    on_text=lambda t: _emit({"t": t}),
+                    cancelled=dead.is_set,
+                )
+                _emit({"stats": stats or {}})
+            except Exception as exc:
+                _emit({"error": f"{type(exc).__name__}: {exc}"})
+        try:
+            self.wfile.write(b"0\r\n\r\n")
+            self.wfile.flush()
+        except Exception:
+            pass
+
+
+def main():
+    global ENGINE, KIND
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--kind", choices=("colibri", "k3"),
+                    default=os.environ.get("CODERAI_MUX_KIND", "colibri"))
+    ap.add_argument("--host", default=os.environ.get("CODERAI_MUX_HOST", "0.0.0.0"))
+    ap.add_argument("--port", type=int, default=int(os.environ.get("CODERAI_MUX_PORT", "8092")))
+    ap.add_argument("--model-dir", default=os.environ.get("CODERAI_MUX_MODEL_DIR", ""),
+                    help="container/checkpoint directory the engine serves")
+    ap.add_argument("--ctx", type=int, default=int(os.environ.get("CODERAI_MUX_CTX", "0") or 0))
+    ap.add_argument("--config-json", default=os.environ.get("CODERAI_MUX_CONFIG", "{}"),
+                    help="engine config fields (the ds4/colibri/k3 config block) as JSON")
+    args = ap.parse_args()
+
+    KIND = args.kind
+    try:
+        cfg = json.loads(args.config_json or "{}")
+    except Exception as exc:
+        print(f"--config-json is not valid JSON: {exc}", file=sys.stderr)
+        return 2
+    print(f"[mux-service] starting {args.kind} engine on {args.model_dir or '(config)'} …",
+          flush=True)
+    ENGINE = _build_engine(args.kind, cfg, args.model_dir, args.ctx)
+    print(f"[mux-service] {args.kind} ready, serving on {args.host}:{args.port}", flush=True)
+    ThreadingHTTPServer((args.host, args.port), Handler).serve_forever()
+
+
+if __name__ == "__main__":
+    sys.exit(main() or 0)