From 2eb435d3a1db12dcc91671e2d8b26fd9c489ec26 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             | 210 +++++++++++++++++++++++++++++++++
 mcp/streamable_test.go         |   2 +-
 4 files changed, 395 insertions(+), 17 deletions(-)

diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go
index 22e6faee..222fefde 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 0893953c..7279b7fb 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 12bf52e7..e6add523 100644
--- a/mcp/client_test.go
+++ b/mcp/client_test.go
@@ -13,6 +13,8 @@ import (
 	"fmt"
 	"log/slog"
 	"maps"
+	"net/http"
+	"net/http/httptest"
 	"slices"
 	"sync"
 	"sync/atomic"
@@ -1511,3 +1513,211 @@ func TestCallToolPreservesSharedParams(t *testing.T) {
 		})
 	}
 }
+
+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 013c77d7..eda97a17 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 9f0c2955e48d0c7ba7e47aa59736775398763f40 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 222fefde..dac196e2 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 7279b7fb..eecaa900 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 b5475a877651e7b3310008b24fa724f29930718d 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 47563523..41694321 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 909ca1c1..d998e1f9 100644
--- a/docs/protocol.md
+++ b/docs/protocol.md
@@ -442,6 +442,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 66f36ff2..591f900d 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 0c005ab0..a3b1a3fa 100644
--- a/internal/docs/protocol.src.md
+++ b/internal/docs/protocol.src.md
@@ -337,6 +337,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 dac196e2..c3d608db 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 eecaa900..0fc8240d 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 2b977394674f24cf8a9b2d183b4577f6b6af78ba 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 41694321..224e7c6a 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 591f900d..11ef82ed 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.