Merge pull request #325 from ag2ai/version0-6-1

Bump version to 0.6.1

autogen/version.py

@@ -2,4 +2,4 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "0.6.0"
+__version__ = "0.6.1"

test/io/test_websockets.py

@@ -9,6 +9,7 @@
 
 import pytest
 from conftest import skip_openai
+from websockets.exceptions import ConnectionClosed
 
 import autogen
 from autogen.cache.cache import Cache
@@ -159,17 +160,20 @@ def on_connect(iostream: IOWebsockets, success_dict: dict[str, bool] = success_d
                 websocket.send("Please write a poem about spring in a city of your choice.")
 
                 while True:
-                    message = websocket.recv()
-                    message = message.decode("utf-8") if isinstance(message, bytes) else message
-                    # drop the newline character
-                    if message.endswith("\n"):
-                        message = message[:-1]
-
-                    print(message, end="", flush=True)
-
-                    if "TERMINATE" in message:
-                        print()
-                        print(" - Received TERMINATE message. Exiting.", flush=True)
+                    try:
+                        message = websocket.recv()
+                        message = message.decode("utf-8") if isinstance(message, bytes) else message
+                        # drop the newline character
+                        if message.endswith("\n"):
+                            message = message[:-1]
+
+                        print(message, end="", flush=True)
+
+                        if "TERMINATE" in message:
+                            print()
+                            print(" - Received TERMINATE message.", flush=True)
+                    except ConnectionClosed as e:
+                        print("Connection closed:", e, flush=True)
                         break
 
         assert success_dict["success"]