cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
from __future__ import annotations
 
import json
import os
import sys
from pathlib import Path
from typing import Any
 
_CONFIG_PATH = Path.home() / ".aquin" / "config.json"
_BASE_URL = "https://api.aquin.app"
 
 
def _load_config() -> dict:
    if _CONFIG_PATH.exists():
        with open(_CONFIG_PATH) as f:
            return json.load(f)
    return {}
 
 
def _save_config(data: dict) -> None:
    _CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(_CONFIG_PATH, "w") as f:
        json.dump(data, f, indent=2)
 
 
def _get_api_key(explicit: str | None = None) -> str:
    """Optional cloud token from env / legacy config — not required for local use."""
    from aquin.auth_config import resolve_api_key
    return resolve_api_key(explicit=explicit)
 
 
def cmd_status(args: list[str]) -> None:
    """aquin status — loaded model and local engine state"""
    plain = False
    as_json = False
    i = 0
    while i < len(args):
        if args[i] == "--plain":
            plain = True
            i += 1
        elif args[i] == "--json":
            as_json = True
            i += 1
        elif args[i] in ("--api-key", "-k", "--live"):
            print(f"Removed flag: {args[i]} (no account auth in this framework).", file=sys.stderr)
            sys.exit(1)
        elif args[i] in ("--help", "-h", "help"):
            print("Usage: aquin status [--plain] [--json]")
            print("")
            print("  Show loaded model and local engine state.")
            print("  --plain     ASCII bullets (no color)")
            print("  --json      Machine-readable rows")
            return
        else:
            print(f"Unknown flag: {args[i]}", file=sys.stderr)
            sys.exit(1)
 
    from aquin.status_display import collect_status_snapshot, render_status, status_payload
 
    snap = collect_status_snapshot(logged_in=False)
 
    if as_json:
        print(json.dumps(status_payload(snap), ensure_ascii=False))
        return
 
    render_status(snap, plain=plain)
 
 
def cmd_version(args: list[str]) -> None:
    """aquin version  — show the installed CLI version"""
    if args and args[0] in ("--help", "-h", "help"):
        print("Usage: aquin version [--json]")
        print("")
        print("  Show the installed CLI build version.")
        print("  --json   Machine-readable")
        return
 
    from aquin.version_display import render_version, version_payload
 
    if "--json" in args:
        print(json.dumps(version_payload(), ensure_ascii=False))
        return
 
    render_version()
 
 
def cmd_license(args: list[str]) -> None:
    """Show the AQIT Apache 2.0 license"""
    if args and args[0] in ("--help", "-h", "help"):
        print("Usage: aqit license [--json] [--plain]")
        print("")
        print("  Display the AQIT Apache License 2.0.")
        print("  --json    Machine-readable text + acceptance status")
        print("  --plain   Full ASCII frame (legacy layout)")
        return
 
    from aquin.auth_config import load_config
    from aquin.license_acceptance import is_license_accepted
    from aquin.license_text import LICENSE_URL, read_license_text
 
    as_json = "--json" in args
    plain = "--plain" in args
    text = read_license_text()
    accepted = is_license_accepted()
    record = (load_config().get("license_accepted") or {}) if accepted else {}
    accepted_at = str(record.get("accepted_at") or "") if isinstance(record, dict) else ""
 
    if as_json:
        try:
            sys.stdout.reconfigure(encoding="utf-8")  # type: ignore[attr-defined]
        except Exception:
            pass
        print(json.dumps({
            "accepted": accepted,
            "acceptedAt": accepted_at or None,
            "url": LICENSE_URL,
            "text": text,
        }, ensure_ascii=False))
        return
 
    if plain:
        from aquin.license_display import render_license_document
 
        render_license_document(context="view")
        return
 
    # Quiet default (npm CLI / scripts)
    try:
        sys.stdout.reconfigure(encoding="utf-8")  # type: ignore[attr-defined]
    except Exception:
        pass
    print(f"license  {LICENSE_URL}")
    if accepted:
        if accepted_at:
            print(f"accepted {accepted_at}")
        else:
            print("accepted")
    else:
        print("not accepted")
    print()
    print("full text: aquin license --plain")
 
 
def cmd_update(args: list[str]) -> None:
    """aquin update  — upgrade aquin to the latest version from PyPI"""
    import subprocess
    from importlib.metadata import PackageNotFoundError, version
 
    if args and args[0] in ("--help", "-h", "help"):
        print("Usage: aquin update [--json]")
        print("")
        print("  Upgrade the Aquin engine package from PyPI.")
        print("  --json   Machine-readable result (npm CLI)")
        return
 
    as_json = "--json" in args
 
    try:
        current = version("aquin")
    except PackageNotFoundError:
        current = "unknown"
 
    if not as_json:
        print(f"Current version: {current}")
        print("Checking for updates...")
 
    result = subprocess.run(
        [sys.executable, "-m", "pip", "install", "--upgrade", "aquin"],
        capture_output=True,
        text=True,
    )
 
    if result.returncode != 0:
        err = (result.stderr or result.stdout or "Update failed.")[-1000:]
        if as_json:
            print(json.dumps({
                "ok": False,
                "current": current,
                "error": err.strip() or "Update failed.",
            }, ensure_ascii=False))
        else:
            print("Update failed:")
            print(err)
        sys.exit(1)
 
    try:
        import importlib
        import importlib.metadata as _meta
 
        importlib.invalidate_caches()
        new = _meta.version("aquin")
    except Exception:
        new = "unknown"
 
    updated = new != current
    if as_json:
        print(json.dumps({
            "ok": True,
            "current": current,
            "new": new,
            "updated": updated,
        }, ensure_ascii=False))
        return
 
    if updated:
        print(f"Updated to {new}")
    else:
        print("Already up to date.")
 
 
def _detect_desktop_platform() -> tuple[str, str]:
    """Return (os_label, arch) for desktop install messaging."""
    import platform
 
    system = platform.system()
    machine = (platform.machine() or "").lower()
    if machine in ("x86_64", "amd64"):
        arch = "x64"
    elif machine in ("aarch64", "arm64"):
        arch = "arm64"
    elif machine in ("i386", "i686", "x86"):
        arch = "x86"
    else:
        arch = machine or "unknown"
 
    if system == "Darwin":
        return "macOS", arch
    if system == "Windows":
        return "Windows", arch
    if system == "Linux":
        return "Linux", arch
    return system or "Unknown OS", arch
 
 
def cmd_desktop(args: list[str]) -> None:
    """aquin desktop install  — install the Aquin desktop app (stub until releases ship)."""
    if not args or args[0] in ("--help", "-h", "help"):
        print("Usage: aquin desktop install")
        print("")
        print("  Detect this machine's OS, confirm, then install Aquin desktop.")
        print("  Download packaging is not published yet; yes only acknowledges for now.")
        return
 
    if args[0] != "install":
        print(f"Unknown desktop subcommand: {args[0]}")
        print("Usage: aquin desktop install")
        sys.exit(1)
 
    rest = args[1:]
    if rest and rest[0] in ("--help", "-h", "help"):
        print("Usage: aquin desktop install")
        print("")
        print("  Detect this machine's OS, confirm, then install Aquin desktop.")
        print("  Download packaging is not published yet; yes only acknowledges for now.")
        return
 
    os_label, arch = _detect_desktop_platform()
    print(f"Detected: {os_label} ({arch})")
    try:
        answer = input("Install Aquin desktop on this machine? [y/N] ").strip().lower()
    except EOFError:
        print("Cancelled.")
        return
 
    if answer not in ("y", "yes"):
        print("Cancelled.")
        return
 
    print("Yes.")
    print(f"Desktop install for {os_label} ({arch}) will download and install here when builds ship.")
 
 
def cmd_list(args: list[str]) -> None:
    """aquin list simulation"""
    if not args:
        print("Usage: aquin list simulation")
        sys.exit(1)
    if args[0] in ("model", "models"):
        print("Removed. Pass a HuggingFace model id to aquin load model <id>.", file=sys.stderr)
        sys.exit(1)
    if args[0] in ("sae", "saes"):
        print("Removed. Use: aquin load sae <model-l{n}> | --user <name> | --path <file>", file=sys.stderr)
        sys.exit(1)
    if args[0] in ("watch", "watches"):
        print("Removed. Training watch is no longer part of the CLI.", file=sys.stderr)
        print("Use the SDK: aquin.init() / run.log() → ./aquin_run/", file=sys.stderr)
        sys.exit(1)
    if args[0] == "simulation":
        _cmd_simulation_catalog("list simulation", "list_simulation_runs", args[1:])
        return
    if args[0] == "simulations":
        print("Use: aquin list simulation", file=sys.stderr)
        sys.exit(1)
    print("Usage: aquin list simulation")
    sys.exit(1)
 
 
