LLeMbas/src/lembas/services/agent/draft.py

"""A chat that does not exist yet, so its panels can.

Chats are created lazily -- there is no endpoint that makes an empty one, and
the row appears together with its first message. That is a rule worth keeping:
an opened-and-abandoned composer should leave nothing behind. But it also meant
the terminal and the canvas were unavailable on the one screen where you are
deciding *which machine to work on*, which is exactly when you want to look
around it first.

A draft is the smallest thing that fixes that: an id, and the three facts the
panels need behind it. It is not a chat and never becomes one -- when the first
prompt is sent, a real chat is created and the draft's shell and tabs are
**adopted** into it, which is a re-key and a copy rather than a promotion.

The id is derived from (owner, connection, directory) rather than invented, so
that returning to the same new-chat screen finds the same shell and the same
tabs instead of quietly starting a second one. It is a hash so that neither the
directory nor the owner is legible in a URL.
"""

from __future__ import annotations

import hashlib
import time
from dataclasses import dataclass, field
from typing import Any

# How long a draft survives without being touched. Generous, because it is
# holding somebody's open files while they decide what to do; bounded, because
# nothing else will ever clean it up -- an abandoned new-chat screen leaves no
# row to cascade from and no chat to delete.
IDLE_TIMEOUT = 3600.0

# The prefix a draft id carries. It has to be distinguishable from a chat id at
# a glance and by code: `Chat.id` is 32 hex characters from `new_id`, so
# nothing here can collide with one by accident.
PREFIX = "draft_"


@dataclass
class Draft:
    """What a draft knows, which is only what the panels ask for."""

    id: str
    owner_id: str
    profile_id: str
    project_dir: str
    # The canvas's tab strip, in the shape `Chat.canvas_json` holds. In memory
    # rather than on a row for the obvious reason, and carried onto the chat at
    # adoption.
    canvas_json: dict = field(default_factory=dict)
    touched_at: float = field(default_factory=time.monotonic)


_DRAFTS: dict[str, Draft] = {}


def is_draft(chat_id: str) -> bool:
    return bool(chat_id) and chat_id.startswith(PREFIX)


def key_for(owner_id: str, profile_id: str, project_dir: str) -> str:
    """The id for one (owner, connection, directory), stably.

    Derived rather than random so that reopening the new-chat screen on the same
    target finds the shell that is already running there. The owner is in the
    hash so that two people pointed at the same directory of the same connection
    do not share a draft -- they would share a *shell*, and the terminal's own
    "one chat, one shell" rule is scoped to a person's chats.
    """
    material = "\0".join((owner_id, profile_id, project_dir or ""))
    digest = hashlib.sha256(material.encode("utf-8")).hexdigest()
    return f"{PREFIX}{digest[:24]}"


def remember(owner_id: str, profile_id: str, project_dir: str) -> Draft:
    """The draft for this target, created if this is the first time."""
    _sweep()
    key = key_for(owner_id, profile_id, project_dir)
    draft = _DRAFTS.get(key)
    if draft is None:
        draft = Draft(
            id=key, owner_id=owner_id, profile_id=profile_id, project_dir=project_dir or ""
        )
        _DRAFTS[key] = draft
    draft.touched_at = time.monotonic()
    return draft


def get(draft_id: str, owner_id: str) -> Draft | None:
    """One draft, if it is this person's.

    The id is a hash of the owner, so a draft belonging to somebody else cannot
    be guessed -- but it is checked rather than assumed, because "unguessable"
    is not an authorisation and the next caller might build the id differently.
    """
    draft = _DRAFTS.get(draft_id or "")
    if draft is None or draft.owner_id != owner_id:
        return None
    draft.touched_at = time.monotonic()
    return draft


def forget(draft_id: str) -> None:
    _DRAFTS.pop(draft_id or "", None)


def clear() -> None:
    _DRAFTS.clear()


def as_chat(draft: Draft) -> Any:
    """A `Chat` the panels can use, constructed and never saved.

    This is the whole trick, and it is worth being precise about why it is safe.
    `canvas.agent_ready`, `canvas._executor`, `_load_agent`/`_save_agent` and
    `agent_session.resolve` read exactly four things off a chat -- `user_id`,
    `kind`, `ssh_profile_id` and `project_dir` -- and none of them passes the
    chat to a query or writes it back. So a transient row satisfies every one of
    them unchanged, and no code that already works has to learn what a draft is.

    `id` and `canvas_json` are set explicitly: both are *column* defaults, which
    SQLAlchemy applies at flush, and this row is never flushed. An unset `id` is
    not a cosmetic problem -- see `SOURCES_NEEDING_A_CHAT`.
    """
    from lembas.db.models import KIND_AGENT, Chat

    return Chat(
        id=draft.id,
        user_id=draft.owner_id,
        kind=KIND_AGENT,
        ssh_profile_id=draft.profile_id,
        project_dir=draft.project_dir,
        canvas_json=dict(draft.canvas_json or {}),
        agent_mode="",
        scope_json={},
    )


# Canvas sources a draft may not open, refused by name.
#
# `scratch` needs a row: `scratch_service.for_chat` would write a `ScratchDoc`
# keyed on a chat that does not exist, which is the lazy-creation rule broken
# outright rather than bent.
#
# `file` is the one that matters. `canvas._load_file` authorises with
# `attachment.chat_id != chat.id`, and an upload made on the new-chat screen is
# stored with `chat_id=None`. If a draft's chat carried no id, `None != None` is
# False and every unclaimed attachment its owner has would open from any draft
# canvas. `as_chat` sets an id, so that comparison already fails -- but relying
# on it would mean the guarantee lives in an id-shaped coincidence. It is stated
# here instead, where it can be read and tested.
SOURCES_NEEDING_A_CHAT = frozenset({"scratch", "file"})


def refuses(source: str) -> bool:
    return source in SOURCES_NEEDING_A_CHAT


def _sweep() -> None:
    """Drop drafts nobody has touched in a long while.

    On write rather than on a timer: a draft holds no connection and no process,
    only a little state, so there is nothing to close and nothing that leaks by
    being late. The shell it points at has its own reaper.
    """
    cutoff = time.monotonic() - IDLE_TIMEOUT
    for key in [k for k, d in _DRAFTS.items() if d.touched_at < cutoff]:
        _DRAFTS.pop(key, None)


__all__ = [
    "SOURCES_NEEDING_A_CHAT",
    "Draft",
    "as_chat",
    "clear",
    "forget",
    "get",
    "is_draft",
    "key_for",
    "refuses",
    "remember",
]