bridge.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.
 
"""
Tool dispatch bridge: call(tool_name, args, ctx) -> dict.
Routes tool names to their compute implementations with no HTTP, no tab_queue,
and no gpu_executor. All GPU imports happen lazily inside handler functions so
this module imports cleanly even without a GPU or heavy ML deps installed.
 
ctx must contain: session_id, api_key, base_url, state (dict with activeModelId, memory).
"""
from __future__ import annotations
 
from typing import Any, Callable
 
# ---------------------------------------------------------------------------
# In-process sandbox state store (dataset/config tools)
# ---------------------------------------------------------------------------
_SANDBOX_STATE: dict[str, Any] = {}
 
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
 
def _get_model_id(ctx: dict, args: dict) -> str:
    from aquin.compute.model_loader import get_active_model_id
 
    state = ctx.get("state", {})
    model_id = args.get("model_id") or state.get("activeModelId") or get_active_model_id()
    if not model_id:
        raise ValueError("No model loaded for this session. Run: aquin load --model <id>")
    return model_id
 
 
def _as_str_list(value: Any, *, fallback: list[str] | None = None) -> list[str]:
    """Normalize CLI/API list flags — never iterate a JSON string character-by-character."""
    import json
 
    if value is None:
        return fallback or []
    if isinstance(value, list):
        return [str(v) for v in value]
    if isinstance(value, str):
        stripped = value.strip()
        if stripped.startswith("["):
            try:
                parsed = json.loads(stripped)
                if isinstance(parsed, list):
                    return [str(v) for v in parsed]
            except json.JSONDecodeError:
                pass
        return [value]
    return [str(value)]
 
 
def _as_json_dict(value: Any) -> dict | None:
    """Parse topics and other object-shaped CLI flags."""
    import json
 
    if value is None:
        return None
    if isinstance(value, dict):
        return value
    if isinstance(value, str):
        stripped = value.strip()
        if stripped.startswith("{"):
            try:
                parsed = json.loads(stripped)
                if isinstance(parsed, dict):
                    return parsed
            except json.JSONDecodeError:
                pass
    return None
 
 
def _as_json_list(value: Any) -> list:
    """Parse JSON array CLI flags (capabilities, evals, etc.)."""
    import json
 
    if value is None:
        return []
    if isinstance(value, list):
        return value
    if isinstance(value, str):
        stripped = value.strip()
        if stripped.startswith("["):
            try:
                parsed = json.loads(stripped)
                if isinstance(parsed, list):
                    return parsed
            except json.JSONDecodeError:
                pass
    return []
 
 
def _get_model(ctx: dict):
    """Return the loaded model, loading it if necessary."""
    from aquin.compute.model_loader import get_active_model_id, get_loaded_model, load_model
 
    model = get_loaded_model()
    if model is None:
        model_id = ctx.get("state", {}).get("activeModelId") or get_active_model_id()
        if not model_id:
            raise ValueError("No model loaded for this session. Run: aquin load --model <id>")
        model = load_model(model_id)
    return model
 
 
def _not_impl(tool_name: str) -> dict:
    return {"status": "not_implemented", "tool": tool_name}
 
 
# ---------------------------------------------------------------------------
# Individual handlers
# ---------------------------------------------------------------------------
 
def _run_steer_and_show(args: dict, ctx: dict) -> dict:
    from aquin.compute.model_loader import resolve_model_id
    from aquin.compute.steer_vector import run_steer_with_vector
 
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
 
    if args.get("eval") and (args.get("save") or args.get("out") or args.get("output_path")):
        return {"error": "Cannot combine --eval with --save. Export the vector first, then steer --eval --vector <path>."}
 
    vector_path = args.get("vector") or args.get("vector_path")
    feature_idx_raw = args.get("feature_idx")
    if vector_path:
        feature_idx = int(feature_idx_raw) if feature_idx_raw is not None else None
    elif feature_idx_raw is not None:
        feature_idx = int(feature_idx_raw)
    else:
        return {"error": "Provide --feature_idx <n> or --vector <path> to a saved LAT file."}
 
    steer_strength = float(args.get("steer_strength", args.get("strength", 20.0)))
    explicit_prompt = args.get("prompt")
    if explicit_prompt is not None:
        prompt = str(explicit_prompt)
    elif args.get("eval"):
        prompt = None
    else:
        prompt = ctx.get("state", {}).get("memory", {}).get("lastPrompt") or "Hello"
    layer = args.get("layer")
    if layer is not None:
        layer = int(layer)
    max_new_tokens = int(args.get("max_new_tokens", 80))
 
    return run_steer_with_vector(
        model_id=model_id,
        prompt=prompt,
        steer_strength=steer_strength,
        layer=layer,
        feature_idx=feature_idx,
        vector_path=vector_path,
        feature_label=args.get("feature_label"),
        max_new_tokens=max_new_tokens,
        ctx=ctx,
        args=args,
    )
 
 
def _extract_steer_vector(args: dict, ctx: dict) -> dict:
    from aquin.compute.model_loader import resolve_model_id
    from aquin.compute.steer_vector import extract_steer_vector
 
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
 
    output = args.get("save") or args.get("out") or args.get("output_path")
    if not output:
        return {"error": "--save <path> is required."}
 
    feature_idx_raw = args.get("feature_idx")
    if feature_idx_raw is None:
        return {"error": "--feature_idx <n> is required."}
 
    layer = args.get("layer")
    try:
        return extract_steer_vector(
            model_id,
            int(feature_idx_raw),
            output,
            layer=int(layer) if layer is not None else None,
            feature_label=args.get("feature_label"),
            probe_id=args.get("probe_id"),
        )
    except (ValueError, FileNotFoundError) as e:
        return {"error": str(e)}
    except Exception as e:
        return {"error": f"steer --save failed: {e}"}
 
 
