address review comments

Author: findleyr

mcp/streamable.go

@@ -491,7 +491,7 @@ type stream struct {
 	// If non-nil, deliver writes data directly to the HTTP response.
 	//
 	// Only one HTTP response may receive messages at a given time. An active
-	// HTTP connection acquires ownership of the stream by setting
+	// HTTP connection acquires ownership of the stream by setting this field.
 	deliver func(data []byte, final bool) error
 
 	// streamRequests is the set of unanswered incoming requests for the stream.
@@ -500,10 +500,10 @@ type stream struct {
 	requests map[jsonrpc.ID]struct{}
 }
 
-// done reports whether the stream is logically complete.
+// doneLocked reports whether the stream is logically complete.
 //
 // s.mu must be held while calling this function.
-func (s *stream) done() bool {
+func (s *stream) doneLocked() bool {
 	return len(s.requests) == 0 && s.id != ""
 }
 
@@ -585,42 +585,73 @@ func (c *streamableServerConn) serveGET(w http.ResponseWriter, req *http.Request
 		}
 	}
 
+	ctx, cancel := context.WithCancel(req.Context())
+	defer cancel()
+
+	stream, done := c.acquireStream(ctx, w, streamID, &lastIdx)
+	if stream == nil {
+		return
+	}
+	// Release the stream when we're done.
+	defer func() {
+		stream.mu.Lock()
+		stream.deliver = nil
+		stream.mu.Unlock()
+	}()
+
+	select {
+	case <-ctx.Done():
+		// request cancelled
+	case <-done:
+		// request complete
+	case <-c.done:
+		// session closed
+	}
+}
+
+// writeEvent writes an SSE event to w corresponding to the given stream, data, and index.
+// lastIdx is incremented before writing, so that it continues to point to the index of the
+// last event written to the stream.
+func (c *streamableServerConn) writeEvent(w http.ResponseWriter, stream *stream, data []byte, lastIdx *int) error {
+	*lastIdx++
+	e := Event{
+		Name: "message",
+		Data: data,
+	}
+	if c.eventStore != nil {
+		e.ID = formatEventID(stream.id, *lastIdx)
+	}
+	if _, err := writeEvent(w, e); err != nil {
+		return err
+	}
+	return nil
+}
+
+// acquireStream acquires the stream and replays all events since lastIdx, if
+// any, updating lastIdx accordingly. If non-nil, the resulting stream will be
+// registered for receiving new messages, and the resulting done channel will
+// be closed when all related messages have been delivered.
+//
+// If any errors occur, they will be written to w and the resulting stream will
+// be nil. The resulting stream may also be nil if the stream is complete.
+//
+// Importantly, this function must hold the stream mutex until done replaying
+// all messages, so that no delivery or storage of new messages occurs while
+// the stream is still replaying.
+func (c *streamableServerConn) acquireStream(ctx context.Context, w http.ResponseWriter, streamID string, lastIdx *int) (*stream, chan struct{}) {
 	c.mu.Lock()
 	stream, ok := c.streams[streamID]
 	c.mu.Unlock()
 	if !ok {
 		http.Error(w, "unknown stream", http.StatusBadRequest)
-		return
+		return nil, nil
 	}
 
-	write := func(data []byte) error {
-		lastIdx++
-		e := Event{
-			Name: "message",
-			Data: data,
-		}
-		if c.eventStore != nil {
-			e.ID = formatEventID(stream.id, lastIdx)
-		}
-		if _, err := writeEvent(w, e); err != nil {
-			return err
-		}
-		return nil
-	}
-
-	ctx, cancel := context.WithCancel(req.Context())
-	defer cancel()
-
-	// Acquire the stream.
-	//
-	// Importantly, hold the mutex until we've replayed all messages, since we
-	// don't want to allow delivery (or storage) of new messages until we've
-	// replayed everything that has been stored thus far.
 	stream.mu.Lock()
+	defer stream.mu.Unlock()
 	if stream.deliver != nil {
-		stream.mu.Unlock()
 		http.Error(w, "stream ID conflicts with ongoing stream", http.StatusConflict)
-		return
+		return nil, nil
 	}
 
 	// Collect events to replay. Collect them all before writing, so that we
@@ -631,9 +662,8 @@ func (c *streamableServerConn) serveGET(w http.ResponseWriter, req *http.Request
 	// messages, and registered our delivery function.
 	var toReplay [][]byte
 	if c.eventStore != nil {
-		for data, err := range c.eventStore.After(ctx, c.SessionID(), stream.id, lastIdx) {
+		for data, err := range c.eventStore.After(ctx, c.SessionID(), stream.id, *lastIdx) {
 			if err != nil {
-				stream.mu.Unlock()
 				// We can't replay events, perhaps because the underlying event store
 				// has garbage collected its storage.
 				//
@@ -643,7 +673,7 @@ func (c *streamableServerConn) serveGET(w http.ResponseWriter, req *http.Request
 				// 400 is not really accurate, but should at least have no side effects.
 				// Other SDKs (typescript) do not have a mechanism for events to be purged.
 				http.Error(w, "failed to replay events", http.StatusBadRequest)
-				return
+				return nil, nil
 			}
 			toReplay = append(toReplay, data)
 		}
@@ -656,27 +686,21 @@ func (c *streamableServerConn) serveGET(w http.ResponseWriter, req *http.Request
 	if stream.id == "" {
 		// Issue #410: the standalone SSE stream is likely not to receive messages
 		// for a long time. Ensure that headers are flushed.
-		//
-		// For other requests, delay the writing of the header in case we
-		// may want to set an error status.
-		// (see the TODO: this probably isn't worth it).
 		w.WriteHeader(http.StatusOK)
 		if f, ok := w.(http.Flusher); ok {
 			f.Flush()
 		}
 	}
 
 	for _, data := range toReplay {
-		if err := write(data); err != nil {
-			stream.mu.Unlock()
-			return
+		if err := c.writeEvent(w, stream, data, lastIdx); err != nil {
+			return nil, nil
 		}
 	}
 
-	if stream.done() {
+	if stream.doneLocked() {
 		// Nothing more to do.
-		stream.mu.Unlock()
-		return
+		return nil, nil
 	}
 
 	// Finally register a delivery function and unlock the stream, allowing the
@@ -686,29 +710,13 @@ func (c *streamableServerConn) serveGET(w http.ResponseWriter, req *http.Request
 		if err := ctx.Err(); err != nil {
 			return err
 		}
-		err := write(data)
+		err := c.writeEvent(w, stream, data, lastIdx)
 		if final {
 			close(done)
 		}
 		return err
 	}
-	stream.mu.Unlock()
-
-	// Release the stream when we're done.
-	defer func() {
-		stream.mu.Lock()
-		stream.deliver = nil
-		stream.mu.Unlock()
-	}()
-
-	select {
-	case <-ctx.Done():
-		// request cancelled
-	case <-done:
-		// request complete
-	case <-c.done:
-		// session closed
-	}
+	return stream, done
 }
 
 // servePOST handles an incoming message, and replies with either an outgoing
@@ -854,23 +862,12 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
 		}
 	} else {
 		// Write events in the order we receive them.
-		lastIndex := 0
+		lastIndex := -1
 		stream.deliver = func(data []byte, final bool) error {
 			if final {
 				defer close(done)
 			}
-			e := Event{
-				Name: "message",
-				Data: data,
-			}
-			if c.eventStore != nil {
-				e.ID = formatEventID(stream.id, lastIndex)
-			}
-			lastIndex++
-			if _, err := writeEvent(w, e); err != nil {
-				return err
-			}
-			return nil
+			return c.writeEvent(w, stream, data, &lastIndex)
 		}
 	}
 
@@ -1024,7 +1021,7 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e
 
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	if s.done() {
+	if s.doneLocked() {
 		return fmt.Errorf("%w: write to closed stream", jsonrpc2.ErrRejected)
 	}
 	if responseTo.IsValid() {
@@ -1040,7 +1037,7 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e
 		}
 	}
 	if s.deliver != nil {
-		if err := s.deliver(data, s.done()); err != nil {
+		if err := s.deliver(data, s.doneLocked()); err != nil {
 			// TODO: report a side-channel error.
 		} else {
 			delivered = true