Fix typecheck errors and improve WebSocket transport implementation

- Fix INotifee interface compliance in WebSocket demo
- Fix handler function signatures to be async (THandler compatibility)
- Fix is_closed method usage with proper type checking
- Fix pytest.raises multiple exception type issue
- Fix line length violations (E501) across multiple files
- Add debugging logging to Noise security module for troubleshooting
- Update WebSocket transport examples and tests
- Improve transport registry error handling

examples/websocket/test_tcp_echo.py

@@ -5,7 +5,6 @@ Simple TCP echo demo to verify basic libp2p functionality.
 
 import argparse
 import logging
-import sys
 import traceback
 
 import multiaddr
@@ -18,10 +17,10 @@ from libp2p.network.swarm import Swarm
 from libp2p.peer.id import ID
 from libp2p.peer.peerinfo import info_from_p2p_addr
 from libp2p.peer.peerstore import PeerStore
-from libp2p.security.insecure.transport import InsecureTransport, PLAINTEXT_PROTOCOL_ID
+from libp2p.security.insecure.transport import PLAINTEXT_PROTOCOL_ID, InsecureTransport
 from libp2p.stream_muxer.yamux.yamux import Yamux
-from libp2p.transport.upgrader import TransportUpgrader
 from libp2p.transport.tcp.tcp import TCP
+from libp2p.transport.upgrader import TransportUpgrader
 
 # Enable debug logging
 logging.basicConfig(level=logging.DEBUG)
@@ -31,12 +30,13 @@ logger = logging.getLogger("libp2p.tcp-example")
 # Simple echo protocol
 ECHO_PROTOCOL_ID = TProtocol("/echo/1.0.0")
 
+
 async def echo_handler(stream):
     """Simple echo handler that echoes back any data received."""
     try:
         data = await stream.read(1024)
         if data:
-            message = data.decode('utf-8', errors='replace')
+            message = data.decode("utf-8", errors="replace")
             print(f"📥 Received: {message}")
             print(f"📤 Echoing back: {message}")
             await stream.write(data)
@@ -45,6 +45,7 @@ async def echo_handler(stream):
         logger.error(f"Echo handler error: {e}")
         await stream.close()
 
+
 def create_tcp_host():
     """Create a host with TCP transport."""
     # Create key pair and peer store
@@ -60,31 +61,35 @@ def create_tcp_host():
         },
         muxer_transports_by_protocol={TProtocol("/yamux/1.0.0"): Yamux},
     )
-    
+
     # Create TCP transport
     transport = TCP()
-    
+
     # Create swarm and host
     swarm = Swarm(peer_id, peer_store, upgrader, transport)
     host = BasicHost(swarm)
-    
+
     return host
 
+
 async def run(port: int, destination: str) -> None:
     localhost_ip = "0.0.0.0"
 
     if not destination:
         # Create first host (listener) with TCP transport
         listen_addr = multiaddr.Multiaddr(f"/ip4/{localhost_ip}/tcp/{port}")
-        
+
         try:
             host = create_tcp_host()
             logger.debug("Created TCP host")
-            
+
             # Set up echo handler
             host.set_stream_handler(ECHO_PROTOCOL_ID, echo_handler)
 
-            async with host.run(listen_addrs=[listen_addr]), trio.open_nursery() as nursery:
+            async with (
+                host.run(listen_addrs=[listen_addr]),
+                trio.open_nursery() as (nursery),
+            ):
                 # Start the peer-store cleanup task
                 nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
 
@@ -95,15 +100,15 @@ async def run(port: int, destination: str) -> None:
                 if not addrs:
                     print("❌ Error: No addresses found for the host")
                     return
-                
+
                 server_addr = str(addrs[0])
                 client_addr = server_addr.replace("/ip4/0.0.0.0/", "/ip4/127.0.0.1/")
 
                 print("🌐 TCP Server Started Successfully!")
                 print("=" * 50)
                 print(f"📍 Server Address: {client_addr}")
