From dc51ca08e1c59e5086803826c0a1c23a4df9d2bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A9s=20Contreras=20Guill=C3=A9n?= Date: Sat, 6 Jun 2026 09:49:02 +0200 Subject: [PATCH] fix(websocket): log on_disconnect errors, gate on accept, correct protocol docs + tests + bump v26.06.19 Surfaced by an audit while validating implement-websocket (skill clean: messages flow, broadcast, disconnect cleanup all proven). The websocket module was untested (322 lines). - adapter: on_disconnect cleanup errors were swallowed by contextlib.suppress(Exception) with no logging -> now logged (warning + traceback), matching the handler-error path. - adapter: on_disconnect ran unconditionally in finally even when the handler raised before accept() -> now gated on the new WebSocketSession.accepted flag. - handler: WebSocketHandler docstrings implied on_connect/on_message are auto-dispatched (they are not; only on_disconnect is) -> corrected to state they are caller-invoked. Implementing on_message expecting framework dispatch was a silent no-op. - decorators: documented that the controller instance is a singleton shared across all connections (per-connection state belongs on the WebSocketSession, not self). Added WebSocketSession.accepted + tests/websocket/ (5 tests; module was untested). Regression guards confirmed to fail without the adapter fix. Gates: mypy --strict (607), ruff + format, full suite 3725 passed. --- CHANGELOG.md | 39 ++++++ README.md | 2 +- pyproject.toml | 2 +- src/pyfly/__init__.py | 2 +- src/pyfly/websocket/adapters/starlette.py | 14 +- src/pyfly/websocket/decorators.py | 4 + src/pyfly/websocket/handler.py | 33 +++-- tests/websocket/__init__.py | 7 + tests/websocket/test_ws_lifecycle.py | 155 ++++++++++++++++++++++ uv.lock | 2 +- 10 files changed, 239 insertions(+), 21 deletions(-) create mode 100644 tests/websocket/__init__.py create mode 100644 tests/websocket/test_ws_lifecycle.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 64e92a2c..a705fcfa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,45 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). --- +## v26.06.19 (2026-06-06) + +### Fixed + +- **WebSocket `on_disconnect` failures are no longer silently swallowed.** The + Starlette adapter wrapped the `on_disconnect` cleanup hook in + `contextlib.suppress(Exception)`, so a failing cleanup (leaked resource, lock) + vanished without a trace. Failures are now logged (`warning` + traceback), + matching the handler-error path. +- **`on_disconnect` runs only when the connection was accepted.** It previously + fired unconditionally in `finally`, so a handler that raised before + `session.accept()` got a spurious disconnect for a never-completed handshake. + `WebSocketSession` now tracks `accepted`, and the adapter gates the hook on it. +- **`WebSocketHandler` protocol docstrings corrected.** They implied `on_connect` + and `on_message` were auto-dispatched by the framework; they are **not** (only + `on_disconnect` is). The `@websocket_mapping` method owns the full lifecycle + (accept + receive loop). Implementing `on_connect`/`on_message` and expecting + the framework to call them was a silent no-op; the docstrings now state they are + caller-invoked. (The `implement-websocket` skill already documented this + correctly.) + +### Added + +- **`WebSocketSession.accepted`** property. +- **`tests/websocket/` suite (5 tests)** — the module was previously untested. + Covers message flow, disconnect cleanup, the accept-gating + error-logging + fixes, the `accepted` flag, and the on_message-not-auto-dispatched contract. + +### Notes + +- Documented that the WebSocket controller instance is a **singleton shared + across all connections** — keep per-connection state on the `WebSocketSession`, + never on `self`. + +These surfaced in an audit while validating the `implement-websocket` skill (which +validated clean — messages flow, broadcast, and disconnect cleanup all proven). + +--- + ## v26.06.18 (2026-06-06) ### Tests diff --git a/README.md b/README.md index 08132ecb..d26c666e 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ Firefly Framework Python 3.12+ License: Apache 2.0 - Version: 26.06.18 + Version: 26.06.19 Type Checked: mypy strict Code Style: Ruff Async First diff --git a/pyproject.toml b/pyproject.toml index a0fbd397..ea1d6635 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "pyfly" # CalVer YY.MM.PATCH — package metadata uses PEP 440 normalized form (26.5.4); # git tag, GitHub release and human-readable display use leading-zero form # (v26.05.04) to match the Java/.NET/Go siblings. -version = "26.6.18" +version = "26.6.19" description = "The official Python implementation of the Firefly Framework — DI, CQRS, EDA, hexagonal architecture, and more." readme = "README.md" license = "Apache-2.0" diff --git a/src/pyfly/__init__.py b/src/pyfly/__init__.py index 0345ec5b..a144e775 100644 --- a/src/pyfly/__init__.py +++ b/src/pyfly/__init__.py @@ -13,4 +13,4 @@ # limitations under the License. """PyFly — Enterprise Python Framework.""" -__version__ = "26.06.18" +__version__ = "26.06.19" diff --git a/src/pyfly/websocket/adapters/starlette.py b/src/pyfly/websocket/adapters/starlette.py index 33188c20..be2c0525 100644 --- a/src/pyfly/websocket/adapters/starlette.py +++ b/src/pyfly/websocket/adapters/starlette.py @@ -21,7 +21,6 @@ from __future__ import annotations -import contextlib import inspect import logging from typing import Any @@ -94,13 +93,18 @@ async def lazy_ws_endpoint(websocket: WebSocket) -> None: # silently (audit #232). _logger.warning("websocket handler '%s' raised", method_name, exc_info=True) finally: - # Invoke an on_disconnect lifecycle hook if the controller - # defines one, so handlers can clean up (audit #232). + # Invoke on_disconnect for cleanup — but only when the connection + # was actually accepted, so a handler that errored before accept() + # does not get a spurious disconnect for a never-completed + # handshake. Log (not silently swallow) cleanup failures so leaked + # resources surface (audit #232). on_disconnect = getattr(_cache["instance"], "on_disconnect", None) - if callable(on_disconnect): - with contextlib.suppress(Exception): + if session.accepted and callable(on_disconnect): + try: result = on_disconnect(session) if inspect.isawaitable(result): await result + except Exception: + _logger.warning("websocket on_disconnect for '%s' raised", method_name, exc_info=True) return lazy_ws_endpoint diff --git a/src/pyfly/websocket/decorators.py b/src/pyfly/websocket/decorators.py index 7d50075a..f89cf731 100644 --- a/src/pyfly/websocket/decorators.py +++ b/src/pyfly/websocket/decorators.py @@ -32,6 +32,10 @@ def websocket_mapping(path: str = "") -> Callable[[F], F]: The decorated method must accept a single ``WebSocketSession`` argument and manage the full connection lifecycle (accept, message loop, close). + The controller instance is a process-wide singleton shared across **all** + connections — keep per-connection state in local variables or on the + ``WebSocketSession``, never on ``self`` (it would leak/race across clients). + Usage:: @rest_controller diff --git a/src/pyfly/websocket/handler.py b/src/pyfly/websocket/handler.py index a346c6d3..cba78a06 100644 --- a/src/pyfly/websocket/handler.py +++ b/src/pyfly/websocket/handler.py @@ -25,27 +25,29 @@ @runtime_checkable class WebSocketHandler(Protocol): - """Protocol for WebSocket handler lifecycle methods. - - Implement any combination of these methods on a controller to handle - WebSocket events. All methods are optional — unimplemented hooks are - simply skipped. + """Optional lifecycle hooks a WebSocket controller may define. + + Only :meth:`on_disconnect` is invoked **automatically** by the framework — + after the ``@websocket_mapping`` handler returns or the socket closes, and + only if the connection was accepted. ``on_connect`` and ``on_message`` are + **not** dispatched by the framework: the ``@websocket_mapping`` method owns + the full lifecycle (accept + receive loop). The two are convenience + signatures you may implement and call yourself from that method. """ async def on_connect(self, session: WebSocketSession) -> None: - """Called when a client initiates a WebSocket connection. - - The connection is *not* yet accepted — call ``await session.accept()`` - to complete the handshake. - """ + """Convenience hook — **not** auto-called. Invoke it yourself from your + ``@websocket_mapping`` method (e.g. around ``await session.accept()``).""" ... async def on_message(self, session: WebSocketSession, data: str) -> None: - """Called when a text message is received from the client.""" + """Convenience hook — **not** auto-called. The framework does not dispatch + incoming messages; run your own receive loop and call this if you want.""" ... async def on_disconnect(self, session: WebSocketSession) -> None: - """Called when the WebSocket connection is closed.""" + """Called automatically by the registrar when the handler returns or the + connection closes — only if the connection was accepted.""" ... @@ -59,6 +61,12 @@ class WebSocketSession: def __init__(self, raw: Any) -> None: self._ws = raw + self._accepted = False + + @property + def accepted(self) -> bool: + """Whether the handshake has been accepted (``accept()`` was called).""" + return self._accepted @property def path_params(self) -> dict[str, Any]: @@ -78,6 +86,7 @@ def headers(self) -> Any: async def accept(self, subprotocol: str | None = None) -> None: """Accept the WebSocket connection handshake.""" await self._ws.accept(subprotocol=subprotocol) + self._accepted = True async def send_text(self, data: str) -> None: """Send a text message to the client.""" diff --git a/tests/websocket/__init__.py b/tests/websocket/__init__.py new file mode 100644 index 00000000..b6e82019 --- /dev/null +++ b/tests/websocket/__init__.py @@ -0,0 +1,7 @@ +# Copyright 2026 Firefly Software Foundation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 diff --git a/tests/websocket/test_ws_lifecycle.py b/tests/websocket/test_ws_lifecycle.py new file mode 100644 index 00000000..305a54f5 --- /dev/null +++ b/tests/websocket/test_ws_lifecycle.py @@ -0,0 +1,155 @@ +# Copyright 2026 Firefly Software Foundation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""WebSocket adapter lifecycle tests (v26.06.19). + +The websocket module previously had no tests. These lock in the registrar's +endpoint lifecycle and the v26.06.19 fixes: ``on_disconnect`` runs only when the +connection was accepted, its failures are logged (not silently swallowed), the +``WebSocketSession.accepted`` flag, message flow, and the documented contract +that ``on_message`` is NOT auto-dispatched by the framework. +""" + +from __future__ import annotations + +import logging +from typing import Any + +import pytest +from starlette.websockets import WebSocketDisconnect + +from pyfly.websocket import WebSocketSession +from pyfly.websocket.adapters.starlette import WebSocketRegistrar + + +class _FakeRawWS: + def __init__(self, incoming: tuple[str, ...] = ()) -> None: + self._incoming = list(incoming) + self.sent: list[str] = [] + self.accepted = False + self.path_params: dict[str, Any] = {} + + async def accept(self, subprotocol: str | None = None) -> None: + self.accepted = True + + async def send_text(self, data: str) -> None: + self.sent.append(data) + + async def receive_text(self) -> str: + if self._incoming: + return self._incoming.pop(0) + raise WebSocketDisconnect(1000) + + async def close(self, code: int = 1000, reason: str | None = None) -> None: + pass + + +class _FakeCtx: + def __init__(self, instance: Any) -> None: + self._instance = instance + self.container = type("_C", (), {"_registrations": {}})() + + def get_bean(self, cls: type) -> Any: + return self._instance + + +def _endpoint(instance: Any, method_name: str = "chat") -> Any: + return WebSocketRegistrar._make_lazy_handler(_FakeCtx(instance), type(instance), method_name) + + +class _EchoController: + def __init__(self) -> None: + self.events: list[str] = [] + + async def chat(self, session: WebSocketSession) -> None: + await session.accept() + self.events.append("accept") + while True: + msg = await session.receive_text() + await session.send_text(f"echo:{msg}") + + async def on_disconnect(self, session: WebSocketSession) -> None: + self.events.append("disconnect") + + +class _NoAcceptController: + def __init__(self) -> None: + self.disconnected = False + + async def chat(self, session: WebSocketSession) -> None: + raise RuntimeError("boom before accept") + + async def on_disconnect(self, session: WebSocketSession) -> None: + self.disconnected = True + + +class _BadCleanupController: + async def chat(self, session: WebSocketSession) -> None: + await session.accept() # returns immediately + + async def on_disconnect(self, session: WebSocketSession) -> None: + raise RuntimeError("cleanup failed") + + +class _OnMessageController: + def __init__(self) -> None: + self.on_message_calls = 0 + + async def chat(self, session: WebSocketSession) -> None: + await session.accept() + try: + while True: + await session.receive_text() + except WebSocketDisconnect: + pass + + async def on_message(self, session: WebSocketSession, data: str) -> None: + self.on_message_calls += 1 # must never be auto-invoked + + +@pytest.mark.asyncio +async def test_message_flow_and_disconnect_cleanup() -> None: + ctrl = _EchoController() + raw = _FakeRawWS(["hi", "there"]) + await _endpoint(ctrl)(raw) + assert raw.sent == ["echo:hi", "echo:there"] # messages flowed + assert ctrl.events == ["accept", "disconnect"] # on_disconnect ran after accept + + +@pytest.mark.asyncio +async def test_on_disconnect_not_called_when_never_accepted() -> None: + ctrl = _NoAcceptController() + await _endpoint(ctrl)(_FakeRawWS()) # handler errors before accept; must not raise + assert ctrl.disconnected is False # gated on session.accepted + + +@pytest.mark.asyncio +async def test_on_disconnect_error_is_logged_not_swallowed(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger="pyfly.websocket.adapters.starlette"): + await _endpoint(_BadCleanupController())(_FakeRawWS()) # must not raise + assert any("on_disconnect" in r.getMessage() for r in caplog.records) + + +@pytest.mark.asyncio +async def test_on_message_is_not_auto_dispatched() -> None: + ctrl = _OnMessageController() + await _endpoint(ctrl)(_FakeRawWS(["a", "b"])) + assert ctrl.on_message_calls == 0 # framework never dispatches to on_message + + +@pytest.mark.asyncio +async def test_session_accepted_flag() -> None: + session = WebSocketSession(_FakeRawWS()) + assert session.accepted is False + await session.accept() + assert session.accepted is True diff --git a/uv.lock b/uv.lock index af9738b6..ed9540c5 100644 --- a/uv.lock +++ b/uv.lock @@ -1967,7 +1967,7 @@ wheels = [ [[package]] name = "pyfly" -version = "26.6.18" +version = "26.6.19" source = { editable = "." } dependencies = [ { name = "pydantic" },