omnigent-ai--omnigent/omnigent/stores/host_store.py

"""
Persistent store for host registrations.

Hosts are machines connected via ``omnigent host``. The store
tracks which hosts have ever connected, their names, owners, and
online/offline status. The ``hosts`` table is the source of truth
for ``GET /v1/hosts`` — all server replicas query it. Live WebSocket
connection state is tracked separately in the in-memory
``HostRegistry`` (one per replica).
"""

from __future__ import annotations

import hashlib
import json
import logging
from dataclasses import dataclass

from sqlalchemy import Engine, select, update
from sqlalchemy import delete as sql_delete
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session

from omnigent.db.db_models import SqlConversation, SqlHost, current_workspace_id
from omnigent.db.enum_codecs import decode_host_status, encode_host_status
from omnigent.db.utils import get_or_create_engine, make_managed_session_maker, now_epoch

# A host is considered live only if its row was touched (connect or
# heartbeat) within this window. The host tunnel's ping loop writes a
# heartbeat every PING_INTERVAL_S (30s); three missed heartbeats means
# the host is gone. This freshness gate is the safety net for every
# path that never runs set_offline — hard crash, OOM, deploy/replica
# restart, silent network drop, or a connect that died after the online
# upsert. It must stay >= the tunnel's ping-miss window
# (PING_INTERVAL_S * PING_MISS_THRESHOLD) so a healthy host that is
# still heart-beating is never falsely aged out.
HOST_LIVENESS_TTL_S = 90
HarnessAvailability = bool | str


@dataclass
class Host:
    """
    A registered host machine.

    :param host_id: Stable identifier from the host's local
        ``~/.omnigent/config.yaml``, e.g. ``"host_a1b2c3d4..."``.
    :param name: Human-readable name, e.g. ``"corey-laptop"``.
    :param owner: User ID from the Databricks auth Bearer token,
        e.g. ``"corey.zumar@databricks.com"``.
    :param status: ``"online"`` or ``"offline"``.
    :param created_at: Unix epoch seconds of first registration.
    :param updated_at: Unix epoch seconds the row was last touched —
        a status change (connect/disconnect) or a tunnel heartbeat.
        Used as the host's last-seen for the liveness freshness gate
        (see :data:`HOST_LIVENESS_TTL_S`).
    :param sandbox_provider: Sandbox provider backing a SERVER-MANAGED
        host (``host_type="managed"`` sessions), e.g. ``"modal"``.
        ``None`` for external (user-connected) hosts — non-``None``
        marks the host as server-managed.
    :param sandbox_id: Provider-assigned id of the sandbox currently
        backing a managed host, e.g. ``"sb-a1b2c3"`` — what
        termination is issued against. ``None`` for external hosts.
    :param configured_harnesses: Per-harness readiness reported in the
        host's last ``host.hello`` frame, e.g.
        ``{"claude-sdk": True, "codex": False}``. ``None`` when the
        host has never reported it (older host build) — unknown, not
        "nothing configured".
    """

    host_id: str
    name: str
    owner: str
    status: str
    created_at: int
    updated_at: int
    sandbox_provider: str | None = None
    sandbox_id: str | None = None
    configured_harnesses: dict[str, HarnessAvailability] | None = None


def host_is_live(host: Host, now: int | None = None) -> bool:
    """
    Return whether a :class:`Host` is online and recently seen.

    Pure helper over an already-loaded entity (no DB access), so
    callers that already hold a :class:`Host` — or a list of them —
    don't re-query per row. A host is live only when its ``status`` is
    ``"online"`` **and** its last-seen (``updated_at``) is within
    :data:`HOST_LIVENESS_TTL_S`; the freshness half is what catches a
    host that died without a graceful disconnect.

    :param host: The host entity to evaluate.
    :param now: Unix epoch seconds to measure freshness against;
        defaults to the current time. Pass an explicit value to
        classify many hosts against one consistent clock.
    :returns: ``True`` when the host is online and fresh.
    """
    ref = now if now is not None else now_epoch()
    return host.status == "online" and host.updated_at >= ref - HOST_LIVENESS_TTL_S