-                print(f"🔧 Protocol: /echo/1.0.0")
-                print(f"🚀 Transport: TCP")
+                print("🔧 Protocol: /echo/1.0.0")
+                print("🚀 Transport: TCP")
                 print()
                 print("📋 To test the connection, run this in another terminal:")
                 print(f"   python test_tcp_echo.py -d {client_addr}")
@@ -112,7 +117,7 @@ async def run(port: int, destination: str) -> None:
                 print("─" * 50)
 
                 await trio.sleep_forever()
-                    
+
         except Exception as e:
             print(f"❌ Error creating TCP server: {e}")
             traceback.print_exc()
@@ -121,13 +126,16 @@ async def run(port: int, destination: str) -> None:
     else:
         # Create second host (dialer) with TCP transport
         listen_addr = multiaddr.Multiaddr(f"/ip4/{localhost_ip}/tcp/{port}")
-        
+
         try:
             # Create a single host for client operations
             host = create_tcp_host()
-            
+
             # Start the host for client operations
-            async with host.run(listen_addrs=[listen_addr]), trio.open_nursery() as nursery:
+            async with (
+                host.run(listen_addrs=[listen_addr]),
+                trio.open_nursery() as (nursery),
+            ):
                 # Start the peer-store cleanup task
                 nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
                 maddr = multiaddr.Multiaddr(destination)
@@ -144,7 +152,7 @@ async def run(port: int, destination: str) -> None:
                     print("✅ Successfully connected to TCP server!")
                 except Exception as e:
                     error_msg = str(e)
-                    print(f"\n❌ Connection Failed!")
+                    print("\n❌ Connection Failed!")
                     print(f"   Peer ID: {info.peer_id}")
                     print(f"   Address: {destination}")
                     print(f"   Error: {error_msg}")
@@ -185,24 +193,28 @@ async def run(port: int, destination: str) -> None:
                     traceback.print_exc()
 
                 print("✅ TCP demo completed successfully!")
-                
+
         except Exception as e:
             print(f"❌ Error creating TCP client: {e}")
             traceback.print_exc()
             return
 
+
 def main() -> None:
     description = "Simple TCP echo demo for libp2p"
     parser = argparse.ArgumentParser(description=description)
     parser.add_argument("-p", "--port", default=0, type=int, help="source port number")
-    parser.add_argument("-d", "--destination", type=str, help="destination multiaddr string")
+    parser.add_argument(
+        "-d", "--destination", type=str, help="destination multiaddr string"
+    )
 
     args = parser.parse_args()
-    
+
     try:
         trio.run(run, args.port, args.destination)
     except KeyboardInterrupt:
         pass
 
+
 if __name__ == "__main__":
     main()

