codec_models.py

"""Runtime LLM model switching.

CODEC serves every local model from ONE `mlx_vlm.server` on :8083. That server
resolves the model per request (`get_cached_model(openai_request.model)`), so a
switch needs no second server — only a different `model` field.

**A switch restarts the PM2 process rather than swapping in-place.** In-place
swapping works, but the server does not give the freed memory back: repeated
switching grew it to 52 GB (measured) against a 19 GB model. A recording is
exactly the wrong moment to discover that, and a fresh process is the only
reliable way to reclaim it. Two things make the leak easy to miss:

- PM2's `max_memory_restart` guard cannot see it. PM2 measures RSS, and MLX
  allocates GPU-backed memory that never appears there — the server reads ~60 MB
  of RSS while holding 20 GB. `footprint -p <pid>` (`phys_footprint`) is the only
  honest measure.
- A restart is nearly free. `scripts/start_model_server.sh` preloads the active
  model during FastAPI's lifespan, and uvicorn does not bind the port until
  lifespan finishes — so "port accepts a connection" already means "model
  loaded". That is what `restart_server()` waits for.

Three further consequences shape this module:

1. The server holds ONE model at a time. A switch costs a load pause (~20-60s
   for a 15-20 GB model) and two models can never be resident at once.
   Attempting to keep both loaded thrashes swap and roughly halves generation
   speed — measured.

2. If the target fails to load, the previous model is already gone and CODEC has
   NO brain. That is why `set_active()` probes after switching and reverts to the
   previous model when the probe fails. A bad switch must never leave the box
   mute.

3. The restart is best-effort, never load-bearing. Where PM2 is absent (tests,
   a hand-started server, CI) `restart_server()` reports failure and the switch
   proceeds in-place exactly as before — `probe()` forces the load either way.
   Losing the memory reclaim is a worse outcome than losing the switch.

The chat handler re-reads ~/.codec/config.json on every request, so writing
`llm_model` there takes effect on the next message.
"""
from __future__ import annotations

import contextlib
import fcntl
import json
import os
import shutil
import socket
import subprocess
import time
import urllib.error
import urllib.parse
import urllib.request
from typing import Any, Dict, List, Optional

import codec_cloud_models
from codec_jsonstore import atomic_write_json

CONFIG_PATH = os.path.expanduser("~/.codec/config.json")
HF_CACHE = os.path.expanduser("~/.cache/huggingface/hub")

# The model a FRESH install actually has. packaging/macos/models.json fetches
# Qwen2.5-7B as the bundled tier; this used to name the 20 GB Qwen3.6-35B, which
# no new install ever downloads — so every first chat on a buyer's Mac failed
# with a model-not-found. The developer's machine had both, so it looked fine.
# tests/test_setup_connect.py::test_defaults_match_reality pins these together.
DEFAULT_MODEL = "mlx-community/Qwen2.5-7B-Instruct-4bit"
DEFAULT_BASE_URL = "http://localhost:8083/v1"

# PM2 process that serves the models. Overridable via config.json so a second
# machine with a different service name does not need a code change.
DEFAULT_PM2_PROCESS = "qwen3.6"

# How long to wait for the restarted server to accept connections. Generous on
# purpose: the port does not open until the model finishes preloading, and a
# 20 GB checkpoint on a cold page cache is the slow case.
_RESTART_READY_TIMEOUT = 300.0

# How long to wait for PM2 to report a NEW pid. Only covers process teardown and
# respawn, not the model load, so it is short.
_RESTART_RESPAWN_TIMEOUT = 30.0

# Local dirs that are not general chat models. Speech, vision-only and
# GUI-grounding checkpoints are served on their own paths and must never appear
# in a chat model picker.
_NON_CHAT = ("whisper", "kokoro", "-tts-", "ui-tars", "nanollava", "embed",
             "flux", "bge", "rerank", "stable-diffusion", "clip", "vae",
             "sdxl", "musicgen", "parakeet")

# CODEC serves MLX checkpoints from one mlx_vlm.server; anything outside the
# mlx-community namespace (torch weights, diffusion models, embedding models)
# cannot be loaded by it and must not be offered as a chat model.
_REQUIRED_PREFIX = "mlx-community/"

# A HF cache dir exists as soon as anything is fetched — including a bare
# metadata stub of a few KB. Only treat a model as usable when real weights are
# on disk, otherwise the picker offers models that cannot load.
_MIN_WEIGHT_BYTES = 500 * 1024 * 1024


def _load_config() -> Dict[str, Any]:
    try:
        with open(CONFIG_PATH) as f:
            return json.load(f)
    except (OSError, json.JSONDecodeError):
        return {}


def _friendly(model_id: str) -> str:
    return model_id.split("/")[-1].replace("-4bit", "").replace("-", " ")


def _snapshot_bytes(cache_dir: str) -> int:
    """Size of the newest snapshot only.

    A HF cache dir can hold several revisions whose files hard/sym-link into a
    shared blobs/ store; summing the whole tree double-counts them (Qwen3.6
    reported 40.8 GB for a 19 GB model). Measure one revision, and resolve each
    file to its blob so a symlink is not counted as 0 bytes.
    """
    snaps = os.path.join(cache_dir, "snapshots")
    try:
        revs = [os.path.join(snaps, d) for d in os.listdir(snaps)]
        revs = [d for d in revs if os.path.isdir(d)]
        newest = max(revs, key=os.path.getmtime)
    except (OSError, ValueError):
        return 0
    total = 0
    for fn in os.listdir(newest):
        if not fn.endswith(".safetensors"):
            continue
        try:
            total += os.stat(os.path.realpath(os.path.join(newest, fn))).st_size
        except OSError:
            pass
    return total


def discover_local() -> List[Dict[str, Any]]:
    """Chat-capable models with real weights in the local HF cache."""
    out: List[Dict[str, Any]] = []
    try:
        entries = sorted(os.listdir(HF_CACHE))
    except OSError:
        return out
    for entry in entries:
        if not entry.startswith("models--"):
            continue
        model_id = entry[len("models--"):].replace("--", "/")
        if not model_id.startswith(_REQUIRED_PREFIX):
            continue
        low = model_id.lower()
        if any(tok in low for tok in _NON_CHAT):
            continue
        size = _snapshot_bytes(os.path.join(HF_CACHE, entry))
        if size < _MIN_WEIGHT_BYTES:
            continue  # metadata-only stub
        out.append({"id": model_id, "label": _friendly(model_id),
                    "size_gb": round(size / 1e9, 1)})
    return out