_logger = logging.getLogger(__name__)


def _parse_configured_harnesses(raw: str | None) -> dict[str, HarnessAvailability] | None:
    """
    Parse the JSON-encoded ``hosts.configured_harnesses`` column.

    Tolerant: ``NULL``, malformed JSON, or a non-object payload all
    map to ``None`` ("unknown") — a corrupt column value must degrade
    to no-warning in the UI, never break host listing. Entries with a
    non-bool/string value are dropped for the same reason.

    :param raw: The raw column value, e.g.
        ``'{"claude-sdk": true, "codex": false}'`` or ``None``.
    :returns: The readiness map, or ``None`` when absent or unparseable.
    """
    if raw is None:
        return None
    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError:
        _logger.warning("Ignoring malformed hosts.configured_harnesses value")
        return None
    if not isinstance(parsed, dict):
        return None
    return {k: v for k, v in parsed.items() if isinstance(k, str) and isinstance(v, (bool, str))}


def _row_to_host(row: SqlHost) -> Host:
    """
    Convert a :class:`SqlHost` ORM row to a :class:`Host` entity.

    :param row: The SQLAlchemy ORM row to convert.
    :returns: A :class:`Host` dataclass instance.
    """
    return Host(
        host_id=row.host_id,
        name=row.name,
        owner=row.owner,
        status=decode_host_status(row.status),
        created_at=row.created_at,
        updated_at=row.updated_at,
        sandbox_provider=row.sandbox_provider,
        sandbox_id=row.sandbox_id,
        configured_harnesses=_parse_configured_harnesses(row.configured_harnesses),
    )


def hash_host_launch_token(token: str) -> str:
    """
    Digest a managed-host launch token for storage / lookup.

    Only the digest is ever persisted (``hosts.token_hash``), so a
    database leak does not leak usable credentials, and the
    tunnel-side lookup is by digest — the raw token never touches a
    query.

    :param token: The raw launch token, e.g. the value of
        ``secrets.token_urlsafe(32)``.
    :returns: Hex SHA-256 digest, e.g. ``"9f86d08..."`` (64 chars).
    """
    return hashlib.sha256(token.encode()).hexdigest()


