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 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)

    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 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
    return profile


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,
        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,
        allow=tuple(values.get("allow_default") or ()),
        deny=tuple(values.get("deny_default") or ()),
        limits=Limits(
            steps=int(values.get("max_steps") or 40),
            wall_seconds=float(values.get("max_wall_seconds") or 900),
            output_bytes=int(values.get("max_total_output_bytes") or 1024 * 1024),
        ),
        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),
        spec=ssh_service.spec_from(profile),
    )


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