diff --git a/src/remote/remote_session_manager.py b/src/remote/remote_session_manager.py new file mode 100644 index 000000000..9083ce4df --- /dev/null +++ b/src/remote/remote_session_manager.py @@ -0,0 +1,293 @@ +"""High-level CCR session manager (claude.ai → CCR viewer/control). + +Ports ``typescript/src/remote/RemoteSessionManager.ts:95-324``. + +Wraps ``SessionsWebSocket`` + adds: + + - Per-``request_id`` tracking of pending permission requests + (so ``control_cancel_request`` can find the right tool_use_id). + - Permission round-trip: ``respond_to_permission_request`` builds + the ``control_response`` envelope (allow→updated_input, + deny→message). + - ``cancel_session`` sends a ``control_request`` with subtype + ``interrupt`` (skipped in viewer-only mode). + - ``viewer_only`` flag — accepts the flag; ``cancel_session`` is a + no-op when set. Other viewer-only behaviors (no 60 s response-stuck + timeout, no title updates) live in the front-end consumer per A16. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field + +from src.bridge.messaging import RemotePermissionResponse +from src.bridge.sdk_types import SDKControlPermissionRequest + +from .sessions_websocket import ( + SessionsWebSocket, + SessionsWebSocketCallbacks, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class RemoteSessionConfig: + """Construction-time config. Mirrors TS at lines 50-62.""" + + session_id: str + get_access_token: Callable[[], str] + org_uuid: str + has_initial_prompt: bool = False + #: Viewer-only mode (``claude assistant``): cancel_session is a no-op. + viewer_only: bool = False + + +@dataclass +class RemoteSessionCallbacks: + """Public callbacks. Mirrors TS at lines 64-85.""" + + on_message: Callable[[dict], None | Awaitable[None]] + on_permission_request: Callable[ + [SDKControlPermissionRequest, str], None | Awaitable[None] + ] + on_permission_cancelled: ( + Callable[[str, str | None], None | Awaitable[None]] | None + ) = None + on_connected: Callable[[], None | Awaitable[None]] | None = None + on_disconnected: Callable[[], None | Awaitable[None]] | None = None + on_reconnecting: Callable[[], None | Awaitable[None]] | None = None + on_error: Callable[[Exception], None | Awaitable[None]] | None = None + + +class RemoteSessionManager: + """High-level façade over ``SessionsWebSocket`` for one CCR session.""" + + def __init__( + self, + config: RemoteSessionConfig, + callbacks: RemoteSessionCallbacks, + *, + base_url: str = 'wss://api.anthropic.com', + ) -> None: + self._config = config + self._callbacks = callbacks + self._base_url = base_url + self._websocket: SessionsWebSocket | None = None + self._pending_permission_requests: dict[str, SDKControlPermissionRequest] = {} + + # ─── Public API ──────────────────────────────────────────────────── + + def connect(self) -> None: + """Open the WS to ``/v1/sessions/ws/{id}/subscribe``. + + Mirrors ``RemoteSessionManager.ts:108-141``. Construction + + connect are split so callers can wire callbacks first. + """ + ws_callbacks = SessionsWebSocketCallbacks( + on_message=self._handle_message, + on_connected=self._callbacks.on_connected, + on_close=self._callbacks.on_disconnected, + on_reconnecting=self._callbacks.on_reconnecting, + on_error=self._callbacks.on_error, + ) + self._websocket = SessionsWebSocket( + self._config.session_id, + self._config.org_uuid, + self._config.get_access_token, + ws_callbacks, + base_url=self._base_url, + ) + # Schedule the connect; SessionsWebSocket.connect is async. + loop = asyncio.get_running_loop() + loop.create_task(self._websocket.connect(), name='remote-session-connect') + + def is_connected(self) -> bool: + return self._websocket is not None and self._websocket.is_connected() + + async def send_message(self, content: object, *, uuid: str | None = None) -> bool: + """Send a user prompt to the remote session. + + TS (``RemoteSessionManager.ts:219-242``) routes via HTTP POST to + ``sendEventToRemoteSession`` (a teleport API). For Phase 4 we + send via the WS itself for protocol simplicity; the wire shape + matches what the agent expects on its stdin (``stream-json``). + Returns False if the WS is not connected. + """ + if self._websocket is None or not self._websocket.is_connected(): + return False + envelope = { + 'type': 'user', + 'message': {'role': 'user', 'content': content}, + 'parent_tool_use_id': None, + 'session_id': self._config.session_id, + } + if uuid is not None: + envelope['uuid'] = uuid + # Re-use the WS' send-control-response method to push raw JSON; + # the WS doesn't distinguish between control envelopes and SDK + # messages on send. + await self._websocket.send_control_response(envelope) + return True + + async def respond_to_permission_request( + self, + request_id: str, + result: RemotePermissionResponse, + ) -> None: + """Reply to a ``can_use_tool`` request. + + Mirrors ``RemoteSessionManager.ts:247-282``. Builds the + ``control_response`` envelope with the allow/deny payload. + """ + if request_id not in self._pending_permission_requests: + logger.warning( + '[RemoteSessionManager] no pending permission request: %s', + request_id, + ) + return + del self._pending_permission_requests[request_id] + + if result.behavior == 'allow': + response_payload: dict[str, object] = { + 'behavior': 'allow', + 'updatedInput': getattr(result, 'updated_input', {}), + } + else: + response_payload = { + 'behavior': 'deny', + 'message': getattr(result, 'message', ''), + } + envelope = { + 'type': 'control_response', + 'response': { + 'subtype': 'success', + 'request_id': request_id, + 'response': response_payload, + }, + } + if self._websocket is not None: + await self._websocket.send_control_response(envelope) + + async def cancel_session(self) -> None: + """Send an interrupt to the remote agent. No-op in viewer-only mode. + + Mirrors ``RemoteSessionManager.ts:294-297``. Viewer-only + gating per A16: viewer_only callers must NOT send interrupts. + """ + if self._config.viewer_only: + logger.debug( + '[RemoteSessionManager] viewer_only: cancel_session is a no-op' + ) + return + if self._websocket is None: + return + await self._websocket.send_control_request({'subtype': 'interrupt'}) + + async def disconnect(self) -> None: + """Close the WS and clear pending permission requests.""" + if self._websocket is not None: + await self._websocket.disconnect() + self._websocket = None + self._pending_permission_requests.clear() + + async def reconnect(self) -> None: + """Force-reconnect the WS (e.g., after stale subscription).""" + if self._websocket is not None: + await self._websocket.reconnect() + + def get_session_id(self) -> str: + return self._config.session_id + + # ─── Internal: WS message routing ──────────────────────────────── + + async def _handle_message(self, message: dict) -> None: + """Route messages from the WS to the right consumer callback.""" + msg_type = message.get('type') + + if msg_type == 'control_request': + await self._handle_control_request(message) + return + if msg_type == 'control_cancel_request': + await self._handle_control_cancel(message) + return + if msg_type == 'control_response': + logger.debug('[RemoteSessionManager] received control_response') + return + # All other SDK message types → on_message. + await self._invoke(self._callbacks.on_message, message) + + async def _handle_control_request(self, message: dict) -> None: + request_id = message.get('request_id') + inner = message.get('request') + if not isinstance(request_id, str) or not isinstance(inner, dict): + return + subtype = inner.get('subtype') + if subtype == 'can_use_tool': + self._pending_permission_requests[request_id] = inner # type: ignore[assignment] + await self._invoke( + self._callbacks.on_permission_request, + inner, # type: ignore[arg-type] + request_id, + ) + return + # Unknown subtype — send error so the server doesn't hang. + logger.debug( + '[RemoteSessionManager] unsupported control request subtype: %s', subtype + ) + if self._websocket is not None: + await self._websocket.send_control_response({ + 'type': 'control_response', + 'response': { + 'subtype': 'error', + 'request_id': request_id, + 'error': f'Unsupported control request subtype: {subtype}', + }, + }) + + async def _handle_control_cancel(self, message: dict) -> None: + request_id = message.get('request_id') + if not isinstance(request_id, str): + return + pending = self._pending_permission_requests.pop(request_id, None) + tool_use_id = pending.get('tool_use_id') if pending else None + await self._invoke( + self._callbacks.on_permission_cancelled, request_id, tool_use_id + ) + + @staticmethod + async def _invoke(callback: Callable[..., object] | None, /, *args: object) -> None: + if callback is None: + return + result = callback(*args) + if asyncio.iscoroutine(result): + await result + + +def create_remote_session_config( + session_id: str, + get_access_token: Callable[[], str], + org_uuid: str, + *, + has_initial_prompt: bool = False, + viewer_only: bool = False, +) -> RemoteSessionConfig: + """Convenience helper mirroring ``createRemoteSessionConfig`` in TS.""" + return RemoteSessionConfig( + session_id=session_id, + get_access_token=get_access_token, + org_uuid=org_uuid, + has_initial_prompt=has_initial_prompt, + viewer_only=viewer_only, + ) + + +__all__ = [ + 'RemoteSessionCallbacks', + 'RemoteSessionConfig', + 'RemoteSessionManager', + 'create_remote_session_config', +] diff --git a/src/remote/sdk_message_adapter.py b/src/remote/sdk_message_adapter.py new file mode 100644 index 000000000..c3b080c64 --- /dev/null +++ b/src/remote/sdk_message_adapter.py @@ -0,0 +1,112 @@ +"""SDK-message format adapter for the remote-session bridge. + +Ports the **functional surface** of +``typescript/src/remote/sdkMessageAdapter.ts (302 lines)``: translate +between the SDK message format used in the local REPL and the bridge +wire format. The TS source is large because it handles many edge cases +(streamlined messages, partial assistant chunks, tool-result back-refs, +etc.); we port the **canonical translation paths** the chapter calls +out, not the long tail of edge cases. + +Two directions: + + - **Wire → SDK**: incoming messages from the WS are normalized into + the SDK shape (bridge wire format may use camelCase from older + server versions; we use ``normalize_control_message_keys`` to fix). + - **SDK → Wire**: outgoing user messages are wrapped in the + ``stream-json`` envelope shape the agent expects. + +Tests pin the canonical shapes; new edge cases land here as they're +discovered. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from src.bridge.messaging import normalize_control_message_keys + +logger = logging.getLogger(__name__) + + +def adapt_wire_to_sdk(wire: dict[str, Any]) -> dict[str, Any]: + """Normalize a wire payload to SDK-canonical shape. + + Only known camelCase keys are converted to snake_case; unknown keys + pass through (matches ``normalize_control_message_keys``). + """ + result = normalize_control_message_keys(wire) + if not isinstance(result, dict): + return wire + return result + + +def adapt_sdk_to_wire_user_message( + content: object, + session_id: str, + *, + parent_tool_use_id: str | None = None, + uuid: str | None = None, +) -> dict[str, Any]: + """Build the stream-json envelope for a user prompt. + + Wire shape: + { + 'type': 'user', + 'message': {'role': 'user', 'content': content}, + 'parent_tool_use_id': parent_tool_use_id, + 'session_id': session_id, + 'uuid': uuid, # optional + } + """ + envelope: dict[str, Any] = { + 'type': 'user', + 'message': {'role': 'user', 'content': content}, + 'parent_tool_use_id': parent_tool_use_id, + 'session_id': session_id, + } + if uuid is not None: + envelope['uuid'] = uuid + return envelope + + +def adapt_permission_response( + request_id: str, + behavior: str, + *, + updated_input: dict | None = None, + message: str | None = None, +) -> dict[str, Any]: + """Build the control_response envelope for a permission decision. + + Mirrors ``RemoteSessionManager.respondToPermissionRequest`` payload + construction. + """ + if behavior == 'allow': + inner_response: dict[str, Any] = { + 'behavior': 'allow', + 'updatedInput': updated_input or {}, + } + elif behavior == 'deny': + inner_response = { + 'behavior': 'deny', + 'message': message or '', + } + else: + raise ValueError(f'unknown permission behavior: {behavior!r}') + return { + 'type': 'control_response', + 'response': { + 'subtype': 'success', + 'request_id': request_id, + 'response': inner_response, + }, + } + + +__all__ = [ + 'adapt_permission_response', + 'adapt_sdk_to_wire_user_message', + 'adapt_wire_to_sdk', +] diff --git a/src/remote/sessions_websocket.py b/src/remote/sessions_websocket.py new file mode 100644 index 000000000..a970c7201 --- /dev/null +++ b/src/remote/sessions_websocket.py @@ -0,0 +1,345 @@ +"""WebSocket client for ``/v1/sessions/ws/{id}/subscribe`` (claude.ai → CCR). + +Ports ``typescript/src/remote/SessionsWebSocket.ts:82-404``. + +Reconnection strategy is **discriminated by close code** (per chapter +§"Remote Session Management"): + + - **4003 (unauthorized)**: stop immediately. Permanent rejection. + - **4001 (session not found)**: max 3 retries with linear backoff + (transient during compaction). + - **other transient**: max 5 attempts with constant-delay retry + (``RECONNECT_DELAY_SECONDS = 2.0``). [chapter, unverified: chapter + says "exponential backoff", but the actual TS code uses a constant + 2 s delay; we match the code.] + +Per Risk #22 in the refactoring plan: the ``websockets`` library does +NOT auto-reconnect (unlike JS WebSocket event handlers). We implement +an explicit reconnect loop with the discriminated retry strategy. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import uuid as _uuid +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field + +import websockets +from websockets.asyncio.client import connect as ws_connect + +from src.bridge.close_codes import ( + WS_CLOSE_PERMANENT_UNAUTHORIZED, + WS_CLOSE_SESSION_NOT_FOUND, +) + +logger = logging.getLogger(__name__) + +#: Constant retry delay matching ``SessionsWebSocket.ts:17``. +RECONNECT_DELAY_SECONDS = 2.0 + +#: Max retries for the generic-transient path +#: (``SessionsWebSocket.ts:18`` ``MAX_RECONNECT_ATTEMPTS``). +MAX_RECONNECT_ATTEMPTS = 5 + +#: Max retries for the 4001 session-not-found path +#: (``SessionsWebSocket.ts:26`` ``MAX_SESSION_NOT_FOUND_RETRIES``). +MAX_SESSION_NOT_FOUND_RETRIES = 3 + +#: Application-level ping interval (websockets library handles +#: protocol-level pings; this constant matches TS for parity). +PING_INTERVAL_SECONDS = 30.0 + + +GetAccessToken = Callable[[], str] + + +@dataclass +class SessionsWebSocketCallbacks: + """Per-session event callbacks. Mirrors TS at lines 57-65. + + All callbacks may be sync or async; sync is just called, async is + awaited via ``asyncio.iscoroutine`` check. + """ + + on_message: Callable[[dict], None | Awaitable[None]] + on_close: Callable[[], None | Awaitable[None]] | None = None + on_error: Callable[[Exception], None | Awaitable[None]] | None = None + on_connected: Callable[[], None | Awaitable[None]] | None = None + on_reconnecting: Callable[[], None | Awaitable[None]] | None = None + + +def _is_sessions_message(value: object) -> bool: + """Permissive type guard: any dict with a string ``type`` field. + + Mirrors ``SessionsWebSocket.ts:46-55``. Deliberately permissive so + new server-side message types don't get silently dropped before the + Python client is updated. + """ + return isinstance(value, dict) and isinstance(value.get('type'), str) + + +class SessionsWebSocket: + """WS client for the claude.ai → CCR session-subscribe stream. + + Lifecycle: + ws = SessionsWebSocket(session_id, org_uuid, get_token, callbacks) + await ws.connect() — opens the WS, spawns reader task + await ws.send_control_response(resp) + await ws.send_control_request(inner) + ws.is_connected() + await ws.disconnect() — stop reconnecting, close the WS + await ws.reconnect() — force-close + immediately retry + """ + + def __init__( + self, + session_id: str, + org_uuid: str, + get_access_token: GetAccessToken, + callbacks: SessionsWebSocketCallbacks, + *, + base_url: str = 'wss://api.anthropic.com', + anthropic_version: str = '2023-06-01', + ) -> None: + self._session_id = session_id + self._org_uuid = org_uuid + self._get_access_token = get_access_token + self._callbacks = callbacks + self._base_url = base_url.rstrip('/') + self._anthropic_version = anthropic_version + + self._ws: websockets.asyncio.client.ClientConnection | None = None + self._reader_task: asyncio.Task[None] | None = None + self._reconnect_task: asyncio.Task[None] | None = None + self._closed = False + self._reconnect_attempts = 0 + self._not_found_retries = 0 + + # ─── State queries ──────────────────────────────────────────────── + + def is_connected(self) -> bool: + return self._ws is not None and not self._closed + + # ─── Lifecycle ──────────────────────────────────────────────────── + + async def connect(self) -> None: + """Open the WS and start the reader. Returns once open OR raises.""" + if self._closed: + return + url = ( + f'{self._base_url}/v1/sessions/ws/{self._session_id}/subscribe' + f'?organization_uuid={self._org_uuid}' + ) + headers = { + 'Authorization': f'Bearer {self._get_access_token()}', + 'anthropic-version': self._anthropic_version, + } + try: + ws = await ws_connect( + url, + additional_headers=headers, + ping_interval=PING_INTERVAL_SECONDS, + ) + except (websockets.exceptions.WebSocketException, OSError) as exc: + await self._invoke(self._callbacks.on_error, exc) + self._schedule_reconnect_or_close() + return + + self._ws = ws + self._reconnect_attempts = 0 + self._not_found_retries = 0 + await self._invoke(self._callbacks.on_connected) + self._reader_task = asyncio.get_running_loop().create_task( + self._read_loop(), name='sessions-ws-reader', + ) + + async def disconnect(self) -> None: + """Permanent close — no reconnects, fire on_close.""" + self._closed = True + if self._reconnect_task is not None and not self._reconnect_task.done(): + self._reconnect_task.cancel() + ws, self._ws = self._ws, None + if self._reader_task is not None and not self._reader_task.done(): + self._reader_task.cancel() + try: + await self._reader_task + except (asyncio.CancelledError, websockets.exceptions.ConnectionClosed): + pass + if ws is not None: + try: + await ws.close() + except (websockets.exceptions.ConnectionClosed, OSError): + pass + + async def reconnect(self) -> None: + """Force close + immediate reconnect (resets retry counters). + + Used when the subscription is known stale (e.g., after the + worker container shutdown the server detected on its side). + """ + self._reconnect_attempts = 0 + self._not_found_retries = 0 + ws, self._ws = self._ws, None + if ws is not None: + try: + await ws.close() + except (websockets.exceptions.ConnectionClosed, OSError): + pass + if self._reader_task is not None and not self._reader_task.done(): + self._reader_task.cancel() + try: + await self._reader_task + except (asyncio.CancelledError, websockets.exceptions.ConnectionClosed): + pass + # 500 ms grace before reconnecting (matches TS). + await asyncio.sleep(0.5) + await self.connect() + + # ─── Send API ───────────────────────────────────────────────────── + + async def send_control_response(self, response: dict) -> None: + """Send a ``control_response`` (e.g., permission decision).""" + if self._ws is None or self._closed: + logger.warning('[SessionsWebSocket] cannot send: not connected') + return + try: + await self._ws.send(json.dumps(response)) + except (websockets.exceptions.ConnectionClosed, OSError) as exc: + logger.warning('[SessionsWebSocket] send failed: %s', exc) + + async def send_control_request(self, inner: dict) -> None: + """Wrap ``inner`` in a ``control_request`` envelope and send.""" + if self._ws is None or self._closed: + logger.warning('[SessionsWebSocket] cannot send: not connected') + return + envelope = { + 'type': 'control_request', + 'request_id': str(_uuid.uuid4()), + 'request': inner, + } + try: + await self._ws.send(json.dumps(envelope)) + except (websockets.exceptions.ConnectionClosed, OSError) as exc: + logger.warning('[SessionsWebSocket] send failed: %s', exc) + + # ─── Read loop ──────────────────────────────────────────────────── + + async def _read_loop(self) -> None: + ws = self._ws + assert ws is not None + try: + async for raw in ws: + text = raw if isinstance(raw, str) else raw.decode('utf-8', errors='replace') + try: + parsed = json.loads(text) + except json.JSONDecodeError as exc: + logger.debug('[SessionsWebSocket] parse error: %s', exc) + continue + if not _is_sessions_message(parsed): + logger.debug('[SessionsWebSocket] ignoring non-message: %s', parsed) + continue + await self._invoke(self._callbacks.on_message, parsed) + except websockets.exceptions.ConnectionClosed as exc: + close_code = exc.rcvd.code if exc.rcvd else None + self._on_ws_close(close_code) + except Exception as exc: # noqa: BLE001 + await self._invoke(self._callbacks.on_error, exc) + self._on_ws_close(None) + + def _on_ws_close(self, close_code: int | None) -> None: + """Discriminated reconnect strategy. + + 4003 → permanent stop, fire on_close. + 4001 → max 3 retries with linear backoff. + other → max 5 attempts with constant delay. + """ + if self._closed: + return + self._ws = None + + if close_code == WS_CLOSE_PERMANENT_UNAUTHORIZED: + logger.debug( + '[SessionsWebSocket] permanent close 4003; not reconnecting' + ) + self._closed = True + asyncio.get_running_loop().create_task( + self._invoke(self._callbacks.on_close) + ) + return + + if close_code == WS_CLOSE_SESSION_NOT_FOUND: + self._not_found_retries += 1 + if self._not_found_retries > MAX_SESSION_NOT_FOUND_RETRIES: + logger.debug( + '[SessionsWebSocket] 4001 retry budget exhausted; not reconnecting' + ) + self._closed = True + asyncio.get_running_loop().create_task( + self._invoke(self._callbacks.on_close) + ) + return + delay = RECONNECT_DELAY_SECONDS * self._not_found_retries + self._schedule_reconnect_with_delay(delay) + return + + # Generic transient. + self._reconnect_attempts += 1 + if self._reconnect_attempts > MAX_RECONNECT_ATTEMPTS: + logger.debug('[SessionsWebSocket] not reconnecting (budget exhausted)') + self._closed = True + asyncio.get_running_loop().create_task( + self._invoke(self._callbacks.on_close) + ) + return + self._schedule_reconnect_with_delay(RECONNECT_DELAY_SECONDS) + + def _schedule_reconnect_or_close(self) -> None: + """Initial-connect failure path: same retry strategy as transient close.""" + if self._closed: + return + self._reconnect_attempts += 1 + if self._reconnect_attempts > MAX_RECONNECT_ATTEMPTS: + self._closed = True + asyncio.get_running_loop().create_task( + self._invoke(self._callbacks.on_close) + ) + return + self._schedule_reconnect_with_delay(RECONNECT_DELAY_SECONDS) + + def _schedule_reconnect_with_delay(self, delay_seconds: float) -> None: + async def _do_reconnect() -> None: + await self._invoke(self._callbacks.on_reconnecting) + try: + await asyncio.sleep(delay_seconds) + except asyncio.CancelledError: + return + if self._closed: + return + await self.connect() + + loop = asyncio.get_running_loop() + self._reconnect_task = loop.create_task(_do_reconnect(), name='sessions-ws-reconnect') + + # ─── Internal: callback helper ─────────────────────────────────── + + @staticmethod + async def _invoke(callback: Callable[..., object] | None, /, *args: object) -> None: + if callback is None: + return + result = callback(*args) + if asyncio.iscoroutine(result): + await result + + +__all__ = [ + 'GetAccessToken', + 'MAX_RECONNECT_ATTEMPTS', + 'MAX_SESSION_NOT_FOUND_RETRIES', + 'PING_INTERVAL_SECONDS', + 'RECONNECT_DELAY_SECONDS', + 'SessionsWebSocket', + 'SessionsWebSocketCallbacks', +] diff --git a/tests/remote/test_remote_session_manager.py b/tests/remote/test_remote_session_manager.py new file mode 100644 index 000000000..60c51d4b7 --- /dev/null +++ b/tests/remote/test_remote_session_manager.py @@ -0,0 +1,292 @@ +"""Tests for ``src.remote.remote_session_manager.RemoteSessionManager``.""" + +from __future__ import annotations + +import asyncio +import json +import socket + +import pytest +import websockets +from websockets.asyncio.server import serve as ws_serve + +from src.bridge.messaging import AllowResponse, DenyResponse +from src.remote.remote_session_manager import ( + RemoteSessionCallbacks, + RemoteSessionConfig, + RemoteSessionManager, + create_remote_session_config, +) + + +pytestmark = pytest.mark.integration + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(('127.0.0.1', 0)) + return s.getsockname()[1] + + +class _TestEnv: + def __init__(self): + self.received_messages: list[dict] = [] + self.permission_requests: list[tuple[dict, str]] = [] + self.permission_cancellations: list[tuple[str, str | None]] = [] + self.connected = asyncio.Event() + self.disconnected = asyncio.Event() + self.from_client: list[str] = [] + self.to_client_queue: asyncio.Queue[str] = asyncio.Queue() + + async def server_handler(self, ws): + async def out(): + try: + while True: + line = await self.to_client_queue.get() + await ws.send(line) + except websockets.exceptions.ConnectionClosed: + return + + async def inn(): + try: + async for raw in ws: + self.from_client.append(raw if isinstance(raw, str) else raw.decode()) + except websockets.exceptions.ConnectionClosed: + return + + out_t = asyncio.get_running_loop().create_task(out()) + in_t = asyncio.get_running_loop().create_task(inn()) + try: + await asyncio.wait([out_t, in_t], return_when=asyncio.FIRST_COMPLETED) + finally: + for t in (out_t, in_t): + if not t.done(): + t.cancel() + try: + await t + except (asyncio.CancelledError, websockets.exceptions.ConnectionClosed): + pass + + def make_callbacks(self) -> RemoteSessionCallbacks: + return RemoteSessionCallbacks( + on_message=lambda m: self.received_messages.append(m), + on_permission_request=lambda req, rid: self.permission_requests.append((req, rid)), + on_permission_cancelled=lambda rid, tuid: self.permission_cancellations.append((rid, tuid)), + on_connected=lambda: self.connected.set(), + on_disconnected=lambda: self.disconnected.set(), + ) + + +@pytest.mark.asyncio +async def test_assistant_message_forwarded(): + env = _TestEnv() + port = _free_port() + ws_server = await ws_serve(env.server_handler, '127.0.0.1', port) + try: + config = RemoteSessionConfig( + session_id='cse_x', get_access_token=lambda: 'tok', org_uuid='org', + ) + mgr = RemoteSessionManager(config, env.make_callbacks(), base_url=f'ws://127.0.0.1:{port}') + mgr.connect() + await asyncio.wait_for(env.connected.wait(), timeout=2.0) + await env.to_client_queue.put(json.dumps({ + 'type': 'assistant', 'message': {'content': 'hello'}, + })) + for _ in range(50): + if env.received_messages: + break + await asyncio.sleep(0.02) + assert env.received_messages[0]['type'] == 'assistant' + await mgr.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed() + + +@pytest.mark.asyncio +async def test_permission_request_routed_and_response_sent(): + env = _TestEnv() + port = _free_port() + ws_server = await ws_serve(env.server_handler, '127.0.0.1', port) + try: + config = RemoteSessionConfig( + session_id='cse_x', get_access_token=lambda: 'tok', org_uuid='org', + ) + mgr = RemoteSessionManager(config, env.make_callbacks(), base_url=f'ws://127.0.0.1:{port}') + mgr.connect() + await asyncio.wait_for(env.connected.wait(), timeout=2.0) + + # Send a can_use_tool request from the server. + await env.to_client_queue.put(json.dumps({ + 'type': 'control_request', + 'request_id': 'r1', + 'request': { + 'subtype': 'can_use_tool', + 'tool_name': 'Bash', + 'input': {'command': 'ls'}, + 'tool_use_id': 'tu1', + }, + })) + for _ in range(50): + if env.permission_requests: + break + await asyncio.sleep(0.02) + assert len(env.permission_requests) == 1 + req, rid = env.permission_requests[0] + assert req['tool_name'] == 'Bash' + assert rid == 'r1' + + # Respond allow. + await mgr.respond_to_permission_request('r1', AllowResponse(updated_input={'command': 'ls -la'})) + await asyncio.sleep(0.1) + sent = json.loads(env.from_client[-1]) + assert sent['type'] == 'control_response' + assert sent['response']['response']['behavior'] == 'allow' + assert sent['response']['response']['updatedInput'] == {'command': 'ls -la'} + await mgr.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed() + + +@pytest.mark.asyncio +async def test_permission_cancel_routed_with_tool_use_id(): + env = _TestEnv() + port = _free_port() + ws_server = await ws_serve(env.server_handler, '127.0.0.1', port) + try: + config = RemoteSessionConfig( + session_id='cse_x', get_access_token=lambda: 'tok', org_uuid='org', + ) + mgr = RemoteSessionManager(config, env.make_callbacks(), base_url=f'ws://127.0.0.1:{port}') + mgr.connect() + await asyncio.wait_for(env.connected.wait(), timeout=2.0) + + # Send the permission request first so the manager has a pending entry. + await env.to_client_queue.put(json.dumps({ + 'type': 'control_request', + 'request_id': 'r1', + 'request': { + 'subtype': 'can_use_tool', + 'tool_name': 'Bash', + 'input': {}, + 'tool_use_id': 'tu_99', + }, + })) + for _ in range(50): + if env.permission_requests: + break + await asyncio.sleep(0.02) + + # Now send a control_cancel_request; manager should fire on_permission_cancelled. + await env.to_client_queue.put(json.dumps({ + 'type': 'control_cancel_request', + 'request_id': 'r1', + })) + for _ in range(50): + if env.permission_cancellations: + break + await asyncio.sleep(0.02) + assert env.permission_cancellations == [('r1', 'tu_99')] + await mgr.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed() + + +@pytest.mark.asyncio +async def test_unknown_subtype_returns_error_response(): + env = _TestEnv() + port = _free_port() + ws_server = await ws_serve(env.server_handler, '127.0.0.1', port) + try: + config = RemoteSessionConfig( + session_id='cse_x', get_access_token=lambda: 'tok', org_uuid='org', + ) + mgr = RemoteSessionManager(config, env.make_callbacks(), base_url=f'ws://127.0.0.1:{port}') + mgr.connect() + await asyncio.wait_for(env.connected.wait(), timeout=2.0) + await env.to_client_queue.put(json.dumps({ + 'type': 'control_request', + 'request_id': 'r2', + 'request': {'subtype': 'set_quantum_flux'}, + })) + await asyncio.sleep(0.1) + # Manager should have sent an error response back. + assert any('error' in line for line in env.from_client) + await mgr.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed() + + +@pytest.mark.asyncio +async def test_cancel_session_no_op_in_viewer_only(): + env = _TestEnv() + port = _free_port() + ws_server = await ws_serve(env.server_handler, '127.0.0.1', port) + try: + config = RemoteSessionConfig( + session_id='cse_x', + get_access_token=lambda: 'tok', + org_uuid='org', + viewer_only=True, + ) + mgr = RemoteSessionManager(config, env.make_callbacks(), base_url=f'ws://127.0.0.1:{port}') + mgr.connect() + await asyncio.wait_for(env.connected.wait(), timeout=2.0) + await mgr.cancel_session() # no-op + await asyncio.sleep(0.1) + # Server received nothing — viewer-only suppressed the interrupt. + assert env.from_client == [] + await mgr.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed() + + +@pytest.mark.asyncio +async def test_cancel_session_sends_interrupt_when_not_viewer(): + env = _TestEnv() + port = _free_port() + ws_server = await ws_serve(env.server_handler, '127.0.0.1', port) + try: + config = RemoteSessionConfig( + session_id='cse_x', + get_access_token=lambda: 'tok', + org_uuid='org', + viewer_only=False, + ) + mgr = RemoteSessionManager(config, env.make_callbacks(), base_url=f'ws://127.0.0.1:{port}') + mgr.connect() + await asyncio.wait_for(env.connected.wait(), timeout=2.0) + await mgr.cancel_session() + await asyncio.sleep(0.1) + assert any('interrupt' in line for line in env.from_client) + await mgr.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed() + + +@pytest.mark.asyncio +async def test_send_message_returns_false_when_not_connected(): + config = RemoteSessionConfig( + session_id='cse_x', get_access_token=lambda: 'tok', org_uuid='org', + ) + callbacks = RemoteSessionCallbacks( + on_message=lambda m: None, + on_permission_request=lambda req, rid: None, + ) + mgr = RemoteSessionManager(config, callbacks) + ok = await mgr.send_message('hi') + assert ok is False + + +def test_create_remote_session_config_helper(): + cfg = create_remote_session_config( + 'cse_x', lambda: 'tok', 'org', viewer_only=True, + ) + assert cfg.session_id == 'cse_x' + assert cfg.viewer_only is True + assert cfg.org_uuid == 'org' diff --git a/tests/remote/test_sdk_message_adapter.py b/tests/remote/test_sdk_message_adapter.py new file mode 100644 index 000000000..21f5c7992 --- /dev/null +++ b/tests/remote/test_sdk_message_adapter.py @@ -0,0 +1,84 @@ +"""Tests for ``src.remote.sdk_message_adapter``.""" + +from __future__ import annotations + +import pytest + +from src.remote.sdk_message_adapter import ( + adapt_permission_response, + adapt_sdk_to_wire_user_message, + adapt_wire_to_sdk, +) + + +class TestAdaptWireToSdk: + def test_normalizes_camel_case_keys(self) -> None: + out = adapt_wire_to_sdk({ + 'type': 'control_request', + 'requestId': 'r1', + 'request': {'subtype': 'set_model', 'model': 'opus'}, + }) + assert out == { + 'type': 'control_request', + 'request_id': 'r1', + 'request': {'subtype': 'set_model', 'model': 'opus'}, + } + + def test_passes_through_snake_case(self) -> None: + inp = {'type': 'user', 'message': {'content': 'hi'}} + assert adapt_wire_to_sdk(inp) == inp + + +class TestAdaptSdkToWireUserMessage: + def test_minimal(self) -> None: + env = adapt_sdk_to_wire_user_message('hello', session_id='cse_x') + assert env == { + 'type': 'user', + 'message': {'role': 'user', 'content': 'hello'}, + 'parent_tool_use_id': None, + 'session_id': 'cse_x', + } + + def test_with_uuid(self) -> None: + env = adapt_sdk_to_wire_user_message('hi', session_id='s', uuid='u1') + assert env['uuid'] == 'u1' + + def test_with_parent_tool_use_id(self) -> None: + env = adapt_sdk_to_wire_user_message( + 'sub-prompt', session_id='s', parent_tool_use_id='tu1', + ) + assert env['parent_tool_use_id'] == 'tu1' + + def test_list_content_passed_through(self) -> None: + content = [{'type': 'text', 'text': 'hi'}] + env = adapt_sdk_to_wire_user_message(content, session_id='s') + assert env['message']['content'] == content + + +class TestAdaptPermissionResponse: + def test_allow_with_updated_input(self) -> None: + env = adapt_permission_response( + 'r1', 'allow', updated_input={'command': 'ls -la'}, + ) + assert env['type'] == 'control_response' + assert env['response']['subtype'] == 'success' + assert env['response']['request_id'] == 'r1' + assert env['response']['response']['behavior'] == 'allow' + assert env['response']['response']['updatedInput'] == {'command': 'ls -la'} + + def test_allow_default_empty_input(self) -> None: + env = adapt_permission_response('r1', 'allow') + assert env['response']['response']['updatedInput'] == {} + + def test_deny_with_message(self) -> None: + env = adapt_permission_response('r1', 'deny', message='too risky') + assert env['response']['response']['behavior'] == 'deny' + assert env['response']['response']['message'] == 'too risky' + + def test_deny_default_empty_message(self) -> None: + env = adapt_permission_response('r1', 'deny') + assert env['response']['response']['message'] == '' + + def test_unknown_behavior_raises(self) -> None: + with pytest.raises(ValueError, match='unknown permission behavior'): + adapt_permission_response('r1', 'maybe') diff --git a/tests/remote/test_sessions_websocket.py b/tests/remote/test_sessions_websocket.py new file mode 100644 index 000000000..c79fafdd4 --- /dev/null +++ b/tests/remote/test_sessions_websocket.py @@ -0,0 +1,305 @@ +"""Tests for ``src.remote.sessions_websocket.SessionsWebSocket``. + +Uses an in-process WS echo/script server with controllable close codes +to exercise the discriminated reconnection strategy. +""" + +from __future__ import annotations + +import asyncio +import json +import socket + +import pytest +import websockets +from websockets.asyncio.server import serve as ws_serve + +from src.remote.sessions_websocket import ( + MAX_RECONNECT_ATTEMPTS, + MAX_SESSION_NOT_FOUND_RETRIES, + SessionsWebSocket, + SessionsWebSocketCallbacks, +) + + +pytestmark = pytest.mark.integration + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(('127.0.0.1', 0)) + return s.getsockname()[1] + + +class _ScriptedServer: + """In-process WS server with per-connection close-code control. + + On each connection, immediately closes with the next ``close_codes`` + code. We do NOT wait for client messages — the test purpose is to + exercise the client's reconnection logic, not message exchange. + """ + + def __init__(self, close_codes: list[int]) -> None: + self.close_codes = list(close_codes) + self.connection_count = 0 + self.received_messages: list[str] = [] + + async def handler(self, ws): + self.connection_count += 1 + # Close immediately with the next code — don't wait for messages. + if self.close_codes: + code = self.close_codes.pop(0) + try: + await ws.close(code=code, reason='scripted close') + except (websockets.exceptions.ConnectionClosed, OSError): + pass + + +@pytest.mark.asyncio +async def test_4003_is_permanent_no_reconnect(): + server = _ScriptedServer(close_codes=[4003]) + port = _free_port() + ws_server = await ws_serve(server.handler, '127.0.0.1', port) + try: + on_close_fired = asyncio.Event() + callbacks = SessionsWebSocketCallbacks( + on_message=lambda m: None, + on_close=lambda: on_close_fired.set(), + ) + ws = SessionsWebSocket( + 'sid', 'org', lambda: 'tok', callbacks, + base_url=f'ws://127.0.0.1:{port}', + ) + await ws.connect() + # Server immediately closes with 4003. + await asyncio.wait_for(on_close_fired.wait(), timeout=2.0) + # Only one connection attempt — 4003 is permanent. + assert server.connection_count == 1 + await ws.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed() + + +@pytest.mark.asyncio +async def test_4001_triggers_reconnect_with_on_reconnecting(): + """A single 4001 close fires on_reconnecting then attempts to reconnect. + + NOTE: matches TS — both ``sessionNotFoundRetries`` and + ``reconnectAttempts`` are reset on every successful WS open, so a + server that keeps sending 4001 forever loops indefinitely. The + realistic test is "verify the 4001 path triggers a reconnect at all" + rather than budget exhaustion. + """ + server = _ScriptedServer(close_codes=[4001, 4001]) + port = _free_port() + ws_server = await ws_serve(server.handler, '127.0.0.1', port) + try: + reconnecting_fired = asyncio.Event() + callbacks = SessionsWebSocketCallbacks( + on_message=lambda m: None, + on_reconnecting=lambda: reconnecting_fired.set(), + ) + from src.remote import sessions_websocket as ws_mod + original_delay = ws_mod.RECONNECT_DELAY_SECONDS + ws_mod.RECONNECT_DELAY_SECONDS = 0.02 + try: + ws = SessionsWebSocket( + 'sid', 'org', lambda: 'tok', callbacks, + base_url=f'ws://127.0.0.1:{port}', + ) + await ws.connect() + await asyncio.wait_for(reconnecting_fired.wait(), timeout=5.0) + # Wait for the second connect to happen. + for _ in range(100): + if server.connection_count >= 2: + break + await asyncio.sleep(0.02) + assert server.connection_count >= 2 + finally: + ws_mod.RECONNECT_DELAY_SECONDS = original_delay + await ws.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed() + + +@pytest.mark.asyncio +async def test_initial_connect_failure_exhausts_budget(): + """When the server is unreachable from the start, the initial-connect + retry budget exhausts at MAX_RECONNECT_ATTEMPTS. This is the realistic + "budget exhaustion" path because no successful connect resets the + counter. + """ + # Use a bound-then-released port so connect() fails with refused. + port = _free_port() # binds + releases — port is now likely refused + + on_close_fired = asyncio.Event() + error_count = 0 + + def on_error(_exc): + nonlocal error_count + error_count += 1 + + callbacks = SessionsWebSocketCallbacks( + on_message=lambda m: None, + on_close=lambda: on_close_fired.set(), + on_error=on_error, + ) + from src.remote import sessions_websocket as ws_mod + original_delay = ws_mod.RECONNECT_DELAY_SECONDS + ws_mod.RECONNECT_DELAY_SECONDS = 0.02 + try: + ws = SessionsWebSocket( + 'sid', 'org', lambda: 'tok', callbacks, + base_url=f'ws://127.0.0.1:{port}', + ) + await ws.connect() + await asyncio.wait_for(on_close_fired.wait(), timeout=10.0) + finally: + ws_mod.RECONNECT_DELAY_SECONDS = original_delay + + # Each failed initial connect goes through _schedule_reconnect_or_close + # which increments _reconnect_attempts; budget exhausts at MAX+1. + assert error_count >= MAX_RECONNECT_ATTEMPTS + await ws.disconnect() + + +@pytest.mark.asyncio +async def test_message_dispatch_to_on_message(): + received: list[dict] = [] + + class _MsgServer: + def __init__(self): + self.connection_count = 0 + + async def handler(self, ws): + self.connection_count += 1 + await ws.send(json.dumps({'type': 'assistant', 'message': {'content': 'hi'}})) + await asyncio.sleep(0.1) # keep connection open briefly + + server = _MsgServer() + port = _free_port() + ws_server = await ws_serve(server.handler, '127.0.0.1', port) + try: + callbacks = SessionsWebSocketCallbacks( + on_message=lambda m: received.append(m), + ) + ws = SessionsWebSocket( + 'sid', 'org', lambda: 'tok', callbacks, + base_url=f'ws://127.0.0.1:{port}', + ) + await ws.connect() + for _ in range(50): + if received: + break + await asyncio.sleep(0.02) + assert any(m.get('type') == 'assistant' for m in received) + await ws.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed() + + +@pytest.mark.asyncio +async def test_invalid_json_is_silently_dropped(): + """Server sends garbage; reader logs + drops, doesn't fire on_message.""" + received: list[dict] = [] + + class _GarbageServer: + async def handler(self, ws): + await ws.send('not json {{{') + await ws.send(json.dumps({'type': 'good'})) + await asyncio.sleep(0.1) + + server = _GarbageServer() + port = _free_port() + ws_server = await ws_serve(server.handler, '127.0.0.1', port) + try: + callbacks = SessionsWebSocketCallbacks( + on_message=lambda m: received.append(m), + ) + ws = SessionsWebSocket( + 'sid', 'org', lambda: 'tok', callbacks, + base_url=f'ws://127.0.0.1:{port}', + ) + await ws.connect() + for _ in range(50): + if received: + break + await asyncio.sleep(0.02) + # Only the good message should make it through. + assert received == [{'type': 'good'}] + await ws.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed() + + +@pytest.mark.asyncio +async def test_send_control_response_round_trip(): + server_received: list[str] = [] + + class _EchoServer: + async def handler(self, ws): + try: + async for raw in ws: + server_received.append(raw if isinstance(raw, str) else raw.decode()) + except websockets.exceptions.ConnectionClosed: + pass + + server = _EchoServer() + port = _free_port() + ws_server = await ws_serve(server.handler, '127.0.0.1', port) + try: + callbacks = SessionsWebSocketCallbacks(on_message=lambda m: None) + ws = SessionsWebSocket( + 'sid', 'org', lambda: 'tok', callbacks, + base_url=f'ws://127.0.0.1:{port}', + ) + await ws.connect() + await ws.send_control_response({ + 'type': 'control_response', + 'response': {'subtype': 'success', 'request_id': 'r1'}, + }) + await asyncio.sleep(0.1) + assert any('control_response' in line for line in server_received) + await ws.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed() + + +@pytest.mark.asyncio +async def test_send_control_request_wraps_in_envelope(): + server_received: list[str] = [] + + class _EchoServer: + async def handler(self, ws): + try: + async for raw in ws: + server_received.append(raw if isinstance(raw, str) else raw.decode()) + except websockets.exceptions.ConnectionClosed: + pass + + server = _EchoServer() + port = _free_port() + ws_server = await ws_serve(server.handler, '127.0.0.1', port) + try: + callbacks = SessionsWebSocketCallbacks(on_message=lambda m: None) + ws = SessionsWebSocket( + 'sid', 'org', lambda: 'tok', callbacks, + base_url=f'ws://127.0.0.1:{port}', + ) + await ws.connect() + await ws.send_control_request({'subtype': 'interrupt'}) + await asyncio.sleep(0.1) + assert len(server_received) == 1 + envelope = json.loads(server_received[0]) + assert envelope['type'] == 'control_request' + assert envelope['request']['subtype'] == 'interrupt' + assert 'request_id' in envelope + await ws.disconnect() + finally: + ws_server.close() + await ws_server.wait_closed()