From a0c37d4c863bfc0e45cbdb158bd5d11145bae160 Mon Sep 17 00:00:00 2001
From: xioshark <151488904+Xio-Shark@users.noreply.github.com>
Date: Sat, 17 Oct 2026 19:47:50 +0800
Subject: [PATCH 1/4] internal/jsonrpc2: handle preceding notifications before
 retiring call responses

In readIncoming, incoming *Request messages (including notifications) are
enqueued onto handlerQueue and processed sequentially by handleAsync,
whereas incoming *Response messages were retired immediately on the
reader goroutine. As a result, a call (such as CallTool) could return its
result before the client had finished running handlers for notifications
(such as notifications/progress or notifications/message) that preceded
the response on the wire.

Deliver incoming *Response messages through handlerQueue when
handleAsync is active so that all synchronous notification handlers
queued ahead of a response complete before the awaiting call returns.
When a synchronous handler itself issues a nested Call on the same
connection, retire its response immediately so the handler does not block
waiting on its own handlerQueue slot.

Fixes #1337

Signed-off-by: xioshark <151488904+Xio-Shark@users.noreply.github.com>
---
 internal/jsonrpc2/conn.go      |  84 ++++++++++---
 internal/jsonrpc2/conn_test.go | 116 ++++++++++++++++++
 mcp/client_test.go             | 214 +++++++++++++++++++++++++++++++++
 mcp/streamable_test.go         |   2 +-
 4 files changed, 399 insertions(+), 17 deletions(-)

diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go
index 22e6faeea..222fefded 100644
--- a/internal/jsonrpc2/conn.go
+++ b/internal/jsonrpc2/conn.go
@@ -86,11 +86,24 @@ type inFlightState struct {
 
 	incomingByID map[ID]*incomingRequest // calls only
 
-	// handlerQueue stores the backlog of calls and notifications that were not
-	// already handled by a preempter.
-	// The queue does not include the request currently being handled (if any).
-	handlerQueue   []*incomingRequest
+	// handlerQueue stores the backlog of requests and responses that must be
+	// delivered sequentially. Calls and notifications not already handled by a
+	// preempter are queued here, along with responses that arrived while the
+	// handler was busy so that notifications preceding a response on the wire
+	// are handled before the call awaiting that response returns.
+	// The queue does not include the item currently being handled (if any).
+	handlerQueue   []queueItem
 	handlerRunning bool
+	activeReleaser *releaser
+}
+
+// queueItem is an entry in [inFlightState.handlerQueue]: either an incoming
+// request to pass to [Handler.Handle], or a response to deliver to an
+// [AsyncCall].
+type queueItem struct {
+	req  *incomingRequest
+	ac   *AsyncCall
+	resp *Response
 }
 
 // updateInFlight locks the state of the connection's in-flight requests, allows
@@ -345,6 +358,10 @@ func (c *Connection) Call(ctx context.Context, method string, params any) *Async
 		if err != nil {
 			return
 		}
+		if r, ok := ctx.Value(asyncKey).(*releaser); ok && r.conn == c && !r.isReleased() {
+			ac.callerReleaser = r
+		}
+		ac.activeAtStart = s.activeSyncReleaser()
 		if s.outgoingCalls == nil {
 			s.outgoingCalls = make(map[ID]*AsyncCall)
 		}
@@ -397,6 +414,7 @@ var asyncKey = asyncKeyType{}
 // request is released when it is allowed to run concurrent with other
 // requests, via a call to [Async].)
 type releaser struct {
+	conn     *Connection
 	mu       sync.Mutex
 	ch       chan struct{}
 	released bool
@@ -418,10 +436,25 @@ func (r *releaser) release(soft bool) {
 	}
 }
 
+func (r *releaser) isReleased() bool {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.released
+}
+
+func (s *inFlightState) activeSyncReleaser() *releaser {
+	if s.activeReleaser != nil && !s.activeReleaser.isReleased() {
+		return s.activeReleaser
+	}
+	return nil
+}
+
 type AsyncCall struct {
-	id       ID
-	ready    chan struct{} // closed after response has been set
-	response *Response
+	id             ID
+	ready          chan struct{} // closed after response has been set
+	response       *Response
+	callerReleaser *releaser
+	activeAtStart  *releaser
 }
 
 // ID used for this call.
