From 522ef33f69e4f7818ff8cfd103a67aa303fff284 Mon Sep 17 00:00:00 2001
From: lc6464 <64722907+lc6464@users.noreply.github.com>
Date: Tue, 24 Mar 2026 21:36:51 +0800
Subject: [PATCH] fix(pico): enforce MaxConnections atomically on registration

---
 pkg/channels/pico/pico.go | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go
index ab93f18b3..c948f9d4a 100644
--- a/pkg/channels/pico/pico.go
+++ b/pkg/channels/pico/pico.go
@@ -99,10 +99,13 @@ func NewPicoChannel(cfg config.PicoConfig, messageBus *bus.MessageBus) (*PicoCha
 	}, nil
 }
 
-// createAndAddConnection generates a unique connID and registers it atomically.
-func (c *PicoChannel) createAndAddConnection(conn *websocket.Conn, sessionID string) *picoConn {
+// createAndAddConnection checks MaxConnections and registers a connection atomically.
+func (c *PicoChannel) createAndAddConnection(conn *websocket.Conn, sessionID string, maxConns int) (*picoConn, error) {
 	c.connsMu.Lock()
 	defer c.connsMu.Unlock()
+	if c.connCount >= maxConns {
+		return nil, channels.ErrSendFailed
+	}
 
 	var connID string
 	for {
@@ -127,7 +130,7 @@ func (c *PicoChannel) createAndAddConnection(conn *websocket.Conn, sessionID str
 	bySession[pc.id] = pc
 	c.connCount++
 
-	return pc
+	return pc, nil
 }
 
 // removeConnection deletes a connection from indexes and returns it when found.
@@ -361,7 +364,16 @@ func (c *PicoChannel) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 		sessionID = uuid.New().String()
 	}
 
-	pc := c.createAndAddConnection(conn, sessionID)
+	pc, err := c.createAndAddConnection(conn, sessionID, maxConns)
+	if err != nil {
+		_ = conn.WriteControl(
+			websocket.CloseMessage,
+			websocket.FormatCloseMessage(websocket.CloseTryAgainLater, "too many connections"),
+			time.Now().Add(2*time.Second),
+		)
+		_ = conn.Close()
+		return
+	}
 
 	logger.InfoCF("pico", "WebSocket client connected", map[string]any{
 		"conn_id":    pc.id,