From 43d720641a94467488838f1e80e17edf053fc559 Mon Sep 17 00:00:00 2001 From: Priya Singh Date: Wed, 29 Jul 2026 23:46:29 +0000 Subject: [PATCH] fix(runtime): update shell session wire protocol --- src/bedrock_agentcore/runtime/shell/config.py | 16 +- .../runtime/shell/protocol.py | 9 - .../runtime/shell/session.py | 151 +------ tests/unit/runtime/test_shell.py | 403 ++++-------------- tests/unit/runtime/test_shell_protocol.py | 4 - 5 files changed, 106 insertions(+), 477 deletions(-) diff --git a/src/bedrock_agentcore/runtime/shell/config.py b/src/bedrock_agentcore/runtime/shell/config.py index 489b03ce..a0336d79 100644 --- a/src/bedrock_agentcore/runtime/shell/config.py +++ b/src/bedrock_agentcore/runtime/shell/config.py @@ -6,7 +6,6 @@ _DEFAULT_MAX_RETRIES = 5 _DEFAULT_BASE_DELAY = 1.0 _DEFAULT_MAX_DELAY = 15.0 -_DEFAULT_METADATA_TIMEOUT = 10.0 _DEFAULT_RECONNECT_WINDOW = 900.0 # ~15 min — matches server-side KARP idle timeout _DEFAULT_OUTER_LOOP_DELAY = 30.0 # wait between inner loop exhaustion and next outer attempt @@ -36,16 +35,11 @@ class ReconnectConfig: outer_loop_delay: Seconds to wait between inner loop exhaustion and the next outer retry cycle. on_reconnect: Optional async or sync callback invoked after each - successful reconnect. Receives ``reconnected: bool`` — ``True`` - means the existing PTY was reattached (buffered output will follow - as STDOUT frames); ``False`` means a fresh shell was started. + successful reconnect. Takes no arguments. - Example — log reconnects and flush buffered output to a file: - async def on_reconnect(reconnected: bool) -> None: - if reconnected: - print("Reattached to existing PTY — replaying buffered output") - else: - print("New shell started") + Example — log reconnects: + async def on_reconnect() -> None: + print("Reconnected to shell") config = ReconnectConfig(reconnect_window=None, on_reconnect=on_reconnect) # None = unlimited async with client.open_shell(arn, reconnect_config=config) as shell: @@ -58,4 +52,4 @@ async def on_reconnect(reconnected: bool) -> None: max_delay: float = _DEFAULT_MAX_DELAY reconnect_window: Optional[float] = _DEFAULT_RECONNECT_WINDOW outer_loop_delay: float = _DEFAULT_OUTER_LOOP_DELAY - on_reconnect: Optional[Callable[[bool], Union[Awaitable[None], None]]] = field(default=None, repr=False) + on_reconnect: Optional[Callable[[], Union[Awaitable[None], None]]] = field(default=None, repr=False) diff --git a/src/bedrock_agentcore/runtime/shell/protocol.py b/src/bedrock_agentcore/runtime/shell/protocol.py index a5dbdedc..88d24c0f 100644 --- a/src/bedrock_agentcore/runtime/shell/protocol.py +++ b/src/bedrock_agentcore/runtime/shell/protocol.py @@ -109,7 +109,6 @@ class ShellFramer: # Encode outbound frames ws.send(framer.encode_stdin("ls /workspace\\n")) ws.send(framer.encode_resize(220, 50)) - ws.send(framer.encode_close()) """ MAX_FRAME_SIZE = 64 * 1024 # matches DP WebSocketFlowController limit @@ -204,11 +203,3 @@ def encode_heartbeat(self) -> bytes: Binary WebSocket frame ready to send. """ return bytes([ShellChannel.HEARTBEAT]) - - def encode_close(self) -> bytes: - """Encode a graceful-shutdown CLOSE frame (empty payload). - - Returns: - Binary WebSocket frame ready to send. - """ - return bytes([ShellChannel.CLOSE]) diff --git a/src/bedrock_agentcore/runtime/shell/session.py b/src/bedrock_agentcore/runtime/shell/session.py index 4a9f1e66..435dc8d9 100644 --- a/src/bedrock_agentcore/runtime/shell/session.py +++ b/src/bedrock_agentcore/runtime/shell/session.py @@ -5,8 +5,7 @@ import logging import random import uuid -from collections import deque -from typing import TYPE_CHECKING, AsyncIterator, Deque, Optional +from typing import TYPE_CHECKING, AsyncIterator, Optional import websockets import websockets.exceptions @@ -14,7 +13,7 @@ from ..models import SESSION_HEADER, SHELL_ID_HEADER from ._validation import parse_runtime_arn, validate_shell_id from .auth import AuthMode, OAuthAuth, PresignedAuth -from .config import _DEFAULT_METADATA_TIMEOUT, ReconnectConfig +from .config import ReconnectConfig from .protocol import ShellChannel, ShellFrame, ShellFramer if TYPE_CHECKING: @@ -26,8 +25,7 @@ class ShellSession: r"""Async context manager wrapping a live interactive shell WebSocket session. - Connects on ``__aenter__``, reads the mandatory metadata frame that carries - ``shellId`` and ``reconnected``, and exposes typed send/resize/iterate/close. + Connects on ``__aenter__`` and exposes typed send/resize/iterate/close. When ``reconnect_config`` is provided, transparently reconnects on unexpected disconnects using the same ``shell_id`` and ``session_id`` so the shell's working directory, environment, background jobs, and up to 256 KB of buffered @@ -86,9 +84,8 @@ class ShellSession: async with client.open_shell( runtime_arn, shell_id=shell_id, session_id=session_id ) as shell: - assert shell.reconnected # True → up to 256 KB buffered output follows async for frame in shell: - ... + ... # resumes from the same PTY Attributes: shell_id: Confirmed shell identifier echoed by the server in @@ -100,18 +97,10 @@ class ShellSession: when reconnecting across process restarts — passing a different (or omitted) session ID may cause the platform to provision a fresh VM where the PTY no longer exists. - reconnected: ``True`` when the session resumed an existing PTY (buffered - output will arrive as STDOUT frames immediately after connect); - ``False`` for a fresh shell. kicked: ``True`` when iteration stopped because another client connected with the same ``shell_id`` (close code 4000). The PTY is still alive — a new ``open_shell`` call with the same ID will reconnect to it. - bytes_dropped: Number of bytes lost from the PTY ring buffer during the - most recent disconnect. Non-zero only when the 256 KB ring buffer - overflowed before reconnection completed. Set after the post-drain - STATUS confirmation frame arrives (which follows the buffered STDOUT - burst). Zero if no overflow occurred or on a fresh connection. exit_code: Exit code of the shell process. ``None`` until the shell exits or if the platform terminated the session without providing an exit code (e.g. an InternalError). ``0`` for a clean exit; @@ -157,24 +146,19 @@ def __init__( self._framer = ShellFramer() self._ws: Optional[object] = None self._closed = False - # Frames received during _read_metadata_frame that arrived before the - # 0x03 confirmation (e.g. first shell prompt on 0x01). - self._pending_frames: Deque[ShellFrame] = deque() if shell_id is not None: validate_shell_id(shell_id) # Auto-generate stable reconnect handles when the caller omits them. # Without a fixed session_id, each _connect() would route to a different - # VM and shell_id would never be found → reconnected=False always. + # VM and shell_id would never be found → reconnect would always fail. self.shell_id: str = shell_id or str(uuid.uuid4()) self.session_id: str = session_id or str(uuid.uuid4()) - self.reconnected: bool = False self.kicked: bool = False - self.bytes_dropped: int = 0 self.exit_code: Optional[int] = None async def __aenter__(self) -> "ShellSession": - """Connect and read the initial metadata frame.""" + """Connect to the shell WebSocket.""" try: await self._connect() except Exception as exc: @@ -191,9 +175,7 @@ async def __aexit__(self, *_: object) -> None: # ── Connection management ───────────────────────────────────────────────── async def _connect(self) -> None: - """Open the WebSocket and consume the initial STATUS metadata frame.""" - self._pending_frames.clear() - self.reconnected = False + """Open the WebSocket and extract shellId from 101 response headers.""" self.kicked = False auth = self._auth @@ -249,75 +231,7 @@ async def _connect(self) -> None: self.session_id = header_sid logger.debug("sessionId from 101 header: %r", header_sid) - await self._read_metadata_frame() - if self._ws is not None: - self._closed = False - - async def _read_metadata_frame(self) -> None: - """Consume the first STATUS frame carrying connection confirmation. - - Both the 0x03 confirmation and the first 0x01 stdout - frame are sent after upgrade but their order is non-deterministic. We - wait for a 0x03 frame with metadata.shellId. Any 0x01 frames - that arrive first are stashed in self._pending_frames so __anext__ can - yield them in order. - """ - loop = asyncio.get_running_loop() - deadline = loop.time() + _DEFAULT_METADATA_TIMEOUT - while True: - remaining = deadline - loop.time() - if remaining <= 0: - logger.warning( - "STATUS confirmation not received within %.1fs (deadline exceeded " - "processing earlier frames). Proceeding with client-generated " - "shell_id=%r — reconnected flag may be incorrect.", - _DEFAULT_METADATA_TIMEOUT, - self.shell_id, - ) - return - try: - raw = await asyncio.wait_for(self._ws.recv(), timeout=remaining) - except asyncio.TimeoutError: - logger.warning( - "Timed out waiting for STATUS confirmation after %.1fs " - "(server did not respond). Proceeding with client-generated " - "shell_id=%r — reconnected flag may be incorrect.", - _DEFAULT_METADATA_TIMEOUT, - self.shell_id, - ) - return - except (websockets.exceptions.ConnectionClosedOK, websockets.exceptions.ConnectionClosedError) as exc: - # Server closed before sending STATUS confirmation — session never became usable. - logger.warning("WebSocket closed before STATUS confirmation (shell_id=%r): %s", self.shell_id, exc) - self._ws = None - self._closed = True - raise - except Exception as exc: - logger.error("Unexpected error waiting for STATUS frame (shell_id=%r): %s", self.shell_id, exc) - self._ws = None - self._closed = True - raise - if not isinstance(raw, bytes): - continue - frame = self._framer.decode(raw) - if frame.channel == ShellChannel.STATUS: - try: - status = frame.json() - meta = status.get("metadata", {}) - if meta.get("shellId"): - # This is the connection confirmation frame. - self.shell_id = meta["shellId"] - self.reconnected = bool(meta.get("reconnected", False)) - return - else: - # Termination status — stash for __anext__. - self._pending_frames.append(frame) - return - except (json.JSONDecodeError, KeyError): - logger.warning("Received malformed STATUS frame during connect; skipping: %r", raw) - else: - # Non-STATUS frame (e.g. first stdout prompt) — stash for __anext__. - self._pending_frames.append(frame) + self._closed = False async def _run_inner_retry_loop(self, cfg: ReconnectConfig) -> bool: """Run one inner retry loop (up to max_retries attempts with exponential backoff). @@ -332,9 +246,9 @@ async def _run_inner_retry_loop(self, cfg: ReconnectConfig) -> bool: logger.info("Reconnect attempt %d (shell_id=%s)", attempt, self.shell_id) try: await self._connect() - logger.info("Reconnected (reconnected=%s)", self.reconnected) + logger.info("Reconnected (shell_id=%s)", self.shell_id) if cfg.on_reconnect is not None: - result = cfg.on_reconnect(self.reconnected) + result = cfg.on_reconnect() if asyncio.iscoroutine(result): await result return True @@ -464,14 +378,9 @@ async def resize(self, width: int, height: int) -> None: raise async def close(self) -> None: - """Send a graceful CLOSE frame and close the underlying WebSocket.""" + """Close the underlying WebSocket (shell detaches, stays alive for reconnect window).""" self._closed = True if self._ws is not None: - try: - await self._ws.send(self._framer.encode_close()) - except Exception as exc: - # Best-effort — the connection may already be gone. - logger.debug("Failed to send CLOSE frame during close() (shell_id=%r): %s", self.shell_id, exc) try: await self._ws.close() except Exception as exc: @@ -531,8 +440,7 @@ async def __anext__(self) -> ShellFrame: automatic reconnect attempt using the same ``shell_id``. The iterator resumes transparently — callers do not need to re-enter the context manager. The ``on_reconnect`` callback fires after each - successful reconnect so callers can react to the ``reconnected`` flag - and the incoming buffered-output burst. + successful reconnect. close code 4000 ("kicked by new connection") MUST NOT trigger auto-reconnect. The iterator stops and sets ``self.kicked = True`` so callers can distinguish this case. @@ -543,29 +451,6 @@ async def __anext__(self) -> ShellFrame: called, or when reconnect attempts are exhausted. """ while True: - # Drain any frames buffered during the metadata handshake first. - if self._pending_frames: - frame = self._pending_frames.popleft() - if frame.channel == ShellChannel.CLOSE: - logger.debug( - "CLOSE frame received from pending queue (shell_id=%r)", - self.shell_id, - ) - raise StopAsyncIteration from None - if frame.channel == ShellChannel.STATUS: - try: - status = frame.json() - if self._is_termination_status(status): - self.exit_code = self._parse_exit_code(status) - self._closed = True - return frame - except (json.JSONDecodeError, KeyError): - logger.warning( - "Received malformed STATUS frame in pending queue; skipping termination check: %r", - frame, - ) - return frame - if self._closed or self._ws is None: logger.debug("Session already closed or disconnected (shell_id=%r)", self.shell_id) raise StopAsyncIteration from None @@ -651,17 +536,7 @@ async def __anext__(self) -> ShellFrame: try: status = frame.json() if self._is_confirmation_status(status): - # Second confirmation frame (post-drain) — carries - # bytesDropped when the 256 KB ring buffer overflowed. - # Swallow it; surface bytesDropped via attribute + warning. - dropped = status.get("metadata", {}).get("bytesDropped", 0) - if dropped: - self.bytes_dropped = dropped - logger.warning( - "%d bytes of PTY output lost during disconnect (ring buffer overflow) (shell_id=%r)", - dropped, - self.shell_id, - ) + # Confirmation frame — silently swallow. continue if self._is_termination_status(status): # Shell exited — mark closed so the next __anext__ call diff --git a/tests/unit/runtime/test_shell.py b/tests/unit/runtime/test_shell.py index 4d4abd96..ab6f7651 100644 --- a/tests/unit/runtime/test_shell.py +++ b/tests/unit/runtime/test_shell.py @@ -1,6 +1,5 @@ """Tests for ShellSession and ReconnectConfig.""" -import asyncio import json from unittest.mock import AsyncMock, MagicMock, patch @@ -97,98 +96,36 @@ def _close_frame() -> bytes: class TestShellSessionConnect: @pytest.mark.asyncio - async def test_connect_reads_metadata_frame(self): + async def test_connect_sets_shell_id_from_101_header(self): client = _make_client() - ws = _make_ws(_metadata_frame("my-shell", reconnected=False)) + ws = _make_ws() + ws.response.headers = {SHELL_ID_HEADER: "server-shell"} with patch("websockets.connect", new=AsyncMock(return_value=ws)): session = ShellSession(client, FAKE_ARN, shell_id="my-shell") await session._connect() - assert session.shell_id == "my-shell" - assert session.reconnected is False + assert session.shell_id == "server-shell" @pytest.mark.asyncio - async def test_connect_sets_reconnected_true(self): + async def test_connect_ready_immediately_after_websocket_open(self): + """Connection is ready immediately after WebSocket opens — no blocking on STATUS frame.""" client = _make_client() - ws = _make_ws(_metadata_frame("my-shell", reconnected=True)) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)): session = ShellSession(client, FAKE_ARN, shell_id="my-shell") await session._connect() - assert session.reconnected is True - - @pytest.mark.asyncio - async def test_connect_tolerates_missing_metadata_frame(self): - client = _make_client() - ws = AsyncMock() - ws.recv = AsyncMock(side_effect=asyncio.TimeoutError) - ws.send = AsyncMock() - ws.close = AsyncMock() - - with patch("websockets.connect", new=AsyncMock(return_value=ws)): - session = ShellSession(client, FAKE_ARN) - await session._connect() # must not raise - - assert session.shell_id is not None - - @pytest.mark.asyncio - async def test_connect_ignores_non_status_first_frame(self): - """STDOUT frames arriving before STATUS are stashed; STATUS is still found.""" - client = _make_client() - ws = _make_ws(_stdout_frame("some output"), _metadata_frame("x")) - - with patch("websockets.connect", new=AsyncMock(return_value=ws)): - session = ShellSession(client, FAKE_ARN, shell_id="x") - await session._connect() - - assert session.reconnected is False - assert session.shell_id == "x" - assert len(session._pending_frames) == 1 # stdout frame was stashed - - @pytest.mark.asyncio - async def test_connect_raises_on_connection_closed_error_before_status(self): - """ConnectionClosedError before STATUS arrives propagates out of __aenter__.""" - import websockets.exceptions - - client = _make_client() - ws = AsyncMock() - ws.recv = AsyncMock(side_effect=websockets.exceptions.ConnectionClosedError(None, None)) - ws.send = AsyncMock() - ws.close = AsyncMock() - - with patch("websockets.connect", new=AsyncMock(return_value=ws)): - session = ShellSession(client, FAKE_ARN) - with pytest.raises(websockets.exceptions.ConnectionClosedError): - await session._connect() - - assert session._ws is None - - @pytest.mark.asyncio - async def test_connect_raises_on_connection_closed_ok_before_status(self): - """ConnectionClosedOK before STATUS arrives propagates out of __aenter__.""" - import websockets.exceptions - - client = _make_client() - ws = AsyncMock() - ws.recv = AsyncMock(side_effect=websockets.exceptions.ConnectionClosedOK(None, None)) - ws.send = AsyncMock() - ws.close = AsyncMock() - - with patch("websockets.connect", new=AsyncMock(return_value=ws)): - session = ShellSession(client, FAKE_ARN) - with pytest.raises(websockets.exceptions.ConnectionClosedOK): - await session._connect() - - assert session._ws is None + assert session._closed is False + assert session._ws is ws @pytest.mark.asyncio async def test_session_id_stable_across_two_connect_calls(self): """session_id must not change between _connect() calls — same value routes to same VM.""" client = _make_client() - ws1 = _make_ws(_metadata_frame("my-shell")) - ws2 = _make_ws(_metadata_frame("my-shell")) + ws1 = _make_ws() + ws2 = _make_ws() with patch("websockets.connect", new=AsyncMock(side_effect=[ws1, ws2])): session = ShellSession(client, FAKE_ARN, shell_id="my-shell") @@ -204,7 +141,7 @@ async def test_session_id_stable_across_two_connect_calls(self): async def test_connect_reads_session_id_from_101_header(self): """X-Amzn-Bedrock-AgentCore-Runtime-Session-Id in 101 response updates session_id.""" client = _make_client() - ws = _make_ws(_metadata_frame("my-shell")) + ws = _make_ws() ws.response.headers = { SHELL_ID_HEADER: "my-shell", SESSION_HEADER: "server-session-99", @@ -216,27 +153,6 @@ async def test_connect_reads_session_id_from_101_header(self): assert session.session_id == "server-session-99" - @pytest.mark.asyncio - async def test_connect_timeout_proceeds_with_warning(self, caplog): - """TimeoutError waiting for STATUS logs a warning but session stays usable.""" - import logging - - client = _make_client() - ws = AsyncMock() - ws.recv = AsyncMock(side_effect=asyncio.TimeoutError) - ws.send = AsyncMock() - ws.close = AsyncMock() - ws.response = MagicMock() - ws.response.headers = {} - - with patch("websockets.connect", new=AsyncMock(return_value=ws)): - with caplog.at_level(logging.WARNING, logger="bedrock_agentcore.runtime.shell"): - session = ShellSession(client, FAKE_ARN) - await session._connect() # must not raise - - assert session._ws is ws # connection still alive - assert "server did not respond" in caplog.text - class TestShellSessionInit: def test_invalid_arn_raises_at_construction(self): @@ -337,7 +253,7 @@ class TestShellSessionContextManager: @pytest.mark.asyncio async def test_context_manager_closes_on_exit(self): client = _make_client() - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: @@ -346,17 +262,17 @@ async def test_context_manager_closes_on_exit(self): ws.close.assert_called_once() @pytest.mark.asyncio - async def test_close_sends_close_frame(self): - framer = ShellFramer() + async def test_close_does_not_send_close_frame(self): + """close() no longer sends 0xFF — shell detaches and stays alive for reconnect window.""" client = _make_client() - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as _: pass - sent_frames = [call.args[0] for call in ws.send.call_args_list] - assert framer.encode_close() in sent_frames + # No frames should have been sent (no encode_close) + ws.send.assert_not_called() class TestShellSessionSend: @@ -364,7 +280,7 @@ class TestShellSessionSend: async def test_send_encodes_stdin_frame(self): framer = ShellFramer() client = _make_client() - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: @@ -377,7 +293,7 @@ async def test_send_encodes_stdin_frame(self): async def test_send_bytes(self): framer = ShellFramer() client = _make_client() - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: @@ -390,7 +306,7 @@ async def test_send_bytes(self): async def test_resize(self): framer = ShellFramer() client = _make_client() - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: @@ -404,7 +320,7 @@ class TestShellSessionSendAfterClose: @pytest.mark.asyncio async def test_send_raises_after_close(self): client = _make_client() - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: pass @@ -414,7 +330,7 @@ async def test_send_raises_after_close(self): @pytest.mark.asyncio async def test_send_bytes_raises_after_close(self): client = _make_client() - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: pass @@ -424,7 +340,7 @@ async def test_send_bytes_raises_after_close(self): @pytest.mark.asyncio async def test_resize_raises_after_close(self): client = _make_client() - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: pass @@ -437,7 +353,6 @@ class TestShellSessionIterate: async def test_iterates_stdout_frames(self): client = _make_client() ws = _make_ws( - _metadata_frame(), _stdout_frame("hello"), _stdout_frame(" world"), _close_frame(), @@ -452,11 +367,30 @@ async def test_iterates_stdout_frames(self): assert output == ["hello", " world"] + @pytest.mark.asyncio + async def test_confirmation_frame_swallowed_during_iteration(self): + """Confirmation STATUS frames (metadata.shellId present) are silently swallowed.""" + client = _make_client() + ws = _make_ws( + _metadata_frame("s"), + _stdout_frame("output"), + _close_frame(), + ) + + frames = [] + with patch("websockets.connect", new=AsyncMock(return_value=ws)): + async with ShellSession(client, FAKE_ARN) as shell: + async for frame in shell: + frames.append(frame) + + # Only STDOUT frame yielded — confirmation frame was swallowed + assert len(frames) == 1 + assert frames[0].channel == ShellChannel.STDOUT + @pytest.mark.asyncio async def test_stops_on_exit_status_frame(self): client = _make_client() ws = _make_ws( - _metadata_frame(), _stdout_frame("output"), _exit_frame(0), ) @@ -478,7 +412,7 @@ async def test_stops_on_exit_status_frame(self): @pytest.mark.asyncio async def test_stops_on_close_frame(self): client = _make_client() - ws = _make_ws(_metadata_frame(), _close_frame()) + ws = _make_ws(_close_frame()) frames = [] with patch("websockets.connect", new=AsyncMock(return_value=ws)): @@ -491,7 +425,7 @@ async def test_stops_on_close_frame(self): @pytest.mark.asyncio async def test_stops_on_connection_closed_without_reconnect(self): client = _make_client() - ws = _make_ws(_metadata_frame()) # ConnectionClosed raised on next recv + ws = _make_ws() # ConnectionClosed raised on next recv frames = [] with patch("websockets.connect", new=AsyncMock(return_value=ws)): @@ -511,19 +445,16 @@ async def test_connection_closed_ok_stops_without_reconnect(self): client = _make_client() reconnect_calls = [] - async def on_reconnect(reconnected: bool) -> None: - reconnect_calls.append(reconnected) + async def on_reconnect() -> None: + reconnect_calls.append(True) ws = AsyncMock() ws.send = AsyncMock() ws.close = AsyncMock() - call_count = 0 + ws.response = MagicMock() + ws.response.headers = {} async def recv(): - nonlocal call_count - call_count += 1 - if call_count == 1: - return _metadata_frame() raise websockets.exceptions.ConnectionClosedOK(None, None) ws.recv = recv @@ -546,20 +477,20 @@ async def test_termination_status_prevents_reconnect(self): client = _make_client() reconnect_calls = [] - async def on_reconnect(reconnected: bool) -> None: - reconnect_calls.append(reconnected) + async def on_reconnect() -> None: + reconnect_calls.append(True) ws = AsyncMock() ws.send = AsyncMock() ws.close = AsyncMock() + ws.response = MagicMock() + ws.response.headers = {} call_count = 0 async def recv(): nonlocal call_count call_count += 1 if call_count == 1: - return _metadata_frame() - if call_count == 2: return _exit_frame(0) # If _closed wasn't set, __anext__ would hit this and try to reconnect raise websockets.exceptions.ConnectionClosedOK(None, None) @@ -588,21 +519,18 @@ async def test_close_code_1001_triggers_reconnect(self): client = _make_client() reconnect_calls = [] - async def on_reconnect(reconnected: bool) -> None: - reconnect_calls.append(reconnected) + async def on_reconnect() -> None: + reconnect_calls.append(True) - # First WebSocket: metadata then close 1001. - # Second WebSocket: metadata then clean close 1000. + # First WebSocket: close 1001 immediately. + # Second WebSocket: stdout then clean close. ws1 = AsyncMock() ws1.send = AsyncMock() ws1.close = AsyncMock() - ws1_count = 0 + ws1.response = MagicMock() + ws1.response.headers = {} async def recv1(): - nonlocal ws1_count - ws1_count += 1 - if ws1_count == 1: - return _metadata_frame("session-1") raise websockets.exceptions.ConnectionClosedOK(Close(1001, "Going Away"), None) ws1.recv = recv1 @@ -610,14 +538,14 @@ async def recv1(): ws2 = AsyncMock() ws2.send = AsyncMock() ws2.close = AsyncMock() + ws2.response = MagicMock() + ws2.response.headers = {} ws2_count = 0 async def recv2(): nonlocal ws2_count ws2_count += 1 if ws2_count == 1: - return _metadata_frame("session-1", reconnected=True) - if ws2_count == 2: return _stdout_frame("hello") raise websockets.exceptions.ConnectionClosedOK(None, None) @@ -648,13 +576,10 @@ async def test_close_code_4000_stops_without_reconnect(self): ws = AsyncMock() ws.send = AsyncMock() ws.close = AsyncMock() - call_count = 0 + ws.response = MagicMock() + ws.response.headers = {} async def recv(): - nonlocal call_count - call_count += 1 - if call_count == 1: - return _metadata_frame() # Simulate kicked close code 4000 close_obj = Close(code=4000, reason="replaced by new connection") raise websockets.exceptions.ConnectionClosedError(close_obj, None) @@ -681,13 +606,10 @@ async def test_close_code_4000_sets_kicked(self): ws = AsyncMock() ws.send = AsyncMock() ws.close = AsyncMock() - call_count = 0 + ws.response = MagicMock() + ws.response.headers = {} async def recv(): - nonlocal call_count - call_count += 1 - if call_count == 1: - return _metadata_frame() close_obj = Close(code=4000, reason="replaced by new connection") raise websockets.exceptions.ConnectionClosedError(close_obj, None) @@ -704,7 +626,7 @@ async def recv(): async def test_clean_close_does_not_set_kicked(self): """shell.kicked remains False when the session ends normally.""" client = _make_client() - ws = _make_ws(_metadata_frame()) # metadata frame then ConnectionClosedOK + ws = _make_ws() # ConnectionClosedOK raised on next recv with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: async for _ in shell: @@ -721,11 +643,10 @@ async def test_reconnects_and_resumes_iteration(self): client = _make_client() ws1 = _make_ws( - _metadata_frame("s", reconnected=False), _stdout_frame("before"), end_with=websockets.exceptions.ConnectionClosedError(None, None), ) - ws2 = _make_ws(_metadata_frame("s", reconnected=True), _stdout_frame("after"), _close_frame()) + ws2 = _make_ws(_stdout_frame("after"), _close_frame()) connect_calls = [ws1, ws2] @@ -734,8 +655,8 @@ async def fake_connect(url, extra_headers=None, additional_headers=None, **_kw): on_reconnect_calls = [] - async def on_reconnect(reconnected: bool) -> None: - on_reconnect_calls.append(reconnected) + async def on_reconnect() -> None: + on_reconnect_calls.append(True) config = ReconnectConfig(max_retries=3, on_reconnect=on_reconnect) output = [] @@ -752,13 +673,7 @@ async def on_reconnect(reconnected: bool) -> None: @pytest.mark.asyncio async def test_exhausts_retries_and_stops(self): client = _make_client() - ws = _make_ws(_metadata_frame()) # drops immediately after metadata - - async def fail_connect(url, extra_headers=None): - raise ConnectionRefusedError("server down") - - async def first_connect(url, extra_headers=None): - return ws + ws = _make_ws() # drops immediately call_count = 0 @@ -780,49 +695,14 @@ async def fake_connect(url, extra_headers=None, additional_headers=None, **_kw): # All reconnect attempts failed — iteration stopped gracefully assert frames == [] - @pytest.mark.asyncio - async def test_pending_frames_cleared_on_reconnect(self): - """_pending_frames must be cleared at the start of _connect() so frames - buffered during a previous metadata handshake cannot bleed into a new session. - - Directly verifies the invariant: after __aenter__ plants a stale frame in - _pending_frames, calling _connect() again must empty it. - """ - client = _make_client() - - from bedrock_agentcore.runtime.shell.protocol import ShellFrame - - ws1 = _make_ws(_metadata_frame("s")) - ws2 = _make_ws(_metadata_frame("s", reconnected=True), _stdout_frame("fresh"), _close_frame()) - connect_seq = [ws1, ws2] - - async def fake_connect(url, extra_headers=None, additional_headers=None, **_kw): - return connect_seq.pop(0) - - with patch("websockets.connect", side_effect=fake_connect): - session = ShellSession(client, FAKE_ARN, shell_id="s") - await session.__aenter__() - - # Plant a stale frame as if left over from a prior metadata handshake. - session._pending_frames.append( - ShellFrame(channel=ShellChannel.STDOUT, raw_channel_byte=ShellChannel.STDOUT, payload=b"stale") - ) - assert len(session._pending_frames) == 1 - - # _connect() must clear it. - await session._connect() - assert len(session._pending_frames) == 0 - - await session.__aexit__(None, None, None) - @pytest.mark.asyncio async def test_sync_on_reconnect_callback_accepted(self): """A synchronous on_reconnect callback must also be accepted.""" import websockets.exceptions client = _make_client() - ws1 = _make_ws(_metadata_frame("s"), end_with=websockets.exceptions.ConnectionClosedError(None, None)) - ws2 = _make_ws(_metadata_frame("s", reconnected=True), _close_frame()) + ws1 = _make_ws(end_with=websockets.exceptions.ConnectionClosedError(None, None)) + ws2 = _make_ws(_close_frame()) connect_calls = [ws1, ws2] async def fake_connect(url, extra_headers=None, additional_headers=None, **_kw): @@ -830,8 +710,8 @@ async def fake_connect(url, extra_headers=None, additional_headers=None, **_kw): sync_calls = [] - def sync_callback(reconnected: bool) -> None: - sync_calls.append(reconnected) + def sync_callback() -> None: + sync_calls.append(True) config = ReconnectConfig(max_retries=1, base_delay=0.0, on_reconnect=sync_callback) @@ -848,8 +728,8 @@ async def test_outer_loop_retries_after_inner_exhaustion(self): import websockets.exceptions client = _make_client() - ws1 = _make_ws(_metadata_frame("s"), end_with=websockets.exceptions.ConnectionClosedError(None, None)) - ws2 = _make_ws(_metadata_frame("s", reconnected=True), _close_frame()) + ws1 = _make_ws(end_with=websockets.exceptions.ConnectionClosedError(None, None)) + ws2 = _make_ws(_close_frame()) call_count = 0 @@ -884,7 +764,7 @@ async def fake_connect(url, extra_headers=None, additional_headers=None, **_kw): async def test_reconnect_window_zero_skips_inner_loop(self): """reconnect_window=0.0 must give up immediately — no inner retry attempts at all.""" client = _make_client() - ws = _make_ws(_metadata_frame()) + ws = _make_ws() connect_count = 0 async def fake_connect(url, extra_headers=None, additional_headers=None, **_kw): @@ -907,7 +787,7 @@ async def test_reconnect_window_expiry_stops_iteration(self): """When reconnect_window=0.0 the outer loop gives up immediately after inner exhaustion.""" client = _make_client() - ws = _make_ws(_metadata_frame()) + ws = _make_ws() call_count = 0 @@ -938,7 +818,7 @@ class TestShellSessionExitCode: async def test_exit_code_zero_on_clean_exit(self): """exit_code is 0 after a clean shell exit (status=Success).""" client = _make_client() - ws = _make_ws(_metadata_frame(), _exit_frame(0)) + ws = _make_ws(_exit_frame(0)) with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: @@ -951,7 +831,7 @@ async def test_exit_code_zero_on_clean_exit(self): async def test_exit_code_nonzero(self): """exit_code reflects non-zero exit status from ExitCode cause.""" client = _make_client() - ws = _make_ws(_metadata_frame(), _exit_frame(42)) + ws = _make_ws(_exit_frame(42)) with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: @@ -964,7 +844,7 @@ async def test_exit_code_nonzero(self): async def test_exit_code_none_before_exit(self): """exit_code is None until the termination STATUS frame is processed.""" client = _make_client() - ws = _make_ws(_metadata_frame(), _stdout_frame("hi"), _exit_frame(1)) + ws = _make_ws(_stdout_frame("hi"), _exit_frame(1)) with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: @@ -976,41 +856,6 @@ async def test_exit_code_none_before_exit(self): assert exit_code_mid_loop is None # not set yet during STDOUT frame assert shell.exit_code == 1 - @pytest.mark.asyncio - async def test_exit_code_set_via_pending_frames_path(self): - """exit_code is set when the termination STATUS is drained from _pending_frames.""" - import websockets.exceptions - - client = _make_client() - ws = AsyncMock() - ws.send = AsyncMock() - ws.close = AsyncMock() - call_count = 0 - - async def recv(): - nonlocal call_count - call_count += 1 - if call_count == 1: - # First recv during _connect() returns a STDOUT frame — stashed as pending - return _stdout_frame("stashed") - if call_count == 2: - # Second recv during _connect() returns the metadata confirmation - return _metadata_frame() - if call_count == 3: - return _exit_frame(5) - raise websockets.exceptions.ConnectionClosedOK(None, None) - - ws.recv = recv - - with patch("websockets.connect", new=AsyncMock(return_value=ws)): - async with ShellSession(client, FAKE_ARN) as shell: - frames = [frame async for frame in shell] - - # stashed STDOUT frame drained first, then the exit STATUS frame - assert frames[0].channel == ShellChannel.STDOUT - assert frames[1].channel == ShellChannel.STATUS - assert shell.exit_code == 5 - @pytest.mark.asyncio async def test_exit_code_platform_error_without_exit_code_cause(self): """exit_code is None when a Failure STATUS has no ExitCode cause (e.g. InternalError).""" @@ -1028,7 +873,7 @@ async def test_exit_code_platform_error_without_exit_code_cause(self): } ).encode() client = _make_client() - ws = _make_ws(_metadata_frame(), bytes([ShellChannel.STATUS]) + platform_error) + ws = _make_ws(bytes([ShellChannel.STATUS]) + platform_error) with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN) as shell: @@ -1038,85 +883,13 @@ async def test_exit_code_platform_error_without_exit_code_cause(self): assert shell.exit_code is None -class TestShellSessionBytesDropped: - def _second_confirmation_frame(self, shell_id: str = "test-session", bytes_dropped: int = 1024) -> bytes: - payload = json.dumps( - { - "kind": "Status", - "apiVersion": "v1", - "metadata": {"shellId": shell_id, "reconnected": True, "bytesDropped": bytes_dropped}, - "status": "Success", - } - ).encode() - return bytes([0x03]) + payload - - @pytest.mark.asyncio - async def test_bytes_dropped_set_on_second_confirmation(self): - """Second confirmation frame with bytesDropped sets shell.bytes_dropped.""" - client = _make_client() - ws = _make_ws( - _metadata_frame("s"), - _stdout_frame("output"), - self._second_confirmation_frame("s", bytes_dropped=512), - _close_frame(), - ) - - with patch("websockets.connect", new=AsyncMock(return_value=ws)): - async with ShellSession(client, FAKE_ARN, shell_id="s") as shell: - frames = [frame async for frame in shell] - - assert shell.bytes_dropped == 512 - # Second confirmation must be swallowed — only stdout frame yielded - assert len(frames) == 1 - assert frames[0].channel == ShellChannel.STDOUT - - @pytest.mark.asyncio - async def test_bytes_dropped_zero_when_no_overflow(self): - """bytes_dropped stays 0 when no second confirmation arrives.""" - client = _make_client() - ws = _make_ws(_metadata_frame(), _close_frame()) - - with patch("websockets.connect", new=AsyncMock(return_value=ws)): - async with ShellSession(client, FAKE_ARN) as shell: - async for _ in shell: - pass - - assert shell.bytes_dropped == 0 - - @pytest.mark.asyncio - async def test_second_confirmation_without_bytes_dropped_swallowed(self): - """Second confirmation with no bytesDropped field is still swallowed.""" - client = _make_client() - # Second confirmation without bytesDropped (single overflow frame edge case) - second_conf = json.dumps( - { - "kind": "Status", - "apiVersion": "v1", - "metadata": {"shellId": "s", "reconnected": True}, - "status": "Success", - } - ).encode() - ws = _make_ws( - _metadata_frame("s"), - bytes([0x03]) + second_conf, - _close_frame(), - ) - - with patch("websockets.connect", new=AsyncMock(return_value=ws)): - async with ShellSession(client, FAKE_ARN, shell_id="s") as shell: - frames = [frame async for frame in shell] - - assert shell.bytes_dropped == 0 - assert frames == [] # second confirmation swallowed, close frame stops iteration - - class TestShellSessionAuthModes: """open_shell routes to the correct auth helper based on the auth= argument.""" @pytest.mark.asyncio async def test_sigv4_default_uses_connect_shell(self): client = _make_client() - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)) as mock_connect: async with ShellSession(client, FAKE_ARN, auth="sigv4") as _: @@ -1134,7 +907,7 @@ async def test_presigned_uses_connect_shell_presigned(self): presigned_url = "wss://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/X/ws/shells?X-Amz-Signature=abc" client = MagicMock() client.connect_shell_presigned.return_value = presigned_url - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)) as mock_connect: async with ShellSession(client, FAKE_ARN, auth=PresignedAuth(expires=120)) as _: @@ -1153,7 +926,7 @@ async def test_presigned_uses_connect_shell_presigned(self): async def test_presigned_forwards_expires(self): client = MagicMock() client.connect_shell_presigned.return_value = FAKE_URL - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN, auth=PresignedAuth(expires=60)) as _: @@ -1170,7 +943,7 @@ async def test_oauth_uses_connect_shell_oauth(self): encoded = base64.urlsafe_b64encode(b"tok").decode().rstrip("=") expected_protos = [f"base64UrlBearerAuthorization.{encoded}", "base64UrlBearerAuthorization"] client.connect_shell_oauth.return_value = (FAKE_URL, expected_protos) - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)) as mock_connect: async with ShellSession(client, FAKE_ARN, auth=OAuthAuth(bearer_token="tok")) as _: @@ -1191,7 +964,7 @@ async def test_oauth_forwards_bearer_token(self): encoded = base64.urlsafe_b64encode(b"my-token").decode().rstrip("=") protos = [f"base64UrlBearerAuthorization.{encoded}", "base64UrlBearerAuthorization"] client.connect_shell_oauth.return_value = (FAKE_URL, protos) - ws = _make_ws(_metadata_frame()) + ws = _make_ws() with patch("websockets.connect", new=AsyncMock(return_value=ws)): async with ShellSession(client, FAKE_ARN, auth=OAuthAuth(bearer_token="my-token")) as _: diff --git a/tests/unit/runtime/test_shell_protocol.py b/tests/unit/runtime/test_shell_protocol.py index 062af26b..3e0958bb 100644 --- a/tests/unit/runtime/test_shell_protocol.py +++ b/tests/unit/runtime/test_shell_protocol.py @@ -154,10 +154,6 @@ def test_encode_heartbeat(self): frame = self.framer.encode_heartbeat() assert frame == bytes([ShellChannel.HEARTBEAT]) - def test_encode_close(self): - frame = self.framer.encode_close() - assert frame == bytes([ShellChannel.CLOSE]) - def test_round_trip_stdin(self): original = "echo hello\n" encoded = self.framer.encode_stdin(original)