def _run_multi_steer(args: dict, ctx: dict) -> dict:
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id, get_config
    from aquin.compute.feature_analysis import load_sae, load_norm
    import torch
 
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
 
    # features: list of {feature_idx, strength}
    features = args.get("features", [])
    if not features:
        feature_idx = int(args.get("feature_idx", 0))
        strength = float(args.get("steer_strength", args.get("strength", 20.0)))
        features = [{"feature_idx": feature_idx, "strength": strength}]
 
    prompt = args.get("prompt") or ctx.get("state", {}).get("memory", {}).get("lastPrompt") or "Hello"
    layer = args.get("layer")
 
    model = get_loaded_model()
    if model is None:
        try:
            model = load_model(model_id)
        except Exception as e:
            return {"error": str(e)}
 
    try:
        sae = load_sae(model_id, layer)
    except Exception as e:
        return {"error": f"SAE load failed: {e}"}
 
    cfg = get_config(model_id)
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
    hook_name = f"blocks.{resolved_layer}.hook_resid_post"
 
    from aquin.compute.causal_trace import run_chat, _format_prompt
 
    try:
        baseline = run_chat(prompt, model_id=model_id, max_new_tokens=80, temperature=0.0)
    except Exception as e:
        return {"error": f"Baseline generation failed: {e}"}
 
    steer_vecs = []
    for f in features:
        fidx = int(f.get("feature_idx", 0))
        s = float(f.get("strength", 20.0))
        steer_vecs.append((
            sae.W_dec[fidx].to(device=model.W_E.device, dtype=model.W_E.dtype),
            s,
        ))
 
    def _multi_steer_hook(value, hook):
        for vec, s in steer_vecs:
            value = value + s * vec.unsqueeze(0).unsqueeze(0)
        return value
 
    try:
        fmt_prompt = _format_prompt(model, prompt)
        tokens = model.to_tokens(fmt_prompt)
        cur = tokens
        steered_token_ids = []
        for _ in range(80):
            with torch.no_grad():
                out = model.run_with_hooks(cur, fwd_hooks=[(hook_name, _multi_steer_hook)])
            next_id = int(out[0, -1].argmax().item())
            if next_id == model.tokenizer.eos_token_id:
                break
            steered_token_ids.append(next_id)
            cur = torch.cat([cur, torch.tensor([[next_id]], device=cur.device)], dim=1)
        steered_response = model.tokenizer.decode(steered_token_ids, skip_special_tokens=True)
    except Exception as e:
        steered_response = f"[steer error: {e}]"
 
    from aquin.compute.feature_analysis import format_feature_ref, resolve_feature_label
    enriched_features = []
    for f in features:
        fidx = int(f.get("feature_idx", 0))
        label = f.get("label") or resolve_feature_label(fidx, ctx=ctx, args=args, layer=layer)
        enriched_features.append({
            **f,
            "feature_idx": fidx,
            "label": label,
            "feature_ref": format_feature_ref(fidx, label),
        })
 
    return {
        "features": enriched_features,
        "layer": resolved_layer,
        "prompt": prompt,
        "baseline_response": baseline,
        "steered_response": steered_response,
    }
 
 
def _run_consistency_eval(args: dict, ctx: dict) -> dict:
    from aquin.compute.evals import consistency_eval
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    model = get_loaded_model()
    if model is None:
        model = load_model(model_id)
    query = args.get("query") or args.get("prompt") or "Hello"
    templates_raw = args.get("templates")
    templates = _as_str_list(templates_raw) or None
    return consistency_eval(query, model, templates=templates or None)
 
 
def _run_suppression_eval(args: dict, ctx: dict) -> dict:
    from aquin.compute.evals import suppression_eval
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    model = get_loaded_model()
    if model is None:
        model = load_model(model_id)
    topics = _as_json_dict(args.get("topics"))
    return suppression_eval(model, topics=topics)
 
 
def _run_boundary_eval(args: dict, ctx: dict) -> dict:
    from aquin.compute.evals import boundary_eval
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    model = get_loaded_model()
    if model is None:
        model = load_model(model_id)
    prompts = _as_str_list(
        args.get("prompts"),
        fallback=[args.get("prompt", "Hello")],
    )
    if not prompts:
        prompts = ["Hello"]
    return boundary_eval(prompts, model)
 
 
def _run_audit(args: dict, ctx: dict) -> dict:
    from aquin.compute.evals import consistency_eval, suppression_eval, boundary_eval
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    model = get_loaded_model()
    if model is None:
        model = load_model(model_id)
    query = args.get("query") or args.get("prompt") or "Hello"
    prompts = _as_str_list(args.get("prompts"), fallback=[query])
    consistency = consistency_eval(query, model)
    suppression = suppression_eval(model)
    boundary = boundary_eval(prompts, model)
    return {
        "consistency": consistency,
        "suppression": suppression,
        "boundary": boundary,
        "model_id": model_id,
    }
 
 
def _run_benchmarks_on_top_feature(args: dict, ctx: dict) -> dict:
    from aquin.compute.interp_score import run_interp_score
    from aquin.compute.feature_analysis import load_sae
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    feature_idx = int(args.get("feature_idx", 0))
    memory = ctx.get("state", {}).get("memory") or {}
    layer_raw = args.get("layer")
    if layer_raw is None and memory.get("lastSaeLayer") is not None:
        layer = int(memory["lastSaeLayer"])
    elif layer_raw is not None:
        layer = int(layer_raw)
    else:
        layer = None
    prompt = args.get("prompt") or memory.get("lastPrompt") or "Hello"
    model = get_loaded_model()
    if model is None:
        try:
            model = load_model(model_id)
        except Exception as e:
            return {"error": str(e)}
    try:
        sae = load_sae(model_id, layer)
    except Exception as e:
        return {"error": f"SAE load failed: {e}"}
    from aquin.compute.openai_client import get_openai_client
    client = get_openai_client(ctx)
    interp = run_interp_score(
        feature_idx, prompt, model, sae, client,
        model_id=model_id, layer=layer,
    )
    from aquin.compute.feature_analysis import format_feature_ref
    label = interp.get("label")
    return {
        "feature_idx": feature_idx,
        "feature_ref": format_feature_ref(feature_idx, label),
        "model_id": model_id,
        "label": label,
        "score": interp.get("score"),
        "purity_score": interp.get("purity_score"),
        "mui_score": interp.get("mui_score"),
        "mui_per_position": interp.get("mui_per_position"),
        "mui_mean_kl": interp.get("mui_mean_kl"),
        "baseline_entropy": interp.get("baseline_entropy"),
        "positive_mean": interp.get("positive_mean"),
        "negative_mean": interp.get("negative_mean"),
        "positive_examples": interp.get("positive_examples", []),
        "negative_examples": interp.get("negative_examples", []),
        "error": interp.get("error"),
    }
 
 
def _check_weights(args: dict, ctx: dict) -> dict:
    from datetime import datetime, timezone
 
    import torch
 
    from aquin.compute.weight_rank import run_weight_rank_analysis
    from aquin.compute.weight_trojans import run_weight_trojan_analysis
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id
 
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    try:
        model = get_loaded_model()
        if model is None:
            model = load_model(model_id)
    except Exception as exc:
        from aquin.user_errors import friendly_message
        return {"error": friendly_message(exc), "model_id": model_id}
 
    layer_range = args.get("layer_range")
    try:
        trojan = run_weight_trojan_analysis(model, layer_range=layer_range)
    except NotImplementedError as exc:
        return {
            "error": (
                "Weight trojan scan is not supported for this model architecture. "
                f"{exc}".strip()
            ),
            "model_id": model_id,
        }
    except Exception as _oom_exc:
        from aquin.compute.device import empty_device_cache, is_oom_error
        if not is_oom_error(_oom_exc):
            raise
        empty_device_cache()
        return {
            "error": (
                "GPU ran out of memory during weight trojan scan. "
                "The model already fills most VRAM — try unloading other processes "
                "or run on a host with more headroom."
            ),
            "model_id": model_id,
        }
 
    if trojan.get("error"):
        return {"error": trojan["error"], "model_id": model_id}
 
    collapse_threshold = float(args.get("collapse_threshold") or 0.1)
    try:
        from aquin.compute.device import empty_device_cache
        empty_device_cache()
        rank = run_weight_rank_analysis(model, collapse_threshold=collapse_threshold)
    except NotImplementedError as exc:
        return {
            "error": (
                "Weight rank analysis is not supported for this model architecture. "
                f"{exc}".strip()
            ),
            "model_id": model_id,
            **trojan,
        }
    except Exception as _oom_exc:
        from aquin.compute.device import empty_device_cache, is_oom_error
        if not is_oom_error(_oom_exc):
            raise
        empty_device_cache()
        return {
            "error": (
                "GPU ran out of memory during weight rank analysis. "
                "Trojan scan completed; rank pass needs more free VRAM."
            ),
            "model_id": model_id,
            **trojan,
        }
    rank["model_id"] = model_id
 
    generated_at = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    return {
        "model_id": model_id,
        "generated_at": generated_at,
        **trojan,
        "rank": rank,
    }
 
 
