distinct_server/ui_state.py

"""Pure helpers for the per-browser state owned by Gradio."""

from __future__ import annotations

import copy
import secrets
import time
from typing import Any, Dict, List, Mapping, Tuple


def _id(prefix: str) -> str:
    return f"{prefix}-{secrets.token_hex(10)}"


#: The title a conversation carries until its first request names it.
#: It used to be "New conversation", the same words as the button directly
#: above the list, so an untitled conversation and the control that creates
#: one sat adjacent and read identically.
UNTITLED = "Untitled"


def _new_conversation() -> Dict[str, Any]:
    return {
        "title": UNTITLED,
        "job_ids": [],
        "prompts": [],
        # The browser owns the transcript.  The server erases a prompt the
        # moment its run reaches a terminal state, and drops the answer once
        # this session has read it, so anything still on screen afterwards is
        # held here and nowhere else.  Keyed by job id.
        "answers": {},
    }


def new_session() -> Dict[str, Any]:
    conversation_id = _id("conversation")
    return {
        "session_id": _id("session"),
        "active_conversation_id": conversation_id,
        "conversations": {conversation_id: _new_conversation()},
        "created_at": time.time(),
    }


#: The longest identifier the control plane will accept. Anything longer is
#: refused there with a `ValidationError`, and a browser holding one would
#: turn every page load into a five-hundred.
_MAX_ID = 128


def _identifier(value: Any) -> str:
    """``value`` if it can be a session or conversation id, else empty."""

    if not isinstance(value, str):
        return ""
    text = value.strip()
    if not text or len(text) > _MAX_ID or any(c in text for c in "\r\n"):
        return ""
    return text


def _string_list(value: Any, *, limit: int) -> List[str]:
    if not isinstance(value, (list, tuple)):
        return []
    return [str(item) for item in value if isinstance(item, str)][:limit]


def _mapping_list(value: Any, *, limit: int) -> List[Dict[str, Any]]:
    if not isinstance(value, (list, tuple)):
        return []
    return [dict(item) for item in value if isinstance(item, Mapping)][:limit]


#: How much of one browser's history is worth carrying forward. A transcript
#: is per-browser and unbounded otherwise, and the whole of it is re-rendered
#: on a two-second timer.
_MAX_CONVERSATIONS = 128
_MAX_TURNS = 512


def _number(value: Any) -> float | None:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        return None
    return float(value)


def _adopt_answer(value: Any) -> Dict[str, Any] | None:
    """One kept answer, rebuilt to the shape `DistinctUI._answer_payload` writes.

    Every field is here, and each is the type the renderer indexes into: the
    receipt walks `trace` and `steps` as mappings, the transcript walks `made`
    as mappings, and the estimator reads `duration_seconds` as a number. An
    older build that spelled any of them differently turned the transcript
    into a `TypeError` on every render — which is to say on every page load
    and every timer tick. Anything not listed here is dropped, because a key
    nothing reads is a key nothing can be broken by.
    """

    if not isinstance(value, Mapping):
        return None
    answer: Dict[str, Any] = {
        "output": str(value.get("output", "")),
        "agent": str(value.get("agent", "")),
        "model": str(value.get("model", "")),
        "mode": str(value.get("mode", "")),
        "tools": [item for item in _mapping_list(value.get("tools"), limit=64) if "id" in item],
        "steps": _mapping_list(value.get("steps"), limit=64),
        "trace": _mapping_list(value.get("trace"), limit=64),
        "made": _mapping_list(value.get("made"), limit=32),
        "artifacts": _string_list(value.get("artifacts"), limit=64),
        "energy": dict(value["energy"]) if isinstance(value.get("energy"), Mapping) else {},
        "guard": dict(value["guard"]) if isinstance(value.get("guard"), Mapping) else {},
        "tool_calls": int(_number(value.get("tool_calls")) or 0),
    }
    duration = _number(value.get("duration_seconds"))
    if duration is not None:
        answer["duration_seconds"] = duration
    return answer