examples/websocket/test_websocket_transport.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+"""
+Simple test script to verify WebSocket transport functionality.
+"""
+
+import asyncio
+import logging
+from pathlib import Path
+import sys
+
+# Add the libp2p directory to the path so we can import it
+sys.path.insert(0, str(Path(__file__).parent))
+
+import multiaddr
+
+from libp2p.transport import create_transport, create_transport_for_multiaddr
+from libp2p.transport.upgrader import TransportUpgrader
+
+# Set up logging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+
+async def test_websocket_transport():
+    """Test basic WebSocket transport functionality."""
+    print("🧪 Testing WebSocket Transport Functionality")
+    print("=" * 50)
+
+    # Create a dummy upgrader
+    upgrader = TransportUpgrader({}, {})
+
+    # Test creating WebSocket transport
+    try:
+        ws_transport = create_transport("ws", upgrader)
+        print(f"✅ WebSocket transport created: {type(ws_transport).__name__}")
+
+        # Test creating transport from multiaddr
+        ws_maddr = multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/8080/ws")
+        ws_transport_from_maddr = create_transport_for_multiaddr(ws_maddr, upgrader)
+        print(
+            f"✅ WebSocket transport from multiaddr: "
+            f"{type(ws_transport_from_maddr).__name__}"
+        )
+
+        # Test creating listener
+        handler_called = False
+
+        async def test_handler(conn):
+            nonlocal handler_called
+            handler_called = True
+            print(f"✅ Connection handler called with: {type(conn).__name__}")
+            await conn.close()
+
+        listener = ws_transport.create_listener(test_handler)
+        print(f"✅ WebSocket listener created: {type(listener).__name__}")
+
+        # Test that the transport can be used
+        print(
+            f"✅ WebSocket transport supports dialing: {hasattr(ws_transport, 'dial')}"
+        )
+        print(
+            f"✅ WebSocket transport supports listening: "
+            f"{hasattr(ws_transport, 'create_listener')}"
+        )
+
+        print("\n🎯 WebSocket Transport Test Results:")
+        print("✅ Transport creation: PASS")
+        print("✅ Multiaddr parsing: PASS")
+        print("✅ Listener creation: PASS")
+        print("✅ Interface compliance: PASS")
+
+    except Exception as e:
+        print(f"❌ WebSocket transport test failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return False
+
+    return True
+
+
+async def test_transport_registry():
+    """Test the transport registry functionality."""
+    print("\n🔧 Testing Transport Registry")
+    print("=" * 30)
+
+    from libp2p.transport import (
+        get_supported_transport_protocols,
+        get_transport_registry,
+    )
+
+    registry = get_transport_registry()
+    supported = get_supported_transport_protocols()
+
+    print(f"Supported protocols: {supported}")
+
+    # Test getting transports
+    for protocol in supported:
+        transport_class = registry.get_transport(protocol)
+        class_name = transport_class.__name__ if transport_class else "None"
+        print(f"  {protocol}: {class_name}")
+
+    # Test creating transports through registry
+    upgrader = TransportUpgrader({}, {})
+
+    for protocol in supported:
+        try:
+            transport = registry.create_transport(protocol, upgrader)
+            if transport:
+                print(f"✅ {protocol}: Created successfully")
+            else:
+                print(f"❌ {protocol}: Failed to create")
+        except Exception as e:
+            print(f"❌ {protocol}: Error - {e}")
+
+
+async def main():
+    """Run all tests."""
+    print("🚀 WebSocket Transport Integration Test Suite")
+    print("=" * 60)
+    print()
+
+    # Run tests
+    success = await test_websocket_transport()
+    await test_transport_registry()
+
+    print("\n" + "=" * 60)
+    if success:
+        print("🎉 All tests passed! WebSocket transport is working correctly.")
+    else:
+        print("❌ Some tests failed. Check the output above for details.")
+
+    print("\n🚀 WebSocket transport is ready for use in py-libp2p!")
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        print("\n👋 Test interrupted by user")
+    except Exception as e:
+        print(f"\n❌ Test failed with error: {e}")
+        import traceback
+
+        traceback.print_exc()

examples/websocket/websocket_demo.py

@@ -1,21 +1,26 @@
 import argparse
 import logging
+import signal
 import sys
 import traceback
 
 import multiaddr
 import trio
 
+from libp2p.abc import INotifee
+from libp2p.crypto.ed25519 import create_new_key_pair as create_ed25519_key_pair
 from libp2p.crypto.secp256k1 import create_new_key_pair
 from libp2p.custom_types import TProtocol
 from libp2p.host.basic_host import BasicHost
 from libp2p.network.swarm import Swarm
 from libp2p.peer.id import ID
-from libp2p.peer.peerinfo import PeerInfo, info_from_p2p_addr
+from libp2p.peer.peerinfo import info_from_p2p_addr
 from libp2p.peer.peerstore import PeerStore
-from libp2p.security.insecure.transport import InsecureTransport, PLAINTEXT_PROTOCOL_ID
-from libp2p.security.noise.transport import Transport as NoiseTransport
-from libp2p.security.noise.transport import PROTOCOL_ID as NOISE_PROTOCOL_ID
+from libp2p.security.insecure.transport import PLAINTEXT_PROTOCOL_ID, InsecureTransport
+from libp2p.security.noise.transport import (
+    PROTOCOL_ID as NOISE_PROTOCOL_ID,
+    Transport as NoiseTransport,
+)
 from libp2p.stream_muxer.yamux.yamux import Yamux
 from libp2p.transport.upgrader import TransportUpgrader
 from libp2p.transport.websocket.transport import WebsocketTransport
@@ -25,6 +30,15 @@ logging.basicConfig(level=logging.DEBUG)
 
 logger = logging.getLogger("libp2p.websocket-example")
 
+
+# Suppress KeyboardInterrupt by handling SIGINT directly
+def signal_handler(signum, frame):
+    print("✅ Clean exit completed.")
+    sys.exit(0)
+
+
+signal.signal(signal.SIGINT, signal_handler)
+
 # Simple echo protocol
 ECHO_PROTOCOL_ID = TProtocol("/echo/1.0.0")
 
@@ -34,7 +48,7 @@ async def echo_handler(stream):
     try:
         data = await stream.read(1024)
         if data:
-            message = data.decode('utf-8', errors='replace')
+            message = data.decode("utf-8", errors="replace")
             print(f"📥 Received: {message}")
             print(f"📤 Echoing back: {message}")
             await stream.write(data)
@@ -44,7 +58,7 @@ async def echo_handler(stream):
         await stream.close()
 
 
-def create_websocket_host(listen_addrs=None, use_noise=False):
+def create_websocket_host(listen_addrs=None, use_plaintext=False):
     """Create a host with WebSocket transport."""
     # Create key pair and peer store
     key_pair = create_new_key_pair()
@@ -52,11 +66,22 @@ def create_websocket_host(listen_addrs=None, use_noise=False):
     peer_store = PeerStore()
     peer_store.add_key_pair(peer_id, key_pair)
 
-    if use_noise:
+    if use_plaintext:
+        # Create transport upgrader with plaintext security
+        upgrader = TransportUpgrader(
+            secure_transports_by_protocol={
+                TProtocol(PLAINTEXT_PROTOCOL_ID): InsecureTransport(key_pair)
+            },
+            muxer_transports_by_protocol={TProtocol("/yamux/1.0.0"): Yamux},
+        )
+    else:
+        # Create separate Ed25519 key for Noise protocol
+        noise_key_pair = create_ed25519_key_pair()
+
         # Create Noise transport
         noise_transport = NoiseTransport(
             libp2p_keypair=key_pair,
-            noise_privkey=key_pair.private_key,
+            noise_privkey=noise_key_pair.private_key,
             early_data=None,
             with_noise_pipes=False,
         )
@@ -68,43 +93,85 @@ def create_websocket_host(listen_addrs=None, use_noise=False):
             },
             muxer_transports_by_protocol={TProtocol("/yamux/1.0.0"): Yamux},
         )
