fix(stub): resolve TCP socket leak on target server disconnect and hanging connection attempt

server/src/fix_stub.py

@@ -0,0 +1,66 @@
+import os
+
+def fix_stub_leaks():
+    filepath = 'stub/src/main.rs'
+    with open(filepath, 'r', encoding='utf-8') as f:
+        content = f.read()
+
+    # 1. Fix connect_async to handle shutdown and timeout
+    old_connect = """            match tokio_tungstenite::connect_async(request).await {"""
+    new_connect = """            let connect_result = tokio::select! {
+                _ = shutdown_rx.recv() => {
+                    tracing::info!("Shutdown received during connect");
+                    break;
+                }
+                res = tokio::time::timeout(
+                    tokio::time::Duration::from_secs(5),
+                    tokio_tungstenite::connect_async(request)
+                ) => res,
+            };
+
+            match connect_result {
+                Ok(Ok((ws_stream, _))) => {"""
+    
+    if old_connect in content:
+        content = content.replace(old_connect, new_connect)
+        print("Replaced connect_async")
+
+    # Fix Err block to match the new match structure
+    old_err = """                Err(e) => {
+                    tracing::error!("Failed to connect via WSS: {}", e);
+                    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+                }"""
+    new_err = """                Ok(Err(e)) => {
+                    tracing::error!("Failed to connect via WSS: {}", e);
+                    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+                }
+                Err(_) => {
+                    tracing::error!("Connection attempt timed out");
+                    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+                }"""
+    
+    if old_err in content:
+        content = content.replace(old_err, new_err)
+        print("Replaced Err branch")
+
+    # 2. Fix the break in send_task that exits the stub instead of reconnecting
+    old_select_send = """                        _ = &mut send_task => {
+                            tracing::error!("Send task exited");
+                            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+                            recv_task.abort();
+                            break;
+                        }"""
+    new_select_send = """                        _ = &mut send_task => {
+                            tracing::error!("Send task exited");
+                            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+                            recv_task.abort();
+                        }"""
+
+    if old_select_send in content:
+        content = content.replace(old_select_send, new_select_send)
+        print("Replaced select send_task")
+
+    with open(filepath, 'w', encoding='utf-8') as f:
+        f.write(content)
+
+fix_stub_leaks()