def _normalize_umap_points(raw: list[dict]) -> list[dict]:
    """Map bridge output to web UMAPExplorer shape (feature_idx, x/y/z in 0–1)."""
    rows: list[dict] = []
    for p in raw:
        idx = p.get("feature_idx", p.get("i"))
        if idx is None:
            continue
        try:
            x, y = float(p["x"]), float(p["y"])
        except (KeyError, TypeError, ValueError):
            continue
        z = float(p.get("z", 0))
        rows.append({
            "feature_idx": int(idx),
            "x": x,
            "y": y,
            "z": z,
            "label": p.get("label"),
        })
    if not rows:
        return []
    xs = [r["x"] for r in rows]
    ys = [r["y"] for r in rows]
    zs = [r["z"] for r in rows]
    xmin, xmax = min(xs), max(xs)
    ymin, ymax = min(ys), max(ys)
    zmin, zmax = min(zs), max(zs)
    xspan = max(xmax - xmin, 1e-9)
    yspan = max(ymax - ymin, 1e-9)
    zspan = max(zmax - zmin, 1e-9)
    flat_z = zspan < 1e-9
    out: list[dict] = []
    for r in rows:
        out.append({
            "feature_idx": r["feature_idx"],
            "x": (r["x"] - xmin) / xspan,
            "y": (r["y"] - ymin) / yspan,
            "z": 0.0 if flat_z else (r["z"] - zmin) / zspan,
            **({"label": r["label"]} if r.get("label") else {}),
        })
    return out
 
 
def _try_fetch_precomputed_umap(model_id: str, layer: int) -> list[dict] | None:
    """Load precomputed 3D UMAP from the public SAE DB when available."""
    try:
        import os
        import requests
        from aquin.saes import list_saes
 
        rows = list_saes(model=model_id)
        match = next(
            (r for r in rows if r.get("layer") == layer and r.get("umap_key")),
            None,
        )
        if not match:
            return None
        base = os.environ.get("AQUIN_BASE_URL", "https://api.aquin.app").rstrip("/")
        resp = requests.get(f"{base}/api/public/saes/{match['id']}/umap", timeout=120)
        resp.raise_for_status()
        data = resp.json()
        raw = data.get("points") if isinstance(data, dict) else data
        if not isinstance(raw, list) or not raw:
            return None
        return raw
    except Exception:
        return None
 
 
def _compute_umap_points(model_id: str, layer: int | None) -> tuple[list[dict], int]:
    """Project SAE decoder weights into 3D (matches /db precomputed UMAPs)."""
    from aquin.compute.feature_analysis import load_sae
    import torch
 
    sae = load_sae(model_id, layer)
    W = sae.W_dec.float().cpu().detach()
    try:
        import umap
        reducer = umap.UMAP(n_components=3, random_state=42, n_neighbors=15)
        embedding = reducer.fit_transform(W.numpy())
        points = [
            {"feature_idx": i, "x": float(r[0]), "y": float(r[1]), "z": float(r[2])}
            for i, r in enumerate(embedding)
        ]
    except ImportError:
        W_centered = W - W.mean(0)
        _, _, Vt = torch.linalg.svd(W_centered, full_matrices=False)
        proj = (W_centered @ Vt[:3].T).numpy()
        points = [
            {"feature_idx": i, "x": float(r[0]), "y": float(r[1]), "z": float(r[2])}
            for i, r in enumerate(proj)
        ]
    except Exception as e:
        raise RuntimeError(f"UMAP failed: {e}") from e
    return points, len(points)
 
 
def _ensure_umap_loaded(args: dict, ctx: dict) -> dict:
    from aquin.compute.sync_slim import downsample_umap_points
    from aquin.compute.model_loader import resolve_model_id, get_config
 
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
 
    cfg = get_config(model_id)
    layer = args.get("layer")
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
 
    source = "precomputed"
    try:
        raw_points = _try_fetch_precomputed_umap(model_id, resolved_layer)
        if raw_points is None:
            source = "computed"
            raw_points, _ = _compute_umap_points(model_id, layer)
    except Exception as e:
        return {"error": str(e)}
 
    normalized = _normalize_umap_points(raw_points)
    sync_points = downsample_umap_points(normalized)
    return {
        "points": sync_points,
        "n_features": len(normalized),
        "n_points": len(sync_points),
        "layer": resolved_layer,
        "model_id": model_id,
        "source": source,
    }
 
 
def _run_layer_analysis(args: dict, ctx: dict) -> dict:
    from aquin.compute.layer_analysis import run_layer_analysis
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id
 
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    model = get_loaded_model()
    if model is None:
        model = load_model(model_id)
    try:
        return run_layer_analysis(model, {**args, "model_id": model_id})
    except Exception as e:
        return {"error": str(e)}
 
 
def _run_sae_stats(args: dict, ctx: dict) -> dict:
    import json
    import os
 
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
    from aquin.compute.sae_stats import run_sae_stats
 
    model_id = (
        args.get("model_id")
        or ctx.get("state", {}).get("activeModelId")
        or get_active_model_id()
    )
    if model_id:
        try:
            model_id = resolve_model_id(str(model_id))
            args = {**args, "model_id": model_id}
        except ValueError:
            pass
 
    result = run_sae_stats(args)
    if result.get("error"):
        return result
 
    save_path = args.get("save")
    if save_path:
        from pathlib import Path
        p = Path(os.path.expanduser(str(save_path)))
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(json.dumps(result, indent=2, default=str), encoding="utf-8")
        result["saved_to"] = str(p)
    return result
 
 
def _run_weight_diff(args: dict, ctx: dict) -> dict:
    import json
    import os
 
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
    from aquin.compute.weight_diff import run_weight_diff_from_args
 
    model_id = (
        args.get("model_id")
        or ctx.get("state", {}).get("activeModelId")
        or get_active_model_id()
    )
    if model_id:
        try:
            model_id = resolve_model_id(str(model_id))
            args = {**args, "model_id": model_id}
        except ValueError:
            pass
 
    result = run_weight_diff_from_args(args)
    if result.get("error"):
        return result
 
    save_path = args.get("save")
    if save_path:
        from pathlib import Path
        p = Path(os.path.expanduser(str(save_path)))
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(json.dumps(result, indent=2, default=str), encoding="utf-8")
        result["saved_to"] = str(p)
    return result
 
 
def _run_merge_analysis(args: dict, ctx: dict) -> dict:
    import json
    import os
 
    from aquin.compute.merge_analysis import run_merge_analysis_from_args
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
 
    model_id = (
        args.get("model_id")
        or ctx.get("state", {}).get("activeModelId")
        or get_active_model_id()
    )
    if model_id:
        try:
            model_id = resolve_model_id(str(model_id))
            args = {**args, "model_id": model_id}
        except ValueError:
            pass
 
    result = run_merge_analysis_from_args(args)
    if result.get("error"):
        return result
 
    save_path = args.get("save")
    if save_path:
        from pathlib import Path
        p = Path(os.path.expanduser(str(save_path)))
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(json.dumps(result, indent=2, default=str), encoding="utf-8")
        result["saved_to"] = str(p)
    return result
 
 
def _run_trajectory_analysis(args: dict, ctx: dict) -> dict:
    import json
    import os
 
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
    from aquin.compute.trajectory_analysis import run_trajectory_analysis_from_args
 
    model_id = (
        args.get("model_id")
        or ctx.get("state", {}).get("activeModelId")
        or get_active_model_id()
    )
    if model_id:
        try:
            model_id = resolve_model_id(str(model_id))
            args = {**args, "model_id": model_id}
        except ValueError:
            pass
 
    result = run_trajectory_analysis_from_args(args)
    if result.get("error"):
        return result
 
    save_path = args.get("save")
    if save_path:
        from pathlib import Path
        p = Path(os.path.expanduser(str(save_path)))
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(json.dumps(result, indent=2, default=str), encoding="utf-8")
        result["saved_to"] = str(p)
    return result
 
 
def _run_residual_drift(args: dict, ctx: dict) -> dict:
    import json
    import os
 
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
    from aquin.compute.residual_drift import run_residual_drift_from_args
 
    model_id = (
        args.get("model_id")
        or ctx.get("state", {}).get("activeModelId")
        or get_active_model_id()
    )
    if model_id:
        try:
            model_id = resolve_model_id(str(model_id))
            args = {**args, "model_id": model_id}
        except ValueError:
            pass
 
    result = run_residual_drift_from_args(args)
    if result.get("error"):
        return result
 
    save_path = args.get("save")
    if save_path:
        from pathlib import Path
        p = Path(os.path.expanduser(str(save_path)))
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(json.dumps(result, indent=2, default=str), encoding="utf-8")
        result["saved_to"] = str(p)
    return result
 
 
def _run_confidence_analysis(args: dict, ctx: dict) -> dict:
    import json
    import os
 
    from aquin.compute.confidence_analysis import run_confidence_analysis_from_args
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
 
    model_id = (
        args.get("model_id")
        or ctx.get("state", {}).get("activeModelId")
        or get_active_model_id()
    )
    if model_id:
        try:
            model_id = resolve_model_id(str(model_id))
            args = {**args, "model_id": model_id}
        except ValueError:
            pass
 
    result = run_confidence_analysis_from_args(args)
    if result.get("error"):
        return result
 
    save_path = args.get("save")
    if save_path:
        from pathlib import Path
        p = Path(os.path.expanduser(str(save_path)))
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(json.dumps(result, indent=2, default=str), encoding="utf-8")
        result["saved_to"] = str(p)
    return result
 
 
def _token_kl_divergence(clean_logits, perturbed_logits) -> float:
    """KL(clean || perturbed) on a single token distribution — stable at zero mass."""
    import torch.nn.functional as F
 
    clean_log = F.log_softmax(clean_logits.float(), dim=-1)
    pert_log = F.log_softmax(perturbed_logits.float(), dim=-1)
    clean_prob = clean_log.exp()
    kl = (clean_prob * (clean_log - pert_log)).sum()
    val = float(kl.item())
    return 0.0 if val != val else max(val, 0.0)
 
 
def _run_perturbation_sensitivity(args: dict, ctx: dict) -> dict:
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id, get_config
    from aquin.compute.feature_analysis import load_sae
    import torch
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    model = get_loaded_model()
    if model is None:
        model = load_model(model_id)
    prompt = args.get("prompt") or "Hello world"
    n_channels = int(args.get("n_channels", 10))
    layer = args.get("layer")
    cfg = get_config(model_id)
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
    hook_name = f"blocks.{resolved_layer}.hook_resid_post"
    tokens = model.to_tokens(prompt)
    with torch.no_grad():
        clean_logits = model(tokens)
    clean_last = clean_logits[0, -1]
    try:
        sae = load_sae(model_id, resolved_layer)
        n_features = sae.W_dec.shape[0]
    except Exception:
        n_features = 64
        sae = None
    results = []
    for ch in range(min(n_channels, n_features)):
        def _zero_hook(value, hook, ch_=ch):
            if sae is not None:
                direction = sae.W_dec[ch_].to(device=value.device, dtype=value.dtype)
                coeff = (value * direction).sum(dim=-1, keepdim=True)
                return value - coeff * direction
            value = value.clone()
            value[..., ch_] = 0.0
            return value
 
        with torch.no_grad():
            perturbed = model.run_with_hooks(tokens, fwd_hooks=[(hook_name, _zero_hook)])
        kl = _token_kl_divergence(clean_last, perturbed[0, -1])
        results.append({"channel": ch, "kl_divergence": round(kl, 6)})
    results.sort(key=lambda x: x["kl_divergence"], reverse=True)
 
    channels = [
        {
            "channel": r["channel"],
            "kl_mean": r["kl_divergence"],
            "kl_max": r["kl_divergence"],
            "kl_std": 0.0,
            "layer_contributions": [{"layer": resolved_layer, "kl": r["kl_divergence"]}],
        }
        for r in results
    ]
    kl_vals = [c["kl_mean"] for c in channels]
    global_mean = sum(kl_vals) / len(kl_vals) if kl_vals else 0.0
    global_max = max(kl_vals) if kl_vals else 0.0
 
    return {
        "prompt": prompt,
        "method": args.get("method") or "dropout",
        "n_channels_sampled": len(channels),
        "d_model": model.cfg.d_model,
        "channels": channels,
        "global_mean_kl": round(global_mean, 6),
        "global_max_kl": round(global_max, 6),
        "threshold_critical": round(global_max * 0.85, 6) if global_max else 0.01,
        "threshold_high": round(global_mean + (global_max - global_mean) * 0.5, 6) if kl_vals else 0.005,
        "perturbation_results": results,
        "layer": resolved_layer,
        "model_id": model_id,
    }
 
 