def cmd_info(args: list[str]) -> None:
    """aquin info sae <model-l{layer}>"""
    if not args or args[0] != "sae" or "--help" in args or "-h" in args:
        print("Usage: aquin info sae <model-l{n}> [--json] [--plain]")
        print("Example: aquin info sae pythia-70m-l3")
        print("")
        print("  --json    Machine-readable fields")
        print("  --plain   Rich table layout")
        if args and args[0] == "sae" and ("--help" in args or "-h" in args):
            return
        sys.exit(1)
 
    as_json = "--json" in args
    plain = "--plain" in args
    rest = [a for a in args[1:] if a not in ("--json", "--plain")]
    if not rest:
        print("Usage: aquin info sae <model-l{n}> [--json] [--plain]")
        sys.exit(1)
 
    from aquin.info_sae_display import (
        info_sae_payload,
        render_info_sae,
        render_info_sae_quiet,
    )
    from aquin.saes import get_sae
 
    r = get_sae(rest[0])
    if as_json:
        try:
            sys.stdout.reconfigure(encoding="utf-8")  # type: ignore[attr-defined]
        except Exception:
            pass
        print(json.dumps(info_sae_payload(r), ensure_ascii=False))
        return
    if plain:
        render_info_sae(r)
        return
    render_info_sae_quiet(r)
 
 
def _default_sae_load_output(sae_id: str) -> Path | None:
    """Default ~/.aquin/sae/... path for aquin load sae <model-l{n}>."""
    import re
 
    from aquin.compute.model_loader import resolve_model_id, resolve_sae_path
 
    m = re.match(r"^(.+)-l(\d+)$", sae_id)
    if not m:
        return None
    model_slug, layer_s = m.group(1), int(m.group(2))
    try:
        model_id = resolve_model_id(model_slug)
    except ValueError:
        return None
    return resolve_sae_path(model_id, layer_s)
 
 
def _layer_from_filename_cli(path: Path) -> int | None:
    stem = path.stem
    if stem.startswith("sae_layer"):
        try:
            return int(stem.replace("sae_layer", ""))
        except ValueError:
            return None
    return None
 
 
def _print_load_sae_help() -> None:
    print("Usage:")
    print("  aquin load sae <model-l{n}>              pull catalog SAE")
    print("  aquin load sae --user <name> [--layer N]  activate user-trained SAE")
    print("  aquin load sae --path <file.pt> [--layer N] [--model <id>]  # self-hosted weights")
    print("")
    print("Examples:")
    print("  aquin load sae gpt2-small-l8")
    print("  aquin sae train --layer 8 --name my-run")
    print("  aquin load sae --user my-run --layer 8")
    print("  aquin info sae <model-l{n}>")
 
 
def _warm_daemon_sae(model_id: str, layer: int) -> None:
    """Best-effort: cache the SAE in the persistent daemon so tools reuse it."""
    try:
        from .engine import model_daemon
 
        if model_daemon.is_running():
            model_daemon.warm_sae(model_id, int(layer))
    except Exception:
        pass
 
 
def _cmd_load_sae(args: list[str]) -> None:
    """aquin load sae <model-l{n}> | --user <name> | --path <file>"""
    if not args or args[0] in ("-h", "--help", "help"):
        _print_load_sae_help()
        return
 
    user_name: str | None = None
    path_arg: str | None = None
    layer: int | None = None
    model_id: str | None = None
    explicit_key: str | None = None
    output = "./sae.pt"
    explicit_output = False
    force = False
    positional: list[str] = []
 
    i = 0
    while i < len(args):
        a = args[i]
        if a in ("--force", "-f"):
            force = True
            i += 1
        elif a == "--user" and i + 1 < len(args):
            user_name = args[i + 1]
            i += 2
        elif a == "--path" and i + 1 < len(args):
            path_arg = args[i + 1]
            i += 2
        elif a == "--layer" and i + 1 < len(args):
            layer = int(args[i + 1])
            i += 2
        elif a == "--model" and i + 1 < len(args):
            model_id = args[i + 1]
            i += 2
        elif a == "--key" and i + 1 < len(args):
            explicit_key = args[i + 1]
            i += 2
        elif a.startswith("--"):
            print(f"Unknown flag: {a}", file=sys.stderr)
            sys.exit(1)
        else:
            positional.append(a)
            i += 1
 
    if user_name:
        locked = model_id or _require_locked_model_id()
        from aquin.compute.user_sae import activate_user_sae, resolve_user_sae_path
 
        try:
            sae_path = resolve_user_sae_path(locked, user_name, layer)
        except (FileNotFoundError, ValueError) as exc:
            print(f"Error: {exc}", file=sys.stderr)
            sys.exit(1)
        resolved_layer = layer if layer is not None else _layer_from_filename_cli(sae_path)
        if resolved_layer is None:
            print("Error: could not infer SAE layer. Pass --layer <n>.", file=sys.stderr)
            sys.exit(1)
        binding = activate_user_sae(
            model_id=locked,
            layer=resolved_layer,
            path=sae_path,
            name=user_name,
        )
        print(f"Active user SAE: {user_name}  layer {resolved_layer}")
        print(f"  model: {binding['model_id']}")
        print(f"  path:  {binding['path']}")
        _warm_daemon_sae(binding["model_id"], resolved_layer)
        return
 
    if path_arg:
        from aquin.compute.user_sae import activate_user_sae, resolve_path_sae
 
        try:
            sae_path, resolved_model, resolved_layer = resolve_path_sae(
                path_arg,
                model_id=model_id,
                layer=layer,
            )
        except (FileNotFoundError, ValueError) as exc:
            print(f"Error: {exc}", file=sys.stderr)
            sys.exit(1)
        binding = activate_user_sae(
            model_id=resolved_model,
            layer=resolved_layer,
            path=sae_path,
            name=sae_path.parent.name,
        )
        print(f"Active user SAE from {sae_path}")
        print(f"  model: {binding['model_id']}  layer: {resolved_layer}")
        _warm_daemon_sae(binding["model_id"], resolved_layer)
        return
 
    if not positional:
        _print_load_sae_help()
        sys.exit(1)
 
    import re
 
    from .saes import pull_sae
 
    sae_id = positional[0]
    if not explicit_output:
        default_out = _default_sae_load_output(sae_id)
        if default_out is not None:
            output = str(default_out)
 
    m = re.match(r"^(.+)-l(\d+)$", sae_id)
    if m:
        model_slug, layer_s = m.group(1), int(m.group(2))
        resolved_catalog: str | None = None
        try:
            from aquin.compute.model_loader import (
                resolve_model_id,
                resolve_sae_path,
            )
 
            resolved_catalog = resolve_model_id(model_slug)
        except ValueError:
            resolved_catalog = None
 
        if resolved_catalog:
            from aquin.compute.model_loader import resolve_sae_path
            from aquin.compute.torch_io import is_valid_sae_checkpoint_path
            from aquin.compute.user_sae import clear_active_binding
 
            cfg_layer = layer_s
            sae_path = resolve_sae_path(resolved_catalog, cfg_layer)
            if sae_path.is_file() and not force:
                if is_valid_sae_checkpoint_path(sae_path):
                    clear_active_binding()
                    print(f"SAE already on disk: {sae_path}")
                    _warm_daemon_sae(resolved_catalog, cfg_layer)
                    return
                print(f"Removing invalid catalog SAE cache: {sae_path}", file=sys.stderr)
                sae_path.unlink(missing_ok=True)
                norm_stale = sae_path.parent / f"norm_layer{cfg_layer}.pt"
                if norm_stale.is_file():
                    norm_stale.unlink(missing_ok=True)
            elif sae_path.is_file() and force:
                print(f"Re-downloading catalog SAE (--force): {sae_path}", file=sys.stderr)
                sae_path.unlink(missing_ok=True)
                norm_stale = sae_path.parent / f"norm_layer{cfg_layer}.pt"
                if norm_stale.is_file():
                    norm_stale.unlink(missing_ok=True)
            api_key = explicit_key or _get_api_key()
            print(f"Pulling SAE {resolved_catalog}-l{cfg_layer} from Aquin catalog...")
            pull_sae(f"{resolved_catalog}-l{cfg_layer}", sae_path, api_key)
            clear_active_binding()
            _warm_daemon_sae(resolved_catalog, cfg_layer)
            return
 
    api_key = explicit_key or _get_api_key()
    print(f"Loading SAE {sae_id}...")
    pull_sae(sae_id, output, api_key)
 
 
 
def _require_loaded_model_id() -> str:
    """Return the loaded model id, or exit."""
    from .compute.model_loader import get_active_model_id, resolve_model_id
 
    active = (get_active_model_id() or "").strip()
    if not active:
        print("Error: no model loaded.")
        print("       Run: aquin load model <model-id>")
        sys.exit(1)
    try:
        return resolve_model_id(active)
    except ValueError:
        return active
 
 
def _require_locked_model_id() -> str:
    """Alias — session lock removed; uses loaded model."""
    return _require_loaded_model_id()
 
 
def _assert_no_model_override(raw_args: list[str]) -> None:
    i = 0
    while i < len(raw_args):
        if raw_args[i] == "--model" and i + 1 < len(raw_args):
            print("Error: --model is not supported here.")
            print("       Run: aquin load model <id>")
            sys.exit(1)
        i += 1
 
 
def _build_tool_ctx(model_id: str | None = None) -> dict:
    """Local tool context for GPU dispatch + command tracking."""
    from .engine.main import _load_state
    from .compute.model_loader import get_active_model_id, resolve_model_id
 
    state = _load_state()
    try:
        from aquin.auth_config import resolve_api_key
        api_key = state.get("api_key") or resolve_api_key(allow_missing=True)
    except Exception:
        api_key = state.get("api_key") or ""
    base_url = (state.get("base_url") or os.environ.get("AQUIN_BASE_URL", _BASE_URL)).rstrip("/")
    active = model_id or get_active_model_id() or ""
    if active:
        try:
            active = resolve_model_id(active)
        except ValueError:
            pass
    from aquin.engine.session_memory_store import load_local_memory
    local_mem = load_local_memory("local")
    return {
        "session_id": "",
        "api_key": api_key,
        "base_url": base_url,
        "state": _inject_session_mode({
            "activeModelId": active,
            "memory": local_mem,
        }),
    }
 
 
def cmd_trace(args: list[str]) -> None:
    """aquin trace --prompt <text> --layer <n> [--check] [--umap]"""
    from aquin.cli_flags import reject_legacy_output_flags
 
    if not args or args[0] in ("-h", "--help", "help") or "-h" in args or "--help" in args:
        print("Full attribution for one prompt: generate, SAE features, causal trace, logit lens.")
        print("")
        print("Usage: aquin trace --prompt <text> --layer <n> [--check] [--umap]")
        print("")
        print("  --prompt  Input prompt (model completes from here).")
        print("  --layer   SAE layer with a downloaded checkpoint.")
        print("  --check   Save trace-check.json and trace-check.png in cwd.")
        print("  --umap    Load SAE UMAP projection after the run (web explorer).")
        print("")
        print("Docs: https://aquin.app/docs/inspection-sae/llm")
        return
 
    reject_legacy_output_flags(args)
    print("[trace] starting...", flush=True)
    _guard_llm_only_cmd("trace")
    _assert_no_model_override(args)
    tool_args = _parse_tool_flags(args)
    prompt = tool_args.get("prompt")
    layer_raw = tool_args.get("layer")
    layer: int | None = int(layer_raw) if layer_raw is not None else None
    do_check = bool(tool_args.get("check")) or "--check" in args
    want_umap = _pop_umap_flag(tool_args, args)
 
    if not prompt:
        print("Usage: aquin trace --prompt <text> --layer <n> [--check] [--umap]")
        sys.exit(1)
    if layer is None:
        from aquin.compute.model_loader import format_sae_layer_choice_message, resolve_model_id
 
        mid = _require_locked_model_id()
        try:
            mid = resolve_model_id(mid)
        except ValueError as e:
            print(f"Error: {e}"); sys.exit(1)
        print(
            format_sae_layer_choice_message(
                mid,
                command="trace",
                example_suffix=f'--prompt "{str(prompt)[:60]}"',
            ),
            file=sys.stderr,
        )
        sys.exit(1)
 
    from .compute.model_loader import resolve_model_id, ComputeNotAvailableError, require_sae_layer
    from .compute.loader_shim import apply as _shim_apply
    from .engine.sync_dispatch import run_dispatch
 
    mid = _require_locked_model_id()
    try:
        mid = resolve_model_id(mid)
    except ValueError as e:
        print(f"Error: {e}"); sys.exit(1)
 
    try:
        layer = require_sae_layer(
            mid,
            layer,
            command="trace",
            example_suffix=f'--prompt "{str(prompt)[:60]}"',
        )
    except ValueError as e:
        print(str(e), file=sys.stderr)
        sys.exit(1)
 
    _shim_apply()
    ctx = _build_tool_ctx()
    ctx["state"]["activeModelId"] = mid
 
    try:
        print(f"[trace] running full inspection (sae layer {layer})...")
        result = run_dispatch(
            "run_full_inspection",
            {"prompt": prompt, "layer": layer},
            ctx,
            command="trace",
            ensure_model=mid,
        )
        if isinstance(result, dict) and result.get("error"):
            print(result["error"]); sys.exit(1)
 
        content = result.get("content", {}) if isinstance(result, dict) else {}
        card_data = (result.get("card") or {}).get("data", {}) if isinstance(result, dict) else {}
        response = content.get("response") or card_data.get("response", "")
        top_features = content.get("top_features") or card_data.get("topFeatures", [])
        prompt_tokens = card_data.get("promptTokens", [])
        response_tokens = card_data.get("responseTokens", [])
        attribution = card_data.get("attribution", [])
        sae_layer = content.get("sae_layer") or card_data.get("saeLayer", "?")
        if sae_layer != layer:
            print(
                f"Error: trace used SAE layer {sae_layer} but --layer {layer} was requested.",
                file=sys.stderr,
            )
            sys.exit(1)
        logit_lens = card_data.get("logitLensResults", [])
 
        print(f"[trace] response: {response!r}\n")
        if prompt_tokens:
            print(f"prompt tokens : {prompt_tokens}")
            print(f"response tokens: {response_tokens}")
        print(f"sae layer      : {sae_layer}")
        print(f"\ntop response features ({len(top_features)}):")
        for f in top_features[:10]:
            ref = f.get("feature_ref") or f["feature_idx"]
            print(f"  [{ref}]  act={f['activation']:.3f}  token={f.get('token', '')!r}")
 
        if attribution:
            print(f"\nattribution ({len(attribution)} response tokens with driven features):")
            for a in attribution[:5]:
                driven = ", ".join(
                    f"{d.get('feature_ref') or d['feature_idx']}({d['activation']:.2f})"
                    for d in a.get("driven_by_features", [])[:3]
                )
                print(f"  {a.get('response_token', ''):>12}  <-  {driven}")
 
        if logit_lens:
            print(f"\nlogit lens ({len(logit_lens)} layers), last 3 layers:")
            for row in logit_lens[-3:]:
                tops = "  ".join(f"{t['token']!r}({t['prob']:.3f})" for t in row.get("top_tokens", [])[:3])
                print(f"  layer {row['layer']:>2}: {tops}")
 
        if do_check:
            from .trace_check import write_trace_check
            try:
                json_path, png_path = write_trace_check(
                    result, tool_name="run_full_inspection", cwd=os.getcwd(),
                )
                print(f"Saved {json_path}")
                print(f"Saved {png_path}")
            except Exception as exc:
                from .user_errors import die
                die(message="Could not save check output files.", exc=exc, label="trace --check")
 
        if want_umap:
            _run_umap_followup(
                ctx,
                result=result,
                tool_args={"layer": layer},
                ensure_model=mid,
                layer=layer,
            )
 
    except ComputeNotAvailableError as e:
        from .user_errors import die
        die(str(e) or "Compute is not available on this machine.")
    except Exception as e:
        from .user_errors import die
        die(exc=e)
 
 