class HostStore:
    """
    Persistent store for host registrations backed by SQLAlchemy.

    :param storage_location: SQLAlchemy database URI, e.g.
        ``"sqlite:///hosts.db"``.
    """

    def __init__(self, storage_location: str) -> None:
        """
        Initialize the host store.

        :param storage_location: SQLAlchemy database URI, e.g.
            ``"sqlite:///hosts.db"``.
        """
        self._engine: Engine = get_or_create_engine(storage_location)
        self._session = make_managed_session_maker(self._engine)

    def upsert_on_connect(
        self,
        host_id: str,
        name: str,
        owner: str,
        *,
        allow_host_id_reown: bool = False,
        configured_harnesses: dict[str, HarnessAvailability] | None = None,
    ) -> Host:
        """
        Register or update a host on WebSocket connect.

        Inserts a new row if ``host_id`` does not exist, otherwise
        updates ``name``, ``owner``, ``status``, and ``updated_at``.
        Called by the host tunnel endpoint when a host sends its
        ``host.hello`` frame.

        The upsert keys on the ``(owner, name)`` primary key, but
        ``host_id`` carries its own UNIQUE constraint. When the same
        physical host re-registers under a *different* owner (e.g. a
        local server respawned with a flipped auth posture changes the
        owner between an accounts user and the reserved ``local`` user),
        the ``(owner, name)`` lookup misses and a plain INSERT would
        collide on ``host_id``. That collision is a deliberate W2-class
        boundary in shared deployments — a different user must not be
        able to claim another user's host_id — so re-owning is gated
        behind *allow_host_id_reown*, which the server sets only for the
        loopback single-user local server. Remote / multi-user servers
        never set it, so the hijack boundary stays intact (the INSERT
        raises ``IntegrityError`` and fails the handshake closed).

        :param host_id: Stable host identifier, e.g.
            ``"host_a1b2c3d4..."``.
        :param name: Human-readable name from ``config.yaml``, e.g.
            ``"corey-laptop"``.
        :param owner: Authenticated user ID from the Bearer token,
            e.g. ``"corey.zumar@databricks.com"``.
        :param allow_host_id_reown: When ``True`` and a row already
            exists for *host_id* under a different ``(owner, name)``,
            re-own that row in place (preserving the ``host_id`` and its
            conversation bindings) instead of inserting. Intended solely
            for the single-user loopback local server.
        :param configured_harnesses: Per-harness readiness from the
            host's ``host.hello`` frame, e.g. ``{"claude-sdk": True}``.
            Written on every connect — including ``None`` from an older
            host that doesn't report it, which correctly resets any
            stale value back to "unknown".
        :returns: The upserted :class:`Host`.
        """
        now = now_epoch()
        harnesses_json = (
            json.dumps(configured_harnesses) if configured_harnesses is not None else None
        )
        with self._session() as session:
            # Primary lookup: by (workspace_id, host_id) — the new PK.
            row = session.get(SqlHost, (current_workspace_id(), host_id))
            if row is not None:
                # W2-class boundary: a different user must not claim another
                # user's host_id. Raise the same IntegrityError the old UNIQUE
                # constraint produced so the tunnel handler rejects the hijack.
                if row.owner != owner and not allow_host_id_reown:
                    raise IntegrityError(
                        "host_id already owned by a different user",
                        params={"host_id": host_id, "owner": owner},
                        orig=Exception("UNIQUE constraint failed: hosts.host_id"),
                    )
                # Known host_id (same owner, or reown opted in): update
                # owner/name in case they changed, then refresh status and timestamp.
                row.owner = owner
                row.name = name
                row.status = encode_host_status("online")
                row.updated_at = now
                row.configured_harnesses = harnesses_json
                return _row_to_host(row)

            # host_id is new — check whether (workspace_id, owner, name)
            # already exists. If it does, the same machine regenerated its
            # identity file: this is a host_id rotation. If allow_host_id_reown
            # is set, also check if any row holds this host_id under a different
            # owner and re-own it instead of inserting.
            if allow_host_id_reown:
                reowned = self._reown_host_id(
                    session,
                    host_id=host_id,
                    name=name,
                    owner=owner,
                    configured_harnesses_json=harnesses_json,
                )
                if reowned is not None:
                    return reowned

            existing_by_name = session.execute(
                select(SqlHost).where(
                    SqlHost.workspace_id == current_workspace_id(),
                    SqlHost.owner == owner,
                    SqlHost.name == name,
                )
            ).scalar_one_or_none()
            if existing_by_name is not None:
                # Same (owner, name), different host_id: identity rotation.
                # host_id is now part of the PK, so we can't UPDATE it via the
                # ORM — delete the old row and insert a fresh one that carries
                # the new host_id while preserving created_at.
                row = self._rotate_host_id(session, existing_by_name, host_id, now, harnesses_json)
                return _row_to_host(row)

            # Genuinely new host: plain INSERT.
            row = SqlHost(
                owner=owner,
                name=name,
                host_id=host_id,
                status=encode_host_status("online"),
                created_at=now,
                updated_at=now,
                configured_harnesses=harnesses_json,
            )
            session.add(row)
            return _row_to_host(row)

    @staticmethod
    def _rotate_host_id(
        session: Session,
        row: SqlHost,
        new_host_id: str,
        now: int,
        harnesses_json: str | None,
    ) -> SqlHost:
        """Replace a host row's host_id while repointing its conversations.

        ``host_id`` is now part of the PK, so an in-place UPDATE is not
        possible via the ORM. The rotation is:

        1. Capture the conversation ids bound to the old host_id.
        2. NULL them so nothing references the old PK value.
        3. DELETE the old row (host_id was the PK member being changed).
        4. INSERT a new row with the new host_id, preserving ``created_at``.
        5. Reattach the captured conversations to the new host_id.

        All steps run inside the caller's transaction so a failure rolls
        the whole upsert back.

        :param session: The active SQLAlchemy session.
        :param row: The existing host row whose ``host_id`` rotates.
        :param new_host_id: The host_id the host reconnected with.
        :param now: Unix epoch seconds for the updated_at timestamp.
        :param harnesses_json: JSON-encoded harness readiness, or None.
        :returns: The newly inserted :class:`SqlHost` row.
        """
        old_host_id = row.host_id
        # Preserve durable fields from the outgoing row before deletion.
        created_at = row.created_at
        owner = row.owner
        name = row.name
        token_hash = row.token_hash
        token_expires_at = row.token_expires_at
        sandbox_provider = row.sandbox_provider
        sandbox_id = row.sandbox_id

        bound_ids = list(
            session.execute(
                select(SqlConversation.id).where(
                    SqlConversation.workspace_id == current_workspace_id(),
                    SqlConversation.host_id == old_host_id,
                )
            ).scalars()
        )
        if bound_ids:
            session.execute(
                update(SqlConversation)
                .where(
                    SqlConversation.workspace_id == current_workspace_id(),
                    SqlConversation.host_id == old_host_id,
                )
                .values(host_id=None)
            )
            session.flush()

        # Delete the old PK row and insert a new one with the rotated host_id.
        session.execute(
            sql_delete(SqlHost).where(
                SqlHost.workspace_id == current_workspace_id(),
                SqlHost.host_id == old_host_id,
            )
        )
        session.flush()

        new_row = SqlHost(
            workspace_id=current_workspace_id(),
            host_id=new_host_id,
            owner=owner,
            name=name,
            status=encode_host_status("online"),
            created_at=created_at,
            updated_at=now,
            token_hash=token_hash,
            token_expires_at=token_expires_at,
            sandbox_provider=sandbox_provider,
            sandbox_id=sandbox_id,
            configured_harnesses=harnesses_json,
        )
        session.add(new_row)
        session.flush()

        if bound_ids:
            session.execute(
                update(SqlConversation)
                .where(
                    SqlConversation.workspace_id == current_workspace_id(),
                    SqlConversation.id.in_(bound_ids),
                )
                .values(host_id=new_host_id)
            )
            session.flush()

        return new_row

    def _reown_host_id(
        self,
        session: Session,
        *,
        host_id: str,
        name: str,
        owner: str,
        configured_harnesses_json: str | None = None,
    ) -> Host | None:
        """Re-own an existing host_id row under a new ``(owner, name)``.

        Used only when ``upsert_on_connect`` opts in via
        ``allow_host_id_reown`` (the single-user loopback local server).
        Updates ``owner``, ``name``, ``status``, and ``updated_at`` on the
        row that already holds *host_id*, leaving ``host_id`` itself
        unchanged so the ``conversations.host_id`` foreign-key bindings
        survive the owner change. ``owner`` / ``name`` are the table's
        primary key, so the change is issued as a Core ``UPDATE`` rather
        than mutating the ORM object's PK in place.

        :param session: The active SQLAlchemy session.
        :param host_id: Host identifier whose row should be re-owned,
            e.g. ``"host_a1b2c3d4..."``.
        :param name: New host name to record, e.g. ``"corey-laptop"``.
        :param owner: New owner to record, e.g. ``"local"`` or
            ``"corey.zumar@databricks.com"``.
        :param configured_harnesses_json: JSON-encoded readiness map from
            the connecting host's hello, e.g.
            ``'{"claude-sdk": true}'``, or ``None`` when unreported.
            Written like the normal connect paths so a re-owned row
            carries fresh (not stale) readiness.
        :returns: The re-owned :class:`Host`, or ``None`` if no row holds
            *host_id* (caller falls through to a normal insert).
        """
        existing = session.execute(
            select(SqlHost).where(
                SqlHost.workspace_id == current_workspace_id(), SqlHost.host_id == host_id
            )
        ).scalar_one_or_none()
        if existing is None:
            return None
        created_at = existing.created_at
        now = now_epoch()
        session.execute(
            update(SqlHost)
            .where(
                SqlHost.workspace_id == current_workspace_id(),
                SqlHost.host_id == host_id,
            )
            .values(
                owner=owner,
                name=name,
                status=encode_host_status("online"),
                updated_at=now,
                configured_harnesses=configured_harnesses_json,
            )
        )
        return Host(
            host_id=host_id,
            name=name,
            owner=owner,
            status="online",
            created_at=created_at,
            updated_at=now,
            sandbox_provider=existing.sandbox_provider,
            sandbox_id=existing.sandbox_id,
            configured_harnesses=_parse_configured_harnesses(configured_harnesses_json),
        )

    def set_offline(self, host_id: str) -> None:
        """
        Mark a host as offline when its WebSocket disconnects.

        No-op if the host does not exist (the disconnect callback
        may fire after a failed registration).

        :param host_id: Host identifier, e.g.
            ``"host_a1b2c3d4..."``.
        """
        with self._session() as session:
            row = session.execute(
                select(SqlHost).where(
                    SqlHost.workspace_id == current_workspace_id(), SqlHost.host_id == host_id
                )
            ).scalar_one_or_none()
            if row is not None:
                row.status = encode_host_status("offline")
                row.updated_at = now_epoch()

    def heartbeat(self, host_id: str) -> None:
        """
        Refresh a host's last-seen timestamp while its tunnel is alive.

        Bumps ``updated_at`` to now so the liveness freshness gate
        (see :data:`HOST_LIVENESS_TTL_S`) keeps treating the host as
        online. Called from the host tunnel's ping loop every
        ``PING_INTERVAL_S``. Does not change ``status`` — a host whose
        ping loop is running is, by construction, still ``"online"``.

        No-op if the host does not exist.

        :param host_id: Host identifier, e.g.
            ``"host_a1b2c3d4..."``.
        """
        # Single UPDATE rather than SELECT-then-mutate: this runs every
        # ping interval for every connected host, so the extra read is
        # pure overhead. A missing host simply matches no rows (a no-op).
        with self._session() as session:
            session.execute(
                update(SqlHost)
                .where(
                    SqlHost.workspace_id == current_workspace_id(),
                    SqlHost.host_id == host_id,
                )
                .values(updated_at=now_epoch())
            )

    def is_online(self, host_id: str) -> bool:
        """
        Return whether a host is currently live, cross-replica.

        A host counts as live only when its row is ``status="online"``
        **and** its last-seen (``updated_at``) is within
        :data:`HOST_LIVENESS_TTL_S`. The freshness check is what
        catches a host that died without a graceful disconnect: the
        ``status`` flag alone stays ``"online"`` forever in that case
        (set_offline only runs on a clean tunnel close), so a stale
        timestamp is the only reliable signal that the host is gone.

        :param host_id: Host identifier, e.g.
            ``"host_a1b2c3d4..."``.
        :returns: ``True`` when the host is online and its last-seen is
            fresh; ``False`` if unknown, offline, or stale.
        """
        row = self.get_host(host_id)
        return row is not None and host_is_live(row)

    def online_host_ids(self, host_ids: list[str]) -> set[str]:
        """
        Return the subset of ``host_ids`` that are currently live.

        Bulk variant of :meth:`is_online` for the sidebar online-dot
        batch path: one ``SELECT ... WHERE host_id IN (...)`` instead
        of a per-host query. Liveness applies the same
        status-plus-freshness gate as :meth:`is_online`, classifying
        every row against one consistent clock.

        :param host_ids: Host identifiers to check, e.g.
            ``["host_abc123", "host_def456"]``. Duplicates are
            tolerated; empty input returns an empty set without
            touching the database.
        :returns: The set of ids whose host row is online and fresh.
            Unknown, offline, or stale ids are absent.
        """
        if not host_ids:
            return set()
        unique_ids = list(set(host_ids))
        ref = now_epoch()
        with self._session() as session:
            rows = session.execute(
                select(SqlHost.host_id, SqlHost.status, SqlHost.updated_at).where(
                    SqlHost.workspace_id == current_workspace_id(),
                    SqlHost.host_id.in_(unique_ids),
                )
            ).all()
        online_code = encode_host_status("online")
        return {
            row.host_id
            for row in rows
            if row.status == online_code and row.updated_at >= ref - HOST_LIVENESS_TTL_S
        }

    def list_hosts(self, owner: str) -> list[Host]:
        """
        List all hosts owned by a specific user.

        Returns both online and offline hosts, ordered by
        ``updated_at`` descending (most recently active first).

        :param owner: User ID to filter by, e.g.
            ``"corey.zumar@databricks.com"``.
        :returns: List of :class:`Host` entities.
        """
        with self._session() as session:
            rows = (
                session.query(SqlHost)
                .filter(
                    SqlHost.workspace_id == current_workspace_id(),
                    SqlHost.owner == owner,
                )
                .order_by(SqlHost.updated_at.desc())
                .all()
            )
            return [_row_to_host(row) for row in rows]

    def get_host(self, host_id: str) -> Host | None:
        """
        Fetch a single host by ID.

        :param host_id: Host identifier, e.g.
            ``"host_a1b2c3d4..."``.
        :returns: The :class:`Host` if found, otherwise ``None``.
        """
        with self._session() as session:
            row = session.execute(
                select(SqlHost).where(
                    SqlHost.workspace_id == current_workspace_id(), SqlHost.host_id == host_id
                )
            ).scalar_one_or_none()
            if row is None:
                return None
            return _row_to_host(row)

    def register_managed_host(
        self,
        *,
        host_id: str,
        name: str,
        owner: str,
        token: str,
        provider: str,
        sandbox_id: str,
        token_expires_at: int,
    ) -> Host:
        """
        Pre-register a server-managed sandbox host with its credential.

        Called by the managed-launch orchestration after the sandbox is
        provisioned and BEFORE the in-sandbox host process starts, so
        the launch token is resolvable by the time the host first dials
        the tunnel. The row is created ``"offline"``; the tunnel's
        normal ``upsert_on_connect`` flips it online when the host
        registers.

        If a row already exists for *host_id* (a RELAUNCH: the host
        identity is durable across sandbox generations so session
        bindings survive a dead sandbox), the credential and sandbox
        columns are overwritten in place — which atomically revokes the
        previous generation's token, since its digest no longer matches
        anything.

        :param host_id: Server-generated host identifier, e.g.
            ``"host_a1b2c3d4..."``.
        :param name: Display name for the host picker, e.g.
            ``"managed-a1b2c3d4"``. Part of the table's
            ``(owner, name)`` primary key.
        :param owner: User the managed host acts for, e.g.
            ``"alice@example.com"``.
        :param token: The RAW launch token (hashed here, never stored),
            e.g. the value of ``secrets.token_urlsafe(32)``.
        :param provider: Sandbox provider name, e.g. ``"modal"``.
        :param sandbox_id: Provider-assigned sandbox id, e.g.
            ``"sb-a1b2c3"``.
        :param token_expires_at: Unix epoch seconds after which the
            token no longer authenticates.
        :returns: The registered :class:`Host`.
        :raises ValueError: If a row for *host_id* exists under a
            DIFFERENT owner — a relaunch may only re-credential a host
            the same user owns.
        """
        now = now_epoch()
        token_hash = hash_host_launch_token(token)
        with self._session() as session:
            existing = session.execute(
                select(SqlHost).where(
                    SqlHost.workspace_id == current_workspace_id(), SqlHost.host_id == host_id
                )
            ).scalar_one_or_none()
            if existing is not None:
                if existing.owner != owner:
                    # Fail closed (W2-class boundary): re-crediting a host
                    # row hands its launch token holder the row owner's
                    # identity, so a cross-owner overwrite would be a host
                    # hijack. host_id is server-generated today (uuid4 per
                    # launch), so this can only fire on a bug or a forged
                    # id — refuse rather than re-own.
                    raise ValueError(
                        f"host {host_id!r} is registered to a different owner; "
                        "refusing to re-credential it"
                    )
                existing.token_hash = token_hash
                existing.token_expires_at = token_expires_at
                existing.sandbox_provider = provider
                existing.sandbox_id = sandbox_id
                existing.updated_at = now
                return _row_to_host(existing)
            row = SqlHost(
                owner=owner,
                name=name,
                host_id=host_id,
                status=encode_host_status("offline"),
                created_at=now,
                updated_at=now,
                token_hash=token_hash,
                token_expires_at=token_expires_at,
                sandbox_provider=provider,
                sandbox_id=sandbox_id,
            )
            session.add(row)
            return _row_to_host(row)

    def resolve_launch_token(self, token: str) -> Host | None:
        """
        Resolve a presented launch token to its managed host, if valid.

        The host tunnel's auth path for managed hosts. Lookup is by
        SHA-256 digest — the comparison happens inside an indexed
        equality query on a uniformly distributed hash, which is not
        byte-by-byte comparable from the network (the standard
        reset-token pattern; no timing oracle on the raw token).
        Expired tokens do not authenticate; the expiry is checked
        atomically with the lookup.

        :param token: The raw token presented by a connecting host.
        :returns: The matching :class:`Host` whose token is unexpired,
            or ``None`` when the token is unknown or expired.
        """
        with self._session() as session:
            row = session.execute(
                select(SqlHost).where(
                    SqlHost.workspace_id == current_workspace_id(),
                    SqlHost.token_hash == hash_host_launch_token(token),
                )
            ).scalar_one_or_none()
            # token_expires_at is written together with token_hash, so a
            # matched row always carries it; the None arm is mypy
            # narrowing that doubles as fail-closed.
            if row is None or row.token_expires_at is None or row.token_expires_at < now_epoch():
                return None
            return _row_to_host(row)

    def delete_host(self, host_id: str) -> None:
        """
        Delete a host row entirely.

        Managed-host teardown: removes the host from the picker AND
        revokes its launch token in one operation (the row IS the
        credential). Explicitly nulls ``conversations.host_id`` for any
        sessions still bound to this host — the DB no longer cascades
        this via FK. No-op when the row does not exist — deletion is
        invoked from best-effort cleanup paths that may race.

        :param host_id: Host identifier, e.g. ``"host_a1b2c3d4..."``.
        """
        with self._session() as session:
            session.execute(
                update(SqlConversation)
                .where(
                    SqlConversation.workspace_id == current_workspace_id(),
                    SqlConversation.host_id == host_id,
                )
                .values(host_id=None)
            )
            session.execute(
                sql_delete(SqlHost).where(
                    SqlHost.workspace_id == current_workspace_id(),
                    SqlHost.host_id == host_id,
                )
            )

    def revoke_launch_token(self, host_id: str) -> None:
        """
        Clear a managed host's launch credential, keeping the row.

        Relaunch-failure cleanup: a failed sandbox RELAUNCH must revoke
        the token it armed (the new sandbox never came up to use it)
        without deleting the durable host row — the session binding
        survives, and the next relaunch attempt re-arms a fresh token
        via :meth:`register_managed_host`. Contrast :meth:`delete_host`,
        which is full teardown. No-op when the row does not exist.

        :param host_id: Host identifier, e.g. ``"host_a1b2c3d4..."``.
        """
        with self._session() as session:
            row = session.execute(
                select(SqlHost).where(
                    SqlHost.workspace_id == current_workspace_id(), SqlHost.host_id == host_id
                )
            ).scalar_one_or_none()
            if row is None:
                return
            row.token_hash = None
            row.token_expires_at = None
            row.updated_at = now_epoch()