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

"""What one agent chat is pointed at, resolved while a session is open.

Everything a runner needs travels in `AgentContext`: the machine, the decrypted
credential, the mode in force, and the two lists that adjust it. Nothing is
looked up later, for the reason `Endpoint` is a frozen copy of a `Connection`
and `ToolContext` carries an owner id rather than a `User` -- a generation
outlives the request that started it, and a detached instance is a trap.

The mode is read **once, at the start of the reply**, and deliberately does not
change under a reply already in flight. Somebody switching to Auto halfway
through must not retroactively approve what is already queued.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field, replace
from typing import Any

from sqlalchemy.orm import Session as DBSession

from lembas.db.models import KIND_AGENT, Chat, SshProfile, User
from lembas.services import settings_store
from lembas.services.agent import hosts, policy
from lembas.services.agent import ssh as ssh_service
from lembas.services.agent.base import Executor
from lembas.services.agent.policy import Limits

log = logging.getLogger(__name__)


@dataclass
class AgentContext:
    """The machine an agent chat acts on, and what it may do there."""

    chat_id: str
    label: str
    project_dir: str
    # The connection's id, carried so a runner can drop the project listing it
    # has just invalidated. `index` is keyed on the connection and the
    # directory, not on the chat -- two chats on one tree share a listing.
    profile_id: str = ""
    mode: str = policy.MODE_MANUAL
    allow: tuple[str, ...] = ()
    deny: tuple[str, ...] = ()
    limits: Limits = field(default_factory=Limits)
    # Per-command bounds, from the instance settings.
    timeout: float = 60.0
    max_timeout: float = 600.0
    max_output: int = 64 * 1024
    # The decrypted credential. Held here and nowhere else, and cleared by
    # `generation` when the reply ends -- a finished Generation lingers five
    # minutes so late followers get the final frames, and a private key should
    # not linger with it.
    spec: dict[str, Any] = field(default_factory=dict)
    # Set only on the per-call copy handed to a runner whose call a person has
    # just allowed. The runners re-check the mode as a backstop, and without
    # this they would refuse the very thing that was approved -- the mode says
    # "ask", and asking is exactly what happened.
    approved: bool = False
    # Absolute paths this reply has read. `file_edit` refuses a file that is not
    # in here, because a patch written from memory against a file the model has
    # not looked at is how a rewrite silently loses somebody's work.
    #
    # Here rather than on `Generation` for two reasons. Runners never see a
    # Generation -- they get a `ToolContext`, which is a session-free snapshot
    # precisely so nothing in a tool holds live state -- and a read path is a
    # fact about the machine, which is what this class is.
    #
    # It is **shared with the approved copy**: `as_approved` is
    # `dataclasses.replace`, which copies field references, so a path read
    # through an approved call is visible here. That is wanted and is not
    # obvious, so there is a test for it.
    #
    # It resets each reply, and that is correct rather than a limitation.
    # `Message.tool_calls_json` is deliberately never replayed as context, so on
    # the next turn the model does not have the file's contents either --
    # requiring a re-read in the reply that edits is asking for something it
    # needs anyway.
    read_paths: set[str] = field(default_factory=set)
    # The plan currently in force, seeded from `chat.plan_message_id` when this
    # is resolved. Mutable and read/written in place by `plan_update`, for a
    # reason that is not obvious: a runner cannot write the message row --
    # `_persist` is the single writer -- so it returns the merged plan on its
    # event and the loop carries it. Two updates in one reply would then both
    # read the same stale plan from the database and the second would lose the
    # first. This snapshot is what they actually merge into.
    plan: dict[str, Any] = field(default_factory=dict)
    # Whether commands may run detached. When off, `shell_run` is byte-for-byte
    # what it always was and the `job_*` tools are not offered -- a command that
    # times out is killed, as before. When on, a command can be launched in the
    # background (or converted to one when it times out) and the model gets the
    # tools to check on it. `on_timeout` is the sub-switch for the auto-convert.
    background: bool = False
    background_on_timeout: bool = True
    # Whether a finished job wakes the model on its own, rather than only being
    # seen when it next runs. Read by the wording here and by the watcher.
    background_notify: bool = True
    # Most jobs watched at once. A watcher is a periodic reconnect, so this is a
    # real resource; past it a job still runs but is not watched or woken for.
    background_max_jobs: int = 5

    def executor(self) -> Executor:
        return ssh_service.SshExecutor(self.spec, self.project_dir)

    def clear(self) -> None:
        self.spec = {}

    def as_approved(self) -> AgentContext:
        """A copy of this context for one call a person has allowed."""
        return replace(self, approved=True)


def _plan_of(db: DBSession, chat: Chat) -> dict[str, Any]:
    """The plan this chat is working to, or an empty dict.

    One `db.get` by primary key -- the column exists to avoid a scan for "the
    newest message carrying a plan", because this runs while a request is
    waiting. The id is validated here rather than constrained in the schema, for
    the reason the column's comment gives.
    """
    from lembas.db.models import Message
    from lembas.services import plans

    if not chat.plan_message_id:
        return {}
    message = db.get(Message, chat.plan_message_id)
    if message is None or message.chat_id != chat.id:
        return {}
    return plans.normalise(message.plan_json)


def profile_for(db: DBSession, chat: Chat, user: User | None) -> SshProfile | None:
    """The connection this chat is pointed at, if it is still usable.

    Ownership is re-checked here rather than trusted from when the chat was
    created: a profile can be deleted, disabled, or moved to a host whose key
    has not been confirmed since, and any of those should stop the chat acting
    rather than be discovered at the first command.
    """
    if chat is None or chat.kind != KIND_AGENT or not chat.ssh_profile_id:
        return None

    profile = db.get(SshProfile, chat.ssh_profile_id)
    if profile is None or not profile.enabled:
        return None
    if user is not None and profile.owner_id != user.id:
        return None
    # A row can predate a setting, so this is asked here rather than trusted
    # from when the profile was saved: an administrator moving the switch to
    # `off` has to stop the chats already pointed at loopback, not only the next
    # one somebody tries to create. See services/agent/hosts.py.
    if not hosts.usable(db, profile):
        return None
    return profile


def _allow_for(chat: Chat) -> tuple[str, ...]:
    """Imported inside `resolve` rather than at module scope.

    `services/tools.py` imports this module's `resolve`, so a top-level import
    back the other way is a cycle.
    """
    from lembas.services import tools as tools_service

    return tools_service.scoped_allow(chat)


def refresh(db: DBSession, agent: AgentContext) -> AgentContext:
    """Re-read the two things a person can change while a reply is running.

    The mode and the chat's own allow list, and nothing else. Everything else on
    the context is fixed for the life of a chat (the connection, the directory)
    or is an instance setting nobody is editing mid-reply.

    Called once per round rather than once per reply. The reply-long snapshot it
    replaces made both controls do nothing until the next turn: switching to
    Auto during a long agent reply went on asking about every call, and
    "Always allow this" was accepted, written to the row, and then ignored for
    the rest of the reply that had just asked. Both look exactly like a control
    that does not work, because for that reply they were.

    Once per *round* and not more often, because a round's calls are authorised
    together: what is already queued was decided under the mode that was in
    force when it was queued, and switching to Auto must not retroactively
    approve it. Mutated in place -- `as_approved` copies field references, so a
    replacement here would leave the approved copy of this round pointing at the
    old one.
    """
    chat = db.get(Chat, agent.chat_id)
    if chat is None:
        return agent

    agent.mode = chat.agent_mode if chat.agent_mode in policy.MODES else policy.MODE_MANUAL
    instance = settings_store.agents(db)
    agent.allow = (*(instance.get("allow_default") or ()), *_allow_for(chat))
    return agent


def resolve(db: DBSession, chat: Chat, user: User | None) -> AgentContext | None:
    """This chat's agent setup, or None if it has none it can use.

    None is the answer to every "no": not an agent chat, the feature switched
    off, the connection gone or disabled, SSH not installed. Each of those means
    the agent tools are not offered at all, which is better than offering a tool
    that fails on its first call.

    A profile whose host key has never been confirmed is deliberately *not* one
    of them. The tools are offered and the failure is explicit, because "check
    the connection and accept its fingerprint" is a thing the reader can act on,
    while a silently missing tool is not.
    """
    profile = profile_for(db, chat, user)
    if profile is None:
        return None

    values = settings_store.agents(db)
    if not values.get("enabled"):
        return None
    if ssh_service.available():
        return None

    return AgentContext(
        chat_id=chat.id,
        label=profile.label,
        plan=_plan_of(db, chat),
        project_dir=chat.project_dir or profile.default_dir or "",
        profile_id=profile.id,
        mode=chat.agent_mode if chat.agent_mode in policy.MODES else policy.MODE_MANUAL,
        # The instance's list, plus whatever this chat's reader has said
        # "always" to on a card. Never the other way round for the deny list:
        # a chat cannot un-deny anything, and `decide` consults deny first
        # regardless.
        allow=(*(values.get("allow_default") or ()), *_allow_for(chat)),
        deny=tuple(values.get("deny_default") or ()),
        limits=Limits(
            steps=int(values.get("max_steps") or 200),
            wall_seconds=float(values.get("max_wall_seconds") or 900),
            output_bytes=int(values.get("max_total_output_bytes") or 1024 * 1024),
            # `or 0` would turn a deliberate 0 into the default, and 0 is how an
            # administrator says "no ceiling". `agents()` has already clamped it.
            completion_tokens=int(values.get("max_completion_tokens", 200_000) or 0),
        ),
        timeout=float(values.get("default_timeout") or 60),
        max_timeout=float(values.get("max_timeout") or 600),
        max_output=int(values.get("max_output_bytes") or 64 * 1024),
        background=bool(values.get("background_enabled")),
        background_on_timeout=bool(values.get("background_on_timeout", True)),
        background_notify=bool(values.get("background_notify", True)),
        background_max_jobs=int(values.get("background_max_jobs") or 5),
        spec=ssh_service.spec_from(profile),
    )


__all__ = ["AgentContext", "profile_for", "resolve"]