From bb87722ac29db0ab48ec5575e3ddc194a3b6355b Mon Sep 17 00:00:00 2001
From: Sam Herrmann <samherrmann.github@gmail.com>
Date: Thu, 14 Aug 2025 11:35:57 -0400
Subject: [PATCH] Close Conn when given context is done

---
 conn.go      | 12 ++++++++++--
 conn_test.go | 25 +++++++++++++++++++++----
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/conn.go b/conn.go
index 33629a8..35c6279 100644
--- a/conn.go
+++ b/conn.go
@@ -44,8 +44,10 @@ var _ JSONRPC2 = (*Conn)(nil)
 // JSON-RPC protocol is symmetric, so a Conn runs on both ends of a
 // client-server connection.
 //
-// NewClient consumes conn, so you should call Close on the returned
-// client not on the given conn.
+// NewConn consumes stream, so you should call Close on the returned
+// Conn not on the given stream or its underlying connection.
+//
+// Conn is closed when the given context's Done channel is closed.
 func NewConn(ctx context.Context, stream ObjectStream, h Handler, opts ...ConnOpt) *Conn {
 
 	ctx, cancel := context.WithCancel(ctx)
@@ -65,6 +67,12 @@ func NewConn(ctx context.Context, stream ObjectStream, h Handler, opts ...ConnOp
 		opt(c)
 	}
 	go c.readMessages(ctx)
+
+	go func() {
+		<-ctx.Done()
+		c.close(nil)
+	}()
+
 	return c
 }
 
diff --git a/conn_test.go b/conn_test.go
index 7313465..5d2a7e4 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -15,6 +15,24 @@ import (
 )
 
 func TestConn(t *testing.T) {
+
+	t.Run("closes when context is done", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+
+		connA, connB := Pipe(ctx, noopHandler{}, noopHandler{})
+		defer connA.Close()
+		defer connB.Close()
+
+		cancel()
+		<-connA.DisconnectNotify()
+
+		got := connA.Close()
+		want := jsonrpc2.ErrClosed
+		if got != want {
+			t.Fatalf("got %v, want %v", got, want)
+		}
+	})
+
 	t.Run("cancels context when closed", func(t *testing.T) {
 		ctxCanceled := make(chan struct{})
 
@@ -24,7 +42,7 @@ func TestConn(t *testing.T) {
 			close(ctxCanceled)
 		})
 
-		connA, connB := Pipe(noopHandler{}, jsonrpc2.AsyncHandler(handler))
+		connA, connB := Pipe(context.Background(), noopHandler{}, jsonrpc2.AsyncHandler(handler))
 		defer connA.Close()
 		defer connB.Close()
 
@@ -232,7 +250,7 @@ func testParams(t *testing.T, want *json.RawMessage, fn func(c *jsonrpc2.Conn) e
 		wg.Done()
 	})
 
-	connA, connB := Pipe(noopHandler{}, handler)
+	connA, connB := Pipe(context.Background(), noopHandler{}, handler)
 	defer connA.Close()
 	defer connB.Close()
 
@@ -278,8 +296,7 @@ func assertRawJSONMessage(t *testing.T, got *json.RawMessage, want *json.RawMess
 
 // 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()
+func Pipe(ctx context.Context, handlerA, handlerB jsonrpc2.Handler) (connA *jsonrpc2.Conn, connB *jsonrpc2.Conn) {
 	a, b := net.Pipe()
 	connA = jsonrpc2.NewConn(ctx, jsonrpc2.NewPlainObjectStream(a), handlerA)
 	connB = jsonrpc2.NewConn(ctx, jsonrpc2.NewPlainObjectStream(b), handlerB)