def cmd_feature_logits(args: list[str]) -> None:
    """aquin feature logit --feature <idx> [--layer <n>] [--topk <n>] [--check] [--umap]"""
    if not args or args[0] in ("-h", "--help", "help") or "-h" in args or "--help" in args:
        print("Usage: aquin feature logit --feature <idx> [--layer <n>] [--topk <n>] [--check] [--umap]")
        return
 
    _guard_llm_only_cmd("feature logit")
    _assert_no_model_override(args)
    feature_idx = None
    layer = None
    top_k = 10
    do_check = False
    want_umap = False
    i = 0
    while i < len(args):
        if args[i] == "--feature" and i + 1 < len(args):
            feature_idx = int(args[i + 1]); i += 2
        elif args[i] == "--layer" and i + 1 < len(args):
            layer = int(args[i + 1]); i += 2
        elif args[i] in ("--topk", "--top-k") and i + 1 < len(args):
            top_k = int(args[i + 1]); i += 2
        elif args[i] == "--check":
            do_check = True; i += 1
        elif args[i] == "--umap":
            want_umap = True; i += 1
        else:
            i += 1
 
    if feature_idx is None:
        print("Usage: aquin feature logit --feature <idx> [--layer <n>] [--topk <n>] [--check] [--umap]")
        sys.exit(1)
 
    from .compute.model_loader import resolve_model_id, ComputeNotAvailableError
    from .engine.sync_dispatch import run_dispatch
 
    mid = _require_locked_model_id()
    ctx = _build_tool_ctx()
    try:
        mid = resolve_model_id(mid)
    except ValueError as e:
        print(f"Error: {e}"); sys.exit(1)
    ctx["state"]["activeModelId"] = mid
 
    tool_args: dict[str, Any] = {"feature_idx": feature_idx, "top_k": top_k}
    if layer is not None:
        tool_args["layer"] = layer
 
    try:
        result = run_dispatch("get_feature_logits", tool_args, ctx, command="feature logit", ensure_model=mid)
        if isinstance(result, dict) and result.get("error"):
            print(f"Error: {result['error']}"); sys.exit(1)
        data = result.get("content", result) if isinstance(result, dict) else result
        ref = data.get("feature_ref") or feature_idx
        print(f"feature {ref}  (model={mid}  layer={layer or 'default'})")
        print(f"\ntop {top_k} boosted tokens:")
        for t in data.get("top", []):
            print(f"  {t['token']!r:<20}  logit={t['logit']:+.4f}")
        print(f"\ntop {top_k} suppressed tokens:")
        for t in data.get("bottom", []):
            print(f"  {t['token']!r:<20}  logit={t['logit']:+.4f}")
 
        if do_check:
            from .feature_logits_check import write_feature_logits_check
            try:
                json_path, png_path = write_feature_logits_check(
                    result, tool_name="get_feature_logits", cwd=os.getcwd(),
                )
                print(f"Saved {json_path}")
                print(f"Saved {png_path}")
            except Exception as exc:
                from .user_errors import die
                die(message="Could not save check output files.", exc=exc, label="feature logit --check")
 
        if want_umap:
            _run_umap_followup(
                ctx,
                result=result,
                tool_args=tool_args,
                ensure_model=mid,
                layer=layer,
                feature_idxs=[feature_idx],
            )
    except ComputeNotAvailableError as e:
        from .user_errors import die
        die(str(e) or "Compute is not available on this machine.")
    except Exception as e:
        from .user_errors import die
        die(exc=e)
 
 
def cmd_feature_neighbors(args: list[str]) -> None:
    """aquin feature neighbor --feature <idx> [--layer <n>] [--topk <n>] [--check] [--umap]"""
    if not args or args[0] in ("-h", "--help", "help") or "-h" in args or "--help" in args:
        print("Usage: aquin feature neighbor --feature <idx> [--layer <n>] [--topk <n>] [--check] [--umap]")
        return
 
    _guard_llm_only_cmd("feature neighbor")
    _assert_no_model_override(args)
    feature_idx = None
    layer = None
    top_k = 8
    do_check = False
    want_umap = False
    i = 0
    while i < len(args):
        if args[i] == "--feature" and i + 1 < len(args):
            feature_idx = int(args[i + 1]); i += 2
        elif args[i] == "--layer" and i + 1 < len(args):
            layer = int(args[i + 1]); i += 2
        elif args[i] in ("--topk", "--top-k") and i + 1 < len(args):
            top_k = int(args[i + 1]); i += 2
        elif args[i] == "--check":
            do_check = True; i += 1
        elif args[i] == "--umap":
            want_umap = True; i += 1
        else:
            i += 1
 
    if feature_idx is None:
        print("Usage: aquin feature neighbor --feature <idx> [--layer <n>] [--topk <n>] [--check] [--umap]")
        sys.exit(1)
 
    from .compute.model_loader import resolve_model_id, ComputeNotAvailableError
    from .engine.sync_dispatch import run_dispatch
 
    mid = _require_locked_model_id()
    ctx = _build_tool_ctx()
    try:
        mid = resolve_model_id(mid)
    except ValueError as e:
        print(f"Error: {e}"); sys.exit(1)
    ctx["state"]["activeModelId"] = mid
 
    tool_args: dict[str, Any] = {"feature_idx": feature_idx, "top_k": top_k}
    if layer is not None:
        tool_args["layer"] = layer
 
    try:
        result = run_dispatch("get_feature_neighbors", tool_args, ctx, command="feature neighbor", ensure_model=mid)
        if isinstance(result, dict) and result.get("error"):
            print(f"Error: {result['error']}"); sys.exit(1)
        data = result.get("content", result) if isinstance(result, dict) else result
        ref = data.get("feature_ref") or feature_idx
        print(f"feature {ref}  (model={mid}  layer={layer or 'default'})")
        print(f"\n{top_k} nearest neighbors:")
        for n in data.get("neighbors", []):
            nref = n.get("feature_ref") or n["feature_idx"]
            print(f"  [{nref}]  similarity={n['similarity']:.4f}")
 
        if do_check:
            from .feature_neighbors_check import write_feature_neighbors_check
            try:
                json_path, png_path = write_feature_neighbors_check(
                    result, tool_name="get_feature_neighbors", cwd=os.getcwd(),
                )
                print(f"Saved {json_path}")
                print(f"Saved {png_path}")
            except Exception as exc:
                from .user_errors import die
                die(message="Could not save check output files.", exc=exc, label="feature neighbor --check")
 
        if want_umap:
            _run_umap_followup(
                ctx,
                result=result,
                tool_args=tool_args,
                ensure_model=mid,
                layer=layer,
                feature_idxs=[feature_idx],
            )
    except ComputeNotAvailableError as e:
        from .user_errors import die
        die(str(e) or "Compute is not available on this machine.")
    except Exception as e:
        from .user_errors import die
        die(exc=e)
 
 
def _push_log(line: str) -> None:
    """Legacy no-op — web session sync removed."""
    _ = line
 
 
