diff --git a/mcp/client.go b/mcp/client.go index a3d36b2..7349ba9 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -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(), } diff --git a/mcp/streamable.go b/mcp/streamable.go index 9ab3303..b48b7c1 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -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) } } diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 9a90d24..9eff74b 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -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", },