diff --git a/agent/src/session/mod.rs b/agent/src/session/mod.rs
index 2339bc2..ce689d0 100644
--- a/agent/src/session/mod.rs
+++ b/agent/src/session/mod.rs
@@ -128,8 +128,13 @@ impl SessionManager {
             tokio::time::sleep(Duration::from_millis(1)).await;
 
             // Check if still connected
-            if !transport.is_connected() {
-                tracing::warn!("Connection lost");
+            if let Some(transport) = self.transport.as_ref() {
+                if !transport.is_connected() {
+                    tracing::warn!("Connection lost");
+                    break;
+                }
+            } else {
+                tracing::warn!("Transport is None");
                 break;
             }
         }