stackpilot/backend/services/exec_service.py

"""Interactive exec (web terminal) into compose-managed containers.

Reuses ``container_service._get_managed`` so a terminal can only be opened on a
container that belongs to a compose-managed stack — never an arbitrary host
container. The raw exec socket is bidirectional (stdin + a TTY-merged
stdout/stderr stream), suitable for piping straight to xterm.js.
"""
from __future__ import annotations

import asyncio
import codecs
import contextlib
import json
import queue
import socket as _socket
import threading

from starlette.websockets import WebSocketDisconnect

from docker_client import get_client, safe_call
from services.container_service import _get_managed

DEFAULT_SHELL = "/bin/sh"


def create_exec(container_id: str, cmd: list[str] | None = None, tty: bool = True) -> str:
    """Create an exec instance on a managed container and return its id."""
    container = _get_managed(container_id)
    client = get_client()
    created = safe_call(
        client.api.exec_create,
        container.id,
        cmd or [DEFAULT_SHELL],
        stdin=True,
        tty=tty,
        stdout=True,
        stderr=True,
    )
    return created["Id"]


def start_exec(exec_id: str, tty: bool = True):
    """Start the exec and return ``(holder, raw_socket)``.

    docker-py 7.x returns a ``socket.SocketIO`` wrapper whose real, recv/sendall
    -capable fd lives at ``._sock``; older versions hand back the socket
    directly. We keep both: ``holder`` is what we close, ``raw`` is what we
    recv/sendall on.
    """
    client = get_client()
    holder = safe_call(client.api.exec_start, exec_id, socket=True, tty=tty, demux=False)
    raw = getattr(holder, "_sock", None) or holder
    return holder, raw


def resize_exec(exec_id: str, height: int, width: int) -> None:
    """Resize the exec's TTY (rows x cols) so the shell wraps correctly."""
    client = get_client()
    safe_call(client.api.exec_resize, exec_id, height=height, width=width)


def exec_exit_code(exec_id: str):
    """Return the exec's ExitCode once it has finished (None while running)."""
    client = get_client()
    info = safe_call(client.api.exec_inspect, exec_id)
    return info.get("ExitCode")


async def pump_exec(websocket, exec_id: str, holder, raw) -> None:
    """Bidirectionally pump an exec socket <-> a WebSocket.

    Browser -> container: JSON ``{"type":"data","data":...}`` keystrokes and
    ``{"type":"resize","rows","cols"}`` control frames (raw text is also
    accepted as keystrokes). Container -> browser: ``{"type":"data","data":...}``
    then a final ``{"type":"exit","code":...}``.

    The blocking ``recv`` and ``sendall`` run on two threads of this session's
    own. They used to run on the event loop's default executor, where an open
    terminal held a worker for its whole lifetime — shared with every
    ``asyncio.to_thread`` in the app. On teardown the socket is shut down, which
    is what unblocks the reader.
    """
    raw.setblocking(True)
    loop = asyncio.get_running_loop()
    closed = asyncio.Event()
    incoming: asyncio.Queue = asyncio.Queue()
    outgoing: queue.Queue = queue.Queue()

    def post(item) -> None:
        try:
            loop.call_soon_threadsafe(incoming.put_nowait, item)
        except RuntimeError:  # loop closed
            pass

    def reader() -> None:
        try:
            while True:
                data = raw.recv(4096)
                if not data:
                    break
                post(data)
        except Exception:  # noqa: BLE001 - shutdown on teardown lands here
            pass
        finally:
            post(None)

    def writer() -> None:
        while True:
            data = outgoing.get()
            if data is None:
                return
            try:
                raw.sendall(data)
            except Exception:  # noqa: BLE001
                return

    name = exec_id[:12]
    threading.Thread(target=reader, name=f"exec-read-{name}", daemon=True).start()
    threading.Thread(target=writer, name=f"exec-write-{name}", daemon=True).start()

    async def to_browser() -> None:
        # Incremental: a multi-byte character (an umlaut, a box-drawing glyph in
        # htop) split across two reads would otherwise turn into two U+FFFD.
        decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
        try:
            while True:
                data = await incoming.get()
                if data is None:
                    break
                text = decoder.decode(data)
                if text:
                    await websocket.send_text(json.dumps({"type": "data", "data": text}))
            if tail := decoder.decode(b"", final=True):
                await websocket.send_text(json.dumps({"type": "data", "data": tail}))
        except Exception:  # noqa: BLE001
            pass
        finally:
            closed.set()

    async def from_browser() -> None:
        try:
            while True:
                msg = await websocket.receive_text()
                obj = None
                try:
                    obj = json.loads(msg)
                except (json.JSONDecodeError, TypeError):
                    obj = None
                if isinstance(obj, dict) and obj.get("type") == "resize":
                    with contextlib.suppress(Exception):
                        resize_exec(exec_id, int(obj.get("rows", 24)), int(obj.get("cols", 80)))
                elif isinstance(obj, dict) and "data" in obj:
                    outgoing.put(str(obj["data"]).encode())
                else:
                    outgoing.put(msg.encode())
        except WebSocketDisconnect:
            pass
        except Exception:  # noqa: BLE001
            pass
        finally:
            closed.set()

    out_task = asyncio.create_task(to_browser())
    in_task = asyncio.create_task(from_browser())
    await closed.wait()

    outgoing.put(None)
    with contextlib.suppress(Exception):
        raw.shutdown(_socket.SHUT_RDWR)
    with contextlib.suppress(Exception):
        holder.close()
    for task in (out_task, in_task):
        task.cancel()
    with contextlib.suppress(Exception):
        await asyncio.gather(out_task, in_task, return_exceptions=True)

    code = None
    with contextlib.suppress(Exception):
        code = exec_exit_code(exec_id)
    with contextlib.suppress(Exception):
        await websocket.send_text(json.dumps({"type": "exit", "code": code}))