-    else:
-        # Create transport upgrader with plaintext security
-        upgrader = TransportUpgrader(
-            secure_transports_by_protocol={
-                TProtocol(PLAINTEXT_PROTOCOL_ID): InsecureTransport(key_pair)
-            },
-            muxer_transports_by_protocol={TProtocol("/yamux/1.0.0"): Yamux},
-        )
-    
+
     # Create WebSocket transport
     transport = WebsocketTransport(upgrader)
-    
+
     # Create swarm and host
     swarm = Swarm(peer_id, peer_store, upgrader, transport)
     host = BasicHost(swarm)
-    
+
     return host
 
 
-async def run(port: int, destination: str, use_noise: bool = False) -> None:
+async def run(port: int, destination: str, use_plaintext: bool = False) -> None:
     localhost_ip = "0.0.0.0"
 
     if not destination:
         # Create first host (listener) with WebSocket transport
         listen_addr = multiaddr.Multiaddr(f"/ip4/{localhost_ip}/tcp/{port}/ws")
-        
+
         try:
-            host = create_websocket_host(use_noise=use_noise)
-            logger.debug(f"Created host with use_noise={use_noise}")
-            
+            host = create_websocket_host(use_plaintext=use_plaintext)
+            logger.debug(f"Created host with use_plaintext={use_plaintext}")
+
             # Set up echo handler
             host.set_stream_handler(ECHO_PROTOCOL_ID, echo_handler)
 
