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:
parent
d6d6edd6d4
commit
87c81260e2
3 changed files with 129 additions and 30 deletions
|
|
@ -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(),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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, ¶ms); 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)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
},
|
||||
|
|
|
|||
Loading…
Reference in a new issue