mcp: gate prime and close events on protocol version 2025-11-25 (#696)

The prime and close SSE events (SEP-1699) were added in protocol version
2025-11-25. Only send these events when the client negotiates that
version or later.

- In servePOST, extract protocol version from InitializeParams for
initialize requests, otherwise use the Mcp-Protocol-Version header
- In serveGET/acquireStream, read protocol version from header
- Only set closeLocked callback when prime/close is supported
- Add unexported protocolVersion field to ClientSessionOptions for
testing

Fixes #686
This commit is contained in:
Robert Findley 2025-12-08 15:24:41 -05:00 • committed by GitHub
commit 87c81260e2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 129 additions and 30 deletions

View file

@ -120,7 +120,11 @@ func (e unsupportedProtocolVersionError) Error() string {
}
// ClientSessionOptions is reserved for future use.
type ClientSessionOptions struct{}
type ClientSessionOptions struct {
// protocolVersion overrides the protocol version sent in the initialize
// request, for testing. If empty, latestProtocolVersion is used.
protocolVersion string
}
func (c *Client) capabilities() *ClientCapabilities {
caps := &ClientCapabilities{}
@ -151,14 +155,18 @@ func (c *Client) capabilities() *ClientCapabilities {
// when it is no longer needed. However, if the connection is closed by the
// server, calls or notifications will return an error wrapping
// [ErrConnectionClosed].
func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptions) (cs *ClientSession, err error) {
func (c *Client) Connect(ctx context.Context, t Transport, opts *ClientSessionOptions) (cs *ClientSession, err error) {
cs, err = connect(ctx, t, c, (*clientSessionState)(nil), nil)
if err != nil {
return nil, err
}
protocolVersion := latestProtocolVersion
if opts != nil && opts.protocolVersion != "" {
protocolVersion = opts.protocolVersion
}
params := &InitializeParams{
ProtocolVersion: latestProtocolVersion,
ProtocolVersion: protocolVersion,
ClientInfo: c.impl,
Capabilities: c.capabilities(),
}

View file

@ -730,7 +730,15 @@ 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)
// Read the protocol version from the header. For GET requests, this should
// always be present since GET only happens after initialization.
protocolVersion := req.Header.Get(protocolVersionHeader)
if protocolVersion == "" {
protocolVersion = protocolVersion20250326
}
supportsPrimeClose := protocolVersion >= protocolVersion20251125
stream, done := c.acquireStream(ctx, w, streamID, &lastIdx, supportsPrimeClose)
if stream == nil {
return
}
@ -792,7 +800,10 @@ func (c *streamableServerConn) writeCloseEvent(w http.ResponseWriter, reconnectA
// 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{}) {
//
// supportsPrimeClose indicates whether the client supports the prime and close
// events (protocol version 2025-11-25 or later).
func (c *streamableServerConn) acquireStream(ctx context.Context, w http.ResponseWriter, streamID string, lastIdx *int, supportsPrimeClose bool) (*stream, chan struct{}) {
// if tempStream is set, the stream is done and we're just replaying messages.
//
// We record a temporary stream to claim exclusive replay rights.
@ -898,16 +909,18 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons
}
return c.writeEvent(w, s.id, Event{Name: "message", Data: data}, lastIdx)
}
s.closeLocked = func(reconnectAfter time.Duration) {
select {
case <-done:
return
default:
if supportsPrimeClose {
s.closeLocked = func(reconnectAfter time.Duration) {
select {
case <-done:
return
default:
}
if reconnectAfter > 0 {
c.writeCloseEvent(w, reconnectAfter)
}
close(done)
}
if reconnectAfter > 0 {
c.writeCloseEvent(w, reconnectAfter)
}
close(done)
}
return s, done
}
@ -959,6 +972,7 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
calls := make(map[jsonrpc.ID]struct{})
tokenInfo := auth.TokenInfoFromContext(req.Context())
isInitialize := false
var initializeProtocolVersion string
for _, msg := range incoming {
if jreq, ok := msg.(*jsonrpc.Request); ok {
// Preemptively check that this is a valid request, so that we can fail
@ -970,6 +984,11 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
}
if jreq.Method == methodInitialize {
isInitialize = true
// Extract the protocol version from InitializeParams.
var params InitializeParams
if err := json.Unmarshal(jreq.Params, &params); err == nil {
initializeProtocolVersion = params.ProtocolVersion
}
}
jreq.Extra = &RequestExtra{
TokenInfo: tokenInfo,
@ -994,6 +1013,15 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
}
}
// The prime and close events were added in protocol version 2025-11-25 (SEP-1699).
// Use the version from InitializeParams if this is an initialize request,
// otherwise use the protocol version header.
effectiveVersion := protocolVersion
if isInitialize && initializeProtocolVersion != "" {
effectiveVersion = initializeProtocolVersion
}
supportsPrimeClose := effectiveVersion >= protocolVersion20251125
// If we don't have any calls, we can just publish the incoming messages and return.
// No need to track a logical stream.
if len(calls) == 0 {
@ -1069,7 +1097,7 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
} else {
// Write events in the order we receive them.
lastIndex := -1
if c.eventStore != nil {
if c.eventStore != nil && supportsPrimeClose {
// Write a priming event.
// We must also write it to the event store in order for indexes to
// align.
@ -1092,16 +1120,18 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
}
return c.writeEvent(w, stream.id, Event{Name: "message", Data: data}, &lastIndex)
}
stream.closeLocked = func(reconnectAfter time.Duration) {
select {
case <-done:
return
default:
if supportsPrimeClose {
stream.closeLocked = func(reconnectAfter time.Duration) {
select {
case <-done:
return
default:
}
if reconnectAfter > 0 {
c.writeCloseEvent(w, reconnectAfter)
}
close(done)
}
if reconnectAfter > 0 {
c.writeCloseEvent(w, reconnectAfter)
}
close(done)
}
}

View file

@ -584,7 +584,7 @@ func TestStreamableServerDisconnect(t *testing.T) {
})
clientSession, err := client.Connect(ctx, &StreamableClientTransport{
Endpoint: httpServer.URL,
}, nil)
}, &ClientSessionOptions{protocolVersion: protocolVersion20251125})
if err != nil {
t.Fatalf("client.Connect() failed: %v", err)
}
@ -752,7 +752,7 @@ func TestStreamableServerTransport(t *testing.T) {
// requests.
// Predefined steps, to avoid repetition below.
initReq := req(1, methodInitialize, &InitializeParams{})
initReq := req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20250618})
initResp := resp(1, &InitializeResult{
Capabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
@ -775,6 +775,30 @@ func TestStreamableServerTransport(t *testing.T) {
wantStatusCode: http.StatusAccepted,
}
// Protocol version 2025-11-25 variants, for testing prime/close events (SEP-1699).
initReq20251125 := req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20251125})
initResp20251125 := resp(1, &InitializeResult{
Capabilities: &ServerCapabilities{
Logging: &LoggingCapabilities{},
Tools: &ToolCapabilities{ListChanged: true},
},
ProtocolVersion: protocolVersion20251125,
ServerInfo: &Implementation{Name: "testServer", Version: "v1.0.0"},
}, nil)
initialize20251125 := streamableRequest{
method: "POST",
messages: []jsonrpc.Message{initReq20251125},
wantStatusCode: http.StatusOK,
wantMessages: []jsonrpc.Message{initResp20251125},
wantSessionID: true,
}
initialized20251125 := streamableRequest{
method: "POST",
headers: http.Header{protocolVersionHeader: {protocolVersion20251125}},
messages: []jsonrpc.Message{initializedMsg},
wantStatusCode: http.StatusAccepted,
}
tests := []struct {
name string
replay bool // if set, use a MemoryEventStore to enable replay
@ -1026,23 +1050,23 @@ func TestStreamableServerTransport(t *testing.T) {
wantSessions: 0, // session deleted
},
{
name: "priming message",
name: "no priming message on old protocol",
replay: true,
requests: []streamableRequest{
initialize,
initialized,
{
method: "POST",
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
wantStatusCode: http.StatusOK,
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
wantBodyContaining: "prime",
method: "POST",
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
wantStatusCode: http.StatusOK,
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
wantBodyNotContaining: "prime",
},
},
wantSessions: 1,
},
{
name: "close message",
name: "no close message on old protocol",
replay: true,
tool: func(t *testing.T, _ context.Context, req *CallToolRequest) {
req.Extra.CloseStream(time.Millisecond)
@ -1055,6 +1079,43 @@ func TestStreamableServerTransport(t *testing.T) {
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
wantStatusCode: http.StatusOK,
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
wantBodyNotContaining: "close",
},
},
wantSessions: 1,
},
{
name: "priming message on 2025-11-25",
replay: true,
requests: []streamableRequest{
initialize20251125,
initialized20251125,
{
method: "POST",
headers: http.Header{protocolVersionHeader: {protocolVersion20251125}},
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
wantStatusCode: http.StatusOK,
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
wantBodyContaining: "prime",
},
},
wantSessions: 1,
},
{
name: "close message on 2025-11-25",
replay: true,
tool: func(t *testing.T, _ context.Context, req *CallToolRequest) {
req.Extra.CloseStream(time.Millisecond)
},
requests: []streamableRequest{
initialize20251125,
initialized20251125,
{
method: "POST",
headers: http.Header{protocolVersionHeader: {protocolVersion20251125}},
messages: []jsonrpc.Message{req(2, "tools/call", &CallToolParams{Name: "tool"})},
wantStatusCode: http.StatusOK,
wantMessages: []jsonrpc.Message{resp(2, &CallToolResult{Content: []Content{}}, nil)},
wantBodyContaining: "close",
wantBodyNotContaining: "result",
},