dev

#!/usr/bin/env bash
#
# Krasis dev helper — single entry point for build, run, test, benchmark.
#
# Usage:
#   ./dev build                     Rebuild Rust extension into this repo only
#   ./dev run <config>              Launch server from a test config
#   ./dev run <config> --benchmark  Launch with benchmark, then serve
#   ./dev test <config>             Run short model test (benchmark + network tests)
#   ./dev test <config> --thorough  Run thorough model test (stress + benchmark + network + large)
#   ./dev network <port>            Run network tests against a running server
#   ./dev network <port> --large    Include large-prompt tests
#   ./dev benchmark <config>        Run standard benchmark (prefill/decode/round trip) and exit
#   ./dev speed-test                Run the fixed standard speed benchmark (QCN INT4 AWQ Polar4)
#   ./dev perplexity <config>       Run perplexity eval (WikiText-2 by default) and exit
#   ./dev release-test <model>      Run full release test (4 configs, produces report)
#   ./dev release-test-all [target]  Run release test on all supported models sequentially
#                                     If a target name is given (e.g. "rc8"), progress is saved
#                                     to a state file. Re-running with the same target resumes
#                                     from the first non-passed model. Without a target, runs
#                                     all models with no resume capability.
#   ./dev reference-test <config>   Compare engine output against BF16 reference data
#   ./dev reference-test --compare <model>  Run 4 quant configs, produce comparative report
#   ./dev reference-inventory       Inventory stored reference artifacts and classify contract state
#   ./dev capture-box <subcommand>  Run HF capture commands through the strict paid-box wrapper
#   ./dev capture-preflight         Verify or bootstrap a host for HF reference capture
#   ./dev reference-prep <model>    Prepare HF capture deps and download a public model
#   ./dev generate-reference <model>  Generate a BF16 HF reference capture
#   ./dev trace-diff <expected> <actual>  Diff two KRASIS_TRACE logs offline
#   ./dev awq-calibrate <config>    Run AWQ calibration, produce attention template
#   ./dev test-kernels [name]       Run CUDA kernel unit tests (or filter by name)
#   ./dev chat [args]               Run krasis chat from source
#   ./dev sanity                    Run sanity test prompts against running server
#   ./dev kill                      Kill all krasis processes and free GPU memory
#   ./dev install                   Sync krasis command to ~/.local/bin (uses source)
#   ./dev shell                     Drop into a Python shell bound to this repo
#   ./dev python <args...>          Run arbitrary Python bound to this repo
#   ./dev verify-imports            Verify imports resolve inside this repo
#
# Timing / profiling:
#   Add --timing to any command to enable per-component decode timing.
#   This inserts GPU synchronization barriers between each decode phase and
#   prints a detailed timing breakdown at the end. Adds ~30-50% overhead so
#   do NOT use for speed benchmarks — only for understanding where time goes.
#
#   Examples:
#     ./dev benchmark qcn --timing    Profile decode components
#     ./dev run qcn --timing          Run server with timing enabled
#     ./dev benchmark qcn             Normal speed benchmark (no timing overhead)
#
# Debug tracing:
#   KRASIS_TRACE=1 enables structured decode tracing with zero normal-mode cost.
#   Optional filters:
#     KRASIS_TRACE_STEPS=0-1
#     KRASIS_TRACE_LAYERS=0-3
#     KRASIS_TRACE_COMPONENTS=embedding,gqa,moe,final
#     KRASIS_TRACE_VALUES=8
#     KRASIS_TRACE_ELEMS=2048
#     KRASIS_TRACE_DUMP_DIR=/tmp/krasis-trace
#     KRASIS_TRACE_PY_COMPARE=1  # one first-step compare via built test endpoint during validate
#
# Configs are short names that map to testconfigs/*.conf:
#   qcn         -> testconfigs/qcn-4-4-a16.conf
#   qcn-a4      -> testconfigs/qcn-4-4-awq.conf
#   q235        -> testconfigs/q235-4-4-a16.conf
#   q235-a4     -> testconfigs/q235-4-4-awq.conf
#   v2lite      -> testconfigs/v2lite-4-4-a16.conf
#   deepseek-vl -> testconfigs/deepseek-vl2-4-4-a16.conf
#   qwen35      -> testconfigs/qwen35-4-4-a16.conf
#   qwen35-a4   -> testconfigs/qwen35-4-4-awq.conf
#   gemma       -> tests/gemma-4-4-a16.conf
#   minimax     -> tests/minimax-m2-4-4-a16.conf
#
set -euo pipefail

# ── Environment ──────────────────────────────────────────────────────
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
infer_repo_home() {
    local path="$1"
    case "$path" in
        /home/*/*)
            local remainder="${path#/home/}"
            local user="${remainder%%/*}"
            [[ -n "$user" ]] && echo "/home/$user"
            ;;
        /Users/*/*)
            local remainder="${path#/Users/}"
            local user="${remainder%%/*}"
            [[ -n "$user" ]] && echo "/Users/$user"
            ;;
        *)
            return 1
            ;;
    esac
}

REFERENCE_CAPTURE_ROOT_SOURCE="home"
if [[ -n "${KRASIS_REFERENCE_CAPTURE_ROOT:-}" ]]; then
    REFERENCE_CAPTURE_ROOT="${KRASIS_REFERENCE_CAPTURE_ROOT}"
    REFERENCE_CAPTURE_ROOT_SOURCE="env"
elif REPO_HOME="$(infer_repo_home "$SCRIPT_DIR")"; then
    REFERENCE_CAPTURE_ROOT="${REPO_HOME}/.krasis"
    REFERENCE_CAPTURE_ROOT_SOURCE="repo_home"
else
    REFERENCE_CAPTURE_ROOT="${HOME}/.krasis"
fi

PYTHON="/home/main/miniconda3/envs/ktransformers/bin/python"
export PYTHONUNBUFFERED=1
# WSL2: add CUDA driver path so Rust cudarc can find libcuda.so
[[ -d /usr/lib/wsl/lib ]] && export LD_LIBRARY_PATH="/usr/lib/wsl/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
PIP="/home/main/miniconda3/envs/ktransformers/bin/pip"
MATURIN="/home/main/.local/bin/maturin"
SO_FILE="$SCRIPT_DIR/python/krasis/krasis.cpython-311-x86_64-linux-gnu.so"
WHEEL_DIR="$SCRIPT_DIR/target/dev-wheels"
REPO_PYTHONPATH="$SCRIPT_DIR/python"
REPO_LIBPATH="$SCRIPT_DIR/python/krasis.libs"
STANDARD_SPEED_TEST_CONFIG="$SCRIPT_DIR/tests/qcn-polar4-awq.conf"
REFERENCE_CAPTURE_MODELS_DIR="${REFERENCE_CAPTURE_ROOT}/models"
REFERENCE_CAPTURE_VENV="${REFERENCE_CAPTURE_ROOT}/reference-capture-venv"
REFERENCE_CAPTURE_READY_JSON="${REFERENCE_CAPTURE_ROOT}/capture-host-ready.json"
REFERENCE_CAPTURE_READY_STAMP="${REFERENCE_CAPTURE_ROOT}/capture-host-ready.stamp"

RED="\033[0;31m"
GREEN="\033[0;32m"
YELLOW="\033[1;33m"
CYAN="\033[0;36m"
BOLD="\033[1m"
NC="\033[0m"

info()  { echo -e "${CYAN}${BOLD}=>${NC} $*"; }
ok()    { echo -e "${GREEN}${BOLD}OK${NC} $*"; }
warn()  { echo -e "${YELLOW}${BOLD}!!${NC} $*"; }
err()   { echo -e "${RED}${BOLD}ERROR${NC} $*" >&2; exit 1; }

build_base_ld_library_path() {
    local current="${LD_LIBRARY_PATH:-}"
    local preferred=""

    # Some container images expose CUDA compat stubs that can override the real
    # host driver and trigger cuda error 803. When both exist, prefer the host
    # driver directories explicitly.
    if [[ -f /usr/lib/x86_64-linux-gnu/libcuda.so.1 ]] && compgen -G "/usr/local/cuda*/compat/libcuda.so*" >/dev/null; then
        preferred="/usr/lib/x86_64-linux-gnu:/lib/x86_64-linux-gnu"
    fi

    if [[ -n "$preferred" ]]; then
        echo "$preferred${current:+:$current}"
    else
        echo "$current"
    fi
}

create_run_dir() {
    local run_type="$1"
    local ts dir counter=1
    mkdir -p "$SCRIPT_DIR/logs"
    ts=$(date +%Y%m%d_%H%M%S)
    dir="$SCRIPT_DIR/logs/${run_type}_${ts}"
    while [[ -e "$dir" ]]; do
        dir="$SCRIPT_DIR/logs/${run_type}_${ts}_$counter"
        counter=$((counter + 1))
    done
    mkdir -p "$dir"
    echo "$dir"
}

# ── Sanity checks ────────────────────────────────────────────────────
[[ -f "$PYTHON" ]] || err "Python not found at $PYTHON
  See DEV.md for environment setup instructions."

[[ -f "$MATURIN" ]] || err "maturin not found at $MATURIN
  Install: pip install maturin"

# ── Config resolver ──────────────────────────────────────────────────
resolve_config() {
    local name="$1"
    case "$name" in
        qcn|QCN)            echo "$SCRIPT_DIR/testconfigs/qcn-4-4-a16.conf" ;;
        qcn-a4)             echo "$SCRIPT_DIR/testconfigs/qcn-4-4-awq.conf" ;;
        q235)               echo "$SCRIPT_DIR/testconfigs/q235-4-4-a16.conf" ;;
        q235-a4)            echo "$SCRIPT_DIR/testconfigs/q235-4-4-awq.conf" ;;
        v2lite|V2-Lite)     echo "$SCRIPT_DIR/testconfigs/v2lite-4-4-a16.conf" ;;
        deepseek-vl|vl2)    echo "$SCRIPT_DIR/testconfigs/deepseek-vl2-4-4-a16.conf" ;;
        qwen35|Qwen3.5)     echo "$SCRIPT_DIR/testconfigs/qwen35-4-4-a16.conf" ;;
        qwen35-a4)          echo "$SCRIPT_DIR/testconfigs/qwen35-4-4-awq.conf" ;;
        gemma|Gemma)        echo "$SCRIPT_DIR/tests/gemma-4-4-a16.conf" ;;
        minimax|MiniMax)    echo "$SCRIPT_DIR/tests/minimax-m2-4-4-a16.conf" ;;
        *)
            # Try as a direct path
            if [[ -f "$name" ]]; then
                echo "$name"
            elif [[ -f "$SCRIPT_DIR/testconfigs/$name" ]]; then
                echo "$SCRIPT_DIR/testconfigs/$name"
            elif [[ -f "$SCRIPT_DIR/tests/$name" ]]; then
                echo "$SCRIPT_DIR/tests/$name"
            else
                err "Unknown config: $name
  Available: qcn, qcn-a4, q235, q235-a4, v2lite, deepseek-vl, qwen35, qwen35-a4, gemma, minimax
  Or pass a path to a .conf file."
            fi
            ;;
    esac
}

# No need for conf_to_args — server.py has native --config support

# Extract port from a .conf file
conf_port() {
    local conf="$1"
    grep '^CFG_PORT=' "$conf" | head -1 | cut -d'"' -f2
}

# ── Auto-rebuild check ──────────────────────────────────────────────
needs_rebuild() {
    # If .so doesn't exist, definitely need to build
    [[ ! -f "$SO_FILE" ]] && return 0

    local so_mtime
    so_mtime=$(stat --format="%Y" "$SO_FILE" 2>/dev/null) || return 0

    # Check if any Rust source or Cargo files are newer than the .so
    while IFS= read -r -d '' f; do
        local f_mtime
        f_mtime=$(stat --format="%Y" "$f" 2>/dev/null) || continue
        if [[ "$f_mtime" -gt "$so_mtime" ]]; then
            return 0
        fi
    done < <(find "$SCRIPT_DIR/src" -name "*.rs" -print0 2>/dev/null)

    # Check Cargo.toml and build.rs
    for f in "$SCRIPT_DIR/Cargo.toml" "$SCRIPT_DIR/build.rs"; do
        if [[ -f "$f" ]]; then
            local f_mtime
            f_mtime=$(stat --format="%Y" "$f" 2>/dev/null) || continue
            if [[ "$f_mtime" -gt "$so_mtime" ]]; then
                return 0
            fi
        fi
    done

    return 1
}

auto_rebuild() {
    if needs_rebuild; then
        warn "Rust source is newer than compiled extension. Rebuilding..."
        do_build
    fi
}

run_repo_python() {
    local ld_path
    ld_path="$(build_base_ld_library_path)"
    if [[ -d "$REPO_LIBPATH" ]]; then
        ld_path="$REPO_LIBPATH${ld_path:+:$ld_path}"
    fi
    PYTHONNOUSERSITE=1 PYTHONPATH="$REPO_PYTHONPATH${PYTHONPATH:+:$PYTHONPATH}" LD_LIBRARY_PATH="$ld_path" "$PYTHON" "$@"
}

exec_repo_python() {
    local ld_path
    ld_path="$(build_base_ld_library_path)"
    if [[ -d "$REPO_LIBPATH" ]]; then
        ld_path="$REPO_LIBPATH${ld_path:+:$ld_path}"
    fi
    exec env PYTHONNOUSERSITE=1 PYTHONPATH="$REPO_PYTHONPATH${PYTHONPATH:+:$PYTHONPATH}" LD_LIBRARY_PATH="$ld_path" "$PYTHON" "$@"
}

reference_capture_python_bin() {
    local capture_python="$REFERENCE_CAPTURE_VENV/bin/python"
    if [[ -x "$capture_python" ]]; then
        echo "$capture_python"
    else
        echo "$PYTHON"
    fi
}

run_reference_capture_python() {
    local ld_path python_bin
    python_bin="$(reference_capture_python_bin)"
    ld_path="$(build_base_ld_library_path)"
    if [[ -d "$REPO_LIBPATH" ]]; then
        ld_path="$REPO_LIBPATH${ld_path:+:$ld_path}"
    fi
    PYTHONNOUSERSITE=1 PYTHONPATH="$REPO_PYTHONPATH${PYTHONPATH:+:$PYTHONPATH}" LD_LIBRARY_PATH="$ld_path" "$python_bin" "$@"
}

print_capture_box_env() {
    info "Capture-box environment"
    echo "repo_root=$SCRIPT_DIR"
    echo "repo_home=$1"
    echo "HOME=$HOME"
    echo "PATH=$PATH"
    echo "capture_root=$REFERENCE_CAPTURE_ROOT"
    echo "capture_root_source=$REFERENCE_CAPTURE_ROOT_SOURCE"
    echo "capture_models_dir=$REFERENCE_CAPTURE_MODELS_DIR"
    echo "capture_venv=$REFERENCE_CAPTURE_VENV"
    echo "ready_json=$REFERENCE_CAPTURE_READY_JSON"
    echo "ready_stamp=$REFERENCE_CAPTURE_READY_STAMP"
}

enter_capture_box_env() {
    local repo_home
    repo_home="$(infer_repo_home "$SCRIPT_DIR" || true)"
    [[ -n "$repo_home" ]] || err "Cannot derive a stable user home from repo path $SCRIPT_DIR.
Run from a checkout under /home/<user>/... or /Users/<user>/..., or set KRASIS_REFERENCE_CAPTURE_ROOT explicitly."

    export HOME="$repo_home"
    export PATH="$repo_home/.cargo/bin:$PATH"
    export KRASIS_CAPTURE_BOX_ACTIVE=1
    export KRASIS_CAPTURE_BOX_REPO_HOME="$repo_home"
    export KRASIS_REFERENCE_CAPTURE_ROOT="$repo_home/.krasis"
    export KRASIS_REFERENCE_CAPTURE_ROOT_SOURCE="capture_box"
    REFERENCE_CAPTURE_ROOT="$KRASIS_REFERENCE_CAPTURE_ROOT"
    REFERENCE_CAPTURE_ROOT_SOURCE="$KRASIS_REFERENCE_CAPTURE_ROOT_SOURCE"
    REFERENCE_CAPTURE_MODELS_DIR="${REFERENCE_CAPTURE_ROOT}/models"
    REFERENCE_CAPTURE_VENV="${REFERENCE_CAPTURE_ROOT}/reference-capture-venv"
    REFERENCE_CAPTURE_READY_JSON="${REFERENCE_CAPTURE_ROOT}/capture-host-ready.json"
    REFERENCE_CAPTURE_READY_STAMP="${REFERENCE_CAPTURE_ROOT}/capture-host-ready.stamp"
    export KRASIS_REFERENCE_CAPTURE_MODELS_DIR="$REFERENCE_CAPTURE_MODELS_DIR"
    export KRASIS_REFERENCE_CAPTURE_VENV="$REFERENCE_CAPTURE_VENV"
    export KRASIS_REFERENCE_CAPTURE_READY_JSON="$REFERENCE_CAPTURE_READY_JSON"
    export KRASIS_REFERENCE_CAPTURE_READY_STAMP="$REFERENCE_CAPTURE_READY_STAMP"

    if [[ ! "$SCRIPT_DIR" =~ ^"$repo_home"/ ]]; then
        err "Repo path $SCRIPT_DIR is not inside the enforced capture-box home $repo_home."
    fi
}

verify_repo_imports() {
    local output
    local ld_path
    ld_path="$(build_base_ld_library_path)"
    if [[ -d "$REPO_LIBPATH" ]]; then
        ld_path="$REPO_LIBPATH${ld_path:+:$ld_path}"
    fi
    output=$(PYTHONNOUSERSITE=1 PYTHONPATH="$REPO_PYTHONPATH" LD_LIBRARY_PATH="$ld_path" SCRIPT_DIR="$SCRIPT_DIR" "$PYTHON" - <<'PY'
import importlib
import os
import pathlib
import sys

repo = pathlib.Path(os.environ["SCRIPT_DIR"]).resolve()
pkg = importlib.import_module("krasis")
ext = importlib.import_module("krasis.krasis")
pkg_path = pathlib.Path(pkg.__file__).resolve()
ext_path = pathlib.Path(ext.__file__).resolve()

def inside_repo(path: pathlib.Path) -> bool:
    try:
        path.relative_to(repo)
        return True
    except ValueError:
        return False

print(f"krasis: {pkg_path}")
print(f"krasis.krasis: {ext_path}")
if not inside_repo(pkg_path) or not inside_repo(ext_path):
    sys.exit(1)
PY
) || {
        echo "$output"
        err "Import verification failed: dev runtime is not bound to this repo."
    }
    info "Import origin check:"
    echo "$output"
}

