atlas-iac/dockerfiles/hermes-webui-hux-context.py

"""Build browser HUX scope only from an authenticated server session."""

from __future__ import annotations

import hashlib
import hmac
import os
from pathlib import Path
import re
import secrets
import stat
import threading
from typing import Any


_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}$")
_KEY_BYTES = 32
_KEY_NAME = ".hux-context-key"
_DEFAULT_PROJECT_SOURCE = "profile:default"
_ID_PURPOSES = {
    "ses": "session",
    "conv": "conversation",
    "prj": "project",
    "run": "run",
    "msg": "message",
}
_KEY_LOCK = threading.Lock()
_KEY_CACHE: tuple[Path, bytes] | None = None


class HuxContextUnavailable(Exception):
    """Signal that trusted context cannot be produced for this response."""


def _safe_value(value: Any, fallback: str | None = None) -> str:
    candidate = value if isinstance(value, str) else ""
    if _RAW_ID.fullmatch(candidate):
        return candidate
    if fallback is not None:
        return fallback
    raise HuxContextUnavailable("WebUI session identity is unavailable")


def _key_path() -> Path:
    from api.config import STATE_DIR

    configured = os.environ.get("HUX_CONTEXT_KEY_FILE", "").strip()
    path = Path(configured) if configured else Path(STATE_DIR) / _KEY_NAME
    if not path.is_absolute() or ".." in path.parts:
        raise HuxContextUnavailable("HUX context key path is invalid")
    root = path.parent
    try:
        info = root.stat(follow_symlinks=False)
    except OSError as exc:
        raise HuxContextUnavailable("HUX context state is unavailable") from exc
    if not stat.S_ISDIR(info.st_mode) or info.st_uid != os.geteuid():
        raise HuxContextUnavailable("HUX context state is unavailable")
    return path


def _read_key(path: Path) -> bytes:
    flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0)
    try:
        descriptor = os.open(path, flags)
        try:
            info = os.fstat(descriptor)
            if (
                not stat.S_ISREG(info.st_mode)
                or info.st_uid != os.geteuid()
                or stat.S_IMODE(info.st_mode) != 0o600
                or info.st_nlink != 1
                or info.st_size != _KEY_BYTES
            ):
                raise HuxContextUnavailable("HUX context key is unsafe")
            value = os.read(descriptor, _KEY_BYTES + 1)
        finally:
            os.close(descriptor)
    except HuxContextUnavailable:
        raise
    except OSError as exc:
        raise HuxContextUnavailable("HUX context key is unavailable") from exc
    if len(value) != _KEY_BYTES:
        raise HuxContextUnavailable("HUX context key is invalid")
    return value


def _create_key(path: Path) -> None:
    flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_CLOEXEC", 0)
    flags |= getattr(os, "O_NOFOLLOW", 0)
    value = secrets.token_bytes(_KEY_BYTES)
    try:
        descriptor = os.open(path, flags, 0o600)
    except FileExistsError:
        return
    except OSError as exc:
        raise HuxContextUnavailable("HUX context key cannot be created") from exc
    try:
        written = os.write(descriptor, value)
        if written != len(value):
            raise HuxContextUnavailable("HUX context key write was incomplete")
        os.fsync(descriptor)
    finally:
        os.close(descriptor)


def _context_key() -> bytes:
    global _KEY_CACHE
    path = _key_path()
    with _KEY_LOCK:
        if _KEY_CACHE and _KEY_CACHE[0] == path:
            return _KEY_CACHE[1]
        _create_key(path)
        value = _read_key(path)
        _KEY_CACHE = (path, value)
        return value


def _project_source() -> str:
    """Return the bounded server-owned source for the default HUX project."""
    configured = os.environ.get("HUX_PROJECT_SOURCE", _DEFAULT_PROJECT_SOURCE).strip()
    return _safe_value(configured)


def _authoritative_scalar(value: Any) -> str | None:
    """Return one bounded scalar already present in a trusted server payload."""
    if isinstance(value, bool):
        return None
    candidate = str(value) if isinstance(value, (str, int)) else ""
    return candidate if _RAW_ID.fullmatch(candidate) else None


def _latest_message_source(session: dict[str, Any]) -> str | None:
    """Find the newest server-stamped message identifier without inventing one."""
    messages = session.get("messages")
    if not isinstance(messages, list):
        return None
    for message in reversed(messages):
        if not isinstance(message, dict):
            continue
        for field in ("id", "message_id"):
            candidate = _authoritative_scalar(message.get(field))
            if candidate is not None:
                return candidate
    return None


def derive_hux_id(key: bytes, prefix: str, purpose: str, slot: str, subject: str, raw: str) -> str:
    """Derive one cross-surface opaque ID from the shared server key."""
    if (
        not isinstance(key, bytes)
        or len(key) != _KEY_BYTES
        or _ID_PURPOSES.get(prefix) != purpose
        or not _SLOT.fullmatch(slot)
        or not _SUBJECT.fullmatch(subject)
        or not _RAW_ID.fullmatch(raw)
    ):
        raise HuxContextUnavailable("HUX context key is invalid")
    message = "\0".join(("hux.context.id.v1", purpose, slot, subject, raw))
    digest = hmac.new(key, message.encode("utf-8"), hashlib.sha256).hexdigest()[:32]
    return f"{prefix}_{digest}"


def build_hux_context(handler: Any, session: dict[str, Any]) -> dict[str, Any]:
    """Return the exact context for one already-authorized server session."""
    from api.hux_bff import _trusted_identity

    if not isinstance(session, dict):
        raise HuxContextUnavailable("WebUI session payload is unavailable")
    raw_session = _safe_value(session.get("session_id"))
    slot, subject, surface = _trusted_identity(handler)
    raw_project = _project_source()
    key = _context_key()
    context = {
        "schema": "hux.webui_context.v1",
        "webui_session_id": raw_session,
        "session_id": derive_hux_id(key, "ses", "session", slot, subject, raw_session),
        "conversation_id": derive_hux_id(
            key, "conv", "conversation", slot, subject, raw_session
        ),
        "project_id": derive_hux_id(key, "prj", "project", slot, subject, raw_project),
        "project_source": raw_project,
        "identity": {
            "tenant_slot": slot,
            "subject": subject,
            "surface": surface,
            "trust": "relay",
        },
    }
    # These optional values exist only when the upstream server has already
    # minted them. A stream id is deliberately never treated as an agent run.
    optional_sources = {
        "run_id": ("run", "run", _authoritative_scalar(session.get("turn_id"))),
        "message_id": ("msg", "message", _latest_message_source(session)),
        "branch_point_message_id": (
            "msg",
            "message",
            _authoritative_scalar(session.get("branch_point_message_id")),
        ),
    }
    for field, (prefix, purpose, source) in optional_sources.items():
        if source is not None:
            context[field] = derive_hux_id(key, prefix, purpose, slot, subject, source)
    notebook = session.get("notebook_id")
    if isinstance(notebook, str) and re.fullmatch(r"nb_[A-Za-z0-9._-]{4,80}", notebook):
        context["notebook_id"] = notebook
    return context


def attach_hux_context(handler: Any, session: Any) -> Any:
    """Attach trusted context or return a clean, feature-disabled payload."""
    if not isinstance(session, dict):
        return session
    clean = dict(session)
    clean.pop("hux_context", None)
    try:
        clean["hux_context"] = build_hux_context(handler, clean)
    except Exception:
        # Chat remains available when auth, state, or the optional HUX path fails.
        clean.pop("hux_context", None)
    return clean