fix(lsp): dispatch notifications off read loop

internal/lsp/client.go

@@ -13,8 +13,14 @@ import (
 )
 
 // NotificationHandler receives server->client notifications (e.g.
-// textDocument/publishDiagnostics). params is the raw JSON payload.
-type NotificationHandler func(method string, params json.RawMessage)
+// textDocument/publishDiagnostics). params is the raw JSON payload. seq is the
+// notification's receipt sequence (see Client.NotificationSeq): the count of
+// notifications read off the wire, including this one, at the moment it was
+// enqueued — NOT when this handler happens to run, which can lag receipt when
+// the queue is backed up. A caller that needs to know whether something newer
+// than a given point has arrived must compare against seq, not against when
+// its own handling code runs.
+type NotificationHandler func(method string, params json.RawMessage, seq int64)
 
 // Client speaks JSON-RPC 2.0 with LSP framing (Content-Length headers) over a
 // reader/writer pair. It is transport-agnostic: server.go wires it to a process's
@@ -32,8 +38,35 @@ type Client struct {
 	closeOnce sync.Once
 	closed    chan struct{}
 	readErr   error
+
+	notifyMu     sync.Mutex
+	notifyQueue  []notification
+	notifyBytes  int
+	notifyReady  chan struct{}
+	notifyClosed bool
+	notifySeq    int64 // count of notifications received (enqueued) so far
+}
+
+type notification struct {
+	method string
+	params json.RawMessage
+	seq    int64
 }
 