def load_model_local(model_id: str, *, as_json: bool = False) -> str:
    """Load a model into VRAM. Returns the resolved model id."""
    from .compute.model_loader import ComputeNotAvailableError
    from .compute.model_runtime import load_weights, resolve_kind, vram_line
    from .engine.main import _load_state, _save_state
    from .load_model_display import LoadProgress, load_model_payload
    from .session_mode import mode_for_model_id, mode_label
 
    try:
        mode = mode_for_model_id(model_id)
    except ValueError as e:
        if as_json:
            print(json.dumps(load_model_payload(
                model=model_id, mode="?", daemon=False, ok=False, error=str(e),
            ), ensure_ascii=False))
        else:
            print(f"Error: {e}")
        sys.exit(1)
 
    label = mode_label(mode)
    progress = None if as_json else LoadProgress(f"Loading {model_id} · {label}")
    _push_log(f"Loading {model_id} ({label})...")
 
    try:
        resolved, _kind = resolve_kind(model_id)
        loaded_via_daemon = False
        vram: str | None = None
        elapsed: float | None = None
 
        try:
            from .engine import model_daemon
 
            if model_daemon.ensure_running():
                if progress:
                    progress.phase("engine ready")
                res = model_daemon.switch_model(resolved, on_tick=progress)
                if res and res.get("ok"):
                    loaded_via_daemon = True
                    vram = res.get("vram")
                    elapsed = progress.elapsed() if progress else None
                elif progress and res and res.get("error"):
                    progress.phase(f"engine failed · {res['error']}")
        except Exception:
            loaded_via_daemon = False
 
        if not loaded_via_daemon:
            if progress:
                progress.phase("loading weights in-process")
 
            def _on_progress(msg: str, _elapsed: float) -> None:
                if progress:
                    progress.tick(msg)
 
            resolved = load_weights(model_id, progress=_on_progress if progress else None)
            elapsed = progress.elapsed() if progress else None
            vram = vram_line()
 
        model_id = resolved
        if as_json:
            print(json.dumps(load_model_payload(
                model=model_id,
                mode=mode,
                daemon=loaded_via_daemon,
                ok=True,
                elapsed_s=elapsed,
                vram=vram,
            ), ensure_ascii=False))
        else:
            tail = f"Model ready · {model_id}"
            if vram:
                tail = f"{tail} · {vram}"
            if progress:
                progress.done(tail)
                if loaded_via_daemon:
                    progress.done("resident in background engine", indent=1)
            else:
                print(tail, flush=True)
        _push_log(f"Model ready: {model_id}  ({label} tools active)")
        state = _load_state()
        _save_state({**state, "session_mode": mode, "active_model_id": model_id})
        return model_id
    except ComputeNotAvailableError as e:
        err = str(e)
        if as_json:
            print(json.dumps(load_model_payload(
                model=model_id, mode=mode, daemon=False, ok=False, error=err,
            ), ensure_ascii=False))
        else:
            print(f"Error: {err}")
        _push_log(f"Error: {err}")
        sys.exit(1)
    except ValueError as e:
        err = str(e)
        if as_json:
            print(json.dumps(load_model_payload(
                model=model_id, mode=mode, daemon=False, ok=False, error=err,
            ), ensure_ascii=False))
        else:
            print(f"Error: {err}")
        _push_log(f"Error: {err}")
        sys.exit(1)
    except Exception as e:
        from .user_errors import friendly_message
 
        err = friendly_message(e)
        if as_json:
            print(json.dumps(load_model_payload(
                model=model_id, mode=mode, daemon=False, ok=False, error=err,
            ), ensure_ascii=False))
        else:
            print(f"Error: {err}", file=sys.stderr)
        _push_log(f"Error: {err}")
        sys.exit(1)
 
 
def cmd_unload(args: list[str]) -> None:
    """aquin unload [--stop] [--json]  — free the resident model's VRAM."""
    if args and args[0] in ("--help", "-h", "help"):
        print("Usage: aquin unload [--stop] [--json]")
        print("")
        print("  Free the loaded model's VRAM (engine stays up by default).")
        print("  --stop   Also stop the background engine")
        print("  --json   Machine-readable result")
        return
 
    as_json = "--json" in args
    stop_daemon = "--stop" in args or "--kill" in args
    from .compute.model_runtime import unload_weights
    from .load_model_display import LoadProgress, unload_payload
 
    progress = None if as_json else LoadProgress("Unloading model")
    freed = False
    daemon_running = False
    message = "Cleared active model."
    try:
        from .engine import model_daemon
 
        daemon_running = model_daemon.is_running()
        if daemon_running:
            if stop_daemon:
                if progress:
                    progress.phase("stopping engine")
                freed = model_daemon.stop()
                message = "Engine stopped · VRAM released"
            else:
                if progress:
                    progress.phase("releasing daemon VRAM")
                freed = model_daemon.unload()
                message = "Model unloaded · engine still running"
        else:
            message = "No background engine · cleared active model"
    except Exception as exc:
        if as_json:
            print(json.dumps(unload_payload(
                ok=False,
                daemon_running=False,
                stopped=stop_daemon,
                message=str(exc),
            ), ensure_ascii=False))
            sys.exit(1)
        print(f"Could not reach background engine: {exc}", file=sys.stderr)
        message = f"Engine unreachable · {exc}"
 
    try:
        if progress:
            progress.phase("releasing local VRAM")
        unload_weights()
    except Exception:
        pass
 
    if as_json:
        print(json.dumps(unload_payload(
            ok=True,
            daemon_running=daemon_running and not stop_daemon,
            stopped=stop_daemon,
            message=message,
        ), ensure_ascii=False))
        return
 
    if progress:
        progress.done(message)
    else:
        print(message, flush=True)
 
 
def cmd_load(args: list[str]) -> None:
    """aquin load model <model-id> | sae <model-l{n}>"""
    if args and args[0] == "sae":
        _cmd_load_sae(args[1:])
        return
 
    rest = list(args)
    if rest and rest[0] == "model":
        rest = rest[1:]
 
    if "--help" in rest or "-h" in rest or (rest and rest[0] == "help") or not rest:
        print("Usage: aquin load model <model-id> [--json]")
        print("       aquin load <model-id>  # legacy shorthand")
        print("       aquin load sae <model-l{n}> | --user <name> | --path <file>")
        print("")
        print("  --json   Machine-readable result")
        print("")
        print("  Env: AQUIN_DEVICE=cuda|mps|cpu|auto  (default: cuda → mps → cpu)")
        print("       AQUIN_ALLOW_CPU=1               slow CPU smoke tests")
        print("")
        print("  Pass a HuggingFace model id (e.g. meta-llama/Llama-3.2-1B-Instruct)")
        if not rest or "--help" in rest or "-h" in rest or rest[0] == "help":
            return
        sys.exit(1)
 
    if rest[0] == "pull":
        print("Use: aquin load sae", file=sys.stderr)
        sys.exit(1)
 
    as_json = "--json" in rest
    rest = [a for a in rest if a != "--json"]
 
    model_id = None
    i = 0
    while i < len(rest):
        if rest[i] == "--model" and i + 1 < len(rest):
            model_id = rest[i + 1]
            i += 2
        elif not rest[i].startswith("-") and model_id is None:
            model_id = rest[i]
            i += 1
        else:
            i += 1
 
    if not model_id:
        print("Usage: aquin load model <model-id> [--json]")
        print("  Example: aquin load model meta-llama/Llama-3.2-1B-Instruct")
        sys.exit(1)
 
    load_model_local(model_id, as_json=as_json)
 
 
def cmd_chat(args: list[str]) -> None:
    """aquin chat — Ink (via npm), Textual detect, or Rich fallback."""
    if "--ink-bridge" in args:
        from aquin.tui.ink_bridge import run_ink_chat_bridge
 
        run_ink_chat_bridge()
        return
 
    if args and args[0] in ("--help", "-h", "help"):
        print("Usage: aquin chat [--plain]")
        print("")
        print("  Multi-turn agent chat (requires login + loaded model).")
        print("  --plain    Force Rich REPL (skip Ink / Textual)")
        return
 
    from aquin.tui.detect import detect_tui_capability
    from aquin.tui.fallback import choose_chat_ui, run_chat_fallback
 
    cap = detect_tui_capability(argv=["aquin", "chat", *args])
    if choose_chat_ui(cap) == "tui":
        # Step 3: launch Textual here. Until then, Rich keeps chat working.
        run_chat_fallback(args, capability=cap)
        return
    run_chat_fallback(args, capability=cap)
 
 
