From 2308db45bbe61c0e0850c7ca6977412fe5ab4dab Mon Sep 17 00:00:00 2001
From: rcourtman <8825017+rcourtman@users.noreply.github.com>
Date: Thu, 24 Sep 2026 18:27:22 +0100
Subject: [PATCH] Wait for slot release in the WebSocket connection limit test

Agent removal occurs inside read-loop cleanup before the handler releases
its per-IP connection slot. The flood test could observe the first state
and retry while the slot was still occupied, causing a spurious rejection.

Wait for the guarded admission count to reach zero before reconnecting,
while preserving rejection at capacity and successful registration after
release.
---
 internal/agentexec/server_websocket_test.go | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/internal/agentexec/server_websocket_test.go b/internal/agentexec/server_websocket_test.go
index 763eb13e07..51559ae970 100644
--- a/internal/agentexec/server_websocket_test.go
+++ b/internal/agentexec/server_websocket_test.go
@@ -672,6 +672,7 @@ func TestHandleWebSocket_RejectsPerIPConnectionFlood(t *testing.T) {
 		t.Fatalf("Dial first connection: %v", err)
 	}
 	defer firstConn.Close()
+	remoteIP := normalizeWebSocketRemoteIP(firstConn.LocalAddr().String())
 
 	wsWriteMessage(t, firstConn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
 		AgentID:  "a1",
@@ -699,6 +700,13 @@ func TestHandleWebSocket_RejectsPerIPConnectionFlood(t *testing.T) {
 
 	firstConn.Close()
 	waitFor(t, 2*time.Second, func() bool { return !s.IsAgentConnected("a1") })
+	// The read loop removes the agent before the handler's deferred slot release.
+	// Wait for admission capacity itself before testing that it can be reused.
+	waitFor(t, 2*time.Second, func() bool {
+		s.mu.RLock()
+		defer s.mu.RUnlock()
+		return s.ipConnCounts[remoteIP] == 0
+	})
 
 	thirdConn, _, err := dialAgentExecWebSocket(t, ts.URL)
 	if err != nil {