fix: gracefull conection drop

Author: krulsaidme0w

acp_test.go

@@ -2,6 +2,7 @@ package acp
 
 import (
 	"context"
+	"encoding/json"
 	"io"
 	"slices"
 	"sync"
@@ -467,6 +468,59 @@ func TestConnectionHandlesNotifications(t *testing.T) {
 	}
 }
 
+func TestConnection_DoesNotCancelInboundContextBeforeDrainingNotificationsOnDisconnect(t *testing.T) {
+	const n = 25
+
+	incomingR, incomingW := io.Pipe()
+
+	var (
+		wg            sync.WaitGroup
+		canceledCount atomic.Int64
+	)
+	wg.Add(n)
+
+	c := NewConnection(func(ctx context.Context, method string, _ json.RawMessage) (any, *RequestError) {
+		defer wg.Done()
+		// Slow down processing so some notifications are handled after the receive
+		// loop observes EOF and signals disconnect.
+		time.Sleep(10 * time.Millisecond)
+		if ctx.Err() != nil {
+			canceledCount.Add(1)
+		}
+		return nil, nil
+	}, io.Discard, incomingR)
+
+	// Write notifications quickly and then close the stream to simulate a peer disconnect.
+	for i := 0; i < n; i++ {
+		if _, err := io.WriteString(incomingW, `{"jsonrpc":"2.0","method":"test/notify","params":{}}`+"\n"); err != nil {
+			t.Fatalf("write notification: %v", err)
+		}
+	}
+	_ = incomingW.Close()
+
+	select {
+	case <-c.Done():
+		// Expected: peer disconnect observed promptly.
+	case <-time.After(2 * time.Second):
+		t.Fatalf("timeout waiting for connection Done()")
+	}
+
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+	select {
+	case <-done:
+	case <-time.After(3 * time.Second):
+		t.Fatalf("timeout waiting for notification handlers")
+	}
+
+	if got := canceledCount.Load(); got != 0 {
+		t.Fatalf("inbound handler context was canceled for %d/%d notifications", got, n)
+	}
+}
+
 // Test initialize method behavior
 func TestConnectionHandlesInitialize(t *testing.T) {
 	c2aR, c2aW := io.Pipe()

connection.go

@@ -10,6 +10,7 @@ import (
 	"log/slog"
 	"sync"
 	"sync/atomic"
+	"time"
 )
 
 type anyMessage struct {
@@ -152,11 +153,26 @@ func (c *Connection) receive() {
 	// First, signal disconnect to callers waiting on responses.
 	c.cancel(cause)
 
-	// Then close and drain notifications so already-received messages are delivered
-	// to handlers with a still-valid inbound context.
+	// Then close the notification queue so already-received messages can drain.
+	// IMPORTANT: Do not block this receive goroutine waiting for the drain to complete;
+	// notification handlers may legitimately block until their context is canceled.
 	close(c.notificationQueue)
-	c.notificationWg.Wait()
-	c.inboundCancel(cause)
+
+	// Cancel inboundCtx after notifications finish, but ensure we don't leak forever if a
+	// handler blocks waiting for cancellation.
+	const drainTimeout = 5 * time.Second
+	go func() {
+		done := make(chan struct{})
+		go func() {
+			c.notificationWg.Wait()
+			close(done)
+		}()
+		select {
+		case <-done:
+		case <-time.After(drainTimeout):
+		}
+		c.inboundCancel(cause)
+	}()
 
 	c.loggerOrDefault().Info("peer connection closed")
 }