@@ -574,7 +607,12 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter
 			c.updateInFlight(func(s *inFlightState) {
 				if ac, ok := s.outgoingCalls[msg.ID]; ok {
 					delete(s.outgoingCalls, msg.ID)
-					ac.retire(msg)
+					active := s.activeSyncReleaser()
+					if !s.handlerRunning || (active != nil && (ac.callerReleaser == active || (len(s.handlerQueue) == 0 && ac.activeAtStart == active))) {
+						ac.retire(msg)
+					} else {
+						s.handlerQueue = append(s.handlerQueue, queueItem{ac: ac, resp: msg})
+					}
 				} else {
 					// TODO: How should we report unexpected responses?
 				}
@@ -674,7 +712,7 @@ func (c *Connection) acceptRequest(ctx context.Context, msg *Request, preempter
 		// (because it may be responding to a request issued by, say, an
 		// asynchronous handler), and in order to get to that response we have
 		// to read all of the requests that came in ahead of it.
-		s.handlerQueue = append(s.handlerQueue, req)
+		s.handlerQueue = append(s.handlerQueue, queueItem{req: req})
 		if !s.handlerRunning {
 			// We start the handleAsync goroutine when it has work to do, and let it
 			// exit when the queue empties.
@@ -698,17 +736,31 @@ func (c *Connection) acceptRequest(ctx context.Context, msg *Request, preempter
 }
 
 // handleAsync invokes the handler on the requests in the handler queue
-// sequentially until the queue is empty.
+// sequentially until the queue is empty, retiring any queued responses in
+// order along the way.
 func (c *Connection) handleAsync() {
 	for {
-		var req *incomingRequest
+		var (
+			item queueItem
+			rel  *releaser
+		)
 		c.updateInFlight(func(s *inFlightState) {
+			s.activeReleaser = nil
 			if len(s.handlerQueue) > 0 {
-				req, s.handlerQueue = s.handlerQueue[0], s.handlerQueue[1:]
+				item, s.handlerQueue = s.handlerQueue[0], s.handlerQueue[1:]
+				if item.req != nil {
+					rel = &releaser{conn: c, ch: make(chan struct{})}
+					s.activeReleaser = rel
+				}
 			} else {
 				s.handlerRunning = false
 			}
 		})
+		if item.ac != nil {
+			item.ac.retire(item.resp)
+			continue
+		}
+		req := item.req
 		if req == nil {
 			return
 		}
@@ -718,6 +770,7 @@ func (c *Connection) handleAsync() {
 		// in-flight request), report that cause rather than the bare
 		// context.Canceled.
 		if err := req.ctx.Err(); err != nil {
+			rel.release(true)
 			if cause := context.Cause(req.ctx); cause != nil {
 				err = cause
 			}
@@ -725,14 +778,13 @@ func (c *Connection) handleAsync() {
 			continue
 		}
 
-		releaser := &releaser{ch: make(chan struct{})}
-		ctx := context.WithValue(req.ctx, asyncKey, releaser)
+		ctx := context.WithValue(req.ctx, asyncKey, rel)
 		go func() {
-			defer releaser.release(true)
+			defer rel.release(true)
 			result, err := c.handler.Handle(ctx, req.Request)
 			c.processResult(c.handler, req, result, err)
 		}()
-		<-releaser.ch
+		<-rel.ch
 	}
 }
 
diff --git a/internal/jsonrpc2/conn_test.go b/internal/jsonrpc2/conn_test.go
index 0893953ca..7279b7fb8 100644
--- a/internal/jsonrpc2/conn_test.go
+++ b/internal/jsonrpc2/conn_test.go
@@ -13,6 +13,7 @@ import (
 	"slices"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 )
 
@@ -259,3 +260,118 @@ func (w *dropWriter) drops() []ID {
 	defer w.mu.Unlock()
 	return slices.Clone(w.dropped)
 }
+
+// TestResponseWaitsForPrecedingNotifications verifies that when notifications
+// arrive on the wire ahead of a call's response, all preceding notification
+// handlers complete before AsyncCall.Await returns (go-sdk#1337).
+func TestResponseWaitsForPrecedingNotifications(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	const numNotifications = 5
+	r := fakeReader{msgs: make(chan Message, numNotifications+1)}
+	w := fakeWriter{out: make(chan Message, 1)}
+
+	var handled atomic.Int32
+	firstStarted := make(chan struct{})
+	releaseFirst := make(chan struct{})
+
+	conn := NewConnection(ctx, ConnectionConfig{
+		Reader: r,
+		Writer: w,
+		Closer: nopCloser{},
+		Bind: func(*Connection) Handler {
+			return HandlerFunc(func(_ context.Context, req *Request) (any, error) {
+				if handled.Add(1) == 1 {
+					close(firstStarted)
+					<-releaseFirst
+				}
+				return nil, nil
+			})
+		},
+	})
+
+	ac := conn.Call(ctx, "work", nil)
+	sent := (<-w.out).(*Request)
+
+	for range numNotifications {
+		n, err := NewNotification("notify", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		r.msgs <- n
+	}
+	resp, err := NewResponse(sent.ID, "ok", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	r.msgs <- resp
+
+	// Wait until the first notification handler is running and blocking handleAsync,
+	// then allow it and all remaining notifications to complete.
+	<-firstStarted
+	close(releaseFirst)
+
+	var got string
+	if err := ac.Await(ctx, &got); err != nil {
+		t.Fatalf("Await failed: %v", err)
+	}
+	if got != "ok" {
+		t.Errorf("Await result = %q, want %q", got, "ok")
+	}
+	if n := handled.Load(); n != numNotifications {
+		t.Errorf("handled %d notifications before Await returned, want %d", n, numNotifications)
+	}
+}
+
+// TestSynchronousHandlerNestedCall verifies that a synchronous handler (one
+// that does not call Async) can make an outgoing Call and await its response
+// without deadlocking on its own handlerQueue slot.
+func TestSynchronousHandlerNestedCall(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	r := fakeReader{msgs: make(chan Message, 4)}
+	w := fakeWriter{out: make(chan Message, 4)}
+	done := make(chan string, 1)
+
+	NewConnection(ctx, ConnectionConfig{
+		Reader: r,
+		Writer: w,
+		Closer: nopCloser{},
+		Bind: func(c *Connection) Handler {
+			return HandlerFunc(func(hctx context.Context, req *Request) (any, error) {
+				if req.Method == "trigger" {
+					var res string
+					if err := c.Call(hctx, "nested", nil).Await(hctx, &res); err != nil {
+						t.Errorf("nested Call failed: %v", err)
+					}
+					done <- res
+				}
+				return nil, nil
+			})
+		},
+	})
+
+	trigger, err := NewNotification("trigger", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	trailing, err := NewNotification("trailing", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	r.msgs <- trigger
+	r.msgs <- trailing
+
+	nestedReq := (<-w.out).(*Request)
+	nestedResp, err := NewResponse(nestedReq.ID, "nested-ok", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	r.msgs <- nestedResp
+
+	if got := <-done; got != "nested-ok" {
+		t.Errorf("nested call result = %q, want %q", got, "nested-ok")
+	}
+}
diff --git a/mcp/client_test.go b/mcp/client_test.go
index bc6dc86e6..e437e058d 100644
--- a/mcp/client_test.go
+++ b/mcp/client_test.go
@@ -12,7 +12,10 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"net/http"
+	"net/http/httptest"
 	"slices"
+	"sync"
 	"sync/atomic"
 	"testing"
 
@@ -1396,3 +1399,214 @@ func TestListRootsCapabilityDisabled(t *testing.T) {
 		})
 	}
 }
+
+// TestProgressBeforeResult verifies that notifications sent by a server before
+// a call's response are handled by the client before CallTool returns
+// (go-sdk#1337).
+func TestProgressBeforeResult(t *testing.T) {
+	const (
+		numCalls        = 100
+		progressPerCall = 3
+	)
+
+	newUpstream := func() *Server {
+		s := NewServer(testImpl, nil)
+		AddTool(s, &Tool{Name: "work"}, func(ctx context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, any, error) {
+			for i := range progressPerCall {
+				if err := req.Session.NotifyProgress(ctx, &ProgressNotificationParams{
+					ProgressToken: req.Params.GetProgressToken(),
+					Progress:      float64(i + 1),
+					Total:         progressPerCall,
+				}); err != nil {
+					return nil, nil, err
+				}
+			}
+			return &CallToolResult{}, nil, nil
+		})
+		return s
+	}
+
+	type progressTracker struct {
+		mu       sync.Mutex
+		seen     map[any]int
+		returned map[any]bool
+		late     int
+	}
+	newTracker := func() *progressTracker {
+		return &progressTracker{seen: make(map[any]int), returned: make(map[any]bool)}
+	}
+	handleProgress := func(tr *progressTracker) func(context.Context, *ProgressNotificationClientRequest) {
+		return func(_ context.Context, req *ProgressNotificationClientRequest) {
+			tr.mu.Lock()
+			defer tr.mu.Unlock()
+			if tr.returned[req.Params.ProgressToken] {
+				tr.late++
+			}
+			tr.seen[req.Params.ProgressToken]++
+		}
+	}
+	callAndCheck := func(t *testing.T, cs *ClientSession, tr *progressTracker) {
+		t.Helper()
+		ctx := context.Background()
+		violations := 0
+		for i := range numCalls {
+			token := fmt.Sprintf("call-%d", i)
+			params := &CallToolParams{Name: "work"}
+			params.SetProgressToken(token)
+			if _, err := cs.CallTool(ctx, params); err != nil {
+				t.Fatalf("CallTool %d: %v", i, err)
+			}
+			tr.mu.Lock()
+			tr.returned[token] = true
+			seen := tr.seen[token]
+			tr.mu.Unlock()
+			if seen < progressPerCall {
+				violations++
+			}
+		}
+		tr.mu.Lock()
+		late := tr.late
+		tr.mu.Unlock()
+		if violations > 0 || late > 0 {
+			t.Errorf("CallTool returned before all progress notifications were handled in %d/%d calls (%d late notifications)", violations, numCalls, late)
+		}
+	}
+
+	t.Run("in_memory", func(t *testing.T) {
+		ctx := context.Background()
+		st, ct := NewInMemoryTransports()
+		ss, err := newUpstream().Connect(ctx, st, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer ss.Close()
+
+		tr := newTracker()
+		cs, err := NewClient(testImpl, &ClientOptions{
+			ProgressNotificationHandler: handleProgress(tr),
+		}).Connect(ctx, ct, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer cs.Close()
+
+		callAndCheck(t, cs, tr)
+	})
+
+	t.Run("streamable_http", func(t *testing.T) {
+		ctx := context.Background()
+		upstream := newUpstream()
+		srv := httptest.NewServer(NewStreamableHTTPHandler(func(*http.Request) *Server { return upstream }, nil))
+		defer srv.Close()
+
+		tr := newTracker()
+		cs, err := NewClient(testImpl, &ClientOptions{
+			ProgressNotificationHandler: handleProgress(tr),
+		}).Connect(ctx, &StreamableClientTransport{Endpoint: srv.URL}, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer cs.Close()
+
+		callAndCheck(t, cs, tr)
+	})
+
+	t.Run("proxy", func(t *testing.T) {
+		ctx := context.Background()
+		upST, upCT := NewInMemoryTransports()
+		upSS, err := newUpstream().Connect(ctx, upST, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer upSS.Close()
+
+		var (
+			forwarders sync.Map
+			dropped    atomic.Int64
+		)
+		upCS, err := NewClient(testImpl, &ClientOptions{
+			ProgressNotificationHandler: func(_ context.Context, req *ProgressNotificationClientRequest) {
+				if f, ok := forwarders.Load(req.Params.ProgressToken); ok {
+					f.(func(*ProgressNotificationParams))(req.Params)
+				} else {
+					dropped.Add(1)
+				}
+			},
+		}).Connect(ctx, upCT, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer upCS.Close()
+
+		proxySrv := NewServer(testImpl, nil)
+		AddTool(proxySrv, &Tool{Name: "work"}, func(ctx context.Context, req *CallToolRequest, _ struct{}) (*CallToolResult, any, error) {
+			token := req.Params.GetProgressToken()
+			forwarders.Store(token, func(n *ProgressNotificationParams) {
+				_ = req.Session.NotifyProgress(ctx, n)
+			})
+			defer forwarders.Delete(token)
+
+			params := &CallToolParams{Name: "work"}
+			params.SetProgressToken(token)
+			res, err := upCS.CallTool(ctx, params)
+			return res, nil, err
+		})
+
+		downST, downCT := NewInMemoryTransports()
+		downSS, err := proxySrv.Connect(ctx, downST, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer downSS.Close()
+
+		tr := newTracker()
+		cs, err := NewClient(testImpl, &ClientOptions{
+			ProgressNotificationHandler: handleProgress(tr),
+		}).Connect(ctx, downCT, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer cs.Close()
+
+		callAndCheck(t, cs, tr)
+		if n := dropped.Load(); n > 0 {
+			t.Errorf("proxy dropped %d upstream progress notifications after upstream CallTool returned", n)
+		}
+	})
+
+	t.Run("notification_handler_nested_call", func(t *testing.T) {
+		ctx := context.Background()
+		s := NewServer(testImpl, &ServerOptions{HasTools: true})
+		st, ct := NewInMemoryTransports()
+		ss, err := s.Connect(ctx, st, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer ss.Close()
+
+		listed := make(chan int, 1)
+		c := NewClient(testImpl, &ClientOptions{
+			ToolListChangedHandler: func(ctx context.Context, req *ToolListChangedRequest) {
+				res, err := req.Session.ListTools(ctx, nil)
+				if err != nil {
+					t.Errorf("ListTools in ToolListChangedHandler: %v", err)
+					listed <- -1
+					return
+				}
+				listed <- len(res.Tools)
+			},
+		})
+		cs, err := c.Connect(ctx, ct, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer cs.Close()
+
+		AddTool(s, &Tool{Name: "greet"}, func(context.Context, *CallToolRequest, struct{}) (*CallToolResult, any, error) {
+			return &CallToolResult{}, nil, nil
+		})
+		if got := <-listed; got != 1 {
+			t.Errorf("ListTools inside ToolListChangedHandler returned %d tools, want 1", got)
+		}
+	})
+}
diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go
index 04c4c814f..4f92558ad 100644
--- a/mcp/streamable_test.go
+++ b/mcp/streamable_test.go
@@ -437,7 +437,7 @@ type clientReplayTest struct {
 }
 
 func testClientReplay(t *testing.T, test clientReplayTest) {
-	notifications := make(chan string)
+	notifications := make(chan string, 4)
 	// Configure the real MCP server.
 	server := NewServer(testImpl, nil)
 

From 004eb9692dbf6c629eea59122b55601a08c341a2 Mon Sep 17 00:00:00 2001
From: xioshark <151488904+Xio-Shark@users.noreply.github.com>
Date: Sun, 18 Oct 2026 18:24:29 +0800
Subject: [PATCH 2/4] internal/jsonrpc2: handle Retire on queued responses and
 non-handler-ctx nested calls

When a response is queued in handlerQueue behind an active handler,
readIncoming has already removed its ID from outgoingCalls. Update
Connection.Retire to also remove and retire a matching entry in
handlerQueue so caller-side retirement (such as context cancellation in
mcp/transport.go) still takes effect immediately and cannot race with
handleAsync draining the queue.

Also compare the monotonic queued-request counter (s.queuedRequests ==
ac.queuedAtStart) instead of requiring len(s.handlerQueue) == 0 when a
synchronous handler issues a nested Call without threading the handler
context, avoiding a deadlock when earlier notifications are already
waiting in handlerQueue.

Signed-off-by: xioshark <151488904+Xio-Shark@users.noreply.github.com>
---
 internal/jsonrpc2/conn.go      |  16 ++++-
 internal/jsonrpc2/conn_test.go | 124 +++++++++++++++++++++++++++------
 2 files changed, 114 insertions(+), 26 deletions(-)

diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go
index 222fefded..dac196e2e 100644
--- a/internal/jsonrpc2/conn.go
+++ b/internal/jsonrpc2/conn.go
@@ -95,6 +95,7 @@ type inFlightState struct {
 	handlerQueue   []queueItem
 	handlerRunning bool
 	activeReleaser *releaser
+	queuedRequests int64
 }
 
 // queueItem is an entry in [inFlightState.handlerQueue]: either an incoming
@@ -362,6 +363,7 @@ func (c *Connection) Call(ctx context.Context, method string, params any) *Async
 			ac.callerReleaser = r
 		}
 		ac.activeAtStart = s.activeSyncReleaser()
+		ac.queuedAtStart = s.queuedRequests
 		if s.outgoingCalls == nil {
 			s.outgoingCalls = make(map[ID]*AsyncCall)
 		}
@@ -389,8 +391,14 @@ func (c *Connection) Retire(ac *AsyncCall, err error) {
 		if s.outgoingCalls[ac.id] == ac {
 			delete(s.outgoingCalls, ac.id)
 			ac.retire(&Response{ID: ac.id, Error: err})
-		} else {
-			// ac was already retired elsewhere.
+			return
+		}
+		for i, item := range s.handlerQueue {
+			if item.ac == ac {
+				s.handlerQueue = append(s.handlerQueue[:i], s.handlerQueue[i+1:]...)
+				ac.retire(&Response{ID: ac.id, Error: err})
+				return
+			}
 		}
 	})
 }
@@ -455,6 +463,7 @@ type AsyncCall struct {
 	response       *Response
 	callerReleaser *releaser
 	activeAtStart  *releaser
+	queuedAtStart  int64
 }
 
 // ID used for this call.
@@ -608,7 +617,7 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter
 				if ac, ok := s.outgoingCalls[msg.ID]; ok {
 					delete(s.outgoingCalls, msg.ID)
 					active := s.activeSyncReleaser()
-					if !s.handlerRunning || (active != nil && (ac.callerReleaser == active || (len(s.handlerQueue) == 0 && ac.activeAtStart == active))) {
+					if !s.handlerRunning || (active != nil && (ac.callerReleaser == active || (s.queuedRequests == ac.queuedAtStart && ac.activeAtStart == active))) {
 						ac.retire(msg)
 					} else {
 						s.handlerQueue = append(s.handlerQueue, queueItem{ac: ac, resp: msg})
@@ -712,6 +721,7 @@ func (c *Connection) acceptRequest(ctx context.Context, msg *Request, preempter
 		// (because it may be responding to a request issued by, say, an
 		// asynchronous handler), and in order to get to that response we have
 		// to read all of the requests that came in ahead of it.
+		s.queuedRequests++
 		s.handlerQueue = append(s.handlerQueue, queueItem{req: req})
 		if !s.handlerRunning {
 			// We start the handleAsync goroutine when it has work to do, and let it
diff --git a/internal/jsonrpc2/conn_test.go b/internal/jsonrpc2/conn_test.go
index 7279b7fb8..eecaa9003 100644
--- a/internal/jsonrpc2/conn_test.go
+++ b/internal/jsonrpc2/conn_test.go
@@ -326,52 +326,130 @@ func TestResponseWaitsForPrecedingNotifications(t *testing.T) {
 
 // TestSynchronousHandlerNestedCall verifies that a synchronous handler (one
 // that does not call Async) can make an outgoing Call and await its response
-// without deadlocking on its own handlerQueue slot.
+// without deadlocking on its own handlerQueue slot, whether it passes the
+// handler context or a background context.
 func TestSynchronousHandlerNestedCall(t *testing.T) {
+	for _, useHandlerCtx := range []bool{true, false} {
+		name := "handler_ctx"
+		if !useHandlerCtx {
+			name = "background_ctx"
+		}
+		t.Run(name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			r := fakeReader{msgs: make(chan Message, 4)}
+			w := fakeWriter{out: make(chan Message, 4)}
+			done := make(chan string, 1)
+
+			NewConnection(ctx, ConnectionConfig{
+				Reader: r,
+				Writer: w,
+				Closer: nopCloser{},
+				Bind: func(c *Connection) Handler {
+					return HandlerFunc(func(hctx context.Context, req *Request) (any, error) {
+						if req.Method == "trigger" {
+							callCtx := hctx
+							if !useHandlerCtx {
+								callCtx = context.Background()
+							}
+							var res string
+							if err := c.Call(callCtx, "nested", nil).Await(callCtx, &res); err != nil {
+								t.Errorf("nested Call failed: %v", err)
+							}
+							done <- res
+						}
+						return nil, nil
+					})
+				},
+			})
+
+			trigger, err := NewNotification("trigger", nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+			trailing, err := NewNotification("trailing", nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+			r.msgs <- trigger
+			r.msgs <- trailing
+
+			nestedReq := (<-w.out).(*Request)
+			nestedResp, err := NewResponse(nestedReq.ID, "nested-ok", nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+			r.msgs <- nestedResp
+
+			if got := <-done; got != "nested-ok" {
+				t.Errorf("nested call result = %q, want %q", got, "nested-ok")
+			}
+		})
+	}
+}
+
+// TestRetireQueuedResponse verifies that Connection.Retire retires an AsyncCall
+// even after its response has been read from the wire and queued behind an
+// in-flight notification handler, without double-retiring when the queue drains.
+func TestRetireQueuedResponse(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	r := fakeReader{msgs: make(chan Message, 4)}
+	r := fakeReader{msgs: make(chan Message)}
 	w := fakeWriter{out: make(chan Message, 4)}
-	done := make(chan string, 1)
+	notifyStarted := make(chan struct{})
+	releaseNotify := make(chan struct{})
+	allHandled := make(chan struct{})
+	var count atomic.Int32
 
-	NewConnection(ctx, ConnectionConfig{
+	conn := NewConnection(ctx, ConnectionConfig{
 		Reader: r,
 		Writer: w,
 		Closer: nopCloser{},
-		Bind: func(c *Connection) Handler {
-			return HandlerFunc(func(hctx context.Context, req *Request) (any, error) {
-				if req.Method == "trigger" {
-					var res string
-					if err := c.Call(hctx, "nested", nil).Await(hctx, &res); err != nil {
-						t.Errorf("nested Call failed: %v", err)
+		Bind: func(*Connection) Handler {
+			return HandlerFunc(func(_ context.Context, req *Request) (any, error) {
+				if req.Method == "notify" {
+					if count.Add(1) == 1 {
+						close(notifyStarted)
+						<-releaseNotify
+					} else {
+						close(allHandled)
 					}
-					done <- res
 				}
 				return nil, nil
 			})
 		},
 	})
 
-	trigger, err := NewNotification("trigger", nil)
-	if err != nil {
-		t.Fatal(err)
-	}
-	trailing, err := NewNotification("trailing", nil)
+	ac := conn.Call(ctx, "work", nil)
+	sent := (<-w.out).(*Request)
+
+	n, err := NewNotification("notify", nil)
 	if err != nil {
 		t.Fatal(err)
 	}
-	r.msgs <- trigger
-	r.msgs <- trailing
+	r.msgs <- n
+	<-notifyStarted
 
-	nestedReq := (<-w.out).(*Request)
-	nestedResp, err := NewResponse(nestedReq.ID, "nested-ok", nil)
+	resp, err := NewResponse(sent.ID, "ok", nil)
 	if err != nil {
 		t.Fatal(err)
 	}
-	r.msgs <- nestedResp
+	r.msgs <- resp
+	// Because r.msgs is unbuffered, sending a second notification only returns
+	// after readIncoming has finished enqueuing resp into handlerQueue and
+	// advanced to the next Read call.
+	r.msgs <- n
+
+	retireErr := errors.New("caller cancelled")
+	conn.Retire(ac, retireErr)
 
-	if got := <-done; got != "nested-ok" {
-		t.Errorf("nested call result = %q, want %q", got, "nested-ok")
+	var got string
+	if err := ac.Await(ctx, &got); !errors.Is(err, retireErr) {
+		t.Fatalf("Await error = %v, want %v", err, retireErr)
 	}
+
+	close(releaseNotify)
+	<-allHandled
 }

From f9fe86917d66bac6da478851e8fe430a6db7b426 Mon Sep 17 00:00:00 2001
From: xioshark <151488904+Xio-Shark@users.noreply.github.com>
Date: Sat, 24 Oct 2026 22:23:10 +0800
Subject: [PATCH 3/4] internal/jsonrpc2: add immediateresponseretire
 compatibility parameter
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The preceding-notification ordering guarantee means a response is
delivered after the notifications ahead of it on the wire have been
handled. A notification handler that waits on the return of the call it
notifies about — for example a progress handler sending on an unbuffered
channel that the caller drains after Call returns — therefore no longer
completes after upgrade, and neither does the call it waits on.

Restore the previous timing behind MCPGODEBUG=immediateresponseretire=1,
following the existing compatibility-parameter pattern. Document the rule
in the Concurrency section of the protocol docs and the parameter in the
MCPGODEBUG history.

Signed-off-by: xioshark <151488904+Xio-Shark@users.noreply.github.com>
---
 docs/mcpgodebug.md              |  9 +++++
 docs/protocol.md                |  5 +++
 internal/docs/mcpgodebug.src.md |  9 +++++
 internal/docs/protocol.src.md   |  5 +++
 internal/jsonrpc2/conn.go       | 19 +++++++++-
 internal/jsonrpc2/conn_test.go  | 64 +++++++++++++++++++++++++++++++++
 6 files changed, 110 insertions(+), 1 deletion(-)

diff --git a/docs/mcpgodebug.md b/docs/mcpgodebug.md
index 475635239..41694321e 100644
--- a/docs/mcpgodebug.md
+++ b/docs/mcpgodebug.md
@@ -79,6 +79,15 @@ Options listed below were added and will be removed in the 1.11.0 version of the
   upgrade from `http` to `https` on the same host with default ports is
   allowed) and keeps the method.
 
+- `immediateresponseretire` added. If set to `1`, the response to a call is
+  retired as soon as it is read, restoring the previous behavior. The default
+  behavior was changed to deliver the response after the notifications that
+  precede it on the wire have been handled, so that a notification is observed
+  before the call it precedes returns. The previous timing is incompatible
+  with a notification handler that waits on the return of the call it notifies
+  about, such as a progress handler sending on an unbuffered channel drained
+  after `Call` returns.
+
 ### 1.8.0
 
 Options listed below were added and will be removed in the 1.9.0 version of the SDK.
diff --git a/docs/protocol.md b/docs/protocol.md
index a77604e8d..cc1db7d03 100644
--- a/docs/protocol.md
+++ b/docs/protocol.md
@@ -437,6 +437,11 @@ implements the following heuristics:
   client goroutine.
 - Calls (such as `tools/call`) are handled asynchronously with respect to
   each other.
+- Notifications that precede a call's response on the wire are handled before
+  the call returns. A notification handler therefore must not wait on the
+  return of the call it is notifying about; set
+  `MCPGODEBUG=immediateresponseretire=1` to restore the previous timing, where
+  the call could return while its notification handler was still running.
 
 See
 [modelcontextprotocol/go-sdk#26](https://github.com/modelcontextprotocol/go-sdk/issues/26)
diff --git a/internal/docs/mcpgodebug.src.md b/internal/docs/mcpgodebug.src.md
index 66f36ff2d..591f900d8 100644
--- a/internal/docs/mcpgodebug.src.md
+++ b/internal/docs/mcpgodebug.src.md
@@ -78,6 +78,15 @@ Options listed below were added and will be removed in the 1.11.0 version of the
   upgrade from `http` to `https` on the same host with default ports is
   allowed) and keeps the method.
 
+- `immediateresponseretire` added. If set to `1`, the response to a call is
+  retired as soon as it is read, restoring the previous behavior. The default
+  behavior was changed to deliver the response after the notifications that
+  precede it on the wire have been handled, so that a notification is observed
+  before the call it precedes returns. The previous timing is incompatible
+  with a notification handler that waits on the return of the call it notifies
+  about, such as a progress handler sending on an unbuffered channel drained
+  after `Call` returns.
+
 ### 1.8.0
 
 Options listed below were added and will be removed in the 1.9.0 version of the SDK.
diff --git a/internal/docs/protocol.src.md b/internal/docs/protocol.src.md
index e1cb37af8..36ff564da 100644
--- a/internal/docs/protocol.src.md
+++ b/internal/docs/protocol.src.md
@@ -332,6 +332,11 @@ implements the following heuristics:
   client goroutine.
 - Calls (such as `tools/call`) are handled asynchronously with respect to
   each other.
+- Notifications that precede a call's response on the wire are handled before
+  the call returns. A notification handler therefore must not wait on the
+  return of the call it is notifying about; set
+  `MCPGODEBUG=immediateresponseretire=1` to restore the previous timing, where
+  the call could return while its notification handler was still running.
 
 See
 [modelcontextprotocol/go-sdk#26](https://github.com/modelcontextprotocol/go-sdk/issues/26)
diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go
index dac196e2e..c3d608dba 100644
--- a/internal/jsonrpc2/conn.go
+++ b/internal/jsonrpc2/conn.go
@@ -33,6 +33,23 @@ import (
 // to use it.
 var nomethodnotfoundcodeinerror = mcpgodebug.Value("nomethodnotfoundcodeinerror")
 
+// immediateresponseretire is a compatibility parameter that restores the
+// previous behavior of retiring a call's response as soon as the read loop sees
+// it, instead of delivering it after the notifications that precede it on the
+// wire. The default behavior keeps the promise made by the Concurrency section
+// of the protocol documentation: a notification is handled before the call
+// awaiting a response that followed it returns.
+//
+// The default is incompatible with a notification handler that waits on the
+// return of the call it is notifying about (for example a progress handler that
+// sends on an unbuffered channel drained after Call returns): the handler is
+// running while the response waits in the queue, so the call cannot return
+// until the handler completes.
+//
+// To restore the previous behavior, set MCPGODEBUG=immediateresponseretire=1.
+// This option will be removed in a future SDK version.
+var immediateresponseretire = mcpgodebug.Value("immediateresponseretire")
+
 // Connection manages the jsonrpc2 protocol, connecting responses back to their
 // calls. Connection is bidirectional; it does not have a designated server or
 // client end.
@@ -617,7 +634,7 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter
 				if ac, ok := s.outgoingCalls[msg.ID]; ok {
 					delete(s.outgoingCalls, msg.ID)
 					active := s.activeSyncReleaser()
-					if !s.handlerRunning || (active != nil && (ac.callerReleaser == active || (s.queuedRequests == ac.queuedAtStart && ac.activeAtStart == active))) {
+					if immediateresponseretire == "1" || !s.handlerRunning || (active != nil && (ac.callerReleaser == active || (s.queuedRequests == ac.queuedAtStart && ac.activeAtStart == active))) {
 						ac.retire(msg)
 					} else {
 						s.handlerQueue = append(s.handlerQueue, queueItem{ac: ac, resp: msg})
diff --git a/internal/jsonrpc2/conn_test.go b/internal/jsonrpc2/conn_test.go
index eecaa9003..0fc8240d3 100644
--- a/internal/jsonrpc2/conn_test.go
+++ b/internal/jsonrpc2/conn_test.go
@@ -15,6 +15,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"testing"
+	"time"
 )
 
 type fakeReader struct {
@@ -324,6 +325,69 @@ func TestResponseWaitsForPrecedingNotifications(t *testing.T) {
 	}
 }
 
+// TestImmediateResponseRetireGODEBUG verifies that
+// MCPGODEBUG=immediateresponseretire=1 restores the previous response timing: a
+// notification handler that waits on the return of the call it notifies about
+// does not deadlock the call.
+func TestImmediateResponseRetireGODEBUG(t *testing.T) {
+	defer func(prev string) { immediateresponseretire = prev }(immediateresponseretire)
+	immediateresponseretire = "1"
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	r := fakeReader{msgs: make(chan Message, 2)}
+	w := fakeWriter{out: make(chan Message, 1)}
+
+	callReturned := make(chan struct{})
+	handlerDone := make(chan struct{})
+
+	conn := NewConnection(ctx, ConnectionConfig{
+		Reader: r,
+		Writer: w,
+		Closer: nopCloser{},
+		Bind: func(*Connection) Handler {
+			return HandlerFunc(func(_ context.Context, req *Request) (any, error) {
+				if req.Method == "notify" {
+					defer close(handlerDone)
+					select {
+					case <-callReturned:
+					case <-time.After(5 * time.Second):
+						t.Error("notification handler: the call did not return while the handler was running")
+					}
+				}
+				return nil, nil
+			})
+		},
+	})
+
+	ac := conn.Call(ctx, "work", nil)
+	sent := (<-w.out).(*Request)
+
+	n, err := NewNotification("notify", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	r.msgs <- n
+	resp, err := NewResponse(sent.ID, "ok", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	r.msgs <- resp
+
+	awaitCtx, cancelAwait := context.WithTimeout(ctx, 5*time.Second)
+	defer cancelAwait()
+	var got string
+	if err := ac.Await(awaitCtx, &got); err != nil {
+		t.Fatalf("Await failed with a notification handler running: %v", err)
+	}
+	if got != "ok" {
+		t.Errorf("Await result = %q, want %q", got, "ok")
+	}
+	close(callReturned)
+	<-handlerDone
+}
+
 // TestSynchronousHandlerNestedCall verifies that a synchronous handler (one
 // that does not call Async) can make an outgoing Call and await its response
 // without deadlocking on its own handlerQueue slot, whether it passes the

From 83fec63a7a4741372e3c5b429aff34728930378f Mon Sep 17 00:00:00 2001
From: xioshark <151488904+Xio-Shark@users.noreply.github.com>
Date: Wed, 28 Oct 2026 16:36:53 +0800
Subject: [PATCH 4/4] internal/docs: the new response-retire timing is the
 incompatible one

The MCPGODEBUG entry credited the incompatibility to the previous timing,
but it is the new default that makes a handler waiting on the return of
the call it notifies about deadlock. Matches conn.go's description of
immediateresponseretire.
---
 docs/mcpgodebug.md              | 2 +-
 internal/docs/mcpgodebug.src.md | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/mcpgodebug.md b/docs/mcpgodebug.md
index 41694321e..224e7c6aa 100644
--- a/docs/mcpgodebug.md
+++ b/docs/mcpgodebug.md
@@ -83,7 +83,7 @@ Options listed below were added and will be removed in the 1.11.0 version of the
   retired as soon as it is read, restoring the previous behavior. The default
   behavior was changed to deliver the response after the notifications that
   precede it on the wire have been handled, so that a notification is observed
-  before the call it precedes returns. The previous timing is incompatible
+  before the call it precedes returns. The new timing is incompatible
   with a notification handler that waits on the return of the call it notifies
   about, such as a progress handler sending on an unbuffered channel drained
   after `Call` returns.
diff --git a/internal/docs/mcpgodebug.src.md b/internal/docs/mcpgodebug.src.md
index 591f900d8..11ef82ed2 100644
--- a/internal/docs/mcpgodebug.src.md
+++ b/internal/docs/mcpgodebug.src.md
@@ -82,7 +82,7 @@ Options listed below were added and will be removed in the 1.11.0 version of the
   retired as soon as it is read, restoring the previous behavior. The default
   behavior was changed to deliver the response after the notifications that
   precede it on the wire have been handled, so that a notification is observed
-  before the call it precedes returns. The previous timing is incompatible
+  before the call it precedes returns. The new timing is incompatible
   with a notification handler that waits on the return of the call it notifies
   about, such as a progress handler sending on an unbuffered channel drained
   after `Call` returns.