From 9ca8fdc65028752165e4e4cdd04cea2bfadd0040 Mon Sep 17 00:00:00 2001
From: Sam Herrmann <samherrmann.github@gmail.com>
Date: Wed, 13 Aug 2025 10:54:12 -0400
Subject: [PATCH] Cancel Handler context when connection closes

---
 conn.go      |  6 +++++
 conn_test.go | 62 ++++++++++++++++++++++++++++++++++++++--------------
 jsonrpc2.go  |  8 +++----
 3 files changed, 55 insertions(+), 21 deletions(-)

diff --git a/conn.go b/conn.go
index 18466cb..33629a8 100644
--- a/conn.go
+++ b/conn.go
@@ -27,6 +27,7 @@ type Conn struct {
 
 	sending sync.Mutex
 
+	cancelCtx  context.CancelFunc
 	disconnect chan struct{}
 
 	logger Logger
@@ -46,10 +47,14 @@ var _ JSONRPC2 = (*Conn)(nil)
 // NewClient consumes conn, so you should call Close on the returned
 // client not on the given conn.
 func NewConn(ctx context.Context, stream ObjectStream, h Handler, opts ...ConnOpt) *Conn {
+
+	ctx, cancel := context.WithCancel(ctx)
+
 	c := &Conn{
 		stream:     stream,
 		h:          h,
 		pending:    map[ID]*call{},
+		cancelCtx:  cancel,
 		disconnect: make(chan struct{}),
 		logger:     log.New(os.Stderr, "", log.LstdFlags),
 	}
@@ -182,6 +187,7 @@ func (c *Conn) close(cause error) error {
 	}
 
 	close(c.disconnect)
+	c.cancelCtx()
 	c.closed = true
 	return c.stream.Close()
 }
diff --git a/conn_test.go b/conn_test.go
index aa32fd9..7313465 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -14,6 +14,40 @@ import (
 	"github.com/sourcegraph/jsonrpc2"
 )
 
+func TestConn(t *testing.T) {
+	t.Run("cancels context when closed", func(t *testing.T) {
+		ctxCanceled := make(chan struct{})
+
+		handler := handlerFunc(func(ctx context.Context, c *jsonrpc2.Conn, r *jsonrpc2.Request) {
+			// Block until the context is canceled.
+			<-ctx.Done()
+			close(ctxCanceled)
+		})
+
+		connA, connB := Pipe(noopHandler{}, jsonrpc2.AsyncHandler(handler))
+		defer connA.Close()
+		defer connB.Close()
+
+		// Send a notification from connA to connB to trigger connB's handler
+		// function.
+		if err := connA.Notify(context.Background(), "foo", nil, nil); err != nil {
+			t.Fatal(err)
+		}
+
+		// Disconnect connA from connB.
+		if err := connA.Close(); err != nil {
+			t.Fatal(err)
+		}
+
+		select {
+		case <-ctxCanceled:
+			// Test passed, the handler's context was canceled.
+		case <-time.After(time.Second):
+			t.Fatal("context not canceled")
+		}
+	})
+}
+
 var paramsTests = []struct {
 	sendParams interface{}
 	wantParams *json.RawMessage
@@ -198,12 +232,12 @@ func testParams(t *testing.T, want *json.RawMessage, fn func(c *jsonrpc2.Conn) e
 		wg.Done()
 	})
 
-	client, server := newClientServer(handler)
-	defer client.Close()
-	defer server.Close()
+	connA, connB := Pipe(noopHandler{}, handler)
+	defer connA.Close()
+	defer connB.Close()
 
 	wg.Add(1)
-	if err := fn(client); err != nil {
+	if err := fn(connA); err != nil {
 		t.Error(err)
 	}
 	wg.Wait()
@@ -242,18 +276,12 @@ func assertRawJSONMessage(t *testing.T, got *json.RawMessage, want *json.RawMess
 	}
 }
 
-func newClientServer(handler jsonrpc2.Handler) (client *jsonrpc2.Conn, server *jsonrpc2.Conn) {
+// Pipe returns two jsonrpc2.Conn, connected via a synchronous, in-memory, full
+// duplex network connection.
+func Pipe(handlerA, handlerB jsonrpc2.Handler) (connA *jsonrpc2.Conn, connB *jsonrpc2.Conn) {
 	ctx := context.Background()
-	connA, connB := net.Pipe()
-	client = jsonrpc2.NewConn(
-		ctx,
-		jsonrpc2.NewPlainObjectStream(connA),
-		noopHandler{},
-	)
-	server = jsonrpc2.NewConn(
-		ctx,
-		jsonrpc2.NewPlainObjectStream(connB),
-		handler,
-	)
-	return client, server
+	a, b := net.Pipe()
+	connA = jsonrpc2.NewConn(ctx, jsonrpc2.NewPlainObjectStream(a), handlerA)
+	connB = jsonrpc2.NewConn(ctx, jsonrpc2.NewPlainObjectStream(b), handlerB)
+	return connA, connB
 }
diff --git a/jsonrpc2.go b/jsonrpc2.go
index 97e26d7..7d3e132 100644
--- a/jsonrpc2.go
+++ b/jsonrpc2.go
@@ -59,10 +59,10 @@ const (
 
 // Handler handles JSON-RPC requests and notifications.
 type Handler interface {
-	// Handle is called to handle a request. No other requests are handled
-	// until it returns. If you do not require strict ordering behavior
-	// of received RPCs, it is suggested to wrap your handler in
-	// AsyncHandler.
+	// Handle is called to handle a request. No other requests are handled until
+	// it returns. If you do not require strict ordering behavior of received
+	// RPCs, it is suggested to wrap your handler in AsyncHandler. The context
+	// is automatically canceled when the connection closes.
 	Handle(context.Context, *Conn, *Request)
 }