def adopt_saved_state(saved: Any) -> Dict[str, Any] | None:
    """Today's shape, rebuilt from whatever an older build left in the browser.

    THE BROWSER'S COPY IS INPUT, NOT STATE.

    It is JSON written by some earlier version of this file, kept in
    localStorage across every deploy since, and editable by anyone with a
    developer console. `load_session` used to check four facts about the top
    level of it — an id, a conversations mapping, an active id inside that
    mapping — and then hand the rest straight to the renderer. Everything one
    level down was trusted: a conversation that was a string, a conversation
    with no `title`, an answer that was a string, an answer whose `made` was
    not a list. Each of those raised out of `load_session`, which runs on
    every single page load, so the app opened on a red error toast and stayed
    that way until the person cleared their site data — and they had no way
    to know that was the fix.

    So the saved copy is rebuilt here field by field instead of inspected.
    What cannot be read is dropped rather than carried, and a transcript with
    nothing readable left in it returns None, which the caller reads as "start
    fresh". Dropping an unreadable turn loses something; refusing to load the
    page loses everything.
    """

    if not isinstance(saved, Mapping):
        return None
    conversations: Dict[str, Any] = {}
    raw = saved.get("conversations")
    if not isinstance(raw, Mapping):
        return None
    for key, value in list(raw.items())[:_MAX_CONVERSATIONS]:
        conversation_id = _identifier(key)
        if not conversation_id or not isinstance(value, Mapping):
            continue
        job_ids = _string_list(value.get("job_ids"), limit=_MAX_TURNS)
        prompts = _string_list(value.get("prompts"), limit=_MAX_TURNS)
        # Padded rather than trimmed to the shorter of the two. A prompt list
        # that fell behind its job list is an older bug's leftovers, and the
        # turn is still worth showing with an empty question attached.
        prompts += [""] * max(0, len(job_ids) - len(prompts))
        answers: Dict[str, Any] = {}
        raw_answers = value.get("answers")
        if isinstance(raw_answers, Mapping):
            for job_id, answer in raw_answers.items():
                if not isinstance(job_id, str) or job_id not in job_ids:
                    continue
                adopted = _adopt_answer(answer)
                if adopted is not None:
                    answers[job_id] = adopted
        title = value.get("title")
        conversations[conversation_id] = {
            "title": str(title) if isinstance(title, str) and title.strip() else UNTITLED,
            "job_ids": job_ids,
            "prompts": prompts[: len(job_ids)],
            "answers": answers,
        }
    if not conversations:
        return None
    active = _identifier(saved.get("active_conversation_id"))
    if active not in conversations:
        active = next(iter(conversations))
    state: Dict[str, Any] = {
        # A session id this server would refuse is worse than no session id:
        # it fails at the control plane, inside the handler, rather than here.
        "session_id": _identifier(saved.get("session_id")) or _id("session"),
        "active_conversation_id": active,
        "conversations": conversations,
        "created_at": saved.get("created_at")
        if isinstance(saved.get("created_at"), (int, float))
        else time.time(),
    }
    defaults = _string_list(saved.get("library_defaults"), limit=256)
    if defaults:
        state["library_defaults"] = defaults
    # The pairing code is browser-owned and worth keeping: dropping it makes
    # the next open of the setup panel mint another, and a person only gets
    # eight live codes before the control plane refuses.
    code = saved.get("pairing_code")
    expires = saved.get("pairing_expires")
    if isinstance(code, str) and code and isinstance(expires, (int, float)):
        state["pairing_code"] = code
        state["pairing_expires"] = float(expires)
    return state


def _validated_copy(state: Mapping[str, Any]) -> Dict[str, Any]:
    value = copy.deepcopy(dict(state))
    if not isinstance(value.get("session_id"), str):
        raise ValueError("session state is missing its id")
    conversations = value.get("conversations")
    if not isinstance(conversations, dict) or not conversations:
        raise ValueError("session state has no conversations")
    if value.get("active_conversation_id") not in conversations:
        raise ValueError("active conversation is unavailable")
    return value


def add_conversation(state: Mapping[str, Any]) -> Tuple[Dict[str, Any], str]:
    value = _validated_copy(state)
    conversation_id = _id("conversation")
    value["conversations"][conversation_id] = _new_conversation()
    value["active_conversation_id"] = conversation_id
    return value, conversation_id


def select_conversation(state: Mapping[str, Any], conversation_id: str) -> Dict[str, Any]:
    value = _validated_copy(state)
    if conversation_id not in value["conversations"]:
        raise ValueError("conversation does not belong to this session")
    value["active_conversation_id"] = conversation_id
    return value


