From 555bdcde847d6418fcf818e9bf8c59c145eeffb5 Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Sat, 29 Aug 2026 21:35:21 +0200 Subject: [PATCH 1/2] Simplify error handling in servers. Rely on the connection's context manager implementation instead of duplicating its logic. --- src/websockets/asyncio/server.py | 6 ++---- src/websockets/sync/server.py | 6 ++---- src/websockets/trio/server.py | 6 ++---- 3 files changed, 6 insertions(+), 12 deletions(-) diff --git a/src/websockets/asyncio/server.py b/src/websockets/asyncio/server.py index 7a168d04..58c1951d 100644 --- a/src/websockets/asyncio/server.py +++ b/src/websockets/asyncio/server.py @@ -744,12 +744,10 @@ async def protocol_handler(connection: ServerConnection) -> None: server.all_connections.add(connection) connection.start_keepalive() try: - await handler(connection) + async with connection: + await handler(connection) except Exception: connection.logger.error("connection handler failed", exc_info=True) - await connection.close(CloseCode.INTERNAL_ERROR) - else: - await connection.close() finally: server.all_connections.discard(connection) diff --git a/src/websockets/sync/server.py b/src/websockets/sync/server.py index f0f3e058..d9ffac45 100644 --- a/src/websockets/sync/server.py +++ b/src/websockets/sync/server.py @@ -770,12 +770,10 @@ def protocol_select_subprotocol( server.all_connections.add(connection) connection.start_keepalive() try: - handler(connection) + with connection: + handler(connection) except Exception: connection.logger.error("connection handler failed", exc_info=True) - connection.close(CloseCode.INTERNAL_ERROR) - else: - connection.close() finally: with server.lock: server.all_connections.discard(connection) diff --git a/src/websockets/trio/server.py b/src/websockets/trio/server.py index 4e104c77..22531ed3 100644 --- a/src/websockets/trio/server.py +++ b/src/websockets/trio/server.py @@ -639,12 +639,10 @@ def protocol_select_subprotocol( server.all_connections.add(connection) connection.start_keepalive() try: - await handler(connection) + async with connection: + await handler(connection) except Exception: connection.logger.error("connection handler failed", exc_info=True) - await connection.aclose(CloseCode.INTERNAL_ERROR) - else: - await connection.aclose() finally: server.all_connections.discard(connection) From 544779cce9ee7a5a05bbbf4da5ae555b6ed94041 Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Sat, 29 Aug 2026 21:42:30 +0200 Subject: [PATCH 2/2] Make clients use close code 1011 on exceptions. According to errata 3227, which is approved, the meaning of code 1011 also applies to clients. It makes sense to use them symetrically. --- src/websockets/asyncio/client.py | 4 ++-- src/websockets/sync/client.py | 4 ++-- src/websockets/trio/client.py | 4 ++-- tests/asyncio/test_client.py | 15 +++++++++++++++ tests/sync/test_client.py | 15 +++++++++++++++ tests/trio/test_client.py | 15 +++++++++++++++ 6 files changed, 51 insertions(+), 6 deletions(-) diff --git a/src/websockets/asyncio/client.py b/src/websockets/asyncio/client.py index b6426230..4e983bbe 100644 --- a/src/websockets/asyncio/client.py +++ b/src/websockets/asyncio/client.py @@ -574,7 +574,7 @@ async def __aenter__(self) -> ClientConnection: if hasattr(self, "connection"): raise RuntimeError("connect() isn't reentrant") self.connection = await self - return self.connection + return await self.connection.__aenter__() async def __aexit__( self, @@ -583,7 +583,7 @@ async def __aexit__( traceback: TracebackType | None, ) -> None: try: - await self.connection.close() + return await self.connection.__aexit__(exc_type, exc_value, traceback) finally: del self.connection diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 95fe53ab..393fc8e2 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -521,7 +521,7 @@ def __enter__(self) -> ClientConnection: raise RuntimeError("connect() isn't reentrant") self.connection = self.connect() self.connection.pending_legacy_warning = False - return self.connection + return self.connection.__enter__() def __exit__( self, @@ -530,7 +530,7 @@ def __exit__( exc_traceback: TracebackType | None, ) -> None: try: - self.connection.close() + return self.connection.__exit__(exc_type, exc_value, exc_traceback) finally: del self.connection diff --git a/src/websockets/trio/client.py b/src/websockets/trio/client.py index 069fe4b2..a18cf7f2 100644 --- a/src/websockets/trio/client.py +++ b/src/websockets/trio/client.py @@ -549,7 +549,7 @@ async def __aenter__(self) -> ClientConnection: await self.__aenter_nursery__() try: self.connection = await self.connect(self.nursery) - return self.connection + return await self.connection.__aenter__() except BaseException as exc: await self.__aexit_nursery__(type(exc), exc, exc.__traceback__) raise AssertionError("expected __aexit_nursery__ to re-raise the exception") @@ -562,7 +562,7 @@ async def __aexit__( ) -> None: try: try: - await self.connection.aclose() + return await self.connection.__aexit__(exc_type, exc_value, traceback) finally: del self.connection finally: diff --git a/tests/asyncio/test_client.py b/tests/asyncio/test_client.py index f8f56308..46bbb59b 100644 --- a/tests/asyncio/test_client.py +++ b/tests/asyncio/test_client.py @@ -60,6 +60,21 @@ async def test_context_manager(self): self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.protocol.state.name, "CLOSED") + async def test_context_manager_normal_exit(self): + """Client closes the connection with code 1000 when exiting normally.""" + async with serve(*args) as server: + async with connect(get_uri(server)) as client: + pass + self.assertEqual(client.close_code, 1000) + + async def test_context_manager_exception(self): + """Client closes the connection with code 1011 when exiting with an error.""" + async with serve(*args) as server: + with self.assertRaises(RuntimeError): + async with connect(get_uri(server)) as client: + raise RuntimeError("BOOM") + self.assertEqual(client.close_code, 1011) + async def test_direct_connection(self): """Client connects to server directly.""" async with serve(*args) as server: diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index e812fcc4..93c81d74 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -67,6 +67,21 @@ def test_context_manager(self): self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.protocol.state.name, "CLOSED") + def test_context_manager_normal_exit(self): + """Client closes the connection with code 1000 when exiting normally.""" + with run_server() as server: + with connect(get_uri(server)) as client: + pass + self.assertEqual(client.close_code, 1000) + + def test_context_manager_exception(self): + """Client closes the connection with code 1011 when exiting with an error.""" + with run_server() as server: + with self.assertRaises(RuntimeError): + with connect(get_uri(server)) as client: + raise RuntimeError("BOOM") + self.assertEqual(client.close_code, 1011) + def test_direct_connection(self): """Client connects to server directly.""" with run_server() as server: diff --git a/tests/trio/test_client.py b/tests/trio/test_client.py index 8d9d3774..b318a83a 100644 --- a/tests/trio/test_client.py +++ b/tests/trio/test_client.py @@ -62,6 +62,21 @@ async def test_context_manager(self): self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.protocol.state.name, "CLOSED") + async def test_context_manager_normal_exit(self): + """Client closes the connection with code 1000 when exiting normally.""" + async with run_server() as server: + async with connect(get_uri(server)) as client: + pass + self.assertEqual(client.close_code, 1000) + + async def test_context_manager_exception(self): + """Client closes the connection with code 1011 when exiting with an error.""" + async with run_server() as server: + with self.assertRaises(RuntimeError): + async with connect(get_uri(server)) as client: + raise RuntimeError("BOOM") + self.assertEqual(client.close_code, 1011) + async def test_explicit_host_port(self): """Client connects using an explicit host / port.""" async with run_server() as server: