mcp: preserve SSE events buffered after the endpoint (#1270)
SSEClientTransport.Connect creates one event scanner to read the endpoint event and a second scanner to read subsequent messages. Each scanner owns a buffered reader, so bytes read ahead by the first scanner are discarded when it is abandoned. A message coalesced with the endpoint can therefore disappear while the connection remains open. Reuse the event scanner across the endpoint handshake and the message-reading goroutine. This preserves unread bytes without changing the public API. The regression test exercises Connect and Read with a controlled HTTP response body, covering both a message coalesced with the endpoint and a message arriving after Connect returns. On unmodified main, the separate-message control passes and the coalesced case times out; both pass with the fix. Validation on linux/amd64 with Go 1.26.8: - go test ./... - go test -race ./... - go vet ./... - staticcheck v0.6.1 ./... - gofmt and git diff --check The HTTP response boundary is simulated; no third-party server or proxy was used. This change concerns the legacy SSE transport. The Streamable HTTP conformance suite was not run.
This commit is contained in:
parent
0c0eb23625
commit
c1918e19ad
2 changed files with 53 additions and 2 deletions
|
|
@ -437,9 +437,11 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) {
|
|||
maxEventSize = DefaultMaxEventSize
|
||||
}
|
||||
|
||||
// Reuse the scanner so events buffered after the endpoint are not lost.
|
||||
events := scanEventsLimited(resp.Body, maxEventSize)
|
||||
msgEndpoint, err := func() (*url.URL, error) {
|
||||
var evt Event
|
||||
for evt, err = range scanEventsLimited(resp.Body, maxEventSize) {
|
||||
for evt, err = range events {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
|
|
@ -468,7 +470,7 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) {
|
|||
go func() {
|
||||
defer s.Close() // close the transport when the GET exits
|
||||
|
||||
for evt, err := range scanEventsLimited(resp.Body, maxEventSize) {
|
||||
for evt, err := range events {
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
|
|
|||
|
|
@ -13,10 +13,13 @@ import (
|
|||
"net/http"
|
||||
"net/http/httptest"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
|
||||
)
|
||||
|
||||
// TestSSEServerTransport_SupportedVersions verifies that the deprecated
|
||||
|
|
@ -432,3 +435,49 @@ func TestSSEServerTransportMaxRequestBodyBytes(t *testing.T) {
|
|||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSSEClientTransportBufferedEvents(t *testing.T) {
|
||||
const endpoint = "event: endpoint\ndata: /messages\n\n"
|
||||
const message = "event: message\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/tools/list_changed\"}\n\n"
|
||||
for _, coalesced := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("coalesced=%t", coalesced), func(t *testing.T) {
|
||||
reader, writer := io.Pipe()
|
||||
defer reader.Close()
|
||||
defer writer.Close()
|
||||
prefix := endpoint
|
||||
if coalesced {
|
||||
prefix += message
|
||||
}
|
||||
transport := &SSEClientTransport{
|
||||
Endpoint: "http://example.com/sse",
|
||||
HTTPClient: &http.Client{Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": {"text/event-stream"}},
|
||||
Body: struct {
|
||||
io.Reader
|
||||
io.Closer
|
||||
}{io.MultiReader(strings.NewReader(prefix), reader), reader},
|
||||
}, nil
|
||||
})},
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
conn, err := transport.Connect(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
if !coalesced {
|
||||
go func() { _, _ = io.WriteString(writer, message) }()
|
||||
}
|
||||
msg, err := conn.Read(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("Read: %v", err)
|
||||
}
|
||||
if req, ok := msg.(*jsonrpc.Request); !ok || req.Method != "notifications/tools/list_changed" {
|
||||
t.Fatalf("Read = %v, want tools/list_changed notification", msg)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue