atlas-iac/services/hermes/plugins/auto-router/hux_mode.py

"""Adopt the conversation's selected HUX friendly mode at the route boundary.

HUX-06: a mode selected in the WebUI must shape the next real Switchyard
request without ever naming a provider for automatic pools.  This adapter
reads the loopback HUX service with worker trust and maps the stored
provider-neutral mode contract onto the public route families the
auto-router already owns.  Every failure is an explicit ``None`` so the
existing routing behaviour is preserved bit-for-bit when HUX is absent,
misconfigured, unreachable, or simply has no selection for the session.
"""

from __future__ import annotations

import hashlib
import hmac
import json
import os
import re
import stat
import time
import urllib.request
from pathlib import Path
from typing import Any

DOMAIN = "hux.context.id.v1"
KEY_BYTES = 32
CACHE_SECONDS = 5.0
TIMEOUT_SECONDS = 2.0
MAX_RESPONSE_BYTES = 256 * 1024
RAW_ID = re.compile(r"^[A-Za-z0-9._:@+-]{1,200}$")
SLOT = re.compile(r"^slot-[0-9]{1,3}$")
SUBJECT = re.compile(r"^usr_[0-9a-f]{64}$")
LOCAL_ROUTE = "atlas/manual/local/qwen-14b"
MANUAL_OVERRIDE = re.compile(r"^atlas/manual/(codex|claude|local)/[a-z0-9][a-z0-9/-]{0,100}$")
# Friendly modes map onto provider-neutral automatic pools; only the explicit
# local-only private mode may pin a route, and it pins the local model.
MODE_ROUTES = {
    "fast": "atlas/auto/fast",
    "thoughtful": "atlas/auto/deep",
    "research": "atlas/auto/deep",
    "create": "atlas/auto/balanced",
    "private": LOCAL_ROUTE,
}

_cache: dict[str, tuple[float, tuple[str, str] | None]] = {}


def _read_secret(path: str, size: int | None = None, modes: frozenset[int] = frozenset({0o400, 0o600})) -> bytes | None:
    """Read one owner-only regular file without following symlinks."""
    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_CLOEXEC", 0)
    descriptor = -1
    try:
        descriptor = os.open(Path(path), flags)
        info = os.fstat(descriptor)
        if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1:
            return None
        if stat.S_IMODE(info.st_mode) not in modes and info.st_uid == os.geteuid():
            return None
        value = os.read(descriptor, 4096)
    except OSError:
        return None
    finally:
        if descriptor >= 0:
            os.close(descriptor)
    if size is not None and len(value) != size:
        return None
    return value


def _derive(key: bytes, prefix: str, purpose: str, slot: str, subject: str, raw: str) -> str | None:
    """Replicate the public ``hux.context.id.v1`` derivation exactly."""
    if not RAW_ID.fullmatch(raw or ""):
        return None
    message = "\0".join((DOMAIN, purpose, slot, subject, raw)).encode("utf-8")
    return f"{prefix}_{hmac.new(key, message, hashlib.sha256).hexdigest()[:32]}"


def _fetch(url: str, headers: dict[str, str]) -> tuple[int, dict[str, Any]] | None:
    """One bounded loopback GET; anything unexpected is None."""
    request = urllib.request.Request(url, headers=headers, method="GET")
    try:
        with urllib.request.urlopen(request, timeout=TIMEOUT_SECONDS) as response:  # noqa: S310 - literal loopback URL built below
            body = response.read(MAX_RESPONSE_BYTES + 1)
            if len(body) > MAX_RESPONSE_BYTES:
                return None
            return response.status, json.loads(body)
    except Exception:
        return None


def _selection_uncached(environ: dict[str, str], raw_session_id: str, fetch: Any) -> tuple[str, str] | None:
    """Resolve ``(route, effort)`` for the session's stored mode, or None."""
    base = environ.get("HUX_BASE_URL", "")
    slot = environ.get("HUX_TENANT_SLOT", "")
    subject_file = environ.get("HUX_SUBJECT_FILE", "")
    key_file = environ.get("HUX_WORKER_KEY_FILE", "")
    context_key_file = environ.get("HUX_CONTEXT_KEY_FILE", "")
    project_source = environ.get("HUX_PROJECT_SOURCE", "profile:default")
    if not (base.startswith("http://127.0.0.1:") and SLOT.fullmatch(slot)
            and subject_file and key_file and context_key_file):
        return None
    context_key = _read_secret(context_key_file, KEY_BYTES)
    subject_raw = _read_secret(subject_file)
    worker_key_raw = _read_secret(key_file)
    if context_key is None or subject_raw is None or worker_key_raw is None:
        return None
    subject = subject_raw.decode("ascii", errors="replace").strip()
    worker_key = worker_key_raw.decode("ascii", errors="replace").strip()
    if not SUBJECT.fullmatch(subject) or not worker_key:
        return None
    conversation = _derive(context_key, "conv", "conversation", slot, subject, raw_session_id)
    project = _derive(context_key, "prj", "project", slot, subject, project_source)
    if not conversation or not project:
        return None
    result = fetch(
        f"{base}/hux/v1/projects/{project}/conversations/{conversation}/mode",
        {
            "X-Hermes-Tenant-Identity": slot,
            "X-Hux-Subject": subject,
            "X-Hux-Surface": "worker",
            "X-Hux-Trust": "worker",
            "X-Hux-Relay-Key": worker_key,
        },
    )
    if result is None or result[0] != 200 or not isinstance(result[1], dict):
        return None
    mode = result[1].get("mode")
    if not isinstance(mode, dict):
        return None
    name = mode.get("mode")
    switchyard = mode.get("switchyard")
    route = MODE_ROUTES.get(name) if isinstance(name, str) else None
    if route is None or not isinstance(switchyard, dict):
        return None
    override = switchyard.get("override_route_id")
    if isinstance(override, str) and MANUAL_OVERRIDE.fullmatch(override):
        if name == "private" and not override.startswith("atlas/manual/local/"):
            return None
        route = override
    return route, ""


def selection(raw_session_id: str, environ: dict[str, str] | None = None, fetch: Any = _fetch) -> tuple[str, str] | None:
    """Cached, fail-open mode adoption for one persisted Hermes session."""
    if not isinstance(raw_session_id, str) or not raw_session_id:
        return None
    now = time.monotonic()
    cached = _cache.get(raw_session_id)
    if cached is not None and cached[0] > now:
        return cached[1]
    try:
        value = _selection_uncached(dict(os.environ if environ is None else environ), raw_session_id, fetch)
    except Exception:
        value = None
    if len(_cache) > 512:
        _cache.clear()
    _cache[raw_session_id] = (now + CACHE_SECONDS, value)
    return value