verify_repo_python_imports() {
    local output
    local ld_path
    ld_path="$(build_base_ld_library_path)"
    if [[ -d "$REPO_LIBPATH" ]]; then
        ld_path="$REPO_LIBPATH${ld_path:+:$ld_path}"
    fi
    local python_bin
    python_bin="$(reference_capture_python_bin)"
    output=$(PYTHONNOUSERSITE=1 PYTHONPATH="$REPO_PYTHONPATH" LD_LIBRARY_PATH="$ld_path" SCRIPT_DIR="$SCRIPT_DIR" "$python_bin" - <<'PY'
import importlib
import os
import pathlib
import sys

repo = pathlib.Path(os.environ["SCRIPT_DIR"]).resolve()
pkg = importlib.import_module("krasis")
contract = importlib.import_module("tests.reference_contract")
pkg_path = pathlib.Path(pkg.__file__).resolve()
contract_path = pathlib.Path(contract.__file__).resolve()
capture_path = (repo / "tests" / "generate_reference.py").resolve()

def inside_repo(path: pathlib.Path) -> bool:
    try:
        path.relative_to(repo)
        return True
    except ValueError:
        return False

print(f"krasis: {pkg_path}")
print(f"tests.reference_contract: {contract_path}")
print(f"tests.generate_reference: {capture_path}")
if not capture_path.is_file():
    sys.exit(1)
if not inside_repo(pkg_path) or not inside_repo(contract_path) or not inside_repo(capture_path):
    sys.exit(1)
PY
) || {
        echo "$output"
        err "Python import verification failed: HF capture runtime is not bound to this repo."
    }
    info "Reference capture python: $python_bin"
    info "Python import origin check:"
    echo "$output"
}

# ── GPU cleanup ─────────────────────────────────────────────────────
# Kill any existing krasis processes and wait for GPUs to be fully clear.
# Called before every run/test to prevent OOMs from stale processes.

cleanup_gpu() {
    # Optional arg: config file path. If provided, only check GPUs listed in CFG_SELECTED_GPUS.
    local conf="${1:-}"
    local selected_gpus=""
    if [[ -n "$conf" && -f "$conf" ]]; then
        selected_gpus=$(grep '^CFG_SELECTED_GPUS=' "$conf" 2>/dev/null | head -1 | cut -d'"' -f2 || true)
    fi

    # Find all krasis server/stress/benchmark Python processes (but not this script)
    local pids
    pids=$(pgrep -f 'python.*krasis\.' 2>/dev/null || true)

    if [[ -n "$pids" ]]; then
        warn "Found existing krasis processes — killing them:"
        for p in $pids; do
            local cmdline
            cmdline=$(ps -p "$p" -o args= 2>/dev/null || echo "unknown")
            warn "  PID $p: $cmdline"
        done

        # SIGTERM first
        for p in $pids; do
            kill -TERM "$p" 2>/dev/null || true
        done

        # Wait up to 10s for graceful exit
        local waited=0
        while [[ $waited -lt 10 ]]; do
            local still_alive=false
            for p in $pids; do
                if kill -0 "$p" 2>/dev/null; then
                    still_alive=true
                    break
                fi
            done
            $still_alive || break
            sleep 1
            waited=$((waited + 1))
        done

        # SIGKILL any survivors
        for p in $pids; do
            if kill -0 "$p" 2>/dev/null; then
                warn "  PID $p still alive after SIGTERM, sending SIGKILL"
                kill -9 "$p" 2>/dev/null || true
            fi
        done

        # Wait for SIGKILL to take effect
        sleep 2
        for p in $pids; do
            wait "$p" 2>/dev/null || true
        done
    fi

    # Now wait for GPU memory to actually be released.
    # CUDA memory isn't freed until the process fully exits.
    if [[ -n "$selected_gpus" ]]; then
        info "Waiting for selected GPUs ($selected_gpus) to clear..."
    else
        info "Waiting for GPU memory to clear..."
    fi
    local gpu_timeout=30
    local gpu_waited=0
    while [[ $gpu_waited -lt $gpu_timeout ]]; do
        local all_clear=true
        while IFS=', ' read -r idx used gpu_mem_total; do
            # If selected_gpus is set, only check those GPUs
            if [[ -n "$selected_gpus" ]]; then
                local check_this=false
                for sg in $(echo "$selected_gpus" | tr ',' ' '); do
                    [[ "$idx" == "$sg" ]] && check_this=true
                done
                $check_this || continue
            fi
            # Baseline is ~200-300 MiB (desktop compositor). 500 MiB threshold.
            if [[ "$used" -gt 500 ]]; then
                all_clear=false
                break
            fi
        done < <(nvidia-smi --query-gpu=index,memory.used,memory.total --format=csv,noheader,nounits 2>/dev/null)

        if $all_clear; then
            if [[ -n "$selected_gpus" ]]; then
                ok "Selected GPUs ($selected_gpus) clear."
            else
                ok "All GPUs clear."
            fi
            return 0
        fi

        sleep 1
        gpu_waited=$((gpu_waited + 1))
    done

    # If we get here, GPUs didn't clear. Show what's using them.
    warn "GPUs still not clear after ${gpu_timeout}s:"
    nvidia-smi 2>/dev/null || true
    err "Cannot proceed — GPU memory not released. Check for non-krasis processes using the GPU."
}

# ── Commands ─────────────────────────────────────────────────────────

do_build_fla() {
    # Cross-compile FLA (Flash Linear Attention) Triton kernels for all
    # target GPU architectures.  The resulting .so files are placed in
    # python/krasis/ so maturin can include them in the wheel.
    local fla_script="$SCRIPT_DIR/src/cuda/fla/compile_kernels.py"
    local fla_src_dir="$SCRIPT_DIR/src/cuda/fla"
    local dest="$SCRIPT_DIR/python/krasis"
    local archs=(80 89 90 120)

    # Check if recompilation is needed: any .so missing or any source newer?
    local needs_rebuild=false
    for arch in "${archs[@]}"; do
        if [[ ! -f "$dest/libkrasis_fla_sm${arch}.so" ]]; then
            needs_rebuild=true
            break
        fi
    done
    if ! $needs_rebuild; then
        # Check if any Python source in fla/ is newer than the oldest .so
        local oldest_so_mtime
        oldest_so_mtime=$(stat --format="%Y" "$dest"/libkrasis_fla_sm*.so 2>/dev/null | sort -n | head -1)
        if [[ -n "$oldest_so_mtime" ]]; then
            while IFS= read -r pyfile; do
                local py_mtime
                py_mtime=$(stat --format="%Y" "$pyfile" 2>/dev/null) || continue
                if [[ "$py_mtime" -gt "$oldest_so_mtime" ]]; then
                    needs_rebuild=true
                    break
                fi
            done < <(find "$fla_src_dir" -name "*.py" -type f 2>/dev/null)
        fi
    fi

    if ! $needs_rebuild; then
        ok "FLA kernels up to date (skipping)"
        return
    fi

    info "Compiling FLA kernels for GPU architectures: ${archs[*]/#/sm_}..."
    local fla_output="/tmp/fla_build"

    KRASIS_DEV_SCRIPT=1 "$PYTHON" "$fla_script" --output-dir "$fla_output" \
        --arch "${archs[@]}" || err "FLA kernel compilation failed"

    mkdir -p "$dest"
    local count=0
    for so in "$fla_output"/libkrasis_fla_sm*.so; do
        [[ -f "$so" ]] || continue
        cp "$so" "$dest/"
        count=$((count + 1))
    done
    ok "FLA kernels: $count architectures compiled"
}

do_build() {
    info "Building Rust extension (repo-local maturin build --release)..."
    export LD_LIBRARY_PATH
    LD_LIBRARY_PATH="$(build_base_ld_library_path)"

    # Compile FLA kernels first so they're in python/krasis/ for maturin
    do_build_fla

    cd "$SCRIPT_DIR"
    mkdir -p "$WHEEL_DIR"
    rm -f "$WHEEL_DIR"/krasis-*.whl
    PATH="$HOME/.cargo/bin:$(dirname "$PYTHON"):$PATH" \
        "$MATURIN" build --release --skip-auditwheel --out "$WHEEL_DIR" 2>&1
    local wheel
    wheel=$(ls -t "$WHEEL_DIR"/krasis-*.whl 2>/dev/null | head -1)
    [[ -n "$wheel" ]] || err "Build succeeded but no wheel was produced in $WHEEL_DIR"
    "$PYTHON" - "$wheel" "$SCRIPT_DIR/python/krasis" "$SCRIPT_DIR/python/krasis.libs" <<'PY'
import pathlib
import shutil
import sys
import zipfile

wheel = pathlib.Path(sys.argv[1])
target = pathlib.Path(sys.argv[2])
lib_target = pathlib.Path(sys.argv[3])
target.mkdir(parents=True, exist_ok=True)
lib_target.mkdir(parents=True, exist_ok=True)
for pattern in ("krasis*.so", "krasis*.pyd", "lib*.so", "*.dylib", "*.dll"):
    for path in target.glob(pattern):
        path.unlink()
for pattern in ("*.so", "*.so.*", "*.dylib", "*.dll"):
    for path in lib_target.glob(pattern):
        path.unlink()
with zipfile.ZipFile(wheel) as zf:
    ext_members = [
        name for name in zf.namelist()
        if name.startswith("krasis/") and name.endswith((".so", ".pyd", ".dll", ".dylib"))
    ]
    lib_members = [
        name for name in zf.namelist()
        if name.startswith("krasis.libs/")
    ]
    if not ext_members:
        raise SystemExit("No compiled artifacts found in wheel")
    for member in ext_members:
        dest = target / pathlib.Path(member).name
        with zf.open(member) as src, open(dest, "wb") as dst:
            shutil.copyfileobj(src, dst)
        print(dest)
    for member in lib_members:
        dest = lib_target / pathlib.Path(member).name
        with zf.open(member) as src, open(dest, "wb") as dst:
            shutil.copyfileobj(src, dst)
        print(dest)
PY
    ok "Build complete. Extension at: $SO_FILE"
    verify_repo_imports
}