-            async with host.run(listen_addrs=[listen_addr]), trio.open_nursery() as nursery:
+            # Add connection event handlers for debugging
+            class DebugNotifee(INotifee):
+                async def opened_stream(self, network, stream):
+                    pass
+
+                async def closed_stream(self, network, stream):
+                    pass
+
+                async def connected(self, network, conn):
+                    print(
+                        f"🔗 New libp2p connection established: "
+                        f"{conn.muxed_conn.peer_id}"
+                    )
+                    if hasattr(conn.muxed_conn, "get_security_protocol"):
+                        security = conn.muxed_conn.get_security_protocol()
+                    else:
+                        security = "Unknown"
+
+                    print(f"   Security: {security}")
+
+                async def disconnected(self, network, conn):
+                    print(f"🔌 libp2p connection closed: {conn.muxed_conn.peer_id}")
+
+                async def listen(self, network, multiaddr):
+                    pass
+
+                async def listen_close(self, network, multiaddr):
+                    pass
+
+            host.get_network().register_notifee(DebugNotifee())
+
+            # Create a cancellation token for clean shutdown
+            cancel_scope = trio.CancelScope()
+
+            async def signal_handler():
+                with trio.open_signal_receiver(signal.SIGINT, signal.SIGTERM) as (
+                    signal_receiver
+                ):
+                    async for sig in signal_receiver:
+                        print(f"\n🛑 Received signal {sig}")
+                        print("✅ Shutting down WebSocket server...")
+                        cancel_scope.cancel()
+                        return
+
+            async with (
+                host.run(listen_addrs=[listen_addr]),
+                trio.open_nursery() as (nursery),
+            ):
                 # Start the peer-store cleanup task
                 nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
 
+                # Start the signal handler
+                nursery.start_soon(signal_handler)
+
                 # Get the actual address and replace 0.0.0.0 with 127.0.0.1 for client
                 # connections
                 addrs = host.get_addrs()
@@ -113,18 +180,19 @@ async def run(port: int, destination: str, use_noise: bool = False) -> None:
                     print("❌ Error: No addresses found for the host")
                     print("Debug: host.get_addrs() returned empty list")
                     return
-                
+
                 server_addr = str(addrs[0])
                 client_addr = server_addr.replace("/ip4/0.0.0.0/", "/ip4/127.0.0.1/")
 
                 print("🌐 WebSocket Server Started Successfully!")
                 print("=" * 50)
                 print(f"📍 Server Address: {client_addr}")
-                print(f"🔧 Protocol: /echo/1.0.0")
-                print(f"🚀 Transport: WebSocket (/ws)")
+                print("🔧 Protocol: /echo/1.0.0")
+                print("🚀 Transport: WebSocket (/ws)")
                 print()
                 print("📋 To test the connection, run this in another terminal:")
-                print(f"   python websocket_demo.py -d {client_addr}")
+                plaintext_flag = " --plaintext" if use_plaintext else ""
+                print(f"   python websocket_demo.py -d {client_addr}{plaintext_flag}")
                 print()
                 print("⏳ Waiting for incoming WebSocket connections...")
                 print("─" * 50)
@@ -132,32 +200,34 @@ async def run(port: int, destination: str, use_noise: bool = False) -> None:
                 # Add a custom handler to show connection events
                 async def custom_echo_handler(stream):
                     peer_id = stream.muxed_conn.peer_id
-                    print(f"\n🔗 New WebSocket Connection!")
+                    print("\n🔗 New WebSocket Connection!")
                     print(f"   Peer ID: {peer_id}")
-                    print(f"   Protocol: /echo/1.0.0")
-                    
+                    print("   Protocol: /echo/1.0.0")
+
                     # Show remote address in multiaddr format
                     try:
                         remote_address = stream.get_remote_address()
                         if remote_address:
                             print(f"   Remote: {remote_address}")
                     except Exception:
-                        print(f"   Remote: Unknown")
-                    
-                    print(f"   ─" * 40)
+                        print("   Remote: Unknown")
+
+                    print("   ─" * 40)
 
                     # Call the original handler
                     await echo_handler(stream)
 
-                    print(f"   ─" * 40)
+                    print("   ─" * 40)
                     print(f"✅ Echo request completed for peer: {peer_id}")
                     print()
 
                 # Replace the handler with our custom one
                 host.set_stream_handler(ECHO_PROTOCOL_ID, custom_echo_handler)
 
-                await trio.sleep_forever()
-                    
+                # Wait indefinitely or until cancelled
+                with cancel_scope:
+                    await trio.sleep_forever()
+
         except Exception as e:
             print(f"❌ Error creating WebSocket server: {e}")
             traceback.print_exc()
@@ -166,15 +236,47 @@ async def run(port: int, destination: str, use_noise: bool = False) -> None:
     else:
         # Create second host (dialer) with WebSocket transport
         listen_addr = multiaddr.Multiaddr(f"/ip4/{localhost_ip}/tcp/{port}/ws")
-        
+
         try:
             # Create a single host for client operations
-            host = create_websocket_host(use_noise=use_noise)
-            
+            host = create_websocket_host(use_plaintext=use_plaintext)
+
             # Start the host for client operations
-            async with host.run(listen_addrs=[listen_addr]), trio.open_nursery() as nursery:
+            async with (
+                host.run(listen_addrs=[listen_addr]),
+                trio.open_nursery() as (nursery),
+            ):
                 # Start the peer-store cleanup task
                 nursery.start_soon(host.get_peerstore().start_cleanup_task, 60)
+
+                # Add connection event handlers for debugging
+                class ClientDebugNotifee(INotifee):
+                    async def opened_stream(self, network, stream):
+                        pass
+
+                    async def closed_stream(self, network, stream):
+                        pass
+
+                    async def connected(self, network, conn):
+                        print(
+                            f"🔗 Client: libp2p connection established: "
+                            f"{conn.muxed_conn.peer_id}"
+                        )
+
+                    async def disconnected(self, network, conn):
+                        print(
+                            f"🔌 Client: libp2p connection closed: "
+                            f"{conn.muxed_conn.peer_id}"
+                        )
+
+                    async def listen(self, network, multiaddr):
+                        pass
+
+                    async def listen_close(self, network, multiaddr):
+                        pass
+
+                host.get_network().register_notifee(ClientDebugNotifee())
+
                 maddr = multiaddr.Multiaddr(destination)
                 info = info_from_p2p_addr(maddr)
                 print("🔌 WebSocket Client Starting...")
@@ -185,21 +287,34 @@ async def run(port: int, destination: str, use_noise: bool = False) -> None:
 
                 try:
                     print("🔗 Connecting to WebSocket server...")
+                    print(f"   Security: {'Plaintext' if use_plaintext else 'Noise'}")
                     await host.connect(info)
                     print("✅ Successfully connected to WebSocket server!")
                 except Exception as e:
                     error_msg = str(e)
-                    if "unable to connect" in error_msg or "SwarmException" in error_msg:
-                        print(f"\n❌ Connection Failed!")
-                        print(f"   Peer ID: {info.peer_id}")
-                        print(f"   Address: {destination}")
-                        print(f"   Error: {error_msg}")
-                        print()
-                        print("💡 Troubleshooting:")
-                        print("   • Make sure the WebSocket server is running")
-                        print("   • Check that the server address is correct")
-                        print("   • Verify the server is listening on the right port")
-                        return
+                    print("\n❌ Connection Failed!")
+                    print(f"   Peer ID: {info.peer_id}")
+                    print(f"   Address: {destination}")
+                    print(f"   Security: {'Plaintext' if use_plaintext else 'Noise'}")
+                    print(f"   Error: {error_msg}")
+                    print(f"   Error type: {type(e).__name__}")
+
+                    # Add more detailed error information for debugging
+                    if hasattr(e, "__cause__") and e.__cause__:
+                        print(f"   Root cause: {e.__cause__}")
+                        print(f"   Root cause type: {type(e.__cause__).__name__}")
+
+                    print()
+                    print("💡 Troubleshooting:")
+                    print("   • Make sure the WebSocket server is running")
+                    print("   • Check that the server address is correct")
+                    print("   • Verify the server is listening on the right port")
+                    print(
+                        "   • Ensure both client and server use the same sec protocol"
+                    )
+                    if not use_plaintext:
+                        print("   • Noise over WebSocket may have compatibility issues")
+                    return
 
                 # Create a stream and send test data
                 try:
