fix(websocket): log on_disconnect errors, gate on accept, correct protocol docs + tests (v26.06.19)

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

README.md

@@ -11,7 +11,7 @@
   <a href="https://github.com/fireflyframework"><img src="https://img.shields.io/badge/Firefly_Framework-official-ff6600?logo=data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCI+PHBhdGggZmlsbD0id2hpdGUiIGQ9Ik0xMiAyQzYuNDggMiAyIDYuNDggMiAxMnM0LjQ4IDEwIDEwIDEwIDEwLTQuNDggMTAtMTBTMTcuNTIgMiAxMiAyeiIvPjwvc3ZnPg==" alt="Firefly Framework"></a>
   <a href="https://www.python.org/"><img src="https://img.shields.io/badge/python-3.12%2B-blue?logo=python&logoColor=white" alt="Python 3.12+"></a>
   <a href="LICENSE"><img src="https://img.shields.io/badge/license-Apache%202.0-green" alt="License: Apache 2.0"></a>
-  <a href="#"><img src="https://img.shields.io/badge/version-26.06.18-brightgreen" alt="Version: 26.06.18"></a>
+  <a href="#"><img src="https://img.shields.io/badge/version-26.06.19-brightgreen" alt="Version: 26.06.19"></a>
   <a href="#"><img src="https://img.shields.io/badge/type--checked-mypy%20strict-blue?logo=python&logoColor=white" alt="Type Checked: mypy strict"></a>
   <a href="#"><img src="https://img.shields.io/badge/code%20style-ruff-purple?logo=ruff&logoColor=white" alt="Code Style: Ruff"></a>
   <a href="#"><img src="https://img.shields.io/badge/async-first-brightgreen" alt="Async First"></a>

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"

src/pyfly/__init__.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 """PyFly — Enterprise Python Framework."""
 
-__version__ = "26.06.18"
+__version__ = "26.06.19"

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

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

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."""

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

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