diff --git a/09-behavioral-design-patterns/02-state-failsafe-socket/offlineState.js b/09-behavioral-design-patterns/02-state-failsafe-socket/offlineState.js
index ac7a8b3..1bdf18c 100644
--- a/09-behavioral-design-patterns/02-state-failsafe-socket/offlineState.js
+++ b/09-behavioral-design-patterns/02-state-failsafe-socket/offlineState.js
@@ -3,15 +3,27 @@ import { createConnection } from 'node:net'
 export class OfflineState {
   constructor(failsafeSocket) {
     this.failsafeSocket = failsafeSocket
+    this._timer = null
   }
 
   send(data) {
     this.failsafeSocket.queue.push(data)
   }
 
+  resetTimer() {
+    if (this._timer) {
+      clearTimeout(this._timer)
+      this._timer = null
+    }
+  }
+
   activate() {
     const retry = () => {
-      setTimeout(() => this.activate(), 1000)
+      this.resetTimer()
+      this._timer = setTimeout(() => {
+        this.resetTimer()
+        this.activate()
+      }, 1000)
     }
 
     console.log(
@@ -22,6 +34,7 @@ export class OfflineState {
       () => {
         console.log('Connection established')
         this.failsafeSocket.socket.removeListener('error', retry)
+        this.failsafeSocket.socket.on('error', retry)
         this.failsafeSocket.changeState('online')
       }
     )
diff --git a/09-behavioral-design-patterns/02-state-failsafe-socket/onlineState.js b/09-behavioral-design-patterns/02-state-failsafe-socket/onlineState.js
index adc41f9..512186f 100644
--- a/09-behavioral-design-patterns/02-state-failsafe-socket/onlineState.js
+++ b/09-behavioral-design-patterns/02-state-failsafe-socket/onlineState.js
@@ -44,6 +44,18 @@ export class OnlineState {
   }
 
   activate() {
+    this.failsafeSocket.socket.once('error', err => {
+      console.error('Connection lost:', err.message)
+      this.failsafeSocket.socket.destroy()
+      this.failsafeSocket.socket = null
+    })
+
+    this.failsafeSocket.socket.once('close', () => {
+      console.log('Connection closed')
+      this.failsafeSocket.socket = null
+      this.failsafeSocket.changeState('offline')
+    })
+
     this._tryFlush()
   }
 }