fix: reconnect on restart instead of exiting

Previously the agent called os.Exit(0) when receiving a server restart
event, requiring Docker or manual intervention to bring it back. Now
it reconnects immediately with a fresh session, which is better for
both dev (no lost process) and production (faster recovery).

agent.go

@@ -68,8 +68,9 @@ func runAgent(ctx context.Context, wsURL, token string) {
 			return
 		}
 		if errors.Is(err, errRestartRequested) {
-			osExit(0)
+			slog.Info("restart requested, reconnecting immediately")
-			return
+			retryDelay = time.Second
+			continue
 		}
 		if err != nil {
 			slog.Error("agent disconnected", "error", err)

agent_test.go

@@ -10,6 +10,7 @@ import (
 	"net"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -1053,35 +1054,36 @@ func TestRunSessionReadErrorDuringJoin(t *testing.T) {
 }
 
 func TestRunAgentReconnectOnError(t *testing.T) {
-	// Server that fails first connection then succeeds, then sends restart
+	// Server that fails first connection then succeeds, then sends restart.
+	// Agent should reconnect (not exit) after both error and restart.
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer func() { _ = ln.Close() }()
 
-	connCount := 0
+	var connCount atomic.Int32
 	go func() {
 		for {
 			conn, err := ln.Accept()
 			if err != nil {
 				return
 			}
-			connCount++
+			count := connCount.Add(1)
 			buf := make([]byte, 4096)
 			n, _ := conn.Read(buf)
 			key := extractWSKey(string(buf[:n]))
 			accept := computeAcceptKey(key)
 
-			if connCount == 1 {
+			if count == 1 {
-				// First connection: upgrade then close immediately
+				// First connection: upgrade then close immediately (triggers error reconnect)
 				resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
 				_, _ = conn.Write([]byte(resp))
 				frameBuf := make([]byte, 4096)
 				_, _ = conn.Read(frameBuf)
 				_ = conn.Close()
-			} else {
+			} else if count == 2 {
-				// Second connection: proper session with restart
+				// Second connection: proper session with restart (triggers restart reconnect)
 				resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
 				_, _ = conn.Write([]byte(resp))
 				frameBuf := make([]byte, 4096)
@@ -1102,15 +1104,13 @@ func TestRunAgentReconnectOnError(t *testing.T) {
 				_, _ = conn.Write(makeTextFrame(restart))
 				time.Sleep(time.Second)
 				_ = conn.Close()
+			} else {
+				// Third connection: agent successfully reconnected after restart
+				_ = conn.Close()
 			}
 		}
 	}()
 
-	origExit := osExit
-	defer func() { osExit = origExit }()
-	exitCalled := make(chan int, 1)
-	osExit = func(code int) { exitCalled <- code }
-
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
@@ -1120,14 +1120,22 @@ func TestRunAgentReconnectOnError(t *testing.T) {
 		close(done)
 	}()
 
+	// Wait for 3rd connection attempt (proves agent reconnected after both error and restart)
+	deadline := time.After(10 * time.Second)
+	for {
+		if connCount.Load() >= 3 {
+			cancel()
+			break
+		}
 		select {
-	case code := <-exitCalled:
+		case <-deadline:
-		if code != 0 {
+			t.Fatalf("expected at least 3 connections, got %d", connCount.Load())
-			t.Errorf("expected exit code 0, got %d", code)
+		default:
+			time.Sleep(50 * time.Millisecond)
 		}
-	case <-time.After(10 * time.Second):
-		t.Error("runAgent did not reconnect and restart")
 	}
+
+	<-done
 }
 
 func TestRunAgentContextCancellation(t *testing.T) {
@@ -1149,41 +1157,32 @@ func TestRunAgentContextCancellation(t *testing.T) {
 }
 
 func TestRunAgentRestart(t *testing.T) {
-	origExit := osExit
+	// After receiving a restart event, the agent should reconnect (not exit).
-	defer func() { osExit = origExit }()
-
-	exitCalled := make(chan int, 1)
-	osExit = func(code int) {
-		exitCalled <- code
-	}
-
-	// Start a fake WebSocket server that accepts the join and sends restart
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer func() { _ = ln.Close() }()
 
+	var connCount atomic.Int32
 	go func() {
+		for {
 			conn, err := ln.Accept()
 			if err != nil {
 				return
 			}
-		defer func() { _ = conn.Close() }()
+			count := connCount.Add(1)
 
 			buf := make([]byte, 4096)
 			n, _ := conn.Read(buf)
-		reqStr := string(buf[:n])
+			key := extractWSKey(string(buf[:n]))
-		key := extractWSKey(reqStr)
 			accept := computeAcceptKey(key)
 			resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
 			_, _ = conn.Write([]byte(resp))
 
-		// Read join frame
 			frameBuf := make([]byte, 4096)
 			_, _ = conn.Read(frameBuf)
 
-		// Send join OK reply
 			reply, _ := json.Marshal(channelMsg{
 				Topic:   "agent:agent-0",
 				Event:   "phx_reply",
@@ -1192,7 +1191,8 @@ func TestRunAgentRestart(t *testing.T) {
 			})
 			_, _ = conn.Write(makeTextFrame(reply))
 
-		// Small delay then send restart event
+			if count == 1 {
+				// First connection: send restart event
 				time.Sleep(50 * time.Millisecond)
 				restart, _ := json.Marshal(channelMsg{
 					Topic:   "agent:agent-0",
@@ -1200,29 +1200,37 @@ func TestRunAgentRestart(t *testing.T) {
 					Payload: json.RawMessage(`{}`),
 				})
 				_, _ = conn.Write(makeTextFrame(restart))
-
+				time.Sleep(time.Second)
-		// Keep connection open briefly
+			}
-		time.Sleep(2 * time.Second)
+			_ = conn.Close()
+		}
 	}()
 
-	addr := ln.Addr().String()
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
 	done := make(chan struct{})
 	go func() {
-		runAgent(ctx, "ws://"+addr, "test-token")
+		runAgent(ctx, "ws://"+ln.Addr().String(), "test-token")
 		close(done)
 	}()
 
+	// Wait for 2nd connection (proves agent reconnected after restart instead of exiting)
+	deadline := time.After(5 * time.Second)
+	for {
+		if connCount.Load() >= 2 {
+			cancel()
+			break
+		}
 		select {
-	case code := <-exitCalled:
+		case <-deadline:
-		if code != 0 {
+			t.Fatalf("expected at least 2 connections, got %d", connCount.Load())
-			t.Errorf("expected exit code 0, got %d", code)
+		default:
+			time.Sleep(50 * time.Millisecond)
 		}
-	case <-time.After(5 * time.Second):
-		t.Error("osExit was not called after restart")
 	}
+
+	<-done
 }
 
 // readMaskedFrame reads a single masked WebSocket frame from the server side.