/omnigent/server/routes/terminal_attach.py

"""WebSocket endpoint exposing an agent's live terminals to the browser.

This module hosts the resource-addressed terminal-attach route:

- ``WS /sessions/{session_id}/resources/terminals/{terminal_id}/attach``
  - open a bidirectional bridge between a browser xterm.js client
  and the runner-side tmux session named by ``terminal_id``.

Runner-aware execution
----------------------

Terminal state lives in whichever process owns the
:class:`TerminalRegistry` that ``sys_terminal_launch`` mutated:

- **In-process runner**: the server passes
  :func:`omnigent.runtime.get_terminal_registry` into
  :func:`omnigent.runner.create_runner_app`, so the registry is
  shared. The server resolves the terminal id locally and bridges
  the PTY in-process.
- **Out-of-process runner** over the WebSocket tunnel: the runner
  owns the tmux socket. The server proxies WebSocket frames over
  the tunnel via a multiplexed WS channel
  (``omnigent/server/_runner_ws_tunnel.py``), and the runner's
  resource-addressed WS route runs ``tmux attach`` and bridges the
  PTY.

The proxy uses a factory configured via
:func:`omnigent.runtime.set_runner_ws_factory` in the server
lifespan. When the factory is unset, the server route falls back
to resolving the terminal in the local registry.

Wire protocol on the WebSocket
------------------------------

- **Server → client**: every PTY read is forwarded as a *binary*
  WebSocket frame. xterm.js's ``term.write()`` accepts ``Uint8Array``
  directly and runs it through its ANSI parser, so colors, cursor
  motion, alternate screen, mouse modes all work transparently.
- **Client → server**:
    - **Text frames** are JSON control messages:
      ``{"type": "resize", "cols": N, "rows": M}``. Parsed and applied
      to the PTY via ``ioctl(TIOCSWINSZ)``. Unknown shapes are ignored
      so future control-message additions don't immediately break
      older servers.
    - **Binary frames** are raw input bytes written verbatim to the
      PTY. xterm.js's ``onData`` callback emits these for keystrokes,
      pasted text, and mouse-mode reports.

Read-only mode
--------------

When the URL has ``?read_only=true``, binary input frames are dropped
silently at the server *and* the runner. The attach process itself
runs ``tmux attach -r`` as defense-in-depth, even if a frame got
past the application filter, tmux would refuse keystrokes from this
client.

Write attach is owner-only
--------------------------

A terminal is a single shared PTY driving one process that runs as the
session owner. Raw keystroke bytes carry no per-user identity, so input
typed by anyone other than the owner would be acted on — and, for the
agent's TUI, persisted into conversation history — as if the owner typed
it. To keep that attribution honest, an *interactive* (write) attach
requires ``LEVEL_OWNER``; non-owners can only attach read-only and drive
the agent through the chat composer, which carries the real sender's
identity. This holds uniformly for the agent's own terminal and for
user-launched shells, since both attach through this route.
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
from typing import Final

from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect, WebSocketException
from starlette import status

from omnigent.errors import OmnigentError
from omnigent.runtime import (
    get_runner_ws_factory,
    get_terminal_registry,
)
from omnigent.server.auth import LEVEL_OWNER, LEVEL_READ, AuthProvider
from omnigent.server.routes._auth_helpers import require_access
from omnigent.stores import ConversationStore
from omnigent.stores.permission_store import PermissionStore
from omnigent.terminals.control_bridge import bridge_tmux_control_to_websocket
from omnigent.terminals.ws_bridge import (
    WS_CLOSE_INTERNAL_ERROR,
    WS_CLOSE_TERMINAL_NOT_FOUND,
    bridge_tmux_pty_to_websocket,
)

_logger = logging.getLogger(__name__)

_WS_CLOSE_TERMINAL_NOT_FOUND: Final[int] = WS_CLOSE_TERMINAL_NOT_FOUND
_WS_CLOSE_INTERNAL_ERROR: Final[int] = WS_CLOSE_INTERNAL_ERROR


def create_terminal_attach_router(
    *,
    auth_provider: AuthProvider | None = None,
    permission_store: PermissionStore | None = None,
    conversation_store: ConversationStore | None = None,
) -> APIRouter:
    """
    Build the router exposing the terminal-attach WebSocket route.

    Wired into the FastAPI app under the ``/v1`` prefix in
    :func:`omnigent.server.app.create_app`. The list/CRUD endpoints
    for terminals live on the sessions router under
    ``/v1/sessions/{id}/resources/terminals``; this module only owns
    the WebSocket bridge.

    :param auth_provider: Optional provider used to authenticate the
        WebSocket handshake. When ``None`` and permissions are disabled,
        attach behaves as before for single-user/local deployments.
    :param permission_store: Optional session permission store. When
        provided, a read-only attach requires read access and an
        interactive (write) attach requires owner access — only the
        session owner may type into the shared PTY.
    :param conversation_store: Conversation store used by permission checks.
    :returns: An :class:`APIRouter` carrying the attach route.
    """
    router = APIRouter()

    @router.websocket("/sessions/{session_id}/resources/terminals/{terminal_id}/attach")
    async def attach_terminal_by_resource_id(
        websocket: WebSocket,
        session_id: str,
        terminal_id: str,
        read_only: bool = Query(default=False),
        transport: str | None = Query(default=None),
    ) -> None:
        """
        Attach to a terminal by resource id via WebSocket.

        Proxies to the runner's resource-addressed WS endpoint when a
        runner tunnel factory is installed, or falls back to resolving
        the terminal id locally and bridging in-process.

        :param websocket: The accepted FastAPI :class:`WebSocket`.
        :param session_id: Session/conversation identifier.
        :param terminal_id: Opaque terminal resource id,
            e.g. ``"terminal_bash_s1"``.
        :param read_only: Pass ``-r`` to tmux when ``True``.
        :param transport: Optional per-attach transport override
            (``"control"`` / ``"pty"``); forwarded to the runner and used to
            pick the control-mode vs PTY bridge in-process. ``None`` defers to
            the terminal spec / global default.
        """
        from omnigent.entities.session_resources import (
            resolve_terminal_entry_by_resource_id,
        )

        await _authorize_terminal_attach(
            websocket,
            session_id=session_id,
            read_only=read_only,
            auth_provider=auth_provider,
            permission_store=permission_store,
            conversation_store=conversation_store,
        )
        await websocket.accept()

        ws_factory = get_runner_ws_factory()
        if ws_factory is not None:
            from urllib.parse import urlencode

            qs_params = {"read_only": "true" if read_only else "false"}
            if transport is not None:
                qs_params["transport"] = transport
            qs = urlencode(qs_params)
            runner_path = (
                f"/v1/sessions/{session_id}/resources/terminals/{terminal_id}/attach?{qs}"
            )
            try:
                runner_cm = ws_factory(runner_path)
            except Exception:  # noqa: BLE001
                await websocket.close(
                    code=_WS_CLOSE_INTERNAL_ERROR,
                    reason="runner attach factory failed",
                )
                return
            try:
                from omnigent.runtime import telemetry

                with telemetry.span(
                    "terminal.attach",
                    attributes={
                        "session.id": session_id,
                        "terminal.id": terminal_id,
                        "terminal.read_only": read_only,
                        "terminal.mode": "runner-proxy",
                    },
                ):
                    async with runner_cm as runner_ws:
                        await _shuttle_ws_frames(websocket, runner_ws)
            except _RunnerWSClosed as closed:
                code = (
                    closed.code
                    if closed.code and closed.code >= 1000
                    else _WS_CLOSE_INTERNAL_ERROR
                )
                with contextlib.suppress(RuntimeError):
                    await websocket.close(
                        code=code,
                        reason=closed.reason or "",
                    )
            except Exception:  # noqa: BLE001
                with contextlib.suppress(RuntimeError):
                    await websocket.close(
                        code=_WS_CLOSE_INTERNAL_ERROR,
                        reason="runner attach proxy failed",
                    )
            return

        try:
            terminal_registry = get_terminal_registry()
        except RuntimeError:
            await websocket.close(
                code=_WS_CLOSE_TERMINAL_NOT_FOUND,
                reason="terminal registry not available",
            )
            return

        entry = resolve_terminal_entry_by_resource_id(
            session_id,
            terminal_id,
            terminal_registry,
        )
        if entry is None or not entry.instance.running or not await entry.instance.is_alive():
            await websocket.close(
                code=_WS_CLOSE_TERMINAL_NOT_FOUND,
                reason="terminal resource not found or not running",
            )
            return

        from omnigent.inner.terminal import (
            TERMINAL_TRANSPORT_CONTROL,
            resolve_terminal_transport,
        )
        from omnigent.runtime import telemetry

        resolved_transport = resolve_terminal_transport(
            override=transport,
            spec_transport=entry.instance.terminal_transport,
        )
        with telemetry.span(
            "terminal.attach",
            attributes={
                "session.id": session_id,
                "terminal.id": terminal_id,
                "terminal.read_only": read_only,
                "terminal.mode": "in-process",
                "terminal.transport": resolved_transport,
            },
        ):
            bridge = (
                bridge_tmux_control_to_websocket
                if resolved_transport == TERMINAL_TRANSPORT_CONTROL
                else bridge_tmux_pty_to_websocket
            )
            await bridge(
                websocket,
                socket_path=str(entry.instance.socket_path),
                tmux_target=entry.instance.tmux_target,
                read_only=read_only,
            )

    return router


async def _authorize_terminal_attach(
    websocket: WebSocket,
    *,
    session_id: str,
    read_only: bool,
    auth_provider: AuthProvider | None,
    permission_store: PermissionStore | None,
    conversation_store: ConversationStore | None,
) -> None:
    """
    Authorize a terminal-attach WebSocket before accepting it.

    Interactive attaches write bytes to the shared PTY, which runs as the
    session owner and (for the agent's TUI) persists input into history
    under the owner's identity. Raw keystrokes carry no per-user
    attribution, so an interactive attach requires ``LEVEL_OWNER`` — only
    the owner can drive the terminal. Read-only attaches still expose
    terminal output, so read access is the minimum; non-owners attach
    read-only and send input through the chat composer (which carries the
    real sender's identity). When permissions are disabled
    (``permission_store is None``), preserve the existing
    single-user/dev behavior.

    :param websocket: The incoming FastAPI :class:`WebSocket`, used to
        resolve the caller's identity via *auth_provider*.
    :param session_id: Session/conversation identifier the attach
        targets, e.g. ``"conv_abc123"``.
    :param read_only: ``True`` for a view-only attach (requires
        ``LEVEL_READ``); ``False`` for an interactive write attach
        (requires ``LEVEL_OWNER``).
    :param auth_provider: Provider used to resolve the caller's user id
        from the WebSocket handshake. Required when *permission_store*
        is set.
    :param permission_store: Session permission store. ``None`` disables
        all checks (single-user/dev mode).
    :param conversation_store: Conversation store consulted by the
        access check. Required when *permission_store* is set.
    :raises WebSocketException: With ``WS_1008_POLICY_VIOLATION`` when
        auth is misconfigured, the caller is unauthenticated, or the
        caller lacks the required level for the requested mode.
    """
    if permission_store is None:
        return
    if auth_provider is None or conversation_store is None:
        raise WebSocketException(
            code=status.WS_1008_POLICY_VIOLATION,
            reason="terminal attach authorization is not configured",
        )

    user_id = auth_provider.get_user_id(websocket)
    if user_id is None:
        raise WebSocketException(
            code=status.WS_1008_POLICY_VIOLATION,
            reason="authentication required",
        )

    required_level = LEVEL_READ if read_only else LEVEL_OWNER
    try:
        await require_access(
            user_id,
            session_id,
            required_level,
            permission_store,
            conversation_store,
        )
    except OmnigentError as exc:
        _logger.info(
            "Rejected terminal attach for session %s as user %s: %s",
            session_id,
            user_id,
            exc,
        )
        raise WebSocketException(
            code=status.WS_1008_POLICY_VIOLATION,
            reason="not authorized",
        ) from exc


class _RunnerWSClosed(Exception):
    """Carries a runner-side close so the browser side mirrors it."""

    def __init__(self, code: int | None, reason: str | None) -> None:
        super().__init__(f"runner WS closed: code={code} reason={reason!r}")
        self.code = code
        self.reason = reason


async def _shuttle_ws_frames(browser_ws: WebSocket, runner_ws: object) -> None:
    """
    Forward frames between *browser_ws* (FastAPI) and *runner_ws*
    (websockets client connection) until either side closes.

    The two libraries use slightly different APIs:

    - FastAPI: ``await ws.receive()`` returns a dict with ``text``
      or ``bytes``. Sends are :meth:`WebSocket.send_text` and
      :meth:`WebSocket.send_bytes`.
    - :mod:`websockets`: ``await ws.recv()`` returns ``str`` or
      ``bytes``; ``await ws.send(data)`` accepts either.

    We translate at the boundary. When either side raises, we
    cancel the other task and let the outer ``finally`` close the
    browser-side WS.

    :param browser_ws: The browser-facing FastAPI WebSocket.
    :param runner_ws: The runner-facing websockets client.
    """
    from websockets.exceptions import ConnectionClosed

    async def _browser_to_runner() -> None:
        try:
            while True:
                msg = await browser_ws.receive()
                msg_type = msg.get("type")
                if msg_type == "websocket.disconnect":
                    return
                text = msg.get("text")
                data = msg.get("bytes")
                if text is not None:
                    await runner_ws.send(text)  # type: ignore[attr-defined]
                elif data is not None:
                    await runner_ws.send(data)  # type: ignore[attr-defined]
        except WebSocketDisconnect:
            return
        except ConnectionClosed:
            return

    async def _runner_to_browser() -> None:
        try:
            while True:
                msg = await runner_ws.recv()  # type: ignore[attr-defined]
                if isinstance(msg, bytes | bytearray | memoryview):
                    await browser_ws.send_bytes(bytes(msg))
                else:
                    await browser_ws.send_text(msg)
        except ConnectionClosed as cc:
            # Surface the runner-side close code so the browser
            # mirrors it (4404 for missing terminal, etc.). Use
            # ``rcvd`` rather than the deprecated ``code`` / ``reason``
            # attributes so we read what the peer actually sent
            # (``sent`` would be the close we transmitted).
            rcvd = cc.rcvd
            code = rcvd.code if rcvd is not None else None
            reason = rcvd.reason if rcvd is not None else None
            raise _RunnerWSClosed(code, reason) from cc

    b2r = asyncio.create_task(_browser_to_runner())
    r2b = asyncio.create_task(_runner_to_browser())

    runner_closed: _RunnerWSClosed | None = None
    try:
        done, pending = await asyncio.wait(
            {b2r, r2b},
            return_when=asyncio.FIRST_COMPLETED,
        )
        for task in pending:
            task.cancel()
            with contextlib.suppress(asyncio.CancelledError, Exception):
                await task
        for task in done:
            exc = task.exception()
            if isinstance(exc, _RunnerWSClosed):
                runner_closed = exc
            elif exc is not None:
                _logger.warning("terminal-attach proxy: task crashed: %r", exc)
    finally:
        if runner_closed is not None:
            raise runner_closed