do_run() {
    [[ $# -lt 1 ]] && err "Usage: ./dev run <config> [--benchmark] [extra args...]"
    local conf
    conf=$(resolve_config "$1")
    shift
    local run_dir
    run_dir=$(create_run_dir "dev-run")
    cleanup_gpu "$conf"
    auto_rebuild
    info "Launching from: $conf"
    info "Run dir: $run_dir"
    verify_repo_imports
    export KRASIS_RUN_DIR="$run_dir"
    export KRASIS_RUN_TYPE="dev-run"
    exec_repo_python -m krasis.server --config "$conf" "$@"
}

do_test() {
    [[ $# -lt 1 ]] && err "Usage: ./dev test <config> [--thorough]"
    local conf
    conf=$(resolve_config "$1")
    local thorough=false
    [[ "${2:-}" == "--thorough" ]] && thorough=true
    cleanup_gpu "$conf"
    auto_rebuild
    local run_dir
    run_dir=$(create_run_dir "dev-test")

    local port
    port=$(conf_port "$conf")
    [[ -z "$port" ]] && port=8012

    info "=== Short Model Test ==="
    info "Config: $conf"
    info "Port: $port"
    info "Run dir: $run_dir"

    # Step 1: Launch server with benchmark
    info "Step 1: Launching server with --benchmark..."
    local logfile
    logfile="$run_dir/test_stdout.log"
    info "Log: $logfile"

    verify_repo_imports
    KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="dev-test" run_repo_python -m krasis.server --config "$conf" --benchmark > >(tee "$logfile") 2>&1 &
    local pid=$!

    # Wait for benchmark to complete AND server to be ready.
    # Must wait for BENCHMARK COMPLETE, not just "Server ready", because
    # the benchmark runs on a separate thread. Sending HTTP requests while
    # the benchmark is still using the GPU causes CUDA_ERROR_ILLEGAL_ADDRESS.
    local ready=false
    local timeout=1200  # 20 minutes for model load + benchmark
    local elapsed=0
    while kill -0 "$pid" 2>/dev/null && [[ $elapsed -lt $timeout ]]; do
        if grep -q "BENCHMARK COMPLETE" "$logfile" 2>/dev/null; then
            ready=true
            break
        fi
        sleep 2
        elapsed=$((elapsed + 2))
    done

    if ! $ready; then
        if kill -0 "$pid" 2>/dev/null; then
            warn "Timed out waiting for server (${timeout}s). Killing..."
            kill -TERM "$pid" 2>/dev/null || true
            sleep 2
            kill -9 "$pid" 2>/dev/null || true
        fi
        err "Server did not start. Check log: $logfile"
    fi

    ok "Server is ready. Benchmark results are in the log."

    # Step 2: Run network multi-prompt tests
    info "Step 2: Running network multi-prompt tests..."
    if $thorough; then
        run_repo_python "$SCRIPT_DIR/tests/test_network.py" --port "$port" --large 2>&1 | tee -a "$logfile"
    else
        run_repo_python "$SCRIPT_DIR/tests/test_network.py" --port "$port" 2>&1 | tee -a "$logfile"
    fi

    local net_result=$?

    # Cleanup: kill the server
    info "Stopping server..."
    kill -TERM "$pid" 2>/dev/null || true
    sleep 2
    kill -9 "$pid" 2>/dev/null || true
    wait "$pid" 2>/dev/null || true

    echo ""
    if [[ $net_result -eq 0 ]]; then
        ok "=== All tests PASSED ==="
    else
        err "Network tests FAILED (exit code $net_result)"
    fi
    info "Full log: $logfile"
}

do_release_test() {
    [[ $# -lt 1 ]] && err "Usage: ./dev release-test <model> [--force-rebuild-cache]
  model = directory name under ~/.krasis/models/
  Example: ./dev release-test Qwen3-Coder-Next"
    local model="$1"
    shift
    cleanup_gpu
    local run_dir
    run_dir=$(create_run_dir "release-test")
    info "Starting release test for model: $model"
    info "Run dir: $run_dir"
    verify_repo_imports
    KRASIS_DEV_SCRIPT=1 KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="release-test" run_repo_python "$SCRIPT_DIR/tests/release_test.py" "$model" "$@"
}

do_release_test_all() {
    local SESSION_WS="Release Test"
    local release_target="${1:-}"
    local MODELS=(
        "Qwen3-Coder-Next"
        "Qwen3.5-35B-A3B"
        "Qwen3.5-122B-A10B"
        "Qwen3-235B-A22B"
        "Qwen3.5-397B-A17B"
    )

    local total=${#MODELS[@]}
    local completed=0
    local skipped=0
    local failed_model=""
    local resuming=false
    local state_file=""
    local start_time
    start_time=$(date +%s)

    # State file for resume support (only when release target is given)
    if [[ -n "$release_target" ]]; then
        local state_dir="$SCRIPT_DIR/logs/release-tests"
        mkdir -p "$state_dir"
        state_file="$state_dir/release-all-${release_target}.state"

        if [[ -f "$state_file" ]]; then
            resuming=true
            # Count previously passed models
            for model in "${MODELS[@]}"; do
                local status
                status=$(grep "^${model}:" "$state_file" 2>/dev/null | cut -d: -f2) || true
                if [[ "$status" == "PASSED" ]]; then
                    skipped=$((skipped + 1))
                    completed=$((completed + 1))
                fi
            done

            if [[ $skipped -eq $total ]]; then
                ok "All models already passed for release target '$release_target'. Nothing to do."
                session-send "$SESSION_WS" "release-test-all '$release_target': all $total models already passed. Nothing to do." 2>/dev/null || true
                return 0
            fi

            info "=== Resuming Release Test All: $release_target ==="
            info "Previously passed: $skipped/$total"
        else
            info "=== Release Test All: $release_target ==="
            # Create initial state file
            for model in "${MODELS[@]}"; do
                echo "${model}:PENDING" >> "$state_file"
            done
        fi
    else
        info "=== Release Test All ==="
    fi

    info "Models: ${MODELS[*]}"
    info "Total: $total"

    # Verify all models exist before starting
    for model in "${MODELS[@]}"; do
        local model_dir="$HOME/.krasis/models/$model"
        if [[ ! -d "$model_dir" ]]; then
            local msg="FAILED: Model $model not found at $model_dir. Cannot run release-test-all."
            err "$msg"
        fi
    done

    # Session start message
    if $resuming; then
        session-send "$SESSION_WS" "Resuming release-test-all '$release_target': $skipped/$total already passed, $(( total - skipped )) remaining.
Models: $(printf '%s, ' "${MODELS[@]}" | sed 's/, $//')" 2>/dev/null || warn "Could not send Session message"
    else
        local target_info=""
        [[ -n "$release_target" ]] && target_info=" ($release_target)"
        session-send "$SESSION_WS" "Starting release-test-all${target_info}: $total models to test.
Models: $(printf '%s, ' "${MODELS[@]}" | sed 's/, $//')" 2>/dev/null || warn "Could not send Session message"
    fi

    for model in "${MODELS[@]}"; do
        # Skip already-passed models on resume
        if [[ -n "$state_file" && -f "$state_file" ]]; then
            local prev_status
            prev_status=$(grep "^${model}:" "$state_file" 2>/dev/null | cut -d: -f2) || true
            if [[ "$prev_status" == "PASSED" ]]; then
                info "=== Skipping $model (already passed) ==="
                continue
            fi
        fi

        local idx=$((completed + 1))
        info "=== [$idx/$total] Starting: $model ==="
        session-send "$SESSION_WS" "[$idx/$total] Starting release test: $model" 2>/dev/null || true

        local model_start
        model_start=$(date +%s)

        # Run release test — cleanup_gpu is called inside do_release_test
        if do_release_test "$model"; then
            completed=$((completed + 1))
            local model_elapsed=$(( $(date +%s) - model_start ))
            local model_mins=$((model_elapsed / 60))
            local model_secs=$((model_elapsed % 60))
            ok "[$completed/$total] PASSED: $model (${model_mins}m ${model_secs}s)"
            session-send "$SESSION_WS" "[$completed/$total] PASSED: $model (${model_mins}m ${model_secs}s)" 2>/dev/null || true

            # Update state file
            if [[ -n "$state_file" ]]; then
                sed -i "s/^${model}:.*/${model}:PASSED/" "$state_file"
            fi
        else
            failed_model="$model"
            local model_elapsed=$(( $(date +%s) - model_start ))
            local model_mins=$((model_elapsed / 60))
            local model_secs=$((model_elapsed % 60))
            local msg="FAILED: $model (${model_mins}m ${model_secs}s). Stopping. $completed/$total models completed."
            warn "$msg"
            session-send "$SESSION_WS" "$msg" 2>/dev/null || true

            # Update state file
            if [[ -n "$state_file" ]]; then
                sed -i "s/^${model}:.*/${model}:FAILED/" "$state_file"
            fi

            err "Release test failed on $model. Aborting."
        fi
    done

    local configs_elapsed=$(( $(date +%s) - start_time ))
    local configs_mins=$((configs_elapsed / 60))
    local configs_secs=$((configs_elapsed % 60))
    local target_suffix=""
    [[ -n "$release_target" ]] && target_suffix=" ($release_target)"
    local configs_msg="All configs passed: $completed/$total models in ${configs_mins}m ${configs_secs}s. Starting perplexity phase.${target_suffix}"
    ok "$configs_msg"
    session-send "$SESSION_WS" "$configs_msg" 2>/dev/null || true

    # ── Perplexity phase: one INT4/AWQ pass per model ──────────────
    info "=== Perplexity Phase: INT4/AWQ for all models ==="
    session-send "$SESSION_WS" "Starting perplexity phase: INT4/AWQ for $total models" 2>/dev/null || true

    local ppl_passed=0
    local ppl_failed=0
    local best_gpu_idx
    best_gpu_idx=$(nvidia-smi --query-gpu=index,memory.total --format=csv,noheader,nounits | sort -t',' -k2 -nr | head -1 | cut -d',' -f1 | tr -d ' ')

    for model in "${MODELS[@]}"; do
        local model_dir="$HOME/.krasis/models/$model"
        local config_json="$model_dir/config.json"

        if [[ ! -f "$config_json" ]]; then
            warn "No config.json for $model — skipping perplexity"
            ppl_failed=$((ppl_failed + 1))
            session-send "$SESSION_WS" "[PPL] SKIP: $model (no config.json)" 2>/dev/null || true
            continue
        fi

        # Get num_layers from model config
        local num_layers
        num_layers=$("$PYTHON" -c "
import json, sys
c = json.load(open('$config_json'))
tc = c.get('text_config', c)
print(tc.get('num_hidden_layers', 0))
") || num_layers=0

        if [[ "$num_layers" -eq 0 ]]; then
            warn "Could not detect num_layers for $model — skipping perplexity"
            ppl_failed=$((ppl_failed + 1))
            session-send "$SESSION_WS" "[PPL] SKIP: $model (unknown layers)" 2>/dev/null || true
            continue
        fi

        # Generate temp INT4/AWQ config
        local ppl_conf
        ppl_conf=$(mktemp /tmp/krasis-ppl-XXXXXX.conf)
        cat > "$ppl_conf" <<PPLCONF
# Perplexity config — INT4/INT4 AWQ — $model
MODEL_PATH="$model_dir"
CFG_SELECTED_GPUS="$best_gpu_idx"
CFG_PP_PARTITION="$num_layers"
CFG_LAYER_GROUP_SIZE="2"
CFG_KV_DTYPE="fp8_e4m3"
CFG_GPU_EXPERT_BITS="4"
CFG_CPU_EXPERT_BITS="4"
CFG_ATTENTION_QUANT="awq"
CFG_SHARED_EXPERT_QUANT="int8"
CFG_DENSE_MLP_QUANT="int8"
CFG_LM_HEAD_QUANT="int8"
PPLCONF

        info "=== [PPL] $model (INT4/AWQ) ==="
        session-send "$SESSION_WS" "[PPL] Starting: $model" 2>/dev/null || true
        local ppl_start
        ppl_start=$(date +%s)

        if do_perplexity "$ppl_conf" --max-tokens 20000; then
            local ppl_elapsed=$(( $(date +%s) - ppl_start ))
            local ppl_mins=$((ppl_elapsed / 60))
            local ppl_secs=$((ppl_elapsed % 60))
            ppl_passed=$((ppl_passed + 1))
            ok "[PPL] PASSED: $model (${ppl_mins}m ${ppl_secs}s)"
            session-send "$SESSION_WS" "[PPL] PASSED: $model (${ppl_mins}m ${ppl_secs}s)" 2>/dev/null || true
        else
            local ppl_elapsed=$(( $(date +%s) - ppl_start ))
            local ppl_mins=$((ppl_elapsed / 60))
            local ppl_secs=$((ppl_elapsed % 60))
            ppl_failed=$((ppl_failed + 1))
            warn "[PPL] FAILED: $model (${ppl_mins}m ${ppl_secs}s)"
            session-send "$SESSION_WS" "[PPL] FAILED: $model (${ppl_mins}m ${ppl_secs}s)" 2>/dev/null || true
        fi

        rm -f "$ppl_conf"
    done

    # ── Final summary ──────────────────────────────────────────────
    local total_elapsed=$(( $(date +%s) - start_time ))
    local total_mins=$((total_elapsed / 60))
    local total_secs=$((total_elapsed % 60))
    local final_msg="COMPLETE: $completed/$total models passed, perplexity $ppl_passed/$total passed, total ${total_mins}m ${total_secs}s.${target_suffix}"
    ok "$final_msg"
    session-send "$SESSION_WS" "$final_msg" 2>/dev/null || true

    # Clean up state file on full success
    if [[ -n "$state_file" ]]; then
        rm -f "$state_file"
        info "State file removed (all passed)."
    fi
}

do_benchmark() {
    [[ $# -lt 1 ]] && err "Usage: ./dev benchmark <config> [extra server args...]"
    local conf
    conf=$(resolve_config "$1")
    shift
    cleanup_gpu "$conf"
    auto_rebuild
    local run_dir
    run_dir=$(create_run_dir "dev-benchmark")
    local logfile
    logfile="$run_dir/benchmark_stdout.log"
    info "Running standard benchmark: $conf"
    info "Run dir: $run_dir"
    info "Log: $logfile"
    verify_repo_imports
    cd "$SCRIPT_DIR" && KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="dev-benchmark" run_repo_python -m krasis.server --config "$conf" --benchmark-only "$@" 2>&1 | tee "$logfile"
    info "Benchmark log saved: $logfile"
}

do_speed_test() {
    [[ -f "$STANDARD_SPEED_TEST_CONFIG" ]] || err "Standard speed-test config not found: $STANDARD_SPEED_TEST_CONFIG"
    info "Standard speed benchmark config: $STANDARD_SPEED_TEST_CONFIG"
    do_benchmark "$STANDARD_SPEED_TEST_CONFIG" "$@"
}

do_decode_harness() {
    [[ $# -lt 1 ]] && err "Usage: ./dev decode-harness <config> [extra args...]"
    local conf
    conf=$(resolve_config "$1")
    shift
    cleanup_gpu "$conf"
    auto_rebuild
    local run_dir
    run_dir=$(create_run_dir "dev-decode-harness")
    local logfile
    logfile="$run_dir/decode_harness_stdout.log"
    info "Running decode harness: $conf"
    info "Run dir: $run_dir"
    info "Log: $logfile"
    verify_repo_imports
    cd "$SCRIPT_DIR" && KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="dev-decode-harness" KRASIS_DEV_SCRIPT="1" run_repo_python "$SCRIPT_DIR/tests/decode_harness.py" --config "$conf" "$@" 2>&1 | tee "$logfile"
    info "Decode harness log saved: $logfile"
}

do_perplexity() {
    [[ $# -lt 1 ]] && err "Usage: ./dev perplexity <config> [--dataset wikitext-2] [--max-tokens N]"
    local conf
    conf=$(resolve_config "$1")
    shift
    cleanup_gpu "$conf"
    auto_rebuild
    local run_dir
    run_dir=$(create_run_dir "dev-perplexity")
    local logfile
    logfile="$run_dir/perplexity_stdout.log"
    info "Running perplexity: $conf"
    info "Run dir: $run_dir"
    info "Log: $logfile"
    verify_repo_imports
    cd "$SCRIPT_DIR" && KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="dev-perplexity" run_repo_python -m perplexity.measure_ppl --config "$conf" "$@" 2>&1 | tee "$logfile"
    info "Perplexity log saved: $logfile"
}

do_awq_calibrate() {
    [[ $# -lt 1 ]] && err "Usage: ./dev awq-calibrate <config> [--tokens 256000] [--dataset wikitext-2]"
    local conf
    conf=$(resolve_config "$1")
    shift
    if [[ "${KRASIS_SKIP_GPU_CLEANUP:-0}" != "1" ]]; then
        cleanup_gpu "$conf"
    else
        info "Skipping GPU cleanup for nested AWQ calibration"
    fi
    auto_rebuild
    local run_dir
    run_dir=$(create_run_dir "dev-awq-calibrate")
    info "Running AWQ calibration: $conf"
    info "Run dir: $run_dir"
    verify_repo_imports
    cd "$SCRIPT_DIR" && KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="dev-awq-calibrate" run_repo_python -m krasis.awq_calibrate --config "$conf" "$@"
}

do_network() {
    [[ $# -lt 1 ]] && err "Usage: ./dev network <port> [--large] [--quick]"
    local port="$1"
    shift
    exec_repo_python "$SCRIPT_DIR/tests/test_network.py" --port "$port" "$@"
}

do_test_kernels() {
    info "Running CUDA kernel unit tests..."
    # First ensure test data exists
    local test_data_dir="$SCRIPT_DIR/tests/kernel_test_data"
    if [[ ! -d "$test_data_dir" ]] || [[ -z "$(ls -A "$test_data_dir" 2>/dev/null)" ]]; then
        info "Generating kernel test data..."
        cd "$SCRIPT_DIR" && run_repo_python tests/generate_kernel_test_data.py
    fi

    # PyO3 extension modules need Python library for test linking.
    # Get the python lib dir from the conda env Python.
    local pyver
    pyver=$("$PYTHON" -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
    local pylibdir
    pylibdir=$("$PYTHON" -c "import sysconfig; print(sysconfig.get_config_var('LIBDIR'))")
    export PYO3_PYTHON="$PYTHON"
    export LD_LIBRARY_PATH="$pylibdir${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
    # Tell the linker where to find libpython and to link it
    export RUSTFLAGS="${RUSTFLAGS:-} -L $pylibdir -l python${pyver}"

    # Build and run kernel tests (lib only, skip other test binaries)
    cd "$SCRIPT_DIR"
    if [[ $# -gt 0 ]]; then
        local filter="kernel_tests::test_$1"
        info "Filter: $filter"
        cargo test --release --lib "$filter" -- --nocapture 2>&1
    else
        cargo test --release --lib kernel_tests -- --nocapture 2>&1
    fi
}

do_validate() {
    local config="${1:?Usage: ./dev validate <config> [--no-server] [--l1-only] [--no-logprobs] [--no-prefill]}"
    shift
    local conf_path
    conf_path=$(resolve_config "$config")
    if [[ ! -f "$conf_path" ]]; then
        die "Config not found: $conf_path"
    fi
    do_build
    info "Running model validation: $conf_path"
    verify_repo_imports
    KRASIS_DEV_SCRIPT=1 exec_repo_python "$SCRIPT_DIR/tests/validate_model.py" "$conf_path" "$@"
}

do_reference_test() {
    [[ $# -lt 1 ]] && err "Usage: ./dev reference-test <config> [--verbose] [--no-server] [--port PORT]
       ./dev reference-test --compare <model> [--verbose]
  Runs reference test comparing engine output against BF16 reference data.
  Produces an HTML report in logs/reference-tests/.
  Single config:  ./dev reference-test qwen35-a4
  Comparison:     ./dev reference-test --compare qcn"

    if [[ "$1" == "--compare" ]]; then
        shift
        [[ $# -lt 1 ]] && err "Usage: ./dev reference-test --compare <model>
  Available models: qcn, q35b, q122b
  Single-config aliases: gemma, minimax"
        local model="$1"
        shift
        cleanup_gpu
        auto_rebuild
        local run_dir
        run_dir=$(create_run_dir "reference-test")
        info "Running comparative reference test: $model"
        info "Run dir: $run_dir"
        verify_repo_imports
        KRASIS_DEV_SCRIPT=1 KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="reference-test" run_repo_python "$SCRIPT_DIR/tests/reference_test.py" --compare "$model" "$@"
    else
        local conf
        conf=$(resolve_config "$1")
        shift
        cleanup_gpu "$conf"
        auto_rebuild
        local run_dir
        run_dir=$(create_run_dir "reference-test")
        info "Running reference test: $conf"
        info "Run dir: $run_dir"
        verify_repo_imports
        KRASIS_DEV_SCRIPT=1 KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="reference-test" run_repo_python "$SCRIPT_DIR/tests/reference_test.py" "$conf" "$@"
    fi
}

do_reference_inventory() {
    cleanup_gpu
    auto_rebuild
    local run_dir
    run_dir=$(create_run_dir "reference-inventory")
    info "Running reference inventory"
    info "Run dir: $run_dir"
    verify_repo_imports
    KRASIS_DEV_SCRIPT=1 KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="reference-inventory" run_repo_python "$SCRIPT_DIR/tests/reference_inventory.py" "$@"
}

do_capture_preflight() {
    info "Running capture host preflight"
    KRASIS_DEV_SCRIPT=1 \
    KRASIS_DEV_PYTHON="$PYTHON" \
    KRASIS_DEV_PIP="$PIP" \
    KRASIS_DEV_MATURIN="$MATURIN" \
    KRASIS_REFERENCE_CAPTURE_ROOT="$REFERENCE_CAPTURE_ROOT" \
    KRASIS_REFERENCE_CAPTURE_ROOT_SOURCE="$REFERENCE_CAPTURE_ROOT_SOURCE" \
    KRASIS_REFERENCE_CAPTURE_MODELS_DIR="$REFERENCE_CAPTURE_MODELS_DIR" \
    KRASIS_REFERENCE_CAPTURE_VENV="$REFERENCE_CAPTURE_VENV" \
    KRASIS_REFERENCE_CAPTURE_READY_JSON="$REFERENCE_CAPTURE_READY_JSON" \
    KRASIS_REFERENCE_CAPTURE_READY_STAMP="$REFERENCE_CAPTURE_READY_STAMP" \
    "$SCRIPT_DIR/scripts/capture_preflight.sh" "$@"
}

do_reference_prep() {
    [[ $# -lt 1 && ! " ${*:-} " =~ [[:space:]](--help|-h|--list-models|--deps-only|--print-deps-json)[[:space:]] ]] && err "Usage: ./dev reference-prep <model> [--detach] [--print-urls|--urls-only] [--deps-only|--download-only|--print-deps-json] [--max-workers N] [--retry-count N] [--retry-backoff-sec N] [--enable-xet|--disable-xet] [--dest-root PATH]
  Prepares the HF reference-capture environment and downloads a public model to ${REFERENCE_CAPTURE_MODELS_DIR}.
  Examples:
    ./dev capture-preflight --bootstrap
    ./dev reference-prep qcn
    ./dev reference-prep qwen35 --detach
    ./dev reference-prep qcn --download-only --max-workers 4 --retry-count 8
    ./dev reference-prep --deps-only
    ./dev reference-prep --print-deps-json
    ./dev reference-prep gemma --urls-only"
    info "Preparing reference host: ${1:-listing models}"
    KRASIS_DEV_SCRIPT=1 KRASIS_DEV_PYTHON="$PYTHON" KRASIS_DEV_PIP="$PIP" \
    KRASIS_REFERENCE_CAPTURE_ROOT="$REFERENCE_CAPTURE_ROOT" \
    KRASIS_REFERENCE_CAPTURE_ROOT_SOURCE="$REFERENCE_CAPTURE_ROOT_SOURCE" \
    KRASIS_REFERENCE_CAPTURE_MODELS_DIR="$REFERENCE_CAPTURE_MODELS_DIR" \
    KRASIS_REFERENCE_CAPTURE_VENV="$REFERENCE_CAPTURE_VENV" \
    "$SCRIPT_DIR/scripts/reference_capture_prep.sh" "$@"
}

do_generate_reference() {
    [[ $# -lt 1 ]] && err "Usage: ./dev generate-reference <model> [--max-tokens N] [--profile PROFILE]
  Generates a BF16 HuggingFace greedy reference with full contract metadata.
  Profiles:
    auto                    thinking_off if supported, otherwise default
    greedy_chat_default     default template invocation
    greedy_chat_thinking_off  pass enable_thinking=False
    greedy_chat_thinking_on   pass enable_thinking=True
  Example:
    ./dev generate-reference Qwen3-Coder-Next --profile greedy_chat_thinking_off"
    cleanup_gpu
    local run_dir logfile dev_cmd
    run_dir=$(create_run_dir "generate-reference")
    logfile="$run_dir/generate_reference_stdout.log"
    printf -v dev_cmd '%q ' ./dev generate-reference "$@"
    dev_cmd="${dev_cmd% }"
    info "Generating reference: $1"
    info "Run dir: $run_dir"
    verify_repo_python_imports
    set +e
    KRASIS_DEV_SCRIPT=1 KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="generate-reference" KRASIS_DEV_CMD="$dev_cmd" \
        KRASIS_REFERENCE_CAPTURE_ROOT="$REFERENCE_CAPTURE_ROOT" \
        KRASIS_REFERENCE_CAPTURE_ROOT_SOURCE="$REFERENCE_CAPTURE_ROOT_SOURCE" \
        KRASIS_REFERENCE_CAPTURE_MODELS_DIR="$REFERENCE_CAPTURE_MODELS_DIR" \
        run_reference_capture_python "$SCRIPT_DIR/tests/generate_reference.py" "$@" 2>&1 | tee "$logfile"
    local status=${PIPESTATUS[0]}
    set -e
    return "$status"
}

do_capture_box() {
    local subcmd="${1:-help}"
    shift || true
    enter_capture_box_env

    case "$subcmd" in
        env)
            print_capture_box_env "$KRASIS_CAPTURE_BOX_REPO_HOME"
            ;;
        preflight|capture-preflight)
            print_capture_box_env "$KRASIS_CAPTURE_BOX_REPO_HOME"
            do_capture_preflight "$@"
            ;;
        prep|reference-prep)
            print_capture_box_env "$KRASIS_CAPTURE_BOX_REPO_HOME"
            do_reference_prep "$@"
            ;;
        stage)
            local dry_run=0
            local skip_bootstrap=0
            local -a stage_models=()
            while [[ $# -gt 0 ]]; do
                case "$1" in
                    --dry-run)
                        dry_run=1
                        shift
                        ;;
                    --no-bootstrap)
                        skip_bootstrap=1
                        shift
                        ;;
                    -h|--help)
                        echo "Usage: ./dev capture-box stage [--dry-run] [--no-bootstrap] <model> [more models...]"
                        echo ""
                        echo "Bootstrap the paid-box capture host once, then queue one or more model"
                        echo "downloads through the enforced capture-box environment."
                        echo ""
                        echo "Options:"
                        echo "  --dry-run       Print the exact staged commands without running them"
                        echo "  --no-bootstrap  Skip capture-preflight --bootstrap and assume the host is already ready"
                        echo ""
                        echo "Example:"
                        echo "  ./dev capture-box stage qcn gemma"
                        echo "  ./dev capture-box stage --dry-run qcn gemma"
                        return 0
                        ;;
                    -*)
                        err "Unknown capture-box stage option: $1"
                        ;;
                    *)
                        stage_models+=("$1")
                        shift
                        ;;
                esac
            done
            [[ "${#stage_models[@]}" -gt 0 ]] || err "Usage: ./dev capture-box stage [--dry-run] [--no-bootstrap] <model> [more models...]"
            print_capture_box_env "$KRASIS_CAPTURE_BOX_REPO_HOME"
            if [[ "$skip_bootstrap" -eq 0 ]]; then
                info "Stage bootstrap target: ${stage_models[0]}"
                if [[ "$dry_run" -eq 1 ]]; then
                    local bootstrap_cmd
                    printf -v bootstrap_cmd '%q ' ./dev capture-box preflight --bootstrap --model "${stage_models[0]}"
                    echo "DRY RUN: ${bootstrap_cmd% }"
                else
                    do_capture_preflight --bootstrap --model "${stage_models[0]}"
                fi
            else
                warn "Skipping capture bootstrap as requested"
            fi
            local model
            for model in "${stage_models[@]}"; do
                info "Queueing staged model: $model"
                if [[ "$dry_run" -eq 1 ]]; then
                    local prep_cmd
                    printf -v prep_cmd '%q ' ./dev capture-box prep "$model" --detach
                    echo "DRY RUN: ${prep_cmd% }"
                else
                    do_reference_prep "$model" --detach
                fi
            done
            if [[ "$dry_run" -eq 1 ]]; then
                info "Dry run complete"
            else
                ok "Staged ${#stage_models[@]} model download(s)"
                echo "Next: run ./dev capture-box generate <model> --profile <profile> once the first model is ready."
            fi
            ;;
        generate|generate-reference)
            print_capture_box_env "$KRASIS_CAPTURE_BOX_REPO_HOME"
            [[ $# -ge 1 ]] || err "Usage: ./dev capture-box generate <model> [--profile PROFILE] [--max-tokens N]"
            info "Verifying paid-box capture host for model: $1"
            do_capture_preflight --bootstrap --model "$1"
            do_generate_reference "$@"
            ;;
        help|-h|--help)
            echo "Usage: ./dev capture-box <subcommand> [args...]"
            echo ""
            echo "Strict paid-box wrapper for HF capture work."
            echo "It enforces one remote identity before any expensive work:"
            echo "  - HOME := inferred repo owner home"
            echo "  - PATH prepends <repo-home>/.cargo/bin"
            echo "  - capture root := <repo-home>/.krasis"
            echo "  - capture root source := capture_box"
            echo ""
            echo "Subcommands:"
            echo "  env                            Print the enforced environment and exit"
            echo "  preflight [args...]            Run ./dev capture-preflight under the enforced environment"
            echo "  prep [args...]                 Run ./dev reference-prep under the enforced environment"
            echo "  stage [args...]                Bootstrap once, then queue one or more model downloads"
            echo "  generate [args...]             Run ./dev generate-reference under the enforced environment"
            echo ""
            echo "Examples:"
            echo "  ./dev capture-box env"
            echo "  ./dev capture-box preflight --bootstrap --model qcn"
            echo "  ./dev capture-box prep qcn --detach"
            echo "  ./dev capture-box stage qcn gemma"
            echo "  ./dev capture-box generate Qwen3-Coder-Next --profile greedy_chat_thinking_off"
            ;;
        *)
            err "Unknown capture-box subcommand: $subcmd
  Expected one of: env, preflight, prep, stage, generate"
            ;;
    esac
}

do_trace_diff() {
    [[ $# -lt 2 ]] && err "Usage: ./dev trace-diff <expected-trace> <actual-trace> [--events mark,tensor,...] [--components decode,...] [--contexts ctx1,ctx2] [--phases artifact,validation] [--show-window N] [--resync-window N] [--family-resync-window N]
  Diffs two KRASIS_TRACE logs offline, reports the first divergence, summarizes matching-prefix plus context/phase counts,
  classifies the divergence, gives bounded event-level and family-region resync hints for short shifted regions,
  adds typed analysis for tensor/logits/route mismatches, and reports stronger family-packet signatures."
    local run_dir
    run_dir=$(create_run_dir "trace-diff")
    info "Running trace diff"
    info "Run dir: $run_dir"
    verify_repo_imports
    KRASIS_DEV_SCRIPT=1 KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="trace-diff" run_repo_python "$SCRIPT_DIR/tests/trace_diff.py" "$@"
}

do_kill() {
    info "Killing all krasis processes and freeing GPU memory..."

    # Kill any krasis Python processes
    local pids
    pids=$(pgrep -f 'python.*krasis\.' 2>/dev/null || true)
    if [[ -n "$pids" ]]; then
        for p in $pids; do
            local cmdline
            cmdline=$(ps -p "$p" -o args= 2>/dev/null || echo "unknown")
            warn "  Killing PID $p: $cmdline"
            kill -9 "$p" 2>/dev/null || true
        done
        sleep 2
    fi

    # Check for zombie processes holding GPU memory
    local zombie_pids
    zombie_pids=$(ps -eo pid,ppid,stat,args 2>/dev/null | awk '$3 ~ /Z/ && /krasis|python/ {print $1 " " $2}' || true)
    if [[ -n "$zombie_pids" ]]; then
        warn "Found zombie processes — killing their parents:"
        while read -r zpid zppid; do
            warn "  Zombie PID $zpid, killing parent PID $zppid"
            kill -9 "$zppid" 2>/dev/null || true
        done <<< "$zombie_pids"
        sleep 2
    fi

    # Kill any process using GPU memory (except desktop compositor type stuff)
    local gpu_pids
    gpu_pids=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader,nounits 2>/dev/null | tr -d ' ' || true)
    if [[ -n "$gpu_pids" ]]; then
        for p in $gpu_pids; do
            [[ -z "$p" ]] && continue
            local cmdline
            cmdline=$(ps -p "$p" -o args= 2>/dev/null || echo "unknown")
            warn "  GPU process PID $p: $cmdline"
            kill -9 "$p" 2>/dev/null || true
        done
        sleep 2
    fi

    # Wait up to 10s for GPU to clear
    local waited=0
    while [[ $waited -lt 10 ]]; do
        local gpu0_used gpu1_used
        gpu0_used=$(nvidia-smi --id=0 --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | tr -d ' ')
        gpu1_used=$(nvidia-smi --id=1 --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | tr -d ' ')
        if [[ "${gpu0_used:-9999}" -lt 500 ]] && [[ "${gpu1_used:-9999}" -lt 500 ]]; then
            ok "All GPUs clear."
            return 0
        fi
        sleep 1
        waited=$((waited + 1))
    done

    # Last resort: nvidia-smi GPU reset
    warn "GPUs still not clear after 10s. Attempting nvidia-smi --gpu-reset..."
    if sudo nvidia-smi --gpu-reset 2>&1; then
        ok "GPU reset successful."
    else
        err "GPU reset failed. You may need to reboot.
  To reboot: sudo reboot"
    fi
}

do_shell() {
    exec_repo_python "$@"
}

do_python() {
    exec_repo_python "$@"
}

do_verify_imports() {
    verify_repo_imports
}

# ── Extract --timing flag ────────────────────────────────────────
# Can appear anywhere in args. Removes it and sets KRASIS_DECODE_TIMING=1.
TIMING_FLAG=false
args=()
for arg in "$@"; do
    if [[ "$arg" == "--timing" ]]; then
        TIMING_FLAG=true
    else
        args+=("$arg")
    fi
done
set -- "${args[@]}"

if $TIMING_FLAG; then
    export KRASIS_DECODE_TIMING=1
    info "Decode timing ENABLED (adds ~30-50% overhead — not for speed benchmarks)"
fi

#── Dispatch ─────────────────────────────────────────────────────────
cmd="${1:-help}"
shift || true

case "$cmd" in
    build)          do_build ;;
    run)            do_run "$@" ;;
    test)           do_test "$@" ;;
    benchmark)      do_benchmark "$@" ;;
    speed-test|speedtest) do_speed_test "$@" ;;
    decode-harness) do_decode_harness "$@" ;;
    release-test)   do_release_test "$@" ;;
    release-test-all) do_release_test_all "$@" ;;
    reference-test) do_reference_test "$@" ;;
    reference-inventory) do_reference_inventory "$@" ;;
    capture-box)    do_capture_box "$@" ;;
    capture-preflight) do_capture_preflight "$@" ;;
    reference-prep) do_reference_prep "$@" ;;
    generate-reference) do_generate_reference "$@" ;;
    trace-diff)     do_trace_diff "$@" ;;
    perplexity)     do_perplexity "$@" ;;
    awq-calibrate)  do_awq_calibrate "$@" ;;
    network)        do_network "$@" ;;
    test-kernels)   do_test_kernels "$@" ;;
    validate)       do_validate "$@" ;;
    shell)          do_shell "$@" ;;
    python)         do_python "$@" ;;
    verify-imports) do_verify_imports ;;
    help|-h|--help)
        echo "Krasis dev helper"
        echo ""
        echo "Usage: ./dev <command> [args...]"
        echo ""
        echo "Core commands:"
        echo "  build                          Rebuild Rust extension into this repo only"
        echo "  run <config> [flags]           Launch server (add --benchmark for benchmark mode)"
        echo "  benchmark <config>             Run standard benchmark and exit"
        echo "  speed-test                     Run the fixed standard speed benchmark"
        echo "  decode-harness <config>        Run synthetic GPU decode harness and exit"
        echo "  perplexity <config>            Run perplexity eval (WikiText-2) and exit"
        echo "  awq-calibrate <config>         Run AWQ attention calibration (produces template)"
        echo ""
        echo "Test and validation commands:"
        echo "  test <config>                  Short model test (benchmark + network tests)"
        echo "  test <config> --thorough       Thorough test (+ stress + large prompts)"
        echo "  network <port> [--large]       Run network tests against a running server"
        echo "  validate <config>              Validate model output against BF16 reference"
        echo "  reference-test <config>        Compare engine output against BF16 reference data"
        echo "  reference-test --compare <model>  Run comparative reference matrix"
        echo "  reference-inventory            Inventory stored reference artifacts"
        echo "  capture-box <subcommand>      Strict paid-box wrapper for HF reference capture"
        echo "  capture-preflight             Verify or bootstrap a host for HF reference capture"
        echo "  reference-prep <model>         Prepare HF capture deps and download a public model"
        echo "  generate-reference <model>     Generate BF16 HF reference capture"
        echo "  trace-diff <expected> <actual> Diff two KRASIS_TRACE logs offline"
        echo "  release-test <model>           Run full release test (4 configs, report)"
        echo "  release-test-all [target]      Run release test on all supported models sequentially"
        echo ""
        echo "Maintenance commands:"
        echo "  test-kernels [name]            Run CUDA kernel unit tests (or filter by name)"
        echo "  chat [args]                    Run krasis chat from source"
        echo "  sanity                         Run sanity test prompts against running server"
        echo "  kill                           Kill all krasis/GPU processes, reset if needed"
        echo "  install                        Sync krasis command to ~/.local/bin (uses source)"
        echo "  shell                          Python shell bound to this repo"
        echo "  python <args>                  Run Python bound to this repo"
        echo "  verify-imports                 Verify imports resolve inside this repo"
        echo ""
        echo "Global flags:"
        echo "  --timing    Enable per-component decode timing (adds ~30-50% overhead)"
        echo "              Prints detailed breakdown of where decode time is spent."
        echo "              NOT for speed benchmarks — only for profiling."
        echo ""
        echo "Reference workflow:"
        echo "  1. Inventory current artifacts:"
        echo "       ./dev reference-inventory"
        echo "  2. On paid boxes, enforce the remote environment first:"
        echo "       ./dev capture-box env"
        echo "  3. Prove the host is ready before queueing downloads/captures:"
        echo "       ./dev capture-box preflight --bootstrap --model <model>"
        echo "  4. Prepare the host and fetch the public model if needed:"
        echo "       ./dev capture-box prep <model>"
        echo "  5. Generate an explicit reference profile:"
        echo "       ./dev capture-box generate <model> --profile <profile>"
        echo "    Or stage multiple downloads at once:"
        echo "       ./dev capture-box stage qcn gemma"
        echo "  6. Validate the capture against a runtime config:"
        echo "       ./dev validate <config> --max-prompts 0 --no-server --port 1"
        echo "  7. Run the real comparison:"
        echo "       ./dev reference-test <config>"
        echo "  8. Diff stored traces offline when needed:"
        echo "       ./dev trace-diff <expected-trace> <actual-trace>"
        echo "  Reference-prep aliases:"
        echo "    qcn, qwen35, q122b, gemma, minimax, q235, q397"
        echo "  Notes:"
        echo "    - Uses public Hugging Face repo IDs only; does not use an HF token."
        echo "    - capture root:"
        echo "        $REFERENCE_CAPTURE_ROOT"
        echo "    - paid-box wrapper:"
        echo "        ./dev capture-box ..."
        echo "    - capture-preflight writes:"
        echo "        $REFERENCE_CAPTURE_READY_JSON"
        echo "        $REFERENCE_CAPTURE_READY_STAMP"
        echo "    - reference-prep installs capture/download deps in the isolated capture venv:"
        echo "        $REFERENCE_CAPTURE_VENV"
        echo "    - reference-prep downloads models to:"
        echo "        $REFERENCE_CAPTURE_MODELS_DIR"
        echo "    - Use --detach to keep a long download running with a log file."
        echo ""
        echo "Reference profiles for ./dev generate-reference:"
        echo "  auto                       Use thinking_off when template supports enable_thinking"
        echo "  greedy_chat_default        Default chat template invocation"
        echo "  greedy_chat_thinking_off   Pass enable_thinking=False"
        echo "  greedy_chat_thinking_on    Pass enable_thinking=True"
        echo "  Output files:"
        echo "    greedy_chat_default.json"
        echo "    greedy_chat_thinking_off.json"
        echo "    greedy_chat_thinking_on.json"
        echo "  Notes:"
        echo "    - Use explicit --profile for any durable reference capture."
        echo "    - Legacy greedy_reference.json is compatibility-only and may fail closed."
        echo "    - New captures store contract metadata and should be regenerated through ./dev only."
        echo ""
        echo "Trace / contract debugging:"
        echo "  KRASIS_TRACE=1 enables structured trace output."
        echo "  KRASIS_TRACE_COMPONENTS=config_contract traces reference inventory/generation/validation state."
        echo "  ./dev trace-diff compares stored KRASIS_TRACE logs, reports the first divergence,"
        echo "  summarizes matching-prefix plus context/phase counts, can suggest a short"
        echo "  resync point when one stream has inserted or dropped a few events, emits"
        echo "  typed mismatch analysis for tensor, values, logits, route, and python_compare events,"
        echo "  adds a top-level divergence classification summary,"
        echo "  adds local decode-family analysis for missing/unexpected sub-events within one"
        echo "  step/layer packet, adds bounded family-region resync hints across neighboring"
        echo "  decode packets, reports stronger family-packet signatures, and reports the"
        echo "  nearest preceding config_contract state for decode divergences."
        echo "  Example:"
        echo "    KRASIS_TRACE=1 KRASIS_TRACE_COMPONENTS=config_contract ./dev reference-inventory"
        echo "    ./dev trace-diff /tmp/expected.log /tmp/actual.log --components config_contract --contexts reference_test_single --phases artifact,validation"
        echo ""
        echo "Configs: qcn, qcn-a4, q235, q235-a4, v2lite, deepseek-vl, qwen35, qwen35-a4"
        echo "         (or path to .conf file)"
        echo "Standard speed benchmark config: tests/qcn-polar4-awq.conf"
        echo ""
        echo "Examples:"
        echo "  ./dev build"
        echo "  ./dev run qcn"
        echo "  ./dev run qcn --benchmark"
        echo "  ./dev speed-test"
        echo "  ./dev benchmark qcn --timing    Profile decode components"
        echo "  ./dev benchmark qcn             Normal speed benchmark"
        echo "  ./dev decode-harness qcn --prompt-len 8192 --steps 64"
        echo "  ./dev release-test Qwen3-Coder-Next    Full release test"
        echo "  ./dev release-test-all                 Test all models (no resume)"
        echo "  ./dev release-test-all rc8             Test all models (resume on re-run with same target)"
        echo "  ./dev capture-box env"
        echo "  ./dev capture-box preflight --bootstrap --model qcn"
        echo "  ./dev capture-box prep qcn --detach"
        echo "  ./dev capture-box stage qcn gemma"
        echo "  ./dev capture-box generate Qwen3-Coder-Next --profile greedy_chat_thinking_off"
        echo "  ./dev awq-calibrate qcn    Build AWQ attention template"
        echo "  ./dev perplexity qcn"
        echo "  ./dev trace-diff /tmp/expected.log /tmp/actual.log --events mark,tensor"
        echo "  ./dev trace-diff /tmp/expected.log /tmp/actual.log --components config_contract --contexts reference_inventory --show-window 3"
        echo "  ./dev trace-diff /tmp/expected.log /tmp/actual.log --components decode --resync-window 6"
        echo "  ./dev trace-diff /tmp/expected.log /tmp/actual.log --components decode --family-resync-window 3"
        echo "  ./dev test qcn"
        echo "  ./dev network 8012"
        echo "  ./dev python -m krasis.launcher"
        ;;
    kill)       do_kill ;;
    chat)       verify_repo_imports; exec_repo_python -m krasis.chat "$@" ;;
    sanity)
        verify_repo_imports
        run_dir=$(create_run_dir "dev-sanity")
        info "Run dir: $run_dir"
        KRASIS_RUN_DIR="$run_dir" KRASIS_RUN_TYPE="dev-sanity" exec_repo_python -m krasis.chat --sanitytest "$@"
        ;;
    install)
        info "Installing krasis command to ~/.local/bin/krasis..."
        mkdir -p "$HOME/.local/bin"
        # Symlink so installed command is always identical to repo script
        rm -f "$HOME/.local/bin/krasis"
        ln -s "$SCRIPT_DIR/krasis" "$HOME/.local/bin/krasis"
        # Clean up any stale pip/venv artifacts
        rm -f "$HOME/.local/bin/krasis-chat" "$HOME/.local/bin/krasis-setup"
        rm -rf "$HOME/.krasis/venv"
        ok "Installed. 'krasis' command symlinked to $SCRIPT_DIR/krasis"
        ;;
    *)
        err "Unknown command: $cmd
  Run ./dev help for usage."
        ;;
esac