@@ -242,8 +357,18 @@ async def run(port: int, destination: str, use_noise: bool = False) -> None:
                 finally:
                     # Ensure stream is closed
                     try:
-                        if stream and not await stream.is_closed():
-                            await stream.close()
+                        if stream:
+                            # Check if stream has is_closed method and use it
+                            has_is_closed = hasattr(stream, "is_closed") and callable(
+                                getattr(stream, "is_closed")
+                            )
+                            if has_is_closed:
+                                # type: ignore[attr-defined]
+                                if not await stream.is_closed():
+                                    await stream.close()
+                            else:
+                                # Fallback: just try to close the stream
+                                await stream.close()
                     except Exception:
                         pass
 
@@ -256,7 +381,10 @@ async def run(port: int, destination: str, use_noise: bool = False) -> None:
                 print("✅ libp2p integration verified!")
                 print()
                 print("🚀 Your WebSocket transport is ready for production use!")
-                
+
+                # Add a small delay to ensure all cleanup is complete
+                await trio.sleep(0.1)
+
         except Exception as e:
             print(f"❌ Error creating WebSocket client: {e}")
             traceback.print_exc()
@@ -266,12 +394,15 @@ async def run(port: int, destination: str, use_noise: bool = False) -> None:
 def main() -> None:
     description = """
     This program demonstrates the libp2p WebSocket transport.
-    First run 'python websocket_demo.py -p <PORT> [--noise]' to start a WebSocket server.
-    Then run 'python websocket_demo.py <ANOTHER_PORT> -d <DESTINATION> [--noise]'
+    First run
+    'python websocket_demo.py -p <PORT> [--plaintext]' to start a WebSocket server.
+    Then run
+    'python websocket_demo.py <ANOTHER_PORT> -d <DESTINATION> [--plaintext]'
     where <DESTINATION> is the multiaddress shown by the server.
 
-    By default, this example uses plaintext security for communication.
-    Use --noise for testing with Noise encryption (experimental).
+    By default, this example uses Noise encryption for secure communication.
+    Use --plaintext for testing with unencrypted communication
+    (not recommended for production).
     """
 
     example_maddr = (
@@ -287,20 +418,30 @@ def main() -> None:
         help=f"destination multiaddr string, e.g. {example_maddr}",
     )
     parser.add_argument(
-        "--noise",
+        "--plaintext",
         action="store_true",
-        help="use Noise encryption instead of plaintext security",
+        help=(
+            "use plaintext security instead of Noise encryption "
+            "(not recommended for production)"
+        ),
     )
 
     args = parser.parse_args()
 
-    # Determine security mode: use plaintext by default, Noise if --noise is specified
-    use_noise = args.noise
-    
+    # Determine security mode: use Noise by default,
+    # plaintext if --plaintext is specified
+    use_plaintext = args.plaintext
+
     try:
-        trio.run(run, args.port, args.destination, use_noise)
+        trio.run(run, args.port, args.destination, use_plaintext)
     except KeyboardInterrupt:
-        pass
+        # This is expected when Ctrl+C is pressed
+        # The signal handler already printed the shutdown message
+        print("✅ Clean exit completed.")
+        return
+    except Exception as e:
+        print(f"❌ Unexpected error: {e}")
+        return
 
 
 if __name__ == "__main__":