+// The notification backlog is bounded by both message count and retained
+// payload bytes. A well-behaved handler drains far faster than any single burst
+// fills either limit; sustained overload — a language server emitting faster
+// than the single handler can consume, or a handler stuck waiting on a
+// re-entrant Call — is a fatal condition for this client, not something to
+// paper over by growing the queue without bound. Hitting either limit fails the
+// client (see enqueueNotification): IsClosed becomes true, and the manager
+// evicts and restarts the session on next use, exactly as it does for any other
+// dead client.
+const (
+	notifyQueueLimit     = 4096
+	notifyQueueByteLimit = 16 << 20 // 16 MiB
+)
+
 type rpcError struct {
 	Code    int             `json:"code"`
 	Message string          `json:"message"`
@@ -81,10 +114,12 @@ type incomingMessage struct {
 // server process exits); call Close to stop using the client.
 func NewClient(r io.Reader, w io.Writer) *Client {
 	client := &Client{
-		writer:  w,
-		pending: make(map[int64]chan rpcResponse),
-		closed:  make(chan struct{}),
+		writer:      w,
+		pending:     make(map[int64]chan rpcResponse),
+		closed:      make(chan struct{}),
+		notifyReady: make(chan struct{}, 1),
 	}
+	go client.notificationLoop()
 	go client.readLoop(bufio.NewReader(r))
 	return client
 }
@@ -166,12 +201,7 @@ func (c *Client) readLoop(reader *bufio.Reader) {
 			// required or the server can block waiting on it (e.g. registerCapability).
 			_ = c.write(outgoingReply{JSONRPC: "2.0", ID: msg.ID, Result: nil})
 		case msg.Method != "":
-			c.mu.Lock()
-			handler := c.handler
-			c.mu.Unlock()
-			if handler != nil {
-				handler(msg.Method, msg.Params)
-			}
+			c.enqueueNotification(notification{method: msg.Method, params: msg.Params})
 		case hasID:
 			var id int64
 			if err := json.Unmarshal(msg.ID, &id); err == nil {
@@ -181,6 +211,117 @@ func (c *Client) readLoop(reader *bufio.Reader) {
 	}
 }
 
+func (c *Client) notificationLoop() {
+	for {
+		select {
+		case <-c.closed:
+			return
+		case <-c.notifyReady:
+			for {
+				notification, ok := c.dequeueNotification()
+				if !ok {
+					break
+				}
+				c.mu.Lock()
+				handler := c.handler
+				c.mu.Unlock()
+				if handler != nil {
+					handler(notification.method, notification.params, notification.seq)
+				}
+			}
+		}
+	}
+}
+
+// enqueueNotification hands a server notification to the worker loop. It never
+// blocks and never silently discards a message the handler could still act on:
+// the queue grows instead, up to its message and byte limits.
+//
+// The alternatives to growing are worse. Blocking the read loop when a buffer
+// fills is the deadlock this dispatch exists to avoid — a handler that calls
+// Client.Call waits for a response frame the blocked reader can no longer
+// deliver. Dropping the oldest queued item instead loses protocol state
+// permanently: a textDocument/publishDiagnostics for one URI is the server's
+// only report for that URI, so discarding it makes session.waitForDiagnostics
+// time out and Manager.Check return nothing even though the server published
+// findings.
+//
+// Growth is bounded in practice by how much the server emits while a handler
+// runs, and the queue is released as soon as it drains. But "in practice" is
+// not a limit: a handler that never returns, or a server that sustains a
+// higher rate than the single handler can drain, would otherwise grow this
+// queue's full json.RawMessage payloads without bound until the heap gives
+// out. The count limit alone is insufficient because notification payloads are
+// peer-controlled and can be large; the byte limit prevents a few large
+// diagnostics publishes from exhausting the heap before the count limit is
+// reached. Either limit turns overload into an explicit, observable failure —
+// the client is failed and closed — rather than unbounded protocol retention.
+func (c *Client) enqueueNotification(item notification) {
+	c.notifyMu.Lock()
+	if c.notifyClosed {
+		// The worker loop has already stopped, so anything queued now would never
+		// be handled. Retaining it would grow the queue for as long as the
+		// transport stays readable after Close — Server.Shutdown closes the client
+		// before closing stdin, so a server emitting notifications while it handles
+		// shutdown/exit keeps the read loop feeding a queue nobody drains.
+		c.notifyMu.Unlock()
+		return
+	}
+	itemBytes := len(item.method) + len(item.params)
+	if len(c.notifyQueue) >= notifyQueueLimit || itemBytes > notifyQueueByteLimit-c.notifyBytes {
+		c.notifyMu.Unlock()
+		c.failPending(fmt.Errorf(
+			"lsp client: notification backlog exceeded %d messages or %d bytes",
+			notifyQueueLimit,
+			notifyQueueByteLimit,
+		))
+		return
+	}
+	c.notifySeq++
+	item.seq = c.notifySeq
+	c.notifyQueue = append(c.notifyQueue, item)
+	c.notifyBytes += itemBytes
+	c.notifyMu.Unlock()
+
+	select {
+	case c.notifyReady <- struct{}{}:
+	default:
+		// A wake-up is already pending; the worker drains the whole queue per wake,
+		// so this item is covered by it.
+	}
+}
+
+// NotificationSeq returns the number of notifications received (read off the
+// wire and enqueued) so far, including any still waiting to be dispatched to
+// the handler. A caller that wants to know whether a notification newer than
+// "now" has arrived should snapshot this before triggering whatever produces
+// it, then require a subsequently-observed seq to be strictly greater: a
+// notification already sitting in the queue at snapshot time has seq <= the
+// snapshot, even if the handler doesn't run for it until afterward.
+func (c *Client) NotificationSeq() int64 {
+	c.notifyMu.Lock()
+	defer c.notifyMu.Unlock()
+	return c.notifySeq
+}
+
+func (c *Client) dequeueNotification() (notification, bool) {
+	c.notifyMu.Lock()
+	defer c.notifyMu.Unlock()
+	if len(c.notifyQueue) == 0 {
+		return notification{}, false
+	}
+	item := c.notifyQueue[0]
+	c.notifyBytes -= len(item.method) + len(item.params)
+	c.notifyQueue[0] = notification{}
+	c.notifyQueue = c.notifyQueue[1:]
+	if len(c.notifyQueue) == 0 {
+		// Release the backing array once drained; re-slicing alone would keep a
+		// burst's worth of capacity (and its already-consumed items) alive.
+		c.notifyQueue = nil
+	}
+	return item, true
+}
+
 func (c *Client) deliver(id int64, resp rpcResponse) {
 	c.mu.Lock()
 	ch, ok := c.pending[id]
@@ -204,9 +345,27 @@ func (c *Client) failPending(err error) {
 			ch <- rpcResponse{Err: &rpcError{Code: -1, Message: err.Error()}}
 		}
 		close(c.closed)
+		c.closeNotifications()
 	})
 }
 
+// closeNotifications stops accepting notifications and releases anything still
+// queued, so a closed client retains nothing. It runs inside closeOnce, after
+// c.closed is signaled: an enqueue racing with shutdown therefore either sees
+// notifyClosed and drops its item, or appends just before the flag is set and has
+// its item cleared here.
+//
+// Queued-but-unhandled notifications are dropped rather than drained: the worker
+// loop is already gone by definition of close, and callers that need diagnostics
+// (Manager.Check) collect them before shutting the client down.
+func (c *Client) closeNotifications() {
+	c.notifyMu.Lock()
+	c.notifyClosed = true
+	c.notifyQueue = nil
+	c.notifyBytes = 0
+	c.notifyMu.Unlock()
+}
+
 func (c *Client) readError() error {
 	c.mu.Lock()
 	defer c.mu.Unlock()