internal/jsonrpc2: handle preceding notifications before retiring call responses

docs/mcpgodebug.md

@@ -70,6 +70,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.

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)

internal/docs/mcpgodebug.src.md

@@ -69,6 +69,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.

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)

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.
@@ -85,11 +102,25 @@ 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
+	queuedRequests int64
+}
+
+// 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
@@ -338,6 +369,11 @@ 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()
+		ac.queuedAtStart = s.queuedRequests
 		if s.outgoingCalls == nil {
 			s.outgoingCalls = make(map[ID]*AsyncCall)
 		}
@@ -365,8 +401,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
+			}
 		}
 	})
 }
@@ -390,6 +432,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
@@ -411,10 +454,26 @@ 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
+	queuedAtStart  int64
 }
 
 // ID used for this call.
@@ -567,7 +626,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 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})
+					}
 				} else {
 					// TODO: How should we report unexpected responses?
 				}
@@ -668,7 +732,8 @@ 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.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
 			// exit when the queue empties.
@@ -692,17 +757,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
 		}
@@ -712,21 +791,21 @@ 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
 			}
 			c.processResult("handleAsync", req, nil, err)
 			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
 	}
 }