def _extra_models(cfg: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Operator-declared models that discovery cannot find on its own.

    `discover_local()` only surfaces `mlx-community/` checkpoints sitting in the
    HF cache. A locally fine-tuned/fused model lives at an arbitrary path and has
    no such namespace, so it would never appear in the picker. `extra_models` in
    config.json lets the operator register one explicitly:

        "extra_models": [
          {"id": "/abs/path/to/model", "label": "Fred (Gutenberg 4B)", "size_gb": 2.1}
        ]

    The `id` is passed verbatim as the OpenAI `model` field, so the mlx server
    loads it by path exactly as it loads a hub id. `label`/`size_gb` are optional.
    Malformed entries are skipped rather than breaking the whole picker.
    """
    raw = cfg.get("extra_models")
    out: List[Dict[str, Any]] = []
    if not isinstance(raw, list):
        return out
    for item in raw:
        if not isinstance(item, dict):
            continue
        mid = item.get("id")
        if not isinstance(mid, str) or not mid:
            continue
        entry: Dict[str, Any] = {"id": mid, "label": item.get("label") or _friendly(mid)}
        size = item.get("size_gb")
        entry["size_gb"] = size if isinstance(size, (int, float)) else None
        if codec_cloud_models.entry_for_id(mid, cfg) is not None:
            # A `base_url` off this Mac makes it a cloud model (codec_cloud_models):
            # switching to it rewrites config, it never touches the local server.
            entry["cloud"] = True
        out.append(entry)
    return out


def model_extras(model_id: Optional[str] = None,
                 config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """Per-model overrides declared alongside an `extra_models` entry.

    A locally fine-tuned model is not a drop-in for the everyday one: CODEC's
    operator prompt turns a 4B persona model into a skill-dispatcher (measured:
    it answered a philosophy question with two [SKILL:web_search] tags and
    nothing else), and MLX sampling knobs that a 35B does not need — notably
    repetition_penalty — are what keep a small model from collapsing into a
    refrain. So an entry may carry:

        {"id": "...", "label": "...",
         "system_prompt_file": "~/path/persona.txt",   # or "system_prompt": "..."
         "sampling": {"repetition_penalty": 1.2, "repetition_context_size": 512}}

    Returns {"system_prompt": str|None, "sampling": dict}. Never raises: a bad
    path or malformed entry degrades to "no override" rather than breaking chat.
    """
    cfg = config if config is not None else _load_config()
    mid = model_id or get_active(cfg)
    out: Dict[str, Any] = {"system_prompt": None, "sampling": {}, "no_think": False}
    for em in (cfg.get("extra_models") or []):
        if not isinstance(em, dict) or em.get("id") != mid:
            continue
        sp = em.get("system_prompt")
        path = em.get("system_prompt_file")
        if not sp and isinstance(path, str) and path:
            try:
                with open(os.path.expanduser(path)) as f:
                    sp = f.read().strip()
            except OSError:
                sp = None
        if isinstance(sp, str) and sp.strip():
            out["system_prompt"] = sp.strip()
        sampling = em.get("sampling")
        if isinstance(sampling, dict):
            out["sampling"] = {k: v for k, v in sampling.items()
                               if isinstance(k, str) and isinstance(v, (int, float, str, bool))}
        # A reformat/persona fine-tune that should answer directly (no <think>
        # block) declares `"no_think": true`. Chat forces enable_thinking=False;
        # the bridge path is already thinking-off via codec_llm.
        out["no_think"] = bool(em.get("no_think", False))
        break
    return out


def _visible_ids(cfg: Dict[str, Any]) -> Optional[set]:
    """Optional operator allowlist: `models_visible` in config.json.

    Discovery finds every loadable chat model on disk, which includes ones the
    operator does not want offered (an older generation, a vision checkpoint kept
    only because skills/screenshot_text.py hardcodes it). Listing it here hides
    it from the picker WITHOUT deleting weights that other code paths still use.
    Absent or empty -> show everything discovered.
    """
    raw = cfg.get("models_visible")
    if isinstance(raw, list) and raw:
        return {str(x) for x in raw}
    return None


def get_active(config: Optional[Dict[str, Any]] = None) -> str:
    cfg = config if config is not None else _load_config()
    return cfg.get("llm_model") or DEFAULT_MODEL


def _base_url(cfg: Dict[str, Any]) -> str:
    return cfg.get("llm_base_url", DEFAULT_BASE_URL).rstrip("/")


def list_models() -> Dict[str, Any]:
    """Registry for the picker: local models, annotated from config, plus the
    active one. `roles` in config.json:model_roles maps id -> human purpose."""
    cfg = _load_config()
    roles = cfg.get("model_roles", {}) if isinstance(cfg.get("model_roles"), dict) else {}
    active = get_active(cfg)
    models = discover_local()
    allow = _visible_ids(cfg)
    if allow:
        # The active model is always shown, even if it was removed from the
        # allowlist — hiding what CODEC is currently running would be a lie.
        models = [m for m in models if m["id"] in allow or m["id"] == active]
    for em in _extra_models(cfg):
        # Operator-declared: bypass the discovery allowlist, exactly as the
        # active model does — hiding a model the operator explicitly registered
        # would be a lie.
        if not any(m["id"] == em["id"] for m in models):
            models.append(em)
    for m in models:
        m["role"] = roles.get(m["id"], "")
        m["active"] = (m["id"] == active)
        if m.get("cloud"):
            m["role"] = _cloud_role(m["id"], m["role"], cfg)
    if active and not any(m["id"] == active for m in models):
        # Active model isn't on disk (or was pruned) — still show it as active.
        models.insert(0, {"id": active, "label": _friendly(active), "size_gb": None,
                          "role": roles.get(active, ""), "active": True})
    return {"active": active, "models": models}


def _cloud_role(model_id: str, role: str, cfg: Dict[str, Any]) -> str:
    """Picker text for a cloud model: its role plus this month's spend, so the
    cost is visible where the switch is made."""
    entry = codec_cloud_models.entry_for_id(model_id, cfg)
    if entry is None:
        return role
    spend = (f"${codec_cloud_models.spent_usd(model_id):.2f} of "
             f"${codec_cloud_models.cap_usd(entry):g} this month")
    return f"{role} · {spend}" if role else f"cloud · {spend}"


def probe(model_id: str, timeout: float = 240.0,
          base_url: Optional[str] = None) -> tuple[bool, str]:
    """Force the server to load `model_id` with a 1-token request.

    Returns (ok, detail). This is what makes a switch safe: it surfaces a load
    failure while we still know which model to fall back to. A cloud model gets
    a short real request with its own key instead (codec_cloud_models.probe).
    """
    cfg = _load_config()
    entry = codec_cloud_models.entry_for_id(model_id, cfg)
    if entry is not None:
        return codec_cloud_models.probe(entry, timeout=min(timeout, 30.0))
    url = (base_url or _base_url(cfg)) + "/chat/completions"
    body = json.dumps({
        "model": model_id,
        "messages": [{"role": "user", "content": "ok"}],
        "max_tokens": 1, "temperature": 0,
    }).encode()
    req = urllib.request.Request(url, data=body,
                                 headers={"Content-Type": "application/json"})
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            resp.read()
        return True, f"loaded in {time.time() - t0:.1f}s"
    except urllib.error.HTTPError as e:
        detail = ""
        try:
            detail = e.read().decode("utf-8", "replace")[:200]
        except Exception:
            pass
        return False, f"HTTP {e.code}: {detail}"
    except Exception as e:  # timeout, connection refused, ...
        return False, f"{type(e).__name__}: {e}"


def _pm2_process_name(cfg: Optional[Dict[str, Any]] = None) -> str:
    c = cfg if cfg is not None else _load_config()
    return str(c.get("model_pm2_process") or DEFAULT_PM2_PROCESS)


def _host_port(cfg: Dict[str, Any]) -> tuple[str, int]:
    parsed = urllib.parse.urlparse(_base_url(cfg))
    return parsed.hostname or "localhost", parsed.port or 8083


def _is_listening(host: str, port: int, timeout: float = 1.5) -> bool:
    """Bare TCP connect — the readiness signal, not an HTTP call.

    Same reasoning as the #308 heartbeat fix: a socket that accepts a connection
    proves the process is up without asking it to do any work. Here it proves
    more, because uvicorn binds only after FastAPI's lifespan (and therefore the
    model preload) has finished.
    """
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True
    except OSError:
        return False


def _pm2_pid(name: str) -> Optional[int]:
    """Current pid of a PM2 process, or None if pm2 or the process is absent."""
    pm2 = shutil.which("pm2")
    if not pm2:
        return None
    try:
        out = subprocess.run([pm2, "jlist"], capture_output=True, text=True,
                             timeout=20).stdout
        for proc in json.loads(out):
            if proc.get("name") == name:
                env = proc.get("pm2_env") or {}
                if env.get("status") != "online":
                    return None
                return proc.get("pid") or None
    except (OSError, ValueError, subprocess.SubprocessError):
        return None
    return None


def restart_server(process: Optional[str] = None,
                   ready_timeout: float = _RESTART_READY_TIMEOUT) -> tuple[bool, str]:
    """Switch the model server STOP-then-START so only one model is ever resident.

    Returns (ok, detail). Never raises: a failure here downgrades the switch to
    the old in-place behaviour rather than aborting it.

    Why not `pm2 restart`: restart overlaps the old and new processes — the old
    MLX server still holds its weights (~17-20 GB of unified memory) while the
    new one allocates the next model, so for a few seconds BOTH the 27B and the
    35B are resident. On a 64 GB Mac Studio that peak is survivable but wrong,
    and it is exactly what another session observed. So: STOP the old process,
    WAIT until it is gone and its port is free (weights released), THEN START
    the new one. One model is loaded at a time, by construction.
    """
    cfg = _load_config()
    name = process or _pm2_process_name(cfg)
    host, port = _host_port(cfg)

    pm2 = shutil.which("pm2")
    if not pm2:
        return False, "pm2 not on PATH — switching in place"

    t0 = time.time()

    # 1. STOP the old server and wait for it to actually exit. The model's
    #    unified-memory footprint is freed only when the process dies, so the
    #    port closing + PM2 reporting it offline is our "weights released" signal.
    try:
        subprocess.run([pm2, "stop", name], capture_output=True, text=True, timeout=60)
    except subprocess.SubprocessError as e:
        return False, f"pm2 stop {name} failed: {type(e).__name__}: {e}"

    gone_deadline = time.time() + _RESTART_RESPAWN_TIMEOUT
    while time.time() < gone_deadline:
        # _pm2_pid returns None once status != online; port must also be free so
        # we do not start the new process while the old still holds the weights.
        if _pm2_pid(name) is None and not _is_listening(host, port):
            break
        time.sleep(0.5)
    else:
        return False, (f"{name} did not release within "
                       f"{_RESTART_RESPAWN_TIMEOUT:.0f}s of stop — not starting a "
                       f"second copy on top of it")

    # 2. START the new server (config already points at the new model, so its
    #    lifespan preload loads only that one). `pm2 restart` is the verb that
    #    accepts a bare process name — on a stopped app it starts it.
    try:
        r = subprocess.run([pm2, "restart", name, "--update-env"],
                           capture_output=True, text=True, timeout=60)
    except subprocess.SubprocessError as e:
        return False, f"pm2 start {name} failed: {type(e).__name__}: {e}"
    if r.returncode != 0:
        err = (r.stderr or r.stdout or "").strip().replace("\n", " ")[:200]
        return False, f"pm2 start {name} exited {r.returncode}: {err}"

    # 3. Wait for the new process to come online and bind the port (uvicorn binds
    #    only after the FastAPI lifespan preload finishes).
    deadline = time.time() + ready_timeout
    while time.time() < deadline:
        if _pm2_pid(name) is not None and _is_listening(host, port):
            return True, f"{name} stop-then-start, serving in {time.time() - t0:.1f}s"
        time.sleep(1.0)
    return False, f"{name} started but {host}:{port} never opened within {ready_timeout:.0f}s"


def _write_active(model_id: str) -> None:
    cfg = _load_config()
    cfg["llm_model"] = model_id
    atomic_write_json(CONFIG_PATH, cfg)  # helper already writes 0600


def _enter_cloud(entry: Dict[str, Any], restore: Dict[str, Any]) -> None:
    """Point config at a cloud model. `llm_local_restore` keeps the local
    base_url and model to come back to (and for callers that must stay local,
    codec_cloud_models.local_only)."""
    cfg = _load_config()
    cfg["llm_local_restore"] = restore
    cfg["llm_base_url"] = entry["base_url"].strip().rstrip("/")
    cfg["llm_model"] = entry["id"]
    atomic_write_json(CONFIG_PATH, cfg)


def _leave_cloud(model_id: str) -> Dict[str, Any]:
    """Point config back at the local server with `model_id` active. Returns
    the `llm_local_restore` record it consumed, to go back if the load fails."""
    cfg = _load_config()
    restore = cfg.pop("llm_local_restore", None)
    restore = restore if isinstance(restore, dict) else {}
    cfg["llm_base_url"] = restore.get("llm_base_url") or DEFAULT_BASE_URL
    cfg["llm_model"] = model_id
    atomic_write_json(CONFIG_PATH, cfg)
    return restore


def _image_busy() -> Optional[str]:
    """Why a local model must not load now (an image job is running), or None."""
    try:
        import codec_image
        return codec_image.busy_message()
    except Exception:
        return None


def _local_server_up(base_url: str) -> bool:
    """Is the local model server answering connections? A non-local previous
    provider (AVA, custom) cannot be checked cheaply and counts as up."""
    if not codec_cloud_models.is_local_url(base_url):
        return True
    parsed = urllib.parse.urlparse(base_url)
    return _is_listening(parsed.hostname or "localhost", parsed.port or 8083)


def _switch_to_cloud(entry: Dict[str, Any], previous: str,
                     verify: bool) -> Dict[str, Any]:
    """Switch to a cloud model, keeping it only if it answers.

    The local server is left exactly as it is: the switch exists for when that
    server is down, and whatever state it is in must still be there for the
    switch back. So neither this switch nor its revert restarts anything.
    """
    cfg = _load_config()
    current = codec_cloud_models.active_entry(cfg)
    if current is None:
        restore = {"llm_base_url": cfg.get("llm_base_url") or DEFAULT_BASE_URL,
                   "llm_model": previous}
    else:  # cloud to cloud: keep the record of the local model
        restore = cfg.get("llm_local_restore")
        restore = restore if isinstance(restore, dict) else {}
    _enter_cloud(entry, restore)
    if not verify:
        return {"ok": True, "active": entry["id"], "previous": previous,
                "changed": True, "cloud": True, "detail": "switched (unverified)"}

    ok, detail = probe(entry["id"])
    if ok:
        _emit_audit(previous, entry["id"], True, detail)
        return {"ok": True, "active": entry["id"], "previous": previous,
                "changed": True, "cloud": True, "detail": detail}

    label = entry.get("label") or entry["id"]
    # A slow cloud model beats no model. When the local server is not running
    # (the reason to switch), keep the cloud model after a failure that waiting
    # can fix. A missing key, a reached cap or a rejected request still revert.
    if (current is None and codec_cloud_models.is_transient_failure(detail)
            and not _local_server_up(restore.get("llm_base_url") or DEFAULT_BASE_URL)):
        _emit_audit(previous, entry["id"], True, f"kept while slow: {detail}")
        return {"ok": True, "active": entry["id"], "previous": previous,
                "changed": True, "cloud": True, "slow": True,
                "detail": (f"{label} is slow to answer right now ({detail}); kept it "
                           f"because the local model is not running")}

    # No answer (no key, cap reached, provider down): put config back exactly
    # as it was. Nothing to reload, because nothing was unloaded.
    if current is None:
        _leave_cloud(previous)
    else:
        _enter_cloud(current, restore)
    _emit_audit(previous, entry["id"], False, detail)
    return {"ok": False, "active": previous, "attempted": entry["id"],
            "changed": False, "reverted": True,
            "error": f"{label} did not answer ({detail}) — kept {_friendly(previous)}"}


def set_active(model_id: str, verify: bool = True) -> Dict[str, Any]:
    """Switch the chat model (the model picker and every explicit switch).

    Holds the switch lock, so an automatic fallback never runs in the middle
    of it, and clears the automatic-fallback flag: the owner's pick wins.
    """
    with _switch_lock():
        r = _set_active_unlocked(model_id, verify)
        if r.get("changed"):
            _set_auto_flag(None)
    return r


def _set_active_unlocked(model_id: str, verify: bool = True) -> Dict[str, Any]:
    """Switch the chat model, verifying it loads and reverting if it does not.

    Only models discovered locally may be selected — an arbitrary string would
    unload the working model in exchange for a 404. A cloud model registered in
    `extra_models` (codec_cloud_models) is switched by config alone; see
    _switch_to_cloud.
    """
    previous = get_active()
    cfg = _load_config()
    allow = _visible_ids(cfg)
    known = {m["id"] for m in discover_local()}
    if allow:
        known &= allow            # hidden models are not switchable either
    # Operator-declared models are switchable by definition — they exist only
    # because the operator registered them (see _extra_models / list_models).
    known |= {em["id"] for em in _extra_models(cfg)}
    known |= {previous}
    if model_id not in known:
        return {"ok": False, "error": f"unknown model: {model_id}",
                "active": previous}
    if model_id == previous:
        return {"ok": True, "active": previous, "changed": False,
                "detail": "already active"}

    target = codec_cloud_models.entry_for_id(model_id, cfg)
    if target is not None:
        return _switch_to_cloud(target, previous, verify)

    # A running Create-image job holds ~35 GB; a local model load next to it
    # overflows memory (docs/CHAT-CREATE-IMAGE-DESIGN.md). Cloud switches pass.
    busy = _image_busy()
    if busy:
        return {"ok": False, "error": busy, "active": previous}

    # Coming back from a cloud model points config home first; the switch away
    # never touched the local server, so the rest is an ordinary local switch.
    from_cloud = codec_cloud_models.active_entry(cfg)
    restore = _leave_cloud(model_id) if from_cloud is not None else None
    if restore is None:
        _write_active(model_id)
    if not verify:
        # No restart either: an unverified switch is the caller asking not to
        # wait, and a restart is the longest part of the wait.
        return {"ok": True, "active": model_id, "previous": previous,
                "changed": True, "detail": "switched (unverified)"}

    # Config is written first so the restarted server preloads the NEW model.
    restarted, restart_detail = restart_server()

    ok, detail = probe(model_id)
    if ok:
        _emit_audit(previous, model_id, True, detail, restarted, restart_detail)
        return {"ok": True, "active": model_id, "previous": previous,
                "changed": True, "detail": detail,
                "restarted": restarted, "restart_detail": restart_detail}

    if from_cloud is not None:
        # The local model did not load: go back to the cloud model that was
        # answering. No restart — the cloud model does not live on this Mac.
        _enter_cloud(from_cloud, restore or {})
        reverted_ok, revert_detail = probe(previous)
        _emit_audit(previous, model_id, False, detail, restarted, restart_detail)
        return {"ok": False, "active": previous, "attempted": model_id,
                "changed": False,
                "error": f"{model_id} failed to load ({detail}) — reverted to {previous}",
                "reverted": reverted_ok, "revert_detail": revert_detail,
                "restarted": restarted, "restart_detail": restart_detail}

    # Failed to load — the old model is already unloaded, so put it back and
    # warm it, otherwise CODEC answers nothing at all. Restart again on the way
    # back: the process that just failed to load a model is the last one we want
    # to keep serving from.
    _write_active(previous)
    restart_server()
    reverted_ok, revert_detail = probe(previous)
    _emit_audit(previous, model_id, False, detail, restarted, restart_detail)
    return {"ok": False, "active": previous, "attempted": model_id,
            "changed": False,
            "error": f"{model_id} failed to load ({detail}) — reverted to {previous}",
            "reverted": reverted_ok, "revert_detail": revert_detail,
            "restarted": restarted, "restart_detail": restart_detail}


def _emit_audit(previous: str, requested: str, ok: bool, detail: str,
                restarted: Optional[bool] = None,
                restart_detail: str = "") -> None:
    try:
        from codec_audit import audit
        extra = {"previous": previous, "requested": requested,
                 "ok": ok, "detail": detail[:200]}
        if restarted is not None:
            # Recorded so a switch that quietly fell back to in-place — and
            # therefore did NOT reclaim memory — is visible after the fact.
            extra["restarted"] = restarted
            extra["restart_detail"] = restart_detail[:200]
        audit(event="model_switched", source="codec-models",
              outcome="ok" if ok else "error",
              level="info" if ok else "warning",
              message=f"model switch {previous} -> {requested}",
              extra=extra)
    except Exception:
        pass


# ── Automatic cloud fallback (docs/CLOUD-AUTO-FALLBACK-DESIGN.md) ────────────
# Opt-in: config.json "llm_auto_fallback" names a registered cloud entry. When
# the local model server is not listening, CODEC switches to that entry before a
# chat turn or a voice session builds its prompt (so the observer privacy gate
# sees the cloud endpoint), and the dashboard checks every 60 s. It switches
# back once the local model answers twice in a row. Any explicit switch
# (set_active) clears the flag, so a model the owner picks is never overridden.

_SWITCH_LOCK_PATH = os.path.expanduser("~/.codec/model_switch.lock")
_LOADING_GRACE_S = 300.0   # a just-started server binds its port only once loaded
_SWITCH_BACK_OKS = 2       # good local probes in a row before switching back
_FALLBACK_ALERT_KEY = "llm_auto_fallback"


@contextlib.contextmanager
def _switch_lock(blocking: bool = True):
    """One model switch at a time across CODEC's processes. Yields True when
    held; with blocking=False, False while another switch holds it."""
    os.makedirs(os.path.dirname(_SWITCH_LOCK_PATH), exist_ok=True)
    fd = os.open(_SWITCH_LOCK_PATH, os.O_RDWR | os.O_CREAT, 0o600)
    try:
        try:
            fcntl.flock(fd, fcntl.LOCK_EX | (0 if blocking else fcntl.LOCK_NB))
        except BlockingIOError:
            yield False
            return
        try:
            yield True
        finally:
            fcntl.flock(fd, fcntl.LOCK_UN)
    finally:
        os.close(fd)


def _set_auto_flag(value: Optional[Dict[str, Any]]) -> None:
    cfg = _load_config()
    if value is None:
        if "llm_auto_fallback_active" not in cfg:
            return
        cfg.pop("llm_auto_fallback_active", None)
    else:
        cfg["llm_auto_fallback_active"] = value
    atomic_write_json(CONFIG_PATH, cfg)


def _fallback_entry(cfg: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """The cloud entry the owner opted in to, or None (the default: no fallback)."""
    return codec_cloud_models.entry_for_id(cfg.get("llm_auto_fallback"), cfg)


def _server_loading(cfg: Dict[str, Any]) -> bool:
    """True while the PM2 model server is online but started less than 5 min
    ago: it binds its port only after the model has loaded."""
    pm2 = shutil.which("pm2")
    if not pm2:
        return False
    try:
        out = subprocess.run([pm2, "jlist"], capture_output=True, text=True,
                             timeout=20).stdout
        for proc in json.loads(out):
            if proc.get("name") == _pm2_process_name(cfg):
                env = proc.get("pm2_env") or {}
                started = float(env.get("pm_uptime") or 0) / 1000.0
                return env.get("status") == "online" and time.time() - started < _LOADING_GRACE_S
    except (OSError, ValueError, subprocess.SubprocessError):
        return False
    return False


def _fallback_alert(message: str, *, resolved: bool = False) -> None:
    try:
        import codec_alerts
        if resolved:
            codec_alerts.alert_resolved(_FALLBACK_ALERT_KEY, message)
        else:
            codec_alerts.alert_once(_FALLBACK_ALERT_KEY, "warning", message)
    except Exception:
        pass  # an alert must never stop the switch


def ensure_llm_available() -> Optional[str]:
    """Switch to the owner's cloud fallback if the local model server is down.

    Returns what it did, or None when it did nothing: no opt-in, already on a
    cloud model, the local server listening or still loading, the cloud entry
    blocked (cap reached, no key), or another process switching right now.
    """
    cfg = _load_config()
    entry = _fallback_entry(cfg)
    if entry is None or codec_cloud_models.active_entry(cfg) is not None:
        return None
    if not codec_cloud_models.is_local_url(_base_url(cfg)):
        return None
    host, port = _host_port(cfg)
    if _is_listening(host, port) or _server_loading(cfg):
        return None
    if codec_cloud_models.block_message(entry):
        return None   # stay put: the local error is the honest answer
    with _switch_lock(blocking=False) as held:
        if not held:
            return None
        cfg = _load_config()
        if codec_cloud_models.active_entry(cfg) is not None:
            return None
        local = get_active(cfg)
        r = _set_active_unlocked(entry["id"], verify=True)
        if not (r.get("ok") and r.get("changed")):
            return None
        _set_auto_flag({"from": local, "at": time.time(), "local_oks": 0})
    label = entry.get("label") or entry["id"]
    _fallback_alert(f"CODEC: the local model ({_friendly(local)}) is not answering. "
                    f"Switched to {label} until it is back.")
    return f"switched to {label}"


def maybe_switch_back() -> Optional[str]:
    """After an automatic fallback, go back to the local model once it answers
    two probes in a row. No restart: the server already serves that model."""
    cfg = _load_config()
    flag = cfg.get("llm_auto_fallback_active")
    if not isinstance(flag, dict):
        return None
    cloud = codec_cloud_models.active_entry(cfg)
    restore = cfg.get("llm_local_restore")
    restore = restore if isinstance(restore, dict) else {}
    local = flag.get("from") or restore.get("llm_model") or DEFAULT_MODEL
    ok = False
    if cloud is not None:
        ok, _ = probe(local, timeout=30.0,
                      base_url=(restore.get("llm_base_url") or DEFAULT_BASE_URL).rstrip("/"))
    with _switch_lock(blocking=False) as held:
        if not held:
            return None
        cfg = _load_config()
        flag = cfg.get("llm_auto_fallback_active")
        if not isinstance(flag, dict):
            return None
        cloud = codec_cloud_models.active_entry(cfg)
        if cloud is None:           # left the cloud some other way
            _set_auto_flag(None)
            return None
        oks = int(flag.get("local_oks") or 0) + 1 if ok else 0
        if oks < _SWITCH_BACK_OKS:
            if oks != flag.get("local_oks"):
                flag["local_oks"] = oks
                _set_auto_flag(flag)
            return None
        _leave_cloud(local)
        _set_auto_flag(None)
    _emit_audit(cloud["id"], local, True, "automatic switch back: local model answers again")
    _fallback_alert(f"CODEC: the local model ({_friendly(local)}) answers again. "
                    f"Switched back to it.", resolved=True)
    return f"back on {_friendly(local)}"


def auto_fallback_tick() -> Optional[str]:
    """One background pass (codec-dashboard, every 60 s): fall back, or come back."""
    return ensure_llm_available() or maybe_switch_back()