def remove_conversation(state: Mapping[str, Any], conversation_id: str) -> Dict[str, Any]:
    """Drop a conversation and everything it held, choosing the next active one.

    THE LAST ONE IS EMPTIED RATHER THAN REMOVED. A session with no conversation
    has no valid ``active_conversation_id``, and every renderer downstream reads
    that key; deleting the only one would leave the state in a shape the rest of
    this module is entitled to assume cannot happen. Emptying it gives the
    person what they asked for — the transcript is gone — without inventing a
    state nothing else can render.

    The jobs are dropped from this state, which is the browser's copy and the
    only durable one: the server erases a prompt when its run ends and the
    answer once it has been collected. The caller is responsible for telling
    the control plane, which owns anything still in flight.
    """

    value = _validated_copy(state)
    if conversation_id not in value["conversations"]:
        raise ValueError("conversation does not belong to this session")

    if len(value["conversations"]) == 1:
        fresh = _new_conversation()
        value["conversations"] = {conversation_id: fresh}
        value["active_conversation_id"] = conversation_id
        return value

    order = list(value["conversations"])
    position = order.index(conversation_id)
    del value["conversations"][conversation_id]
    if value["active_conversation_id"] == conversation_id:
        # The neighbour, preferring the one above, which is where the eye
        # already is after a row disappears.
        remaining = list(value["conversations"])
        value["active_conversation_id"] = remaining[max(0, position - 1)]
    return value


def append_job(
    state: Mapping[str, Any],
    *,
    job_id: str,
    prompt: str,
    conversation_id: str | None = None,
) -> Dict[str, Any]:
    value = _validated_copy(state)
    selected = conversation_id or value["active_conversation_id"]
    if selected not in value["conversations"]:
        raise ValueError("conversation does not belong to this session")
    conversation = value["conversations"][selected]
    if job_id in conversation["job_ids"]:
        raise ValueError("job is already present in the conversation")
    conversation["job_ids"].append(job_id)
    conversation["prompts"].append(prompt)
    if conversation["title"] == UNTITLED:
        title = " ".join(prompt.split())
        conversation["title"] = title[:48] + ("…" if len(title) > 48 else "")
    return value


def record_answer(
    state: Mapping[str, Any],
    *,
    job_id: str,
    answer: Mapping[str, Any],
    conversation_id: str | None = None,
) -> Dict[str, Any]:
    """Store a delivered answer in browser-owned state.

    This is the counterpart to ``ControlPlane.claim_result``: the server hands
    the answer over exactly once and then erases it, so if this does not
    persist it, the reply is gone.
    """

    value = _validated_copy(state)
    selected = conversation_id or value["active_conversation_id"]
    if selected not in value["conversations"]:
        raise ValueError("conversation does not belong to this session")
    conversation = value["conversations"][selected]
    if job_id not in conversation["job_ids"]:
        raise ValueError("job does not belong to this conversation")
    conversation.setdefault("answers", {})[job_id] = dict(answer)
    return value


def conversation_turns(state: Mapping[str, Any]) -> List[Dict[str, Any]]:
    """The browser's own transcript: one entry per submitted run, in order."""

    value = _validated_copy(state)
    conversation = value["conversations"][value["active_conversation_id"]]
    answers = conversation.get("answers") or {}
    prompts = conversation.get("prompts") or []
    turns: List[Dict[str, Any]] = []
    for index, job_id in enumerate(conversation.get("job_ids") or []):
        turns.append(
            {
                "job_id": job_id,
                "prompt": prompts[index] if index < len(prompts) else "",
                "answer": answers.get(job_id),
            }
        )
    return turns


def active_conversation(state: Mapping[str, Any]) -> Dict[str, Any]:
    value = _validated_copy(state)
    return value["conversations"][value["active_conversation_id"]]


#: How each conversation state is marked in the rail. The glyph is real text
#: in the option's own label, so the state is carried by a shape and a word
#: and never by the tone alone: the tones are pastel by design and pastels are
#: the first thing to go on a poor screen or to a reader who cannot separate
#: them. The word is omitted for the two resting states, where it would appear
#: on almost every row and say nothing.
CONVERSATION_STATES: Dict[str, tuple[str, str]] = {
    "idle": ("◇", ""),
    "waiting": ("○", "waiting"),
    "generating": ("◐", "working"),
    "completed": ("●", ""),
    "failed": ("■", "failed"),
}


def conversation_choices(
    state: Mapping[str, Any], states: Mapping[str, str] | None = None
) -> List[tuple[str, str]]:
    """The rail's options, each marked with the state of its conversation.

    ``states`` maps a conversation id to one of :data:`CONVERSATION_STATES`.
    It is optional because several call sites rebuild the list after an event
    that cannot have changed any run's state, and an unmarked option is a
    correct option rather than a broken one.
    """

    value = _validated_copy(state)
    marks = states or {}
    options: List[tuple[str, str]] = []
    for conversation_id, conversation in value["conversations"].items():
        glyph, word = CONVERSATION_STATES.get(
            str(marks.get(conversation_id, "idle")), CONVERSATION_STATES["idle"]
        )
        title = str(conversation["title"])
        suffix = f" · {word}" if word else ""
        options.append((f"{glyph} {title}{suffix}", conversation_id))
    return options