def _run_attention_routing(args: dict, ctx: dict) -> dict:
    from aquin.compute.attention_routing import run_attention_routing
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id
 
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    model = get_loaded_model()
    if model is None:
        model = load_model(model_id)
    prompt = args.get("prompt") or args.get("text")
    if not prompt or not str(prompt).strip():
        return {"error": "Missing --prompt. Usage: aquin check attention --prompt <text>"}
    prompt = str(prompt).strip()
    top_k = int(args.get("top_k") or args.get("topk") or 5)
    try:
        return run_attention_routing(model, prompt, top_k=top_k, model_id=model_id)
    except Exception as e:
        return {"error": str(e)}
 
 
def _run_find_feature(args: dict, ctx: dict) -> dict:
    from aquin.compute.find_feature import run_find_feature
    from aquin.compute.model_loader import resolve_model_id
 
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
 
    benchmark_top = int(args.get("benchmark_top") or args.get("benchmark_top_k") or 0)
    openai_client = None
    if benchmark_top > 0:
        try:
            from aquin.compute.openai_client import get_openai_client
            openai_client = get_openai_client(ctx)
        except Exception:
            benchmark_top = 0
 
    session_id = ctx.get("session_id") or (ctx.get("state") or {}).get("session_id")
    checkpoint = args.get("checkpoint") or None
    if checkpoint == "":
        checkpoint = None
 
    try:
        return run_find_feature(
            model_id,
            scorer=str(args.get("scorer") or "deception"),
            prompts_path=args.get("prompts") or args.get("prompts_path"),
            layer=int(args["layer"]) if args.get("layer") is not None else None,
            checkpoint_path=checkpoint,
            top_k=int(args.get("top_k") or args.get("top") or 20),
            direction=str(args.get("direction") or "both"),
            conditioning=str(args.get("conditioning") or "behavior"),
            benchmark_top=benchmark_top,
            persist_key=args.get("persist") or None,
            session_id=str(session_id) if session_id else None,
            openai_client=openai_client,
        )
    except Exception as e:
        return {"error": str(e)}
 
# ---------------------------------------------------------------------------
# Dataset / config state tools
# ---------------------------------------------------------------------------
 
def _dataset_generate(args: dict, ctx: dict) -> dict:
    from pathlib import Path
 
    from aquin.compute.dataset_generate import run_dataset_generate
 
    topic = str(args.get("topic") or "").strip()
    if not topic:
        return {"error": "Pass --topic <text> (e.g. aquin simulate --topic flowers)"}
 
    try:
        count = int(args.get("count") or 5)
    except (TypeError, ValueError):
        return {"error": "count must be an integer"}
 
    cwd = args.get("cwd") or ctx.get("cwd") or Path.cwd()
    output = args.get("save")
 
    try:
        result = run_dataset_generate(
            topic=topic,
            count=count,
            cwd=cwd,
            output=str(output) if output else None,
        )
    except ValueError as e:
        return {"error": str(e)}
    except OSError as e:
        return {"error": f"Could not write dataset file: {e}"}
 
    ds = _SANDBOX_STATE.setdefault("dataset", {})
    ds["rows"] = result["rows"]
    ds["topic"] = result["topic"]
    ds["path"] = result["path"]
 
    return result
 
 
def _session_sim_kind_filter(ctx: dict) -> str | None:
    """LLM sessions list LLM sims."""
    return "simulate"
 
def _assemble_simulation_result(events: list[dict]) -> dict:
    """Fold simulation SSE events into one result dict for compare/load."""
    skip = {"log", "state"}
    result: dict = {}
    for ev in events:
        key = ev.get("type")
        if not key or key in skip:
            continue
        payload = {k: v for k, v in ev.items() if k != "type"}
        if key == "meta":
            result.setdefault("meta", {}).update(payload)
        elif key == "step":
            loss = ev.get("loss")
            if loss is not None:
                result.setdefault("lossHistory", []).append(loss)
        elif key == "gradHeatmap":
            result.setdefault("gradHeatmap", []).append(payload)
        elif key == "signal":
            result.setdefault("signals", []).append({
                "signalType": ev.get("signalType", ""),
                "severity": ev.get("severity", ""),
                "message": ev.get("message", ""),
                "step": ev.get("step", 0),
            })
        else:
            result[key] = payload
    return result
 
 
def _persist_simulation_run(events: list[dict], meta: dict) -> str:
    """Save simulation output under ~/.aquin/runs/<id> for list/compare."""
    import json
    import uuid
    from datetime import datetime, timezone
    from pathlib import Path
 
    run_id = uuid.uuid4().hex[:12]
    run_dir = Path.home() / ".aquin" / "runs" / run_id
    run_dir.mkdir(parents=True, exist_ok=True)
    result = _assemble_simulation_result(events)
    result.update({k: v for k, v in meta.items() if v is not None})
    (run_dir / "result.json").write_text(json.dumps(result, default=str))
    run_meta = {
        "run_id": run_id,
        "saved_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        **meta,
    }
    (run_dir / "meta.json").write_text(json.dumps(run_meta, default=str))
    return run_id
 
 
def _load_simulation_result(run_id: str) -> dict:
    import json
    from pathlib import Path
 
    run_dir = Path.home() / ".aquin" / "runs" / run_id
    result_file = run_dir / "result.json"
    if result_file.exists():
        return json.loads(result_file.read_text())
    meta_file = run_dir / "meta.json"
    if meta_file.exists():
        return json.loads(meta_file.read_text())
    return {"status": "not_found", "run_id": run_id}
 
 
def _list_simulation_runs(args: dict, ctx: dict) -> dict:
    import json
    from datetime import datetime, timezone
    from pathlib import Path
 
    from aquin.compute.card_mapper import _simulation_method_from_meta
 
    runs_dir = Path.home() / ".aquin" / "runs"
    if not runs_dir.exists():
        return {"runs": [], "count": 0}
 
    entries: list[dict] = []
    kind_filter = _session_sim_kind_filter(ctx)
    for run_dir in runs_dir.iterdir():
        if not run_dir.is_dir():
            continue
        run_id = run_dir.name
        meta: dict = {}
        meta_file = run_dir / "meta.json"
        if meta_file.exists():
            try:
                meta = json.loads(meta_file.read_text())
            except json.JSONDecodeError:
                meta = {}
 
        entry_kind = str(meta.get("kind") or "simulate")
        if kind_filter and entry_kind != kind_filter:
            continue
 
        n_samples = None
        harmful_count = None
        method = None
        result_file = run_dir / "result.json"
        if result_file.exists():
            try:
                result = json.loads(result_file.read_text())
                dq = result.get("datasetQuality")
                if isinstance(dq, dict):
                    n_samples = dq.get("nSamples")
                    harmful_count = dq.get("harmfulCount")
                meta_block = result.get("meta")
                if isinstance(meta_block, dict) and meta_block.get("algoConfig"):
                    method = _simulation_method_from_meta(meta_block)
            except json.JSONDecodeError:
                pass
 
        saved_at = str(meta.get("saved_at") or "")
        if not saved_at:
            try:
                saved_at = datetime.fromtimestamp(
                    run_dir.stat().st_mtime, tz=timezone.utc
                ).strftime("%Y-%m-%dT%H:%M:%SZ")
            except OSError:
                saved_at = ""
 
        entries.append({
            "run_id": run_id,
            "saved_at": saved_at,
            "model_id": str(meta.get("model_id") or meta.get("modelId") or ""),
            "topic": meta.get("topic"),
            "dataset_path": meta.get("dataset_path") or meta.get("pairs_path"),
            "kind": entry_kind,
            "n_samples": n_samples,
            "harmful_count": harmful_count,
            "method": method or "LoRA",
        })
 
    entries.sort(key=lambda e: e.get("saved_at") or "", reverse=True)
    return {"runs": entries, "count": len(entries)}
 
 
def _load_simulation_run(args: dict, ctx: dict) -> dict:
    run_id = args.get("run_id", "")
    data = _load_simulation_result(run_id)
    if data.get("status") == "not_found":
        return data
    return {"run_id": run_id, **data}
 
 
def _compare_simulations(args: dict, ctx: dict) -> dict:
    """Compare two simulation runs."""
    run_id_a = args.get("run_id_a") or args.get("run_id_before") or ""
    run_id_b = args.get("run_id_b") or args.get("run_id_after") or ""
    label_a = args.get("label_a", "Before")
    label_b = args.get("label_b", "After")
 
    result_a = args.get("result_a") or (_load_simulation_result(run_id_a) if run_id_a else {})
    result_b = args.get("result_b") or (_load_simulation_result(run_id_b) if run_id_b else {})
 
    if result_a.get("status") == "not_found" or result_b.get("status") == "not_found":
        missing = [
            rid for rid, data in ((run_id_a, result_a), (run_id_b, result_b))
            if data.get("status") == "not_found" and rid
        ]
        return {"error": f"Simulation run not found: {', '.join(missing)}"}
 
    if not result_a or not result_b:
        return {
            "error": "Two simulation results required. Run simulate twice, then: "
            "aquin compare simulation --run_id_a <id> --run_id_b <id>"
        }
 
    from aquin.compute.train_simulate import compare_simulation_results
    comparison = compare_simulation_results(
        result_a, result_b, label_a=label_a, label_b=label_b,
        run_id_a=run_id_a, run_id_b=run_id_b,
    )
    return {
        "run_id_a": run_id_a,
        "run_id_b": run_id_b,
        "label_a": label_a,
        "label_b": label_b,
        "comparison": comparison,
        "attack_surface": comparison.get("modelScores", {}),
    }
 
 
def _simulation_error(events: list) -> str | None:
    for ev in events:
        if ev.get("type") == "error" and ev.get("message"):
            return str(ev["message"])
        if ev.get("type") == "log" and str(ev.get("line", "")).startswith("ERROR:"):
            return str(ev["line"])[6:].strip()
    return None
 
 
 
def _wrap_simulation_result(events: list, final: dict, meta: dict) -> dict:
    run_id = _persist_simulation_run(events, meta)
    assembled = _assemble_simulation_result(events)
    err = _simulation_error(events)
    return {
        "events": events,
        "result": assembled or final,
        "run_id": run_id,
        "status": "error" if err else "complete",
        **({"error": err} if err else {}),
        **{k: v for k, v in meta.items() if k not in ("events", "result", "run_id", "status", "error")},
    }
 
 
# ---------------------------------------------------------------------------
# Custom eval
# ---------------------------------------------------------------------------
 
def _run_custom_eval(args: dict, ctx: dict) -> dict:
    from aquin.compute.evals import custom_eval
    from aquin.compute.model_loader import resolve_model_id
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    scorer_type = args.get("scorer_type") or "semantic_similarity"
    if scorer_type != "semantic_similarity":
        return {
            "error": (
                "run_custom_eval only supports scorer_type=semantic_similarity. "
                "Use built-in evals or run custom logic outside Aquin."
            ),
        }
    prompts = _as_str_list(args.get("prompts"))
    reference_answers = _as_str_list(args.get("reference_answers"))
    return custom_eval(
        args.get("name") or "Custom eval",
        prompts,
        model_id,
        reference_answers=reference_answers,
        threshold=float(args.get("threshold", 0.5)),
        max_tokens=int(args.get("max_tokens", 40)),
        temperature=float(args.get("temperature", 0.0)),
        description=args.get("description"),
    )
 
 
# ---------------------------------------------------------------------------
# Card builders
# ---------------------------------------------------------------------------
 
def _run_red_team(args: dict, ctx: dict) -> dict:
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id
    from aquin.compute.red_team import run_red_team
 
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
 
    # Product: always the full six-vector suite. Ignore partial vectors args
    # so agents/users never get "pick a vector" / unknown-vector loops.
    model = get_loaded_model()
    if model is None:
        try:
            model = load_model(model_id)
        except Exception as e:
            return {"error": str(e)}
 
    return run_red_team(model, model_id, vectors=None)
 
 
def _run_logit_lens(args: dict, ctx: dict) -> dict:
    from aquin.compute.causal_trace import run_logit_lens
    from aquin.compute.model_loader import resolve_model_id
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    prompt = args.get("prompt") or ctx.get("state", {}).get("memory", {}).get("lastPrompt") or "Hello"
    top_k = int(args.get("top_k", 5))
    try:
        results = run_logit_lens(prompt, model_id=model_id, top_k=top_k)
        return {"results": results}
    except Exception as e:
        return {"error": str(e)}
 
 
def _run_trace(args: dict, ctx: dict) -> dict:
    from aquin.compute.causal_trace import run_trace
    from aquin.compute.model_loader import resolve_model_id
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    prompt = args.get("prompt") or ctx.get("state", {}).get("memory", {}).get("lastPrompt") or "Hello"
    target = args.get("target") or ""
    try:
        results = run_trace(prompt, target, model_id=model_id)
        return {"results": results}
    except Exception as e:
        return {"error": str(e)}
 
 
def _run_prompt_attribution(args: dict, ctx: dict) -> dict:
    from aquin.compute.causal_trace import run_prompt_attribution
    from aquin.compute.model_loader import resolve_model_id
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    try:
        result = run_prompt_attribution(
            args.get("prompt", ""),
            args.get("response", ""),
            args.get("prompt_tokens", []),
            args.get("response_tokens", []),
            args.get("sig_prompt_tis", []),
            args.get("sig_response_tis", []),
            model_id=model_id,
            noise_scale=float(args.get("noise_scale", 3.0)),
            n_noise_runs=int(args.get("n_noise_runs", 5)),
        )
        return result
    except Exception as e:
        return {"error": str(e)}
 
 
def _run_full_inspection_bridge(args: dict, ctx: dict) -> dict:
    """Delegates to pipelines.run_full_inspection (same logic as inspect.py handler)."""
    from aquin.compute.pipelines import run_full_inspection as _pipeline
    from aquin.compute.model_loader import resolve_model_id
    state = ctx.get("state", {})
    prompt = args.get("prompt") or state.get("lastPrompt") or "Hello"
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    layer_raw = args.get("layer")
    layer = int(layer_raw) if layer_raw is not None else None
    result = _pipeline(prompt, model_id, ctx, layer=layer)
    if "error" in result:
        return result
    return {
        "content": {
            "prompt": result["prompt"],
            "response": result["response"],
            "model_id": result["model_id"],
            "top_features": result["top_features"],
            "sae_layer": result["sae_layer"],
        },
        "card": {
            "type": "inspectionFull",
            "data": result,
        },
    }
 
 
def _get_feature_logits_bridge(args: dict, ctx: dict) -> dict:
    from aquin.compute.feature_analysis import get_feature_logits as _get_logits
    from aquin.compute.model_loader import get_loaded_model, load_model, resolve_model_id
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    feature_idx = int(args.get("feature_idx", 0))
    layer = args.get("layer")
    model = get_loaded_model()
    if model is None:
        try:
            model = load_model(model_id)
        except Exception as e:
            return {"error": str(e)}
    try:
        from aquin.compute.feature_analysis import enrich_feature_tool_result, prompt_for_labeling
        from aquin.compute.openai_client import get_openai_client
        result = _get_logits(feature_idx, model, model_id=model_id, layer=layer, top_k=int(args.get("top_k", 10)))
        enrich_feature_tool_result(
            result,
            prompt=prompt_for_labeling(ctx, args),
            model=model,
            client=get_openai_client(ctx),
            model_id=model_id,
            layer=layer,
        )
        return {"content": result}
    except Exception as e:
        return {"error": str(e)}
 
 
def _get_feature_neighbors_bridge(args: dict, ctx: dict) -> dict:
    from aquin.compute.feature_analysis import get_feature_neighbors as _get_neighbors
    from aquin.compute.model_loader import resolve_model_id
    model_id = _get_model_id(ctx, args)
    try:
        model_id = resolve_model_id(model_id)
    except Exception as e:
        return {"error": str(e)}
    feature_idx = int(args.get("feature_idx", 0))
    layer = args.get("layer")
    try:
        from aquin.compute.feature_analysis import enrich_feature_tool_result, prompt_for_labeling
        from aquin.compute.model_loader import get_loaded_model, load_model
        from aquin.compute.openai_client import get_openai_client
        result = _get_neighbors(feature_idx, model_id=model_id, layer=layer, top_k=int(args.get("top_k", 8)))
        model = get_loaded_model()
        if model is None:
            model = load_model(model_id)
        enrich_feature_tool_result(
            result,
            prompt=prompt_for_labeling(ctx, args),
            model=model,
            client=get_openai_client(ctx),
            model_id=model_id,
            layer=layer,
            label_neighbors=True,
        )
        return {"content": result}
    except Exception as e:
        return {"error": str(e)}
 
 
# ---------------------------------------------------------------------------
# SSE training launchers — synchronous wrappers for thread-based workers
# ---------------------------------------------------------------------------
# Each worker was designed for FastAPI SSE streaming. Here we run it in a
# thread with a real asyncio loop+queue, drain all emitted payloads, then
# return the final result dict. The caller (executor.py) pushes events to
# cloud via push_events after the call returns.
 