def cmd_prompt(args: list[str]) -> None:
    """aquin prompt <text> — quick try-out generation against the loaded model."""
    if not args or args[0] in ("--help", "-h", "help"):
        print("Usage: aquin prompt <text>")
        print("       aquin prompting <text>")
        print("")
        print("  Generate a response from the loaded LLM (try the model).")
        print("")
        print("  --model <id>         Model to use (default: last aquin load)")
        print("  --max-tokens <n>     Max new tokens (default: 200)")
        print("  --temperature <t>    Sampling temperature (default: 0.7)")
        print("  --json               Machine-readable result")
        print("")
        print("  Examples:")
        print('    aquin prompt "Why is the sky blue?"')
        print("    aquin prompting hello")
        print("    aquin prompt --model gpt2-small Tell me a joke")
        return
 
    as_json = "--json" in args
    model_flag: str | None = None
    max_tokens = 200
    temperature = 0.7
    text_parts: list[str] = []
    i = 0
    while i < len(args):
        a = args[i]
        if a == "--json":
            i += 1
            continue
        if a == "--model" and i + 1 < len(args):
            model_flag = args[i + 1]
            i += 2
            continue
        if a in ("--max-tokens", "--max_new_tokens") and i + 1 < len(args):
            max_tokens = int(args[i + 1])
            i += 2
            continue
        if a == "--temperature" and i + 1 < len(args):
            temperature = float(args[i + 1])
            i += 2
            continue
        if a.startswith("-"):
            print(f"Unknown flag: {a}", file=sys.stderr)
            sys.exit(1)
        text_parts.append(a)
        i += 1
 
    prompt_text = " ".join(text_parts).strip()
    if not prompt_text:
        print("Usage: aquin prompt <text>")
        sys.exit(1)
 
    from .compute.model_loader import get_active_model_id
    from .prompt_display import print_prompt_result, prompt_payload
 
    model_id = (model_flag or get_active_model_id() or "").strip()
    if not model_id:
        err = "No model loaded. Run: aquin load model <id>"
        if as_json:
            print(json.dumps(prompt_payload(
                model="", prompt=prompt_text, response="", ok=False, error=err,
            ), ensure_ascii=False))
        else:
            print(f"Error: {err}", file=sys.stderr)
        sys.exit(1)
 
    used_daemon = False
    response = ""
    try:
        from .engine import model_daemon
 
        if model_daemon.is_running():
            res = model_daemon.prompt(
                prompt_text,
                model_id=model_id,
                max_new_tokens=max_tokens,
                temperature=temperature,
            )
            if res and res.get("ok"):
                response = str(res.get("response") or "")
                model_id = str(res.get("model_id") or model_id)
                used_daemon = True
            elif res and res.get("error"):
                err = str(res["error"])
                if as_json:
                    print(json.dumps(prompt_payload(
                        model=model_id, prompt=prompt_text, response="", ok=False, error=err,
                    ), ensure_ascii=False))
                else:
                    print(f"Error: {err}", file=sys.stderr)
                sys.exit(1)
    except Exception:
        used_daemon = False
 
    if not used_daemon:
        try:
            from .compute.causal_trace import run_chat
 
            if not as_json:
                from .load_model_display import print_phase
                print_phase(f"Generating · {model_id}")
            response = run_chat(
                prompt_text,
                model_id=model_id,
                max_new_tokens=max_tokens,
                temperature=temperature,
            )
        except Exception as exc:
            from .user_errors import friendly_message
 
            err = friendly_message(exc)
            if as_json:
                print(json.dumps(prompt_payload(
                    model=model_id, prompt=prompt_text, response="", ok=False, error=err,
                ), ensure_ascii=False))
            else:
                print(f"Error: {err}", file=sys.stderr)
            sys.exit(1)
 
    if as_json:
        print(json.dumps(prompt_payload(
            model=model_id,
            prompt=prompt_text,
            response=response,
            ok=True,
            daemon=used_daemon,
        ), ensure_ascii=False))
        return
 
    print_prompt_result(model=model_id, prompt=prompt_text, response=response)
 
 
# ── Raw CLI tool dispatch (Step 57) ──────────────────────────────────────────
# Verb maps live in session_mode.py.
 
from .session_mode import (
    LLM_CLI_VERBS,
    SHARED_CLI_VERBS,
    cli_verbs_for_mode,
    resolve_legacy_verb,
    verb_known,
)
 
_CLI_TOOL_MAP: dict[str, str] = cli_verbs_for_mode(None)
 
 
def _connected_state() -> dict:
    from .engine.main import _load_state
    return _load_state()
 
 
def _has_loaded_model() -> bool:
    from aquin.compute.model_loader import get_active_model_id
    return bool(get_active_model_id())
 
 
def _session_mode() -> str | None:
    """Return session mode from the loaded model."""
    from .session_mode import mode_for_active_model
    if not _has_loaded_model():
        return None
    return mode_for_active_model()
 
 
def _active_cli_tool_map() -> dict[str, str]:
    from .session_mode import SHARED_TOOLS, tools_for_mode
 
    mode = _session_mode()
    m = cli_verbs_for_mode(mode)
    if mode is not None:
        allowed = tools_for_mode(mode)
        return {verb: tool for verb, tool in m.items() if tool in allowed}
    if _has_loaded_model():
        return {verb: tool for verb, tool in m.items() if tool in SHARED_TOOLS}
    return m
 
 
def _inject_session_mode(state: dict) -> dict:
    from .session_mode import normalize_mode
 
    mode = _session_mode()
    if mode is not None:
        state = {**state, "session_mode": mode}
    return state
 
 
def _guard_llm_only_cmd(cmd: str) -> None:
    mode = _session_mode()
    if mode is None:
        print(f"[{cmd}] requires a loaded model.")
        print("  aquin load model pythia-70m")
        sys.exit(1)
 
 
def _model_switch_hint(mode: str | None) -> None:
    print("  aquin load model pythia-70m")
 
 
def _guard_tool_verb(verb: str) -> None:
    from .session_mode import mode_label
 
    verb = resolve_legacy_verb(verb)
    active = _active_cli_tool_map()
    if verb in active:
        return
    mode = _session_mode()
    if mode is None and _has_loaded_model() is False:
        print(f"[{verb}] no model loaded — run:")
        print("  aquin load model llama-3.2-1b")
        sys.exit(1)
    if mode is not None and verb_known(verb):
        print(f"[{verb}] not available for the loaded {mode_label(mode)} model.")
        _model_switch_hint(mode)
        sys.exit(1)
    print(f"Unknown command: {verb}")
    print("Run  aquin help  to see available commands.")
    sys.exit(1)
 
 
def _coerce_flag_value(val: str) -> Any:
    """Parse CLI flag values: JSON arrays/objects, bools, numbers, or raw strings."""
    import json as _json
 
    if not isinstance(val, str):
        return val
    stripped = val.strip()
    if stripped.startswith(("[", "{")):
        try:
            return _json.loads(stripped)
        except _json.JSONDecodeError:
            pass
    if stripped.lower() == "true":
        return True
    if stripped.lower() == "false":
        return False
    try:
        return int(stripped)
    except ValueError:
        pass
    try:
        return float(stripped)
    except ValueError:
        pass
    return val
 
 
def _normalize_topk_flags(tool_args: dict[str, Any]) -> None:
    if tool_args.get("topk") is not None and tool_args.get("top_k") is None:
        tool_args["top_k"] = tool_args.pop("topk")
    if tool_args.get("topk_retrieval") is not None and tool_args.get("top_k_retrieval") is None:
        tool_args["top_k_retrieval"] = tool_args.pop("topk_retrieval")
 
 
def _parse_tool_flags(raw: list[str]) -> dict[str, Any]:
    """Parse --key value flags into an args dict."""
    from aquin.cli_flags import reject_legacy_output_flags
 
    reject_legacy_output_flags(raw)
    tool_args: dict[str, Any] = {}
    i = 0
    while i < len(raw):
        a = raw[i]
        if a.startswith("--"):
            if "=" in a:
                key, _, val = a[2:].partition("=")
                tool_args[key.replace("-", "_")] = _coerce_flag_value(val)
                i += 1
            elif i + 1 < len(raw) and not raw[i + 1].startswith("--"):
                key = a[2:].replace("-", "_")
                tool_args[key] = _coerce_flag_value(raw[i + 1])
                i += 2
            else:
                tool_args[a[2:].replace("-", "_")] = True
                i += 1
        else:
            i += 1
    _normalize_topk_flags(tool_args)
    return tool_args
 
 
# Commands where --umap loads the SAE feature map after the primary result.
_UMAP_FOLLOWUP_VERBS: frozenset[str] = frozenset({
    "trace",
    "feature locate",
    "feature logit",
    "feature neighbor",
    "steer",
    "multi-steer",
    "benchmark",
    "sae-stats",
    "features",
    "contrastive",
    "interp",
    "browser",
    "graph",
    "circuit",
    "steer",
    "absorption",
    "polysemy",
    "faithfulness",
    "decomp",
})
 
 
def _pop_umap_flag(tool_args: dict[str, Any], raw_args: list[str] | None = None) -> bool:
    wanted = bool(tool_args.pop("umap", False))
    if raw_args and "--umap" in raw_args:
        wanted = True
    return wanted
 
 
def _coerce_feature_idx(value: Any) -> int | None:
    try:
        if value is None or isinstance(value, bool):
            return None
        return int(value)
    except (TypeError, ValueError):
        return None
 
 
def _feature_idxs_from_payload(result: Any, tool_args: dict[str, Any] | None = None) -> list[int]:
    """Collect feature indices from a tool payload / args for UMAP highlighting."""
    idxs: list[int] = []
    seen: set[int] = set()
 
    def _add(raw: Any) -> None:
        idx = _coerce_feature_idx(raw)
        if idx is None or idx in seen:
            return
        seen.add(idx)
        idxs.append(idx)
 
    args = tool_args or {}
    for key in ("feature_idx", "feature", "target_feature_idx", "chosen_feature_idx"):
        _add(args.get(key))
 
    if not isinstance(result, dict):
        return idxs
 
    payload = result.get("content", result) if isinstance(result.get("content"), dict) else result
    if not isinstance(payload, dict):
        return idxs
 
    for key in ("feature_idx", "feature", "target_feature_idx", "chosen_feature_idx"):
        _add(payload.get(key))
 
    for key in ("top_features", "rankings", "neighbors", "features"):
        rows = payload.get(key)
        if not isinstance(rows, list):
            continue
        for row in rows:
            if isinstance(row, dict):
                _add(row.get("feature_idx") or row.get("feature") or row.get("idx"))
 
    card = result.get("card") if isinstance(result.get("card"), dict) else None
    if card:
        data = card.get("data") if isinstance(card.get("data"), dict) else {}
        _add(data.get("featureIdx") or data.get("chosenFeatureIdx"))
        for key in ("topFeatures", "rankings", "neighbors", "features"):
            rows = data.get(key)
            if not isinstance(rows, list):
                continue
            for row in rows:
                if isinstance(row, dict):
                    _add(row.get("feature_idx") or row.get("featureIdx") or row.get("feature"))
 
    return idxs
 
 
def _layer_from_payload(result: Any, tool_args: dict[str, Any] | None = None) -> int | None:
    args = tool_args or {}
    for key in ("layer", "sae_layer"):
        idx = _coerce_feature_idx(args.get(key))
        if idx is not None:
            return idx
    if not isinstance(result, dict):
        return None
    payload = result.get("content", result) if isinstance(result.get("content"), dict) else result
    if isinstance(payload, dict):
        for key in ("layer", "sae_layer"):
            idx = _coerce_feature_idx(payload.get(key))
            if idx is not None:
                return idx
    card = result.get("card") if isinstance(result.get("card"), dict) else None
    if isinstance(card, dict):
        data = card.get("data") if isinstance(card.get("data"), dict) else {}
        for key in ("layer", "saeLayer", "sae_layer"):
            idx = _coerce_feature_idx(data.get(key))
            if idx is not None:
                return idx
    return None
 
 
def _run_umap_followup(
    ctx: dict[str, Any],
    *,
    result: Any = None,
    tool_args: dict[str, Any] | None = None,
    ensure_model: str | None = None,
    layer: int | None = None,
    feature_idxs: list[int] | None = None,
) -> None:
    """After a feature-producing command, load UMAP and sync the explorer card."""
    from .cli_output import print_tool_result
    from .engine.sync_dispatch import run_dispatch
 
    resolved_layer = layer if layer is not None else _layer_from_payload(result, tool_args)
    idxs = feature_idxs if feature_idxs is not None else _feature_idxs_from_payload(result, tool_args)
    umap_args: dict[str, Any] = {}
    if resolved_layer is not None:
        umap_args["layer"] = resolved_layer
    if ensure_model:
        umap_args["model_id"] = ensure_model
 
    print("\n[--umap] loading UMAP projection...", flush=True)
    try:
        umap_result = run_dispatch(
            "ensure_umap_loaded",
            umap_args,
            ctx,
            command="umap",
            ensure_model=ensure_model,
        )
    except Exception as exc:
        print(f"[--umap] failed: {exc}", file=sys.stderr)
        return
 
    print_tool_result("umap", umap_result, tool_name="ensure_umap_loaded")
    if isinstance(umap_result, dict) and umap_result.get("error"):
        return
    if idxs:
        shown = ", ".join(f"f{i}" for i in idxs[:8])
        more = f" (+{len(idxs) - 8} more)" if len(idxs) > 8 else ""
        print(f"[--umap] features to inspect: {shown}{more}")
 
 
def _cmd_simulation_catalog(verb: str, tool_name: str, extra_args: list[str]) -> None:
    """List/load/compare saved simulation runs — no GPU model required."""
    from .engine.sync_dispatch import dispatch_model_free
    from .engine.tools.registry import _load_stubs  # noqa: ensure registered
 
    _load_stubs()
    tool_args = _parse_tool_flags(extra_args)
    ctx = {**_build_tool_ctx(), "cwd": os.getcwd()}
 
    try:
        result = dispatch_model_free(tool_name, tool_args, ctx, command=verb)
    except Exception as exc:
        from .user_errors import die
        die(exc=exc, label=verb)
 
    from .cli_output import print_tool_result
    print_tool_result(verb, result, tool_name=tool_name)
    if isinstance(result, dict) and result.get("error"):
        sys.exit(1)
 
 
def _require_attention_args(verb: str, tool_name: str, tool_args: dict[str, Any]) -> None:
    """`aquin check attention` must include --prompt."""
    if verb != "attention":
        return
 
    if tool_name == "run_attention_routing":
        prompt = tool_args.get("prompt") or tool_args.get("text")
        if not prompt or not str(prompt).strip():
            print("Error: --prompt is required.", file=sys.stderr)
            print("Usage: aquin check attention --prompt <string> [--check]", file=sys.stderr)
            sys.exit(1)
        tool_args["prompt"] = str(prompt).strip()
 
 
def _save_tool_check(verb: str, result: dict[str, Any], tool_name: str | None) -> None:
    """Write JSON + PNG for verbs that support --check."""
    if not isinstance(result, dict):
        return
 
    from .benchmark_check import has_benchmark_payload
    from .audit_check import has_audit_payload
    from .consistency_eval_check import has_consistency_eval_payload
    from .suppression_eval_check import has_suppression_eval_payload
    from .boundary_eval_check import has_boundary_eval_payload
    from .red_team_check import has_red_team_payload
    from .eval_check import has_eval_payload
 
    if result.get("error"):
        if verb in ("benchmark", "benchmarks") and has_benchmark_payload(result):
            pass
        elif verb == "audit" and has_audit_payload(result):
            pass
        elif verb == "consistency-eval" and has_consistency_eval_payload(result):
            pass
        elif verb == "suppression-eval" and has_suppression_eval_payload(result):
            pass
        elif verb == "boundary-eval" and has_boundary_eval_payload(result):
            pass
        elif verb == "red-team" and has_red_team_payload(result):
            pass
        elif verb == "eval" and has_eval_payload(result):
            pass
        else:
            print(f"[{verb} --check] skipped save: {result['error']}", file=sys.stderr)
            return
 
    writers: dict[str, tuple[str, str]] = {
        "attention": ("aquin.attention_check", "write_attention_check"),
        "layer-analysis": ("aquin.layer_analysis_check", "write_layer_analysis_check"),
        "perturbation": ("aquin.perturbation_check", "write_perturbation_check"),
        "check-weights": ("aquin.check_weights_check", "write_check_weights_check"),
        "benchmark": ("aquin.benchmark_check", "write_benchmark_check"),
        "benchmarks": ("aquin.benchmark_check", "write_benchmark_check"),
        "audit": ("aquin.audit_check", "write_audit_check"),
        "consistency-eval": ("aquin.consistency_eval_check", "write_consistency_eval_check"),
        "suppression-eval": ("aquin.suppression_eval_check", "write_suppression_eval_check"),
        "boundary-eval": ("aquin.boundary_eval_check", "write_boundary_eval_check"),
        "red-team": ("aquin.red_team_check", "write_red_team_check"),
        "eval": ("aquin.eval_check", "write_eval_check"),
    }
    spec = writers.get(verb)
    if not spec:
        return
 
    import importlib
 
    mod = importlib.import_module(spec[0])
    write_fn = getattr(mod, spec[1])
    try:
        json_path, png_path = write_fn(result, tool_name=tool_name, cwd=os.getcwd())
        print(f"Saved {json_path}")
        print(f"Saved {png_path}")
    except Exception as exc:
        from .user_errors import die
        die(message="Could not save check output files.", exc=exc, label=f"{verb} --check")
 
 
def cmd_tool(verb: str, raw_args: list[str]) -> None:
    """Generic handler for raw CLI tool verbs — thin argv layer over ``aquin.sdk``."""
    verb = resolve_legacy_verb(verb)
    _guard_tool_verb(verb)
    _assert_no_model_override(raw_args)
 
    from .engine.tools.registry import _load_stubs  # noqa: ensure registered
    from .session_mode import tool_requires_model
 
    active_map = _active_cli_tool_map()
    tool_name = active_map[verb]
    needs_model = tool_requires_model(tool_name)
    tool_args = _parse_tool_flags(raw_args)
    tool_check = bool(tool_args.pop("check", False)) or "--check" in raw_args
    want_umap = _pop_umap_flag(tool_args, raw_args)
    tool_args.pop("model_id", None)
 
    if want_umap and verb not in _UMAP_FOLLOWUP_VERBS:
        print(
            f"Error: --umap is for SAE feature commands "
            f"(trace, feature locate/logit/neighbor, steer, benchmark).",
            file=sys.stderr,
        )
        sys.exit(1)
 
    if verb == "steer" and tool_args.get("save") and tool_name == "run_steer_and_show":
        if tool_args.get("vector") or tool_args.get("vector_path"):
            print("Error: --save exports a feature vector; use --vector to apply a saved vector.", file=sys.stderr)
            sys.exit(1)
        if tool_args.get("eval"):
            print("Error: cannot combine --eval with --save. Export first, then: aquin steer --eval --vector <path>", file=sys.stderr)
            sys.exit(1)
        if tool_args.get("feature_idx") is None:
            print("Usage: aquin steer --feature_idx <n> --save <path> [--layer N]", file=sys.stderr)
            sys.exit(1)
        tool_name = "extract_steer_vector"
        needs_model = tool_requires_model(tool_name)
 
    _require_attention_args(verb, tool_name, tool_args)
 
    from .compute.model_loader import get_active_model_id
    from aquin.sdk._runtime import build_ctx, invoke
 
    active_model = get_active_model_id() or ""
    if active_model and needs_model:
        tool_args.setdefault("model_id", active_model)
 
    try:
        result = invoke(
            tool_name,
            tool_args,
            command=verb,
            needs_model=needs_model,
            model_id=active_model or None,
            raise_on_error=False,
        )
    except NotImplementedError as e:
        from .user_errors import die
        die(f"This command is not available yet.", exc=e, label=verb)
    except Exception as e:
        from .user_errors import die
        die(exc=e, label=verb)
 
    from .cli_output import print_tool_result
    print_tool_result(verb, result, tool_name=tool_name)
 
    if tool_check:
        _save_tool_check(verb, result, tool_name)
 
    if want_umap and isinstance(result, dict) and not result.get("error"):
        _run_umap_followup(
            build_ctx(model_id=active_model or None),
            result=result,
            tool_args=tool_args,
            ensure_model=active_model or None,
        )
 
    if isinstance(result, dict) and result.get("error"):
        sys.exit(1)
 
 
def _try_simulation_phrase(args: list[str]) -> bool:
    """Dispatch aquin list simulation | compare simulation; redirect load simulation."""
    if len(args) < 2:
        return False
    prefix, noun = args[0], args[1]
    if prefix == "list" and noun in ("simulation", "simulations"):
        cmd_list(["simulation"] + args[2:])
        return True
    if prefix == "load" and noun == "simulation":
        print("Use: aquin replay simulation --run_id <id>", file=sys.stderr)
        sys.exit(1)
    if prefix == "compare" and noun == "simulation":
        _guard_tool_verb("compare simulation")
        _cmd_simulation_catalog("compare simulation", "compare_simulations", args[2:])
        return True
    return False
 
 
def cmd_help(args: list[str]) -> None:
    if args and args[0] == "commands":
        from .commands_cli import cmd_commands
 
        cmd_commands(["--help"])
        return
 
    if args and args[0] in ("--help", "-h", "help"):
        print("Usage: aquin help [commands] [--json] [--plain]")
        print("")
        print("  Show the CLI command map.")
        print("  commands   Tracked-command subcommands")
        print("  --json     Machine-readable sections")
        print("  --plain    Full ASCII banner layout")
        return
 
    as_json = "--json" in args
    plain = "--plain" in args
 
    from aquin.compute.model_loader import get_active_model_id
 
    model_id = get_active_model_id()
    mode = _session_mode()
    show_llm = True
 
    from .help_display_ascii import help_payload, render_help, render_help_quiet
 
    if as_json:
        try:
            sys.stdout.reconfigure(encoding="utf-8")  # type: ignore[attr-defined]
        except Exception:
            pass
        print(json.dumps(help_payload(
            connected=True,
            mode=mode,  # type: ignore[arg-type]
            show_llm=show_llm,
            model_id=model_id,
        ), ensure_ascii=False))
        return
 
    if plain:
        render_help(
            connected=True,
            mode=mode,  # type: ignore[arg-type]
            show_llm=show_llm,
            model_id=model_id,
        )
        return
 
    render_help_quiet(
        connected=True,
        mode=mode,  # type: ignore[arg-type]
        show_llm=show_llm,
        model_id=model_id,
    )
 
 
 
def main() -> None:
    from aquin.quiet import install_cli_quiet_mode
 
    install_cli_quiet_mode()
    try:
        _cli_main()
    except KeyboardInterrupt:
        sys.exit(130)
    except SystemExit:
        raise
    except Exception as exc:
        from .user_errors import die
        die(exc=exc)
 
 
def _cli_main() -> None:
    args = sys.argv[1:]
 
    if args and args[0] == "--cli-version":
        from aquin.version_display import version_payload
 
        print(version_payload()["cli"])
        return
 
    if "--accept-license" in args:
        from aquin.license_acceptance import accept_license_cli
 
        accept_license_cli()
        return
 
    if not args:
        from aquin.ascii_layout import print_logo
 
        print_logo()
        return
 
    cmd = resolve_legacy_verb(args[0])
    # Common typos
    _CMD_TYPOS = {"stauts": "status"}
    cmd = _CMD_TYPOS.get(cmd, cmd)
    rest = args[1:]
 
    from aquin.license_acceptance import ensure_license_accepted
 
    ensure_license_accepted(cmd=cmd)
 
    if _try_simulation_phrase(args):
        return
 
    if cmd == "engine":
        print("Removed. Use: aquin load model <id>")
        sys.exit(1)
    elif cmd in ("login", "logout", "switch"):
        print(f"Removed. This framework has no account auth ({cmd}).", file=sys.stderr)
        sys.exit(1)
    elif cmd == "chat":
        cmd_chat(rest)
    elif cmd in ("prompt", "prompting"):
        cmd_prompt(rest)
    elif cmd == "load":
        cmd_load(rest)
    elif cmd == "unload":
        cmd_unload(rest)
    elif cmd == "commands":
        from .commands_cli import cmd_commands
        cmd_commands(rest)
    elif cmd == "watch":
        print("Removed. Training watch is no longer part of the CLI.", file=sys.stderr)
        print("Use the SDK: aquin.init() / run.log() / run.checkpoint() → ./aquin_run/", file=sys.stderr)
        sys.exit(1)
    elif cmd == "dataset-generate":
        print("Removed. Pass --dataset <path> or --topic to aqit simulate.", file=sys.stderr)
        sys.exit(1)
    elif cmd in ("pairs-generate", "embed-pairs-generate"):
        print("Removed. Pass --dataset <path> or --topic to aqit simulate.", file=sys.stderr)
        sys.exit(1)
    elif cmd == "dataset-analyze":
        print("Removed. Dataset quality check is no longer available.", file=sys.stderr)
        sys.exit(1)
    elif cmd == "space-decomp":
        print("Removed. Embedding space decomposition is no longer available.", file=sys.stderr)
        sys.exit(1)
    elif cmd in ("mem", "mem-write", "mem-read"):
        print("Removed.", file=sys.stderr)
        sys.exit(1)
    elif cmd in ("run-code", "save-artifact"):
        print("Removed.", file=sys.stderr)
        sys.exit(1)
    elif cmd == "status":
        cmd_status(rest)
    elif cmd == "update":
        cmd_update(rest)
    elif cmd == "desktop":
        cmd_desktop(rest)
    elif cmd == "setup":
        from .setup_cli import cmd_setup
        cmd_setup(rest)
    elif cmd == "version" or cmd in ("--version", "-V"):
        cmd_version(rest)
    elif cmd == "license":
        cmd_license(rest)
    elif cmd in ("help", "--help", "-h"):
        cmd_help(rest)
    else:
        from aqit.cli import dispatch as aqit_dispatch, is_aqit_verb
 
        if is_aqit_verb(cmd):
            aqit_dispatch([cmd, *rest])
            return
        print(f"Unknown command: {cmd}")
        print("Run  aqit help  for Recipe/train/inspect. aquin help for the shell.")
        sys.exit(1)