websocket_resilient.py

"""Auto-generated by scripts/sdk-facades.ts. Do not edit.

Auto-reconnect decorator (Decorator pattern) for WebSocket sessions.
Wraps :class:`TerminalSession` / :class:`VncSession` with retry-and-
reconnect: on a non-user-initiated close, the wrapper re-issues the
ticket and re-opens the WebSocket transparently.
"""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass
from typing import Callable, Optional

from .configuration import Configuration
from .websocket import (
    TerminalSession,
    TerminalTarget,
    VncSession,
    VncTarget,
    connect_terminal,
    connect_vnc,
)


@dataclass(frozen=True)
class RetryOptions:
    max_retries: int = 10
    initial_delay_s: float = 0.5
    max_delay_s: float = 30.0
    backoff_multiplier: float = 2.0


def _delay_for(attempt: int, opts: RetryOptions) -> float:
    return min(opts.initial_delay_s * (opts.backoff_multiplier ** attempt), opts.max_delay_s)


class ResilientTerminalSession:
    """Same surface as :class:`TerminalSession`, transparently reconnects."""

    def __init__(
        self,
        config: Configuration,
        target: TerminalTarget,
        on_message: Optional[Callable[[str], None]] = None,
        on_close: Optional[Callable[[int, str], None]] = None,
        on_error: Optional[Callable[[Exception], None]] = None,
        on_reconnect: Optional[Callable[[int], None]] = None,
        on_give_up: Optional[Callable[[Optional[Exception]], None]] = None,
        retry: RetryOptions = RetryOptions(),
    ) -> None:
        self._config = config
        self._target = target
        self._on_message = on_message
        self._on_close = on_close
        self._on_error = on_error
        self._on_reconnect = on_reconnect
        self._on_give_up = on_give_up
        self._retry = retry
        self._closed = False
        self._attempt = 0
        self._session = self._connect()

    def _connect(self) -> TerminalSession:
        return connect_terminal(
            self._config, self._target,
            on_message=self._on_message,
            on_close=self._intercept_close,
            on_error=self._on_error,
        )

    def _intercept_close(self, code: int, reason: str) -> None:
        if self._closed:
            if self._on_close:
                self._on_close(code, reason)
            return
        threading.Thread(target=self._reconnect_loop, name="pve-resilient-reconnect", daemon=True).start()

    def _reconnect_loop(self) -> None:
        last_error: Optional[Exception] = None
        for i in range(self._retry.max_retries):
            self._attempt += 1
            time.sleep(_delay_for(i, self._retry))
            try:
                self._session = self._connect()
                if self._on_reconnect:
                    self._on_reconnect(self._attempt)
                return
            except Exception as exc:  # noqa: BLE001
                last_error = exc
        if self._on_give_up:
            self._on_give_up(last_error)
        if self._on_close:
            self._on_close(1006, "retries exhausted")

    def send(self, text: str) -> None: self._session.send(text)
    def resize(self, cols: int, rows: int) -> None: self._session.resize(cols, rows)

    def close(self) -> None:
        self._closed = True
        self._session.close()


class ResilientVncSession:
    """Same surface as :class:`VncSession`, transparently reconnects."""

    def __init__(
        self,
        config: Configuration,
        target: VncTarget,
        on_frame: Optional[Callable[[bytes], None]] = None,
        on_close: Optional[Callable[[int, str], None]] = None,
        on_error: Optional[Callable[[Exception], None]] = None,
        on_reconnect: Optional[Callable[[int], None]] = None,
        on_give_up: Optional[Callable[[Optional[Exception]], None]] = None,
        retry: RetryOptions = RetryOptions(),
    ) -> None:
        self._config = config
        self._target = target
        self._on_frame = on_frame
        self._on_close = on_close
        self._on_error = on_error
        self._on_reconnect = on_reconnect
        self._on_give_up = on_give_up
        self._retry = retry
        self._closed = False
        self._attempt = 0
        self._session = self._connect()

    def _connect(self) -> VncSession:
        return connect_vnc(
            self._config, self._target,
            on_frame=self._on_frame,
            on_close=self._intercept_close,
            on_error=self._on_error,
        )

    def _intercept_close(self, code: int, reason: str) -> None:
        if self._closed:
            if self._on_close:
                self._on_close(code, reason)
            return
        threading.Thread(target=self._reconnect_loop, name="pve-resilient-reconnect", daemon=True).start()

    def _reconnect_loop(self) -> None:
        last_error: Optional[Exception] = None
        for i in range(self._retry.max_retries):
            self._attempt += 1
            time.sleep(_delay_for(i, self._retry))
            try:
                self._session = self._connect()
                if self._on_reconnect:
                    self._on_reconnect(self._attempt)
                return
            except Exception as exc:  # noqa: BLE001
                last_error = exc
        if self._on_give_up:
            self._on_give_up(last_error)
        if self._on_close:
            self._on_close(1006, "retries exhausted")

    def send(self, data: bytes) -> None: self._session.send(data)

    def close(self) -> None:
        self._closed = True
        self._session.close()


def connect_terminal_resilient(
    config: Configuration,
    target: TerminalTarget,
    **kwargs,
) -> ResilientTerminalSession:
    return ResilientTerminalSession(config, target, **kwargs)


def connect_vnc_resilient(
    config: Configuration,
    target: VncTarget,
    **kwargs,
) -> ResilientVncSession:
    return ResilientVncSession(config, target, **kwargs)