api/state.py

"""Server-side state for the API layer.

Jacobina keeps miru-tracer's deliberate single-user design (one loaded model
per process, see ``ModelManager``): the Logging-mode tracer, the Lens
analysis, and the active-interventions registry are process-global. Interactive
mode stays session-id based via the core ``SessionManager``.
"""

from __future__ import annotations

import threading
from dataclasses import dataclass, field
from typing import Any

from miru_tracer.core.interventions import Intervention


@dataclass
class LoggingSession:
    """The Logging-mode run: a live tracer plus the inputs it started from."""

    lock: threading.Lock = field(default_factory=threading.Lock)
    tracer: Any = None
    # (mode, prompt, chat_json, raw_text, thinking, think_prefill)
    originals: tuple | None = None


@dataclass
class LensSession:
    """The Lens tab's analysis of the last generated sequence.

    ``analysis`` holds input_ids (tensor), model_name, iset, n_layers,
    prompt_len, position_texts, and lazily-cached activations. ``bundle``
    holds the last computed {slice, rows, intervened} the result views render
    from.
    """

    lock: threading.Lock = field(default_factory=threading.Lock)
    analysis: dict | None = None
    bundle: dict | None = None


_logging_session = LoggingSession()
_lens_session = LensSession()

# Active interventions registry: the Lens view edits it, Interactive mode's
# "apply Lens interventions" reads it. Rows are {"enabled": bool,
# "intervention": Intervention}; a lock guards concurrent API calls.
_interventions_lock = threading.Lock()
_intervention_rows: list[dict] = []


def get_logging_session() -> LoggingSession:
    return _logging_session


def get_lens_session() -> LensSession:
    return _lens_session


def get_intervention_rows() -> list[dict]:
    with _interventions_lock:
        return list(_intervention_rows)


def set_intervention_rows(rows: list[dict]) -> None:
    with _interventions_lock:
        _intervention_rows.clear()
        _intervention_rows.extend(rows)


def get_active_interventions() -> list[Intervention]:
    with _interventions_lock:
        return [
            row["intervention"]
            for row in _intervention_rows
            if row.get("enabled", True)
        ]