def _drain_queue(worker_fn, req, extra_kwargs=None, on_event=None):
    """
    Run worker_fn(req, queue, loop) in a thread. Drain the queue until a
    sentinel {"__done__": True} payload arrives. Return (events, final_result).
    """
    import asyncio
    import threading
    import math
 
    def _sanitize(obj):
        if isinstance(obj, float):
            return None if (math.isnan(obj) or math.isinf(obj)) else obj
        if isinstance(obj, dict):
            return {k: _sanitize(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [_sanitize(v) for v in obj]
        return obj
 
    loop = asyncio.new_event_loop()
    queue: asyncio.Queue = asyncio.Queue()
 
    def _run():
        if extra_kwargs:
            worker_fn(req, queue, loop, **extra_kwargs)
        else:
            worker_fn(req, queue, loop)
 
    t = threading.Thread(target=_run, daemon=True)
    t.start()
 
    events = []
    final = {}
 
    async def _collect():
        nonlocal final
        while True:
            item = await queue.get()
            if item.get("__done__"):
                final = {k: v for k, v in item.items() if k != "__done__"}
                break
            events.append(_sanitize(item))
            if on_event:
                on_event(item)
 
    loop.run_until_complete(_collect())
    t.join(timeout=5)
    loop.close()
    return events, final
 
 
def _build_simulate_req(args: dict, ctx: dict):
    from aquin.compute.train_simulate import SimulateRequest
    from aquin.compute.simulate_inputs import prepare_simulation_args
 
    merged = prepare_simulation_args(args)
    sandbox = _SANDBOX_STATE
    rows = (
        merged.get("rows")
        or sandbox.get("dataset", {}).get("rows", [])
    )
 
    model_id = _get_model_id(ctx, merged)
    try:
        from aquin.compute.model_loader import resolve_model_id
        model_id = resolve_model_id(model_id)
    except ValueError:
        pass
    return SimulateRequest(
        model_id=model_id,
        rank=int(merged.get("rank", 8)),
        alpha=int(merged.get("alpha", 16)),
        lr=float(merged.get("lr", 2e-4)),
        epochs=int(merged.get("epochs", 3)),
        dropout=float(merged.get("dropout", 0.05)),
        target_modules=merged.get("targetModules", []),
        warmup_steps=int(merged.get("warmupSteps", 0)),
        grad_clip=float(merged.get("gradClip", 1.0)),
        weight_decay=float(merged.get("weightDecay", 0.01)),
        grad_accum_steps=int(merged.get("gradAccumSteps", 1)),
        optimizer=merged.get("optimizer", "adamw"),
        scheduler=merged.get("scheduler", "cosine"),
        max_seq_len=int(merged.get("maxSeqLen", 128)),
        dataset=rows or [],
        use_qlora=bool(merged.get("useQlora", False)),
        use_rlhf=bool(merged.get("use_rlhf", False)),
        rlhf_beta=float(merged.get("rlhf_beta", 0.1)),
    )
 
 
def _print_simulate_progress(ev: dict) -> None:
    """Print key simulation metrics as they arrive (CLI live output)."""
    t = ev.get("type")
    if t == "datasetQuality":
        print(
            f"[simulate] dataset: {ev.get('nSamples', '?')} samples · "
            f"diversity={float(ev.get('diversityScore', 0)):.2f} · "
            f"harmful={ev.get('harmfulCount', 0)}",
            flush=True,
        )
    elif t == "step" and ev.get("loss") is not None:
        print(f"[simulate] step {ev.get('step', '?')}  loss={float(ev['loss']):.4f}", flush=True)
    elif t == "saePrediction":
        feats = ev.get("topFeatures") or []
        strengthen = sum(1 for f in feats if f.get("direction") == "strengthen")
        suppress = sum(1 for f in feats if f.get("direction") == "suppress")
        print(
            f"[simulate] SAE prediction: {strengthen} strengthen · {suppress} suppress "
            f"({len(feats)} top features)",
            flush=True,
        )
    elif t == "saeDiff":
        print(
            f"[simulate] SAE diff: {ev.get('nChanged', '?')}/{ev.get('nFeatures', '?')} changed  "
            f"mean|Δ|={ev.get('meanAbsDelta', '?')}",
            flush=True,
        )
    elif t == "influenceScores":
        method = ev.get("method", "lissa")
        print(
            f"[simulate] influence ({method}): {ev.get('nHelpful', 0)} helpful · "
            f"{ev.get('nHarmful', 0)} harmful samples",
            flush=True,
        )
    elif t == "lossSharpness":
        lam = ev.get("maxEigenvalue")
        lam_s = f"{float(lam):.2e}" if isinstance(lam, (int, float)) else "?"
        print(
            f"[simulate] loss landscape: {ev.get('sharpnessLabel', '?')} · λ={lam_s}",
            flush=True,
        )
    elif t == "modelDiff":
        parts: list[str] = []
        for label, key in (
            ("consistency", "consistencyScore"),
            ("suppression", "suppressionScore"),
            ("robustness", "robustnessScore"),
        ):
            val = ev.get(key)
            if isinstance(val, (int, float)):
                parts.append(f"{label}={val:.2f}")
        if parts:
            print(f"[simulate] attack surface: {' · '.join(parts)}", flush=True)
    elif t == "calibration":
        print(
            f"[simulate] calibration: base_ece={ev.get('base_ece')}  "
            f"ft_ece={ev.get('ft_ece')}  "
            f"low_conf={len(ev.get('low_confidence_rows') or [])}",
            flush=True,
        )
    elif t == "signal":
        print(f"[simulate] ⚠ {ev.get('message', ev.get('signalType', 'signal'))}", flush=True)
 
 
 
 
 
 
def _run_simulation(args: dict, ctx: dict) -> dict:
    """Full training simulation. Pass --dataset /path, --algo /path, or --topic."""
    session_model_id: str | None = None
    try:
        from aquin.compute.train_simulate import _run_simulation as _worker
        from aquin.compute.simulate_inputs import prepare_simulation_args
 
        merged = prepare_simulation_args(args)
        rows = (
            merged.get("rows")
            or _SANDBOX_STATE.get("dataset", {}).get("rows", [])
        )
        topic = merged.get("topic")
        if not rows:
            if not topic:
                return {
                    "error": "Pass --dataset /path/to/file, or aquin simulate --topic for a quick probe."
                }
            rows = [
                {"instruction": f"What is {topic}?", "response": f"A question about {topic}."},
                {"instruction": f"Explain {topic} simply.", "response": f"{topic} is a topic."},
            ]
        session_model_id = _get_model_id(ctx, merged)
        req = _build_simulate_req({**merged, "rows": rows}, ctx)
        session_model_id = req.model_id
        meta = {"model_id": req.model_id, "topic": topic, "kind": "simulate"}
        if merged.get("dataset_path"):
            meta["dataset_path"] = merged["dataset_path"]
        if merged.get("algo_path"):
            meta["algo_path"] = merged["algo_path"]
        events, final = _drain_queue(_worker, req, on_event=_print_simulate_progress)
        return _wrap_simulation_result(events, final, meta)
    except (FileNotFoundError, ValueError) as e:
        return {"error": str(e)}
    except Exception as e:
        return {"error": f"run_simulation failed: {e}"}
    finally:
        if session_model_id:
            from aquin.compute.vram_guard import restore_session_model
 
            restore_session_model(session_model_id)
 
 
# ---------------------------------------------------------------------------
# TOOL_ROUTES
# ---------------------------------------------------------------------------
# run_full_inspection, get_feature_logits, get_feature_neighbors
# are also registered directly in inspect.py — bridge entries here are the
# canonical compute-only path used by bridge.call() and raw CLI (Step 57).
 
TOOL_ROUTES: dict[str, Callable[[dict, dict], dict]] = {
    # Full inspection pipeline
    "run_full_inspection": _run_full_inspection_bridge,
    # Direct compute routes (server.py route index)
    "run_logit_lens": _run_logit_lens,
    "run_trace": _run_trace,
    "run_prompt_attribution": _run_prompt_attribution,
    # Feature drill-down
    "get_feature_logits": _get_feature_logits_bridge,
    "get_feature_neighbors": _get_feature_neighbors_bridge,
    # Feature analysis
    "run_steer_and_show": _run_steer_and_show,
    "extract_steer_vector": _extract_steer_vector,
    "run_multi_steer": _run_multi_steer,
    # Evals
    "run_consistency_eval": _run_consistency_eval,
    "run_suppression_eval": _run_suppression_eval,
    "run_boundary_eval": _run_boundary_eval,
    "run_audit": _run_audit,
    # Benchmarks / interp
    "run_benchmarks_on_top_feature": _run_benchmarks_on_top_feature,
    "run_find_feature": _run_find_feature,
    # Weights
    "check_weights": _check_weights,
    # UMAP
    "ensure_umap_loaded": _ensure_umap_loaded,
    # Model internals
    "run_layer_analysis": _run_layer_analysis,
    "run_sae_stats": _run_sae_stats,
    "run_confidence_analysis": _run_confidence_analysis,
    "run_weight_diff": _run_weight_diff,
    "run_merge_analysis": _run_merge_analysis,
    "run_trajectory_analysis": _run_trajectory_analysis,
    "run_residual_drift": _run_residual_drift,
    "run_perturbation_sensitivity": _run_perturbation_sensitivity,
    "run_attention_routing": _run_attention_routing,
    # Dataset / config state
    "dataset_generate": _dataset_generate,
    # SSE streaming simulation
    "run_simulation": _run_simulation,
    # Simulation run list
    "list_simulation_runs": _list_simulation_runs,
    "load_simulation_run": _load_simulation_run,
    "compare_simulations": _compare_simulations,
    # UI-only tools
    "run_red_team": _run_red_team,
    # Card builders
    "run_custom_eval": _run_custom_eval,
}
 
 
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
 
def call(tool_name: str, args: dict, ctx: dict) -> dict:
    """
    Dispatch a tool call to its compute implementation.
    Returns a result dict. Never raises — errors are returned as {"error": ...}.
    """
    handler = TOOL_ROUTES.get(tool_name)
    if handler is None:
        return {"error": f"Unknown tool: {tool_name}"}
    try:
        return handler(args, ctx)
    except Exception as exc:
        from aquin.user_errors import friendly_message, log_debug
        log_debug(tool_name, exc)
        return {"error": friendly_message(exc)}