405 Method Not Allowed responses MUST include an Allow header listing supported methods, per RFC 9110 Section 15.5.6. This fixes issues with strict HTTP gateways (like Apigee) that treat 405 responses without an Allow header as malformed, returning 502 Bad Gateway errors. ## Changes - **SSEHandler**: Add `Allow: GET, POST` header for unsupported methods - **StreamableHTTPHandler**: Add Allow header for GET-without-session case: - Stateless mode: `Allow: POST, DELETE` (GET is never valid) - Stateful mode: `Allow: GET, POST, DELETE` (GET is valid once you have a session) - Add tests to verify Allow header presence in all 405 responses ## RFC 9110 Reference [Section 15.5.6 (405 Method Not Allowed)](https://httpwg.org/specs/rfc9110.html#status.405): > The origin server MUST generate an Allow header field in a 405 response containing a list of the target resource's currently supported methods. ## Testing All existing tests pass, plus new tests added: - `TestSSE405AllowHeader` - verifies SSE handler compliance - `TestStreamable405AllowHeader` - verifies Streamable handler compliance in both stateful and stateless modes --------- Co-authored-by: omgitsads <4619+omgitsads@users.noreply.github.com>
2090 lines
70 KiB
Go
2090 lines
70 KiB
Go
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
|
// Use of this source code is governed by an MIT-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
// NOTE: see streamable_server.go and streamable_client.go for detailed
|
|
// documentation of the streamable server design.
|
|
// TODO: move the client and server logic into those files.
|
|
|
|
package mcp
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"maps"
|
|
"math"
|
|
"math/rand/v2"
|
|
"net/http"
|
|
"slices"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/modelcontextprotocol/go-sdk/auth"
|
|
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
|
"github.com/modelcontextprotocol/go-sdk/internal/xcontext"
|
|
"github.com/modelcontextprotocol/go-sdk/jsonrpc"
|
|
)
|
|
|
|
const (
|
|
protocolVersionHeader = "Mcp-Protocol-Version"
|
|
sessionIDHeader = "Mcp-Session-Id"
|
|
lastEventIDHeader = "Last-Event-ID"
|
|
)
|
|
|
|
// A StreamableHTTPHandler is an http.Handler that serves streamable MCP
|
|
// sessions, as defined by the [MCP spec].
|
|
//
|
|
// [MCP spec]: https://modelcontextprotocol.io/2025/03/26/streamable-http-transport.html
|
|
type StreamableHTTPHandler struct {
|
|
getServer func(*http.Request) *Server
|
|
opts StreamableHTTPOptions
|
|
|
|
onTransportDeletion func(sessionID string) // for testing
|
|
|
|
mu sync.Mutex
|
|
sessions map[string]*sessionInfo // keyed by session ID
|
|
}
|
|
|
|
type sessionInfo struct {
|
|
session *ServerSession
|
|
transport *StreamableServerTransport
|
|
// userID is the user ID from the TokenInfo when the session was created.
|
|
// If non-empty, subsequent requests must have the same user ID to prevent
|
|
// session hijacking.
|
|
userID string
|
|
|
|
// If timeout is set, automatically close the session after an idle period.
|
|
timeout time.Duration
|
|
timerMu sync.Mutex
|
|
refs int // reference count
|
|
timer *time.Timer
|
|
}
|
|
|
|
// startPOST signals that a POST request for this session is starting (which
|
|
// carries a client->server message), pausing the session timeout if it was
|
|
// running.
|
|
//
|
|
// TODO: we may want to also pause the timer when resuming non-standalone SSE
|
|
// streams, but that is tricy to implement. Clients should generally make
|
|
// keepalive pings if they want to keep the session live.
|
|
func (i *sessionInfo) startPOST() {
|
|
if i.timeout <= 0 {
|
|
return
|
|
}
|
|
|
|
i.timerMu.Lock()
|
|
defer i.timerMu.Unlock()
|
|
|
|
if i.timer == nil {
|
|
return // timer stopped permanently
|
|
}
|
|
if i.refs == 0 {
|
|
i.timer.Stop()
|
|
}
|
|
i.refs++
|
|
}
|
|
|
|
// endPOST sigals that a request for this session is ending, starting the
|
|
// timeout if there are no other requests running.
|
|
func (i *sessionInfo) endPOST() {
|
|
if i.timeout <= 0 {
|
|
return
|
|
}
|
|
|
|
i.timerMu.Lock()
|
|
defer i.timerMu.Unlock()
|
|
|
|
if i.timer == nil {
|
|
return // timer stopped permanently
|
|
}
|
|
|
|
i.refs--
|
|
assert(i.refs >= 0, "negative ref count")
|
|
if i.refs == 0 {
|
|
i.timer.Reset(i.timeout)
|
|
}
|
|
}
|
|
|
|
// stopTimer stops the inactivity timer permanently.
|
|
func (i *sessionInfo) stopTimer() {
|
|
i.timerMu.Lock()
|
|
defer i.timerMu.Unlock()
|
|
if i.timer != nil {
|
|
i.timer.Stop()
|
|
i.timer = nil
|
|
}
|
|
}
|
|
|
|
// StreamableHTTPOptions configures the StreamableHTTPHandler.
|
|
type StreamableHTTPOptions struct {
|
|
// Stateless controls whether the session is 'stateless'.
|
|
//
|
|
// A stateless server does not validate the Mcp-Session-Id header, and uses a
|
|
// temporary session with default initialization parameters. Any
|
|
// server->client request is rejected immediately as there's no way for the
|
|
// client to respond. Server->Client notifications may reach the client if
|
|
// they are made in the context of an incoming request, as described in the
|
|
// documentation for [StreamableServerTransport].
|
|
Stateless bool
|
|
|
|
// TODO(#148): support session retention (?)
|
|
|
|
// JSONResponse causes streamable responses to return application/json rather
|
|
// than text/event-stream ([§2.1.5] of the spec).
|
|
//
|
|
// [§2.1.5]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#sending-messages-to-the-server
|
|
JSONResponse bool
|
|
|
|
// Logger specifies the logger to use.
|
|
// If nil, do not log.
|
|
Logger *slog.Logger
|
|
|
|
// EventStore enables stream resumption.
|
|
//
|
|
// If set, EventStore will be used to persist stream events and replay them
|
|
// upon stream resumption.
|
|
EventStore EventStore
|
|
|
|
// SessionTimeout configures a timeout for idle sessions.
|
|
//
|
|
// When sessions receive no new HTTP requests from the client for this
|
|
// duration, they are automatically closed.
|
|
//
|
|
// If SessionTimeout is the zero value, idle sessions are never closed.
|
|
SessionTimeout time.Duration
|
|
}
|
|
|
|
// NewStreamableHTTPHandler returns a new [StreamableHTTPHandler].
|
|
//
|
|
// The getServer function is used to create or look up servers for new
|
|
// sessions. It is OK for getServer to return the same server multiple times.
|
|
// If getServer returns nil, a 400 Bad Request will be served.
|
|
func NewStreamableHTTPHandler(getServer func(*http.Request) *Server, opts *StreamableHTTPOptions) *StreamableHTTPHandler {
|
|
h := &StreamableHTTPHandler{
|
|
getServer: getServer,
|
|
sessions: make(map[string]*sessionInfo),
|
|
}
|
|
if opts != nil {
|
|
h.opts = *opts
|
|
}
|
|
|
|
if h.opts.Logger == nil { // ensure we have a logger
|
|
h.opts.Logger = ensureLogger(nil)
|
|
}
|
|
|
|
return h
|
|
}
|
|
|
|
// closeAll closes all ongoing sessions, for tests.
|
|
//
|
|
// TODO(rfindley): investigate the best API for callers to configure their
|
|
// session lifecycle. (?)
|
|
//
|
|
// Should we allow passing in a session store? That would allow the handler to
|
|
// be stateless.
|
|
func (h *StreamableHTTPHandler) closeAll() {
|
|
// TODO: if we ever expose this outside of tests, we'll need to do better
|
|
// than simply collecting sessions while holding the lock: we need to prevent
|
|
// new sessions from being added.
|
|
//
|
|
// Currently, sessions remove themselves from h.sessions when closed, so we
|
|
// can't call Close while holding the lock.
|
|
h.mu.Lock()
|
|
sessionInfos := slices.Collect(maps.Values(h.sessions))
|
|
h.sessions = nil
|
|
h.mu.Unlock()
|
|
for _, s := range sessionInfos {
|
|
s.session.Close()
|
|
}
|
|
}
|
|
|
|
func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
|
|
// Allow multiple 'Accept' headers.
|
|
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/Accept#syntax
|
|
accept := strings.Split(strings.Join(req.Header.Values("Accept"), ","), ",")
|
|
var jsonOK, streamOK bool
|
|
for _, c := range accept {
|
|
switch strings.TrimSpace(c) {
|
|
case "application/json", "application/*":
|
|
jsonOK = true
|
|
case "text/event-stream", "text/*":
|
|
streamOK = true
|
|
case "*/*":
|
|
jsonOK = true
|
|
streamOK = true
|
|
}
|
|
}
|
|
|
|
if req.Method == http.MethodGet {
|
|
if !streamOK {
|
|
http.Error(w, "Accept must contain 'text/event-stream' for GET requests", http.StatusBadRequest)
|
|
return
|
|
}
|
|
} else if (!jsonOK || !streamOK) && req.Method != http.MethodDelete { // TODO: consolidate with handling of http method below.
|
|
http.Error(w, "Accept must contain both 'application/json' and 'text/event-stream'", http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
sessionID := req.Header.Get(sessionIDHeader)
|
|
var sessInfo *sessionInfo
|
|
if sessionID != "" {
|
|
h.mu.Lock()
|
|
sessInfo = h.sessions[sessionID]
|
|
h.mu.Unlock()
|
|
if sessInfo == nil && !h.opts.Stateless {
|
|
// Unless we're in 'stateless' mode, which doesn't perform any Session-ID
|
|
// validation, we require that the session ID matches a known session.
|
|
//
|
|
// In stateless mode, a temporary transport is be created below.
|
|
http.Error(w, "session not found", http.StatusNotFound)
|
|
return
|
|
}
|
|
// Prevent session hijacking: if the session was created with a user ID,
|
|
// verify that subsequent requests come from the same user.
|
|
if sessInfo != nil && sessInfo.userID != "" {
|
|
tokenInfo := auth.TokenInfoFromContext(req.Context())
|
|
if tokenInfo == nil || tokenInfo.UserID != sessInfo.userID {
|
|
http.Error(w, "session user mismatch", http.StatusForbidden)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
if req.Method == http.MethodDelete {
|
|
if sessionID == "" {
|
|
http.Error(w, "Bad Request: DELETE requires an Mcp-Session-Id header", http.StatusBadRequest)
|
|
return
|
|
}
|
|
if sessInfo != nil { // sessInfo may be nil in stateless mode
|
|
// Closing the session also removes it from h.sessions, due to the
|
|
// onClose callback.
|
|
sessInfo.session.Close()
|
|
}
|
|
w.WriteHeader(http.StatusNoContent)
|
|
return
|
|
}
|
|
|
|
switch req.Method {
|
|
case http.MethodPost, http.MethodGet:
|
|
if req.Method == http.MethodGet && (h.opts.Stateless || sessionID == "") {
|
|
if h.opts.Stateless {
|
|
// Per MCP spec: server MUST return 405 if it doesn't offer SSE stream.
|
|
// In stateless mode, GET (SSE streaming) is not supported.
|
|
// RFC 9110 §15.5.6: 405 responses MUST include Allow header.
|
|
w.Header().Set("Allow", "POST")
|
|
http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
|
|
} else {
|
|
// In stateful mode, GET is supported but requires a session ID.
|
|
// This is a precondition error, similar to DELETE without session.
|
|
http.Error(w, "Bad Request: GET requires an Mcp-Session-Id header", http.StatusBadRequest)
|
|
}
|
|
return
|
|
}
|
|
default:
|
|
// RFC 9110 §15.5.6: 405 responses MUST include Allow header.
|
|
if h.opts.Stateless {
|
|
w.Header().Set("Allow", "POST")
|
|
} else {
|
|
w.Header().Set("Allow", "GET, POST, DELETE")
|
|
}
|
|
http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
|
|
// [§2.7] of the spec (2025-06-18) states:
|
|
//
|
|
// "If using HTTP, the client MUST include the MCP-Protocol-Version:
|
|
// <protocol-version> HTTP header on all subsequent requests to the MCP
|
|
// server, allowing the MCP server to respond based on the MCP protocol
|
|
// version.
|
|
//
|
|
// For example: MCP-Protocol-Version: 2025-06-18
|
|
// The protocol version sent by the client SHOULD be the one negotiated during
|
|
// initialization.
|
|
//
|
|
// For backwards compatibility, if the server does not receive an
|
|
// MCP-Protocol-Version header, and has no other way to identify the version -
|
|
// for example, by relying on the protocol version negotiated during
|
|
// initialization - the server SHOULD assume protocol version 2025-03-26.
|
|
//
|
|
// If the server receives a request with an invalid or unsupported
|
|
// MCP-Protocol-Version, it MUST respond with 400 Bad Request."
|
|
//
|
|
// Since this wasn't present in the 2025-03-26 version of the spec, this
|
|
// effectively means:
|
|
// 1. IF the client provides a version header, it must be a supported
|
|
// version.
|
|
// 2. In stateless mode, where we've lost the state of the initialize
|
|
// request, we assume that whatever the client tells us is the truth (or
|
|
// assume 2025-03-26 if the client doesn't say anything).
|
|
//
|
|
// This logic matches the typescript SDK.
|
|
//
|
|
// [§2.7]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#protocol-version-header
|
|
protocolVersion := req.Header.Get(protocolVersionHeader)
|
|
if protocolVersion == "" {
|
|
protocolVersion = protocolVersion20250326
|
|
}
|
|
if !slices.Contains(supportedProtocolVersions, protocolVersion) {
|
|
http.Error(w, fmt.Sprintf("Bad Request: Unsupported protocol version (supported versions: %s)", strings.Join(supportedProtocolVersions, ",")), http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
if sessInfo == nil {
|
|
server := h.getServer(req)
|
|
if server == nil {
|
|
// The getServer argument to NewStreamableHTTPHandler returned nil.
|
|
http.Error(w, "no server available", http.StatusBadRequest)
|
|
return
|
|
}
|
|
if sessionID == "" {
|
|
// In stateless mode, sessionID may be nonempty even if there's no
|
|
// existing transport.
|
|
sessionID = server.opts.GetSessionID()
|
|
}
|
|
transport := &StreamableServerTransport{
|
|
SessionID: sessionID,
|
|
Stateless: h.opts.Stateless,
|
|
EventStore: h.opts.EventStore,
|
|
jsonResponse: h.opts.JSONResponse,
|
|
logger: h.opts.Logger,
|
|
}
|
|
|
|
// Sessions without a session ID are also stateless: there's no way to
|
|
// address them.
|
|
stateless := h.opts.Stateless || sessionID == ""
|
|
// To support stateless mode, we initialize the session with a default
|
|
// state, so that it doesn't reject subsequent requests.
|
|
var connectOpts *ServerSessionOptions
|
|
if stateless {
|
|
// Peek at the body to see if it is initialize or initialized.
|
|
// We want those to be handled as usual.
|
|
var hasInitialize, hasInitialized bool
|
|
{
|
|
// TODO: verify that this allows protocol version negotiation for
|
|
// stateless servers.
|
|
body, err := io.ReadAll(req.Body)
|
|
if err != nil {
|
|
http.Error(w, "failed to read body", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
req.Body.Close()
|
|
|
|
// Reset the body so that it can be read later.
|
|
req.Body = io.NopCloser(bytes.NewBuffer(body))
|
|
|
|
msgs, _, err := readBatch(body)
|
|
if err == nil {
|
|
for _, msg := range msgs {
|
|
if req, ok := msg.(*jsonrpc.Request); ok {
|
|
switch req.Method {
|
|
case methodInitialize:
|
|
hasInitialize = true
|
|
case notificationInitialized:
|
|
hasInitialized = true
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// If we don't have InitializeParams or InitializedParams in the request,
|
|
// set the initial state to a default value.
|
|
state := new(ServerSessionState)
|
|
if !hasInitialize {
|
|
state.InitializeParams = &InitializeParams{
|
|
ProtocolVersion: protocolVersion,
|
|
}
|
|
}
|
|
if !hasInitialized {
|
|
state.InitializedParams = new(InitializedParams)
|
|
}
|
|
state.LogLevel = "info"
|
|
connectOpts = &ServerSessionOptions{
|
|
State: state,
|
|
}
|
|
} else {
|
|
// Cleanup is only required in stateful mode, as transportation is
|
|
// not stored in the map otherwise.
|
|
connectOpts = &ServerSessionOptions{
|
|
onClose: func() {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
if info, ok := h.sessions[transport.SessionID]; ok {
|
|
info.stopTimer()
|
|
delete(h.sessions, transport.SessionID)
|
|
if h.onTransportDeletion != nil {
|
|
h.onTransportDeletion(transport.SessionID)
|
|
}
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
// Pass req.Context() here, to allow middleware to add context values.
|
|
// The context is detached in the jsonrpc2 library when handling the
|
|
// long-running stream.
|
|
session, err := server.Connect(req.Context(), transport, connectOpts)
|
|
if err != nil {
|
|
http.Error(w, "failed connection", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
// Capture the user ID from the token info to enable session hijacking
|
|
// prevention on subsequent requests.
|
|
var userID string
|
|
if tokenInfo := auth.TokenInfoFromContext(req.Context()); tokenInfo != nil {
|
|
userID = tokenInfo.UserID
|
|
}
|
|
sessInfo = &sessionInfo{
|
|
session: session,
|
|
transport: transport,
|
|
userID: userID,
|
|
}
|
|
|
|
if stateless {
|
|
// Stateless mode: close the session when the request exits.
|
|
defer session.Close() // close the fake session after handling the request
|
|
} else {
|
|
// Otherwise, save the transport so that it can be reused
|
|
|
|
// Clean up the session when it times out.
|
|
//
|
|
// Note that the timer here may fire multiple times, but
|
|
// sessInfo.session.Close is idempotent.
|
|
if h.opts.SessionTimeout > 0 {
|
|
sessInfo.timeout = h.opts.SessionTimeout
|
|
sessInfo.timer = time.AfterFunc(sessInfo.timeout, func() {
|
|
sessInfo.session.Close()
|
|
})
|
|
}
|
|
h.mu.Lock()
|
|
h.sessions[transport.SessionID] = sessInfo
|
|
h.mu.Unlock()
|
|
defer func() {
|
|
// If initialization failed, clean up the session (#578).
|
|
if session.InitializeParams() == nil {
|
|
// Initialization failed.
|
|
session.Close()
|
|
}
|
|
}()
|
|
}
|
|
}
|
|
|
|
if req.Method == http.MethodPost {
|
|
sessInfo.startPOST()
|
|
defer sessInfo.endPOST()
|
|
}
|
|
|
|
sessInfo.transport.ServeHTTP(w, req)
|
|
}
|
|
|
|
// A StreamableServerTransport implements the server side of the MCP streamable
|
|
// transport.
|
|
//
|
|
// Each StreamableServerTransport must be connected (via [Server.Connect]) at
|
|
// most once, since [StreamableServerTransport.ServeHTTP] serves messages to
|
|
// the connected session.
|
|
//
|
|
// Reads from the streamable server connection receive messages from http POST
|
|
// requests from the client. Writes to the streamable server connection are
|
|
// sent either to the related stream, or to the standalone SSE stream,
|
|
// according to the following rules:
|
|
// - JSON-RPC responses to incoming requests are always routed to the
|
|
// appropriate HTTP response.
|
|
// - Requests or notifications made with a context.Context value derived from
|
|
// an incoming request handler, are routed to the HTTP response
|
|
// corresponding to that request, unless it has already terminated, in
|
|
// which case they are routed to the standalone SSE stream.
|
|
// - Requests or notifications made with a detached context.Context value are
|
|
// routed to the standalone SSE stream.
|
|
type StreamableServerTransport struct {
|
|
// SessionID is the ID of this session.
|
|
//
|
|
// If SessionID is the empty string, this is a 'stateless' session, which has
|
|
// limited ability to communicate with the client. Otherwise, the session ID
|
|
// must be globally unique, that is, different from any other session ID
|
|
// anywhere, past and future. (We recommend using a crypto random number
|
|
// generator to produce one, as with [crypto/rand.Text].)
|
|
SessionID string
|
|
|
|
// Stateless controls whether the eventstore is 'Stateless'. Server sessions
|
|
// connected to a stateless transport are disallowed from making outgoing
|
|
// requests.
|
|
//
|
|
// See also [StreamableHTTPOptions.Stateless].
|
|
Stateless bool
|
|
|
|
// EventStore enables stream resumption.
|
|
//
|
|
// If set, EventStore will be used to persist stream events and replay them
|
|
// upon stream resumption.
|
|
EventStore EventStore
|
|
|
|
// jsonResponse, if set, tells the server to prefer to respond to requests
|
|
// using application/json responses rather than text/event-stream.
|
|
//
|
|
// Specifically, responses will be application/json whenever incoming POST
|
|
// request contain only a single message. In this case, notifications or
|
|
// requests made within the context of a server request will be sent to the
|
|
// standalone SSE stream, if any.
|
|
//
|
|
// TODO(rfindley): jsonResponse should be exported, since
|
|
// StreamableHTTPOptions.JSONResponse is exported, and we want to allow users
|
|
// to write their own streamable HTTP handler.
|
|
jsonResponse bool
|
|
|
|
// optional logger provided through the [StreamableHTTPOptions.Logger].
|
|
//
|
|
// TODO(rfindley): logger should be exported, since we want to allow users
|
|
// to write their own streamable HTTP handler.
|
|
logger *slog.Logger
|
|
|
|
// connection is non-nil if and only if the transport has been connected.
|
|
connection *streamableServerConn
|
|
}
|
|
|
|
// Connect implements the [Transport] interface.
|
|
func (t *StreamableServerTransport) Connect(ctx context.Context) (Connection, error) {
|
|
if t.connection != nil {
|
|
return nil, fmt.Errorf("transport already connected")
|
|
}
|
|
t.connection = &streamableServerConn{
|
|
sessionID: t.SessionID,
|
|
stateless: t.Stateless,
|
|
eventStore: t.EventStore,
|
|
jsonResponse: t.jsonResponse,
|
|
logger: ensureLogger(t.logger), // see #556: must be non-nil
|
|
incoming: make(chan jsonrpc.Message, 10),
|
|
done: make(chan struct{}),
|
|
streams: make(map[string]*stream),
|
|
requestStreams: make(map[jsonrpc.ID]string),
|
|
}
|
|
// Stream 0 corresponds to the standalone SSE stream.
|
|
//
|
|
// It is always text/event-stream, since it must carry arbitrarily many
|
|
// messages.
|
|
var err error
|
|
t.connection.streams[""], err = t.connection.newStream(ctx, nil, "")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return t.connection, nil
|
|
}
|
|
|
|
type streamableServerConn struct {
|
|
sessionID string
|
|
stateless bool
|
|
jsonResponse bool
|
|
eventStore EventStore
|
|
|
|
logger *slog.Logger
|
|
|
|
incoming chan jsonrpc.Message // messages from the client to the server
|
|
|
|
mu sync.Mutex // guards all fields below
|
|
|
|
// Sessions are closed exactly once.
|
|
isDone bool
|
|
done chan struct{}
|
|
|
|
// Sessions can have multiple logical connections (which we call streams),
|
|
// corresponding to HTTP requests. Additionally, streams may be resumed by
|
|
// subsequent HTTP requests, when the HTTP connection is terminated
|
|
// unexpectedly.
|
|
//
|
|
// Therefore, we use a logical stream ID to key the stream state, and
|
|
// perform the accounting described below when incoming HTTP requests are
|
|
// handled.
|
|
|
|
// streams holds the logical streams for this session, keyed by their ID.
|
|
//
|
|
// Lifecycle: streams persist until all of their responses are received from
|
|
// the server.
|
|
streams map[string]*stream
|
|
|
|
// requestStreams maps incoming requests to their logical stream ID.
|
|
//
|
|
// Lifecycle: requestStreams persist until their response is received.
|
|
requestStreams map[jsonrpc.ID]string
|
|
}
|
|
|
|
func (c *streamableServerConn) SessionID() string {
|
|
return c.sessionID
|
|
}
|
|
|
|
// A stream is a single logical stream of SSE events within a server session.
|
|
// A stream begins with a client request, or with a client GET that has
|
|
// no Last-Event-ID header.
|
|
//
|
|
// A stream ends only when its session ends; we cannot determine its end otherwise,
|
|
// since a client may send a GET with a Last-Event-ID that references the stream
|
|
// at any time.
|
|
type stream struct {
|
|
// id is the logical ID for the stream, unique within a session.
|
|
//
|
|
// The standalone SSE stream has id "".
|
|
id string
|
|
|
|
// logger is used for logging errors during stream operations.
|
|
logger *slog.Logger
|
|
|
|
// mu guards the fields below, as well as storage of new messages in the
|
|
// connection's event store (if any).
|
|
mu sync.Mutex
|
|
|
|
// If pendingJSONMessages is non-nil, this is a JSON stream and messages are
|
|
// collected here until the stream is complete, at which point they are
|
|
// flushed as a single JSON response. Note that the non-nilness of this field
|
|
// is significant, as it signals the expected content type.
|
|
//
|
|
// Note: if we remove support for batching, this could just be a bool.
|
|
pendingJSONMessages []json.RawMessage
|
|
|
|
// w is the HTTP response writer for this stream. A non-nil w indicates
|
|
// that the stream is claimed by an HTTP request (the hanging POST or GET);
|
|
// it is set to nil when the request completes.
|
|
w http.ResponseWriter
|
|
|
|
// done is closed to release the hanging HTTP request.
|
|
//
|
|
// Invariant: a non-nil done implies w is also non-nil, though the converse
|
|
// is not necessarily true: done is set to nil when it is closed, to avoid
|
|
// duplicate closure.
|
|
done chan struct{}
|
|
|
|
// lastIdx is the index of the last written SSE event, for event ID generation.
|
|
// It starts at -1 since indices start at 0.
|
|
lastIdx int
|
|
|
|
// protocolVersion is the protocol version for this stream.
|
|
protocolVersion string
|
|
|
|
// requests is the set of unanswered incoming requests for the stream.
|
|
//
|
|
// Requests are removed when their response has been received.
|
|
// In practice, there is only one request, but in the 2025-03-26 version of
|
|
// the spec and earlier there was a concept of batching, in which POST
|
|
// payloads could hold multiple requests or responses.
|
|
requests map[jsonrpc.ID]struct{}
|
|
}
|
|
|
|
// close sends a 'close' event to the client (if protocolVersion >= 2025-11-25
|
|
// and reconnectAfter > 0) and closes the done channel.
|
|
//
|
|
// The done channel is set to nil after closing, so that done != nil implies
|
|
// the stream is active and done is open. This simplifies checks elsewhere.
|
|
func (s *stream) close(reconnectAfter time.Duration) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if s.done == nil {
|
|
return // stream not connected or already closed
|
|
}
|
|
if s.protocolVersion >= protocolVersion20251125 && reconnectAfter > 0 {
|
|
reconnectStr := strconv.FormatInt(reconnectAfter.Milliseconds(), 10)
|
|
if _, err := writeEvent(s.w, Event{
|
|
Name: "close",
|
|
Retry: reconnectStr,
|
|
}); err != nil {
|
|
s.logger.Warn(fmt.Sprintf("Writing close event: %v", err))
|
|
}
|
|
}
|
|
close(s.done)
|
|
s.done = nil
|
|
}
|
|
|
|
// release releases the stream from its HTTP request, allowing it to be
|
|
// claimed by another request (e.g., for resumption).
|
|
func (s *stream) release() {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.w = nil
|
|
s.done = nil // may already be nil, if the stream is done or closed
|
|
}
|
|
|
|
// deliverLocked writes data to the stream (for SSE) or stores it in
|
|
// pendingJSONMessages (for JSON mode). The eventID is used for SSE event ID;
|
|
// pass "" to omit.
|
|
//
|
|
// If responseTo is valid, it is removed from the requests map. When all
|
|
// requests have been responded to, the done channel is closed and set to nil.
|
|
//
|
|
// Returns true if the stream is now done (all requests have been responded to).
|
|
// The done value is always accurate, even if an error is returned.
|
|
//
|
|
// s.mu must be held when calling this method.
|
|
func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.ID) (done bool, err error) {
|
|
// First, record the response. We must do this *before* returning an error
|
|
// below, as even if the stream is disconnected we want to update our
|
|
// accounting.
|
|
if responseTo.IsValid() {
|
|
delete(s.requests, responseTo)
|
|
}
|
|
// Now, try to deliver the message to the client.
|
|
done = len(s.requests) == 0 && s.id != ""
|
|
if s.done == nil {
|
|
return done, fmt.Errorf("stream not connected or already closed")
|
|
}
|
|
if done {
|
|
defer func() { close(s.done); s.done = nil }()
|
|
}
|
|
// Try to write to the response.
|
|
//
|
|
// If we get here, the request is still hanging (because s.done != nil
|
|
// implies s.w != nil), but may have been cancelled by the client/http layer:
|
|
// there's a brief race between request cancellation and releasing the
|
|
// stream.
|
|
if s.pendingJSONMessages != nil {
|
|
s.pendingJSONMessages = append(s.pendingJSONMessages, data)
|
|
if done {
|
|
// Flush all pending messages as JSON response.
|
|
var toWrite []byte
|
|
if len(s.pendingJSONMessages) == 1 {
|
|
toWrite = s.pendingJSONMessages[0]
|
|
} else {
|
|
toWrite, err = json.Marshal(s.pendingJSONMessages)
|
|
if err != nil {
|
|
return done, err
|
|
}
|
|
}
|
|
if _, err := s.w.Write(toWrite); err != nil {
|
|
return done, err
|
|
}
|
|
}
|
|
} else {
|
|
// SSE mode: write event to response writer.
|
|
s.lastIdx++
|
|
if _, err := writeEvent(s.w, Event{Name: "message", Data: data, ID: eventID}); err != nil {
|
|
return done, err
|
|
}
|
|
}
|
|
return done, nil
|
|
}
|
|
|
|
// doneLocked reports whether the stream is logically complete.
|
|
//
|
|
// s.requests was populated when reading the POST body, requests are deleted as
|
|
// they are responded to. Once all requests have been responded to, the stream
|
|
// is done.
|
|
//
|
|
// s.mu must be held while calling this function.
|
|
func (s *stream) doneLocked() bool {
|
|
return len(s.requests) == 0 && s.id != ""
|
|
}
|
|
|
|
func (c *streamableServerConn) newStream(ctx context.Context, requests map[jsonrpc.ID]struct{}, id string) (*stream, error) {
|
|
if c.eventStore != nil {
|
|
if err := c.eventStore.Open(ctx, c.sessionID, id); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
return &stream{
|
|
id: id,
|
|
requests: requests,
|
|
lastIdx: -1, // indices start at 0, incremented before each write
|
|
logger: c.logger,
|
|
}, nil
|
|
}
|
|
|
|
// We track the incoming request ID inside the handler context using
|
|
// idContextValue, so that notifications and server->client calls that occur in
|
|
// the course of handling incoming requests are correlated with the incoming
|
|
// request that caused them, and can be dispatched as server-sent events to the
|
|
// correct HTTP request.
|
|
//
|
|
// Currently, this is implemented in [ServerSession.handle]. This is not ideal,
|
|
// because it means that a user of the MCP package couldn't implement the
|
|
// streamable transport, as they'd lack this privileged access.
|
|
//
|
|
// If we ever wanted to expose this mechanism, we have a few options:
|
|
// 1. Make ServerSession an interface, and provide an implementation of
|
|
// ServerSession to handlers that closes over the incoming request ID.
|
|
// 2. Expose a 'HandlerTransport' interface that allows transports to provide
|
|
// a handler middleware, so that we don't hard-code this behavior in
|
|
// ServerSession.handle.
|
|
// 3. Add a `func ForRequest(context.Context) jsonrpc.ID` accessor that lets
|
|
// any transport access the incoming request ID.
|
|
//
|
|
// For now, by giving only the StreamableServerTransport access to the request
|
|
// ID, we avoid having to make this API decision.
|
|
type idContextKey struct{}
|
|
|
|
// ServeHTTP handles a single HTTP request for the session.
|
|
func (t *StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request) {
|
|
if t.connection == nil {
|
|
http.Error(w, "transport not connected", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
switch req.Method {
|
|
case http.MethodGet:
|
|
t.connection.serveGET(w, req)
|
|
case http.MethodPost:
|
|
t.connection.servePOST(w, req)
|
|
default:
|
|
// Should not be reached, as this is checked in StreamableHTTPHandler.ServeHTTP.
|
|
w.Header().Set("Allow", "GET, POST")
|
|
http.Error(w, "unsupported method", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
}
|
|
|
|
// serveGET streams messages to a hanging http GET, with stream ID and last
|
|
// message parsed from the Last-Event-ID header.
|
|
//
|
|
// It returns an HTTP status code and error message.
|
|
func (c *streamableServerConn) serveGET(w http.ResponseWriter, req *http.Request) {
|
|
// streamID "" corresponds to the default GET request.
|
|
streamID := ""
|
|
// By default, we haven't seen a last index. Since indices start at 0, we represent
|
|
// that by -1. This is incremented just before each event is written.
|
|
lastIdx := -1
|
|
if len(req.Header.Values(lastEventIDHeader)) > 0 {
|
|
eid := req.Header.Get(lastEventIDHeader)
|
|
var ok bool
|
|
streamID, lastIdx, ok = parseEventID(eid)
|
|
if !ok {
|
|
http.Error(w, fmt.Sprintf("malformed Last-Event-ID %q", eid), http.StatusBadRequest)
|
|
return
|
|
}
|
|
if c.eventStore == nil {
|
|
http.Error(w, "stream replay unsupported", http.StatusBadRequest)
|
|
return
|
|
}
|
|
}
|
|
|
|
ctx := req.Context()
|
|
|
|
// 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
|
|
}
|
|
|
|
stream, done := c.acquireStream(ctx, w, streamID, lastIdx, protocolVersion)
|
|
if stream == nil {
|
|
return
|
|
}
|
|
defer stream.release()
|
|
c.hangResponse(ctx, done)
|
|
}
|
|
|
|
// hangResponse blocks the HTTP response until one of three conditions is met:
|
|
// - ctx is cancelled (the client disconnected or the request timed out)
|
|
// - done is closed (all responses have been sent, or the stream was explicitly closed)
|
|
// - the session is closed
|
|
//
|
|
// This keeps the HTTP connection open so that server-sent events can be
|
|
// written to the response.
|
|
func (c *streamableServerConn) hangResponse(ctx context.Context, done <-chan struct{}) {
|
|
select {
|
|
case <-ctx.Done():
|
|
case <-done:
|
|
case <-c.done:
|
|
}
|
|
}
|
|
|
|
// acquireStream replays all events since lastIdx, and acquires the ongoing
|
|
// stream, if any. If non-nil, the resulting stream will be registered for
|
|
// receiving new messages, and the stream's 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.
|
|
//
|
|
// protocolVersion is the protocol version for this stream, used to determine
|
|
// feature support (e.g. prime and close events were added in 2025-11-25).
|
|
func (c *streamableServerConn) acquireStream(ctx context.Context, w http.ResponseWriter, streamID string, lastIdx int, protocolVersion string) (*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. The spec
|
|
// (https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#resumability-and-redelivery)
|
|
// does not explicitly require exclusive replay, but we enforce it defensively.
|
|
tempStream := false
|
|
c.mu.Lock()
|
|
s, ok := c.streams[streamID]
|
|
if !ok {
|
|
// The stream is logically done, but claim exclusive rights to replay it by
|
|
// adding a temporary entry in the streams map.
|
|
//
|
|
// We create this entry with a non-nil w, to ensure it isn't claimed by
|
|
// another request before we lock it below.
|
|
tempStream = true
|
|
s = &stream{
|
|
id: streamID,
|
|
w: w,
|
|
}
|
|
c.streams[streamID] = s
|
|
|
|
// Since this stream is transient, we must clean up after replaying.
|
|
defer func() {
|
|
c.mu.Lock()
|
|
delete(c.streams, streamID)
|
|
c.mu.Unlock()
|
|
}()
|
|
}
|
|
c.mu.Unlock()
|
|
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
// Check that this stream wasn't claimed by another request.
|
|
if !tempStream && s.w != nil {
|
|
http.Error(w, "stream ID conflicts with ongoing stream", http.StatusConflict)
|
|
return nil, nil
|
|
}
|
|
|
|
// Collect events to replay. Collect them all before writing, so that we
|
|
// have an opportunity to set the HTTP status code on an error.
|
|
//
|
|
// As indicated above, we must do that while holding stream.mu, so that no
|
|
// new messages are added to the eventstore until we've replayed all previous
|
|
// messages, and registered our delivery function.
|
|
var toReplay [][]byte
|
|
if c.eventStore != nil {
|
|
for data, err := range c.eventStore.After(ctx, c.SessionID(), s.id, lastIdx) {
|
|
if err != nil {
|
|
// We can't replay events, perhaps because the underlying event store
|
|
// has garbage collected its storage.
|
|
//
|
|
// We must be careful here: any 404 will signal to the client that the
|
|
// *session* is not found, rather than the stream.
|
|
//
|
|
// 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 nil, nil
|
|
}
|
|
if len(data) > 0 {
|
|
toReplay = append(toReplay, data)
|
|
}
|
|
}
|
|
}
|
|
|
|
w.Header().Set("Cache-Control", "no-cache, no-transform")
|
|
w.Header().Set("Content-Type", "text/event-stream") // Accept checked in [StreamableHTTPHandler]
|
|
w.Header().Set("Connection", "keep-alive")
|
|
|
|
if s.id == "" {
|
|
// Issue #410: the standalone SSE stream is likely not to receive messages
|
|
// for a long time. Ensure that headers are flushed.
|
|
w.WriteHeader(http.StatusOK)
|
|
if f, ok := w.(http.Flusher); ok {
|
|
f.Flush()
|
|
}
|
|
}
|
|
|
|
for _, data := range toReplay {
|
|
lastIdx++
|
|
e := Event{Name: "message", Data: data}
|
|
if c.eventStore != nil {
|
|
e.ID = formatEventID(s.id, lastIdx)
|
|
}
|
|
if _, err := writeEvent(w, e); err != nil {
|
|
return nil, nil
|
|
}
|
|
}
|
|
|
|
if tempStream || s.doneLocked() {
|
|
// Nothing more to do.
|
|
return nil, nil
|
|
}
|
|
|
|
// The stream is not done: set up delivery state before the stream is
|
|
// unlocked, allowing the connection to write new events.
|
|
s.w = w
|
|
s.done = make(chan struct{})
|
|
s.lastIdx = lastIdx
|
|
s.protocolVersion = protocolVersion
|
|
return s, s.done
|
|
}
|
|
|
|
// servePOST handles an incoming message, and replies with either an outgoing
|
|
// message stream or single response object, depending on whether the
|
|
// jsonResponse option is set.
|
|
//
|
|
// It returns an HTTP status code and error message.
|
|
func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Request) {
|
|
if len(req.Header.Values(lastEventIDHeader)) > 0 {
|
|
http.Error(w, "can't send Last-Event-ID for POST request", http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
// Read incoming messages.
|
|
body, err := io.ReadAll(req.Body)
|
|
if err != nil {
|
|
http.Error(w, "failed to read body", http.StatusBadRequest)
|
|
return
|
|
}
|
|
if len(body) == 0 {
|
|
http.Error(w, "POST requires a non-empty body", http.StatusBadRequest)
|
|
return
|
|
}
|
|
// TODO(#674): once we've documented the support matrix for 2025-03-26 and
|
|
// earlier, drop support for matching entirely; that will simplify this
|
|
// logic.
|
|
incoming, isBatch, err := readBatch(body)
|
|
if err != nil {
|
|
http.Error(w, fmt.Sprintf("malformed payload: %v", err), http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
protocolVersion := req.Header.Get(protocolVersionHeader)
|
|
if protocolVersion == "" {
|
|
protocolVersion = protocolVersion20250326
|
|
}
|
|
|
|
if isBatch && protocolVersion >= protocolVersion20250618 {
|
|
http.Error(w, fmt.Sprintf("JSON-RPC batching is not supported in %s and later (request version: %s)", protocolVersion20250618, protocolVersion), http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
// TODO(rfindley): no tests fail if we reject batch JSON requests entirely.
|
|
// We need to test this with older protocol versions.
|
|
// if isBatch && c.jsonResponse {
|
|
// http.Error(w, "server does not support batch requests", http.StatusBadRequest)
|
|
// return
|
|
// }
|
|
|
|
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
|
|
// the HTTP request. If we didn't do this, a request with a bad method or
|
|
// missing ID could be silently swallowed.
|
|
if _, err := checkRequest(jreq, serverMethodInfos); err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
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
|
|
}
|
|
}
|
|
// Include metadata for all requests (including notifications).
|
|
jreq.Extra = &RequestExtra{
|
|
TokenInfo: tokenInfo,
|
|
Header: req.Header,
|
|
}
|
|
if jreq.IsCall() {
|
|
calls[jreq.ID] = struct{}{}
|
|
// See the doc for CloseSSEStream: allow the request handler to
|
|
// explicitly close the ongoing stream.
|
|
jreq.Extra.(*RequestExtra).CloseSSEStream = func(args CloseSSEStreamArgs) {
|
|
c.mu.Lock()
|
|
streamID, ok := c.requestStreams[jreq.ID]
|
|
var stream *stream
|
|
if ok {
|
|
stream = c.streams[streamID]
|
|
}
|
|
c.mu.Unlock()
|
|
|
|
if stream != nil {
|
|
stream.close(args.RetryAfter)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
// If we don't have any calls, we can just publish the incoming messages and return.
|
|
// No need to track a logical stream.
|
|
//
|
|
// See section [§2.1.4] of the spec: "If the server accepts the input, the
|
|
// server MUST return HTTP status code 202 Accepted with no body."
|
|
//
|
|
// [§2.1.4]: https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#sending-messages-to-the-server
|
|
if len(calls) == 0 {
|
|
for _, msg := range incoming {
|
|
select {
|
|
case c.incoming <- msg:
|
|
case <-c.done:
|
|
// The session is closing. Since we haven't yet written any data to the
|
|
// response, we can signal to the client that the session is gone.
|
|
http.Error(w, "session is closing", http.StatusNotFound)
|
|
return
|
|
}
|
|
}
|
|
w.WriteHeader(http.StatusAccepted)
|
|
return
|
|
}
|
|
|
|
// Invariant: we have at least one call.
|
|
//
|
|
// Create a logical stream to track its responses.
|
|
// Important: don't publish the incoming messages until the stream is
|
|
// registered, as the server may attempt to respond to imcoming messages as
|
|
// soon as they're published.
|
|
stream, err := c.newStream(req.Context(), calls, randText())
|
|
if err != nil {
|
|
http.Error(w, fmt.Sprintf("storing stream: %v", err), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
// Set response headers. Accept was checked in [StreamableHTTPHandler].
|
|
w.Header().Set("Cache-Control", "no-cache, no-transform")
|
|
if c.jsonResponse {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
} else {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Header().Set("Connection", "keep-alive")
|
|
}
|
|
if c.sessionID != "" && isInitialize {
|
|
w.Header().Set(sessionIDHeader, c.sessionID)
|
|
}
|
|
|
|
// Set up stream delivery state.
|
|
stream.w = w
|
|
done := make(chan struct{})
|
|
stream.done = done
|
|
stream.protocolVersion = effectiveVersion
|
|
if c.jsonResponse {
|
|
// JSON mode: collect messages in pendingJSONMessages until done.
|
|
// Set pendingJSONMessages to a non-nil value to signal that this is an
|
|
// application/json stream.
|
|
stream.pendingJSONMessages = []json.RawMessage{}
|
|
} else {
|
|
// SSE mode: write a priming event if supported.
|
|
if c.eventStore != nil && effectiveVersion >= protocolVersion20251125 {
|
|
// Write a priming event, as defined by [§2.1.6] of the spec.
|
|
//
|
|
// [§2.1.6]: https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#sending-messages-to-the-server
|
|
//
|
|
// We must also write it to the event store in order for indexes to
|
|
// align.
|
|
if err := c.eventStore.Append(req.Context(), c.sessionID, stream.id, nil); err != nil {
|
|
c.logger.Warn(fmt.Sprintf("Storing priming event: %v", err))
|
|
}
|
|
stream.lastIdx++
|
|
e := Event{Name: "prime", ID: formatEventID(stream.id, stream.lastIdx)}
|
|
if _, err := writeEvent(w, e); err != nil {
|
|
c.logger.Warn(fmt.Sprintf("Writing priming event: %v", err))
|
|
}
|
|
}
|
|
}
|
|
|
|
// TODO(rfindley): if we have no event store, we should really cancel all
|
|
// remaining requests here, since the client will never get the results.
|
|
defer stream.release()
|
|
|
|
// The stream is now set up to deliver messages.
|
|
//
|
|
// Register it before publishing incoming messages.
|
|
c.mu.Lock()
|
|
c.streams[stream.id] = stream
|
|
for reqID := range calls {
|
|
c.requestStreams[reqID] = stream.id
|
|
}
|
|
c.mu.Unlock()
|
|
|
|
// Publish incoming messages.
|
|
for _, msg := range incoming {
|
|
select {
|
|
case c.incoming <- msg:
|
|
// Note: don't select on req.Context().Done() here, since we've already
|
|
// received the requests and may have already published a response message
|
|
// or notification. The client could resume the stream.
|
|
//
|
|
// In fact, this send could be in a separate goroutine.
|
|
case <-c.done:
|
|
// Session closed: we don't know if any data has been written, so it's
|
|
// too late to write a status code here.
|
|
return
|
|
}
|
|
}
|
|
|
|
c.hangResponse(req.Context(), done)
|
|
}
|
|
|
|
// Event IDs: encode both the logical connection ID and the index, as
|
|
// <streamID>_<idx>, to be consistent with the typescript implementation.
|
|
|
|
// formatEventID returns the event ID to use for the logical connection ID
|
|
// streamID and message index idx.
|
|
//
|
|
// See also [parseEventID].
|
|
func formatEventID(sid string, idx int) string {
|
|
return fmt.Sprintf("%s_%d", sid, idx)
|
|
}
|
|
|
|
// parseEventID parses a Last-Event-ID value into a logical stream id and
|
|
// index.
|
|
//
|
|
// See also [formatEventID].
|
|
func parseEventID(eventID string) (streamID string, idx int, ok bool) {
|
|
parts := strings.Split(eventID, "_")
|
|
if len(parts) != 2 {
|
|
return "", 0, false
|
|
}
|
|
streamID = parts[0]
|
|
idx, err := strconv.Atoi(parts[1])
|
|
if err != nil || idx < 0 {
|
|
return "", 0, false
|
|
}
|
|
return streamID, idx, true
|
|
}
|
|
|
|
// Read implements the [Connection] interface.
|
|
func (c *streamableServerConn) Read(ctx context.Context) (jsonrpc.Message, error) {
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case msg, ok := <-c.incoming:
|
|
if !ok {
|
|
return nil, io.EOF
|
|
}
|
|
return msg, nil
|
|
case <-c.done:
|
|
return nil, io.EOF
|
|
}
|
|
}
|
|
|
|
// Write implements the [Connection] interface.
|
|
func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) error {
|
|
// Throughout this function, note that any error that wraps ErrRejected
|
|
// indicates a does not cause the connection to break.
|
|
//
|
|
// Most errors don't break the connection: unlike a true bidirectional
|
|
// stream, a failure to deliver to a stream is not an indication that the
|
|
// logical session is broken.
|
|
data, err := jsonrpc2.EncodeMessage(msg)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if req, ok := msg.(*jsonrpc.Request); ok && req.IsCall() && (c.stateless || c.sessionID == "") {
|
|
// Requests aren't possible with stateless servers, or when there's no session ID.
|
|
return fmt.Errorf("%w: stateless servers cannot make requests", jsonrpc2.ErrRejected)
|
|
}
|
|
|
|
// Find the incoming request that this write relates to, if any.
|
|
var (
|
|
relatedRequest jsonrpc.ID
|
|
responseTo jsonrpc.ID // if valid, the message is a response to this request
|
|
)
|
|
if resp, ok := msg.(*jsonrpc.Response); ok {
|
|
// If the message is a response, it relates to its request (of course).
|
|
relatedRequest = resp.ID
|
|
responseTo = resp.ID
|
|
} else {
|
|
// Otherwise, we check to see if it request was made in the context of an
|
|
// ongoing request. This may not be the case if the request was made with
|
|
// an unrelated context.
|
|
if v := ctx.Value(idContextKey{}); v != nil {
|
|
relatedRequest = v.(jsonrpc.ID)
|
|
}
|
|
}
|
|
|
|
// If the stream is application/json, but the message is not a response, we
|
|
// must send it out of band to the standalone SSE stream.
|
|
if c.jsonResponse && !responseTo.IsValid() {
|
|
relatedRequest = jsonrpc.ID{}
|
|
}
|
|
|
|
// Write the message to the stream.
|
|
var s *stream
|
|
c.mu.Lock()
|
|
if relatedRequest.IsValid() {
|
|
if streamID, ok := c.requestStreams[relatedRequest]; ok {
|
|
s = c.streams[streamID]
|
|
}
|
|
} else {
|
|
s = c.streams[""] // standalone SSE stream
|
|
}
|
|
if responseTo.IsValid() {
|
|
// Once we've responded to a request, disallow related messages by removing
|
|
// the stream association. This also releases memory.
|
|
delete(c.requestStreams, responseTo)
|
|
}
|
|
sessionClosed := c.isDone
|
|
c.mu.Unlock()
|
|
|
|
if s == nil {
|
|
// The request was made in the context of an ongoing request, but that
|
|
// request is complete.
|
|
//
|
|
// In the future, we could be less strict and allow the request to land on
|
|
// the standalone SSE stream.
|
|
return fmt.Errorf("%w: write to closed stream", jsonrpc2.ErrRejected)
|
|
}
|
|
if sessionClosed {
|
|
return errors.New("session is closed")
|
|
}
|
|
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
// Store in eventStore before delivering.
|
|
// TODO(rfindley): we should only append if the response is SSE, not JSON, by
|
|
// pushing down into the delivery layer.
|
|
delivered := false
|
|
var errs []error
|
|
if c.eventStore != nil {
|
|
if err := c.eventStore.Append(ctx, c.sessionID, s.id, data); err != nil {
|
|
errs = append(errs, err)
|
|
} else {
|
|
delivered = true
|
|
}
|
|
}
|
|
|
|
// Compute eventID for SSE streams with event store.
|
|
// Use s.lastIdx + 1 because deliverLocked increments before writing.
|
|
var eventID string
|
|
if c.eventStore != nil {
|
|
eventID = formatEventID(s.id, s.lastIdx+1)
|
|
}
|
|
|
|
done, err := s.deliverLocked(data, eventID, responseTo)
|
|
if err != nil {
|
|
errs = append(errs, err)
|
|
} else {
|
|
delivered = true
|
|
}
|
|
|
|
if done {
|
|
c.mu.Lock()
|
|
delete(c.streams, s.id)
|
|
c.mu.Unlock()
|
|
}
|
|
|
|
if !delivered {
|
|
return fmt.Errorf("%w: undelivered message: %v", jsonrpc2.ErrRejected, errors.Join(errs...))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Close implements the [Connection] interface.
|
|
func (c *streamableServerConn) Close() error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if !c.isDone {
|
|
c.isDone = true
|
|
close(c.done)
|
|
if c.eventStore != nil {
|
|
// TODO: find a way to plumb a context here, or an event store with a long-running
|
|
// close operation can take arbitrary time. Alternative: impose a fixed timeout here.
|
|
return c.eventStore.SessionClosed(context.TODO(), c.sessionID)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// A StreamableClientTransport is a [Transport] that can communicate with an MCP
|
|
// endpoint serving the streamable HTTP transport defined by the 2025-03-26
|
|
// version of the spec.
|
|
type StreamableClientTransport struct {
|
|
Endpoint string
|
|
HTTPClient *http.Client
|
|
// MaxRetries is the maximum number of times to attempt a reconnect before giving up.
|
|
// It defaults to 5. To disable retries, use a negative number.
|
|
MaxRetries int
|
|
|
|
// DisableStandaloneSSE controls whether the client establishes a standalone SSE stream
|
|
// for receiving server-initiated messages.
|
|
//
|
|
// When false (the default), after initialization the client sends an HTTP GET request
|
|
// to establish a persistent server-sent events (SSE) connection. This allows the server
|
|
// to send messages to the client at any time, such as ToolListChangedNotification or
|
|
// other server-initiated requests and notifications. The connection persists for the
|
|
// lifetime of the session and automatically reconnects if interrupted.
|
|
//
|
|
// When true, the client does not establish the standalone SSE stream. The client will
|
|
// only receive responses to its own POST requests. Server-initiated messages will not
|
|
// be received.
|
|
//
|
|
// According to the MCP specification, the standalone SSE stream is optional.
|
|
// Setting DisableStandaloneSSE to true is useful when:
|
|
// - You only need request-response communication and don't need server-initiated notifications
|
|
// - The server doesn't properly handle GET requests for SSE streams
|
|
// - You want to avoid maintaining a persistent connection
|
|
DisableStandaloneSSE bool
|
|
|
|
// TODO(rfindley): propose exporting these.
|
|
// If strict is set, the transport is in 'strict mode', where any violation
|
|
// of the MCP spec causes a failure.
|
|
strict bool
|
|
// If logger is set, it is used to log aspects of the transport, such as spec
|
|
// violations that were ignored.
|
|
logger *slog.Logger
|
|
}
|
|
|
|
// These settings are not (yet) exposed to the user in
|
|
// StreamableClientTransport.
|
|
const (
|
|
// reconnectGrowFactor is the multiplicative factor by which the delay increases after each attempt.
|
|
// A value of 1.0 results in a constant delay, while a value of 2.0 would double it each time.
|
|
// It must be 1.0 or greater if MaxRetries is greater than 0.
|
|
reconnectGrowFactor = 1.5
|
|
// reconnectMaxDelay caps the backoff delay, preventing it from growing indefinitely.
|
|
reconnectMaxDelay = 30 * time.Second
|
|
)
|
|
|
|
var (
|
|
// reconnectInitialDelay is the base delay for the first reconnect attempt.
|
|
//
|
|
// Mutable for testing.
|
|
reconnectInitialDelay = 1 * time.Second
|
|
)
|
|
|
|
// Connect implements the [Transport] interface.
|
|
//
|
|
// The resulting [Connection] writes messages via POST requests to the
|
|
// transport URL with the Mcp-Session-Id header set, and reads messages from
|
|
// hanging requests.
|
|
//
|
|
// When closed, the connection issues a DELETE request to terminate the logical
|
|
// session.
|
|
func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, error) {
|
|
client := t.HTTPClient
|
|
if client == nil {
|
|
client = http.DefaultClient
|
|
}
|
|
maxRetries := t.MaxRetries
|
|
if maxRetries == 0 {
|
|
maxRetries = 5
|
|
} else if maxRetries < 0 {
|
|
maxRetries = 0
|
|
}
|
|
// Create a new cancellable context that will manage the connection's lifecycle.
|
|
// This is crucial for cleanly shutting down the background SSE listener by
|
|
// cancelling its blocking network operations, which prevents hangs on exit.
|
|
//
|
|
// This context should be detached from the incoming context: the standalone
|
|
// SSE request should not break when the connection context is done.
|
|
//
|
|
// For example, consider that the user may want to wait at most 5s to connect
|
|
// to the server, and therefore uses a context with a 5s timeout when calling
|
|
// client.Connect. Let's suppose that Connect returns after 1s, and the user
|
|
// starts using the resulting session. If we didn't detach here, the session
|
|
// would break after 4s, when the background SSE stream is terminated.
|
|
//
|
|
// Instead, creating a cancellable context detached from the incoming context
|
|
// allows us to preserve context values (which may be necessary for auth
|
|
// middleware), yet only cancel the standalone stream when the connection is closed.
|
|
connCtx, cancel := context.WithCancel(xcontext.Detach(ctx))
|
|
conn := &streamableClientConn{
|
|
url: t.Endpoint,
|
|
client: client,
|
|
incoming: make(chan jsonrpc.Message, 10),
|
|
done: make(chan struct{}),
|
|
maxRetries: maxRetries,
|
|
strict: t.strict,
|
|
logger: ensureLogger(t.logger), // must be non-nil for safe logging
|
|
ctx: connCtx,
|
|
cancel: cancel,
|
|
failed: make(chan struct{}),
|
|
disableStandaloneSSE: t.DisableStandaloneSSE,
|
|
}
|
|
return conn, nil
|
|
}
|
|
|
|
type streamableClientConn struct {
|
|
url string
|
|
client *http.Client
|
|
ctx context.Context // connection context, detached from Connect
|
|
cancel context.CancelFunc // cancels ctx
|
|
incoming chan jsonrpc.Message
|
|
maxRetries int
|
|
strict bool // from [StreamableClientTransport.strict]
|
|
logger *slog.Logger // from [StreamableClientTransport.logger]
|
|
|
|
// disableStandaloneSSE controls whether to disable the standalone SSE stream
|
|
// for receiving server-to-client notifications when no request is in flight.
|
|
disableStandaloneSSE bool // from [StreamableClientTransport.DisableStandaloneSSE]
|
|
|
|
// Guard calls to Close, as it may be called multiple times.
|
|
closeOnce sync.Once
|
|
closeErr error
|
|
done chan struct{} // signal graceful termination
|
|
|
|
// Logical reads are distributed across multiple http requests. Whenever any
|
|
// of them fails to process their response, we must break the connection, by
|
|
// failing the pending Read.
|
|
//
|
|
// Achieve this by storing the failure message, and signalling when reads are
|
|
// broken. See also [streamableClientConn.fail] and
|
|
// [streamableClientConn.failure].
|
|
failOnce sync.Once
|
|
_failure error
|
|
failed chan struct{} // signal failure
|
|
|
|
// Guard the initialization state.
|
|
mu sync.Mutex
|
|
initializedResult *InitializeResult
|
|
sessionID string
|
|
}
|
|
|
|
// errSessionMissing distinguishes if the session is known to not be present on
|
|
// the server (see [streamableClientConn.fail]).
|
|
//
|
|
// TODO(rfindley): should we expose this error value (and its corresponding
|
|
// API) to the user?
|
|
//
|
|
// The spec says that if the server returns 404, clients should reestablish
|
|
// a session. For now, we delegate that to the user, but do they need a way to
|
|
// differentiate a 'NotFound' error from other errors?
|
|
var errSessionMissing = errors.New("session not found")
|
|
|
|
var _ clientConnection = (*streamableClientConn)(nil)
|
|
|
|
func (c *streamableClientConn) sessionUpdated(state clientSessionState) {
|
|
c.mu.Lock()
|
|
c.initializedResult = state.InitializeResult
|
|
c.mu.Unlock()
|
|
|
|
// Start the standalone SSE stream as soon as we have the initialized
|
|
// result, if continuous listening is enabled.
|
|
//
|
|
// § 2.2: The client MAY issue an HTTP GET to the MCP endpoint. This can be
|
|
// used to open an SSE stream, allowing the server to communicate to the
|
|
// client, without the client first sending data via HTTP POST.
|
|
//
|
|
// We have to wait for initialized, because until we've received
|
|
// initialized, we don't know whether the server requires a sessionID.
|
|
//
|
|
// § 2.5: A server using the Streamable HTTP transport MAY assign a session
|
|
// ID at initialization time, by including it in a Mcp-Session-Id header
|
|
// on the HTTP response containing the InitializeResult.
|
|
if !c.disableStandaloneSSE {
|
|
c.connectStandaloneSSE()
|
|
}
|
|
}
|
|
|
|
func (c *streamableClientConn) connectStandaloneSSE() {
|
|
resp, err := c.connectSSE(c.ctx, "", 0, true)
|
|
if err != nil {
|
|
// If the client didn't cancel the request, and failure breaks the logical
|
|
// session.
|
|
if c.ctx.Err() == nil {
|
|
c.fail(fmt.Errorf("standalone SSE request failed (session ID: %v): %v", c.sessionID, err))
|
|
}
|
|
return
|
|
}
|
|
|
|
// [§2.2.3]: "The server MUST either return Content-Type:
|
|
// text/event-stream in response to this HTTP GET, or else return HTTP
|
|
// 405 Method Not Allowed, indicating that the server does not offer an
|
|
// SSE stream at this endpoint."
|
|
//
|
|
// [§2.2.3]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#listening-for-messages-from-the-server
|
|
if resp.StatusCode == http.StatusMethodNotAllowed {
|
|
// The server doesn't support the standalone SSE stream.
|
|
resp.Body.Close()
|
|
return
|
|
}
|
|
if resp.Header.Get("Content-Type") != "text/event-stream" {
|
|
// modelcontextprotocol/go-sdk#736: some servers return 200 OK or redirect with
|
|
// non-SSE content type instead of text/event-stream for the standalone
|
|
// SSE stream.
|
|
c.logger.Warn(fmt.Sprintf("got Content-Type %s instead of text/event-stream for standalone SSE stream", resp.Header.Get("Content-Type")))
|
|
resp.Body.Close()
|
|
return
|
|
}
|
|
if resp.StatusCode >= 400 && resp.StatusCode < 500 && !c.strict {
|
|
// modelcontextprotocol/go-sdk#393,#610: some servers return NotFound or
|
|
// other status codes instead of MethodNotAllowed for the standalone SSE
|
|
// stream.
|
|
//
|
|
// Treat this like MethodNotAllowed in non-strict mode.
|
|
c.logger.Warn(fmt.Sprintf("got %d instead of 405 for standalone SSE stream", resp.StatusCode))
|
|
resp.Body.Close()
|
|
return
|
|
}
|
|
summary := "standalone SSE stream"
|
|
if err := c.checkResponse(summary, resp); err != nil {
|
|
c.fail(err)
|
|
return
|
|
}
|
|
go c.handleSSE(c.ctx, summary, resp, nil)
|
|
}
|
|
|
|
// fail handles an asynchronous error while reading.
|
|
//
|
|
// If err is non-nil, it is terminal, and subsequent (or pending) Reads will
|
|
// fail.
|
|
//
|
|
// If err wraps errSessionMissing, the failure indicates that the session is no
|
|
// longer present on the server, and no final DELETE will be performed when
|
|
// closing the connection.
|
|
func (c *streamableClientConn) fail(err error) {
|
|
if err != nil {
|
|
c.failOnce.Do(func() {
|
|
c._failure = err
|
|
close(c.failed)
|
|
})
|
|
}
|
|
}
|
|
|
|
func (c *streamableClientConn) failure() error {
|
|
select {
|
|
case <-c.failed:
|
|
return c._failure
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (c *streamableClientConn) SessionID() string {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
return c.sessionID
|
|
}
|
|
|
|
// Read implements the [Connection] interface.
|
|
func (c *streamableClientConn) Read(ctx context.Context) (jsonrpc.Message, error) {
|
|
if err := c.failure(); err != nil {
|
|
return nil, err
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, ctx.Err()
|
|
case <-c.failed:
|
|
return nil, c.failure()
|
|
case <-c.done:
|
|
return nil, io.EOF
|
|
case msg := <-c.incoming:
|
|
return msg, nil
|
|
}
|
|
}
|
|
|
|
// Write implements the [Connection] interface.
|
|
func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) error {
|
|
if err := c.failure(); err != nil {
|
|
return err
|
|
}
|
|
|
|
var requestSummary string
|
|
var forCall *jsonrpc.Request
|
|
switch msg := msg.(type) {
|
|
case *jsonrpc.Request:
|
|
requestSummary = fmt.Sprintf("sending %q", msg.Method)
|
|
if msg.IsCall() {
|
|
forCall = msg
|
|
}
|
|
case *jsonrpc.Response:
|
|
requestSummary = fmt.Sprintf("sending jsonrpc response #%d", msg.ID)
|
|
default:
|
|
panic("unreachable")
|
|
}
|
|
|
|
data, err := jsonrpc.EncodeMessage(msg)
|
|
if err != nil {
|
|
return fmt.Errorf("%s: %v", requestSummary, err)
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(data))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "application/json, text/event-stream")
|
|
c.setMCPHeaders(req)
|
|
|
|
resp, err := c.client.Do(req)
|
|
if err != nil {
|
|
// Any error from client.Do means the request didn't reach the server.
|
|
// Wrap with ErrRejected so the jsonrpc2 connection doesn't set writeErr
|
|
// and permanently break the connection.
|
|
return fmt.Errorf("%w: %s: %v", jsonrpc2.ErrRejected, requestSummary, err)
|
|
}
|
|
|
|
if err := c.checkResponse(requestSummary, resp); err != nil {
|
|
// Only fail the connection for non-transient errors.
|
|
// Transient errors (wrapped with ErrRejected) should not break the connection.
|
|
if !errors.Is(err, jsonrpc2.ErrRejected) {
|
|
c.fail(err)
|
|
}
|
|
return err
|
|
}
|
|
|
|
if sessionID := resp.Header.Get(sessionIDHeader); sessionID != "" {
|
|
c.mu.Lock()
|
|
hadSessionID := c.sessionID
|
|
if hadSessionID == "" {
|
|
c.sessionID = sessionID
|
|
}
|
|
c.mu.Unlock()
|
|
if hadSessionID != "" && hadSessionID != sessionID {
|
|
resp.Body.Close()
|
|
return fmt.Errorf("mismatching session IDs %q and %q", hadSessionID, sessionID)
|
|
}
|
|
}
|
|
|
|
if forCall == nil {
|
|
resp.Body.Close()
|
|
|
|
// [§2.1.4]: "If the input is a JSON-RPC response or notification:
|
|
// If the server accepts the input, the server MUST return HTTP status code 202 Accepted with no body."
|
|
//
|
|
// [§2.1.4]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#listening-for-messages-from-the-server
|
|
if resp.StatusCode != http.StatusNoContent && resp.StatusCode != http.StatusAccepted {
|
|
errMsg := fmt.Sprintf("unexpected status code %d from non-call", resp.StatusCode)
|
|
// Some servers return 200, even with an empty json body.
|
|
//
|
|
// In strict mode, return an error to the caller.
|
|
c.logger.Warn(errMsg)
|
|
if c.strict {
|
|
return errors.New(errMsg)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
contentType := strings.TrimSpace(strings.SplitN(resp.Header.Get("Content-Type"), ";", 2)[0])
|
|
switch contentType {
|
|
case "application/json":
|
|
go c.handleJSON(requestSummary, resp)
|
|
|
|
case "text/event-stream":
|
|
var forCall *jsonrpc.Request
|
|
if jsonReq, ok := msg.(*jsonrpc.Request); ok && jsonReq.IsCall() {
|
|
forCall = jsonReq
|
|
}
|
|
// Handle the resulting stream. Note that ctx comes from the call, and
|
|
// therefore is already cancelled when the JSON-RPC request is cancelled
|
|
// (or rather, context cancellation is what *triggers* JSON-RPC
|
|
// cancellation)
|
|
go c.handleSSE(ctx, requestSummary, resp, forCall)
|
|
|
|
default:
|
|
resp.Body.Close()
|
|
return fmt.Errorf("%s: unsupported content type %q", requestSummary, contentType)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// testAuth controls whether a fake Authorization header is added to outgoing requests.
|
|
// TODO: replace with a better mechanism when client-side auth is in place.
|
|
var testAuth atomic.Bool
|
|
|
|
func (c *streamableClientConn) setMCPHeaders(req *http.Request) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
|
|
if c.initializedResult != nil {
|
|
req.Header.Set(protocolVersionHeader, c.initializedResult.ProtocolVersion)
|
|
}
|
|
if c.sessionID != "" {
|
|
req.Header.Set(sessionIDHeader, c.sessionID)
|
|
}
|
|
if testAuth.Load() {
|
|
req.Header.Set("Authorization", "Bearer foo")
|
|
}
|
|
}
|
|
|
|
func (c *streamableClientConn) handleJSON(requestSummary string, resp *http.Response) {
|
|
body, err := io.ReadAll(resp.Body)
|
|
resp.Body.Close()
|
|
if err != nil {
|
|
c.fail(fmt.Errorf("%s: failed to read body: %v", requestSummary, err))
|
|
return
|
|
}
|
|
msg, err := jsonrpc.DecodeMessage(body)
|
|
if err != nil {
|
|
c.fail(fmt.Errorf("%s: failed to decode response: %v", requestSummary, err))
|
|
return
|
|
}
|
|
select {
|
|
case c.incoming <- msg:
|
|
case <-c.done:
|
|
// The connection was closed by the client; exit gracefully.
|
|
}
|
|
}
|
|
|
|
// handleSSE manages the lifecycle of an SSE connection. It can be either
|
|
// persistent (for the main GET listener) or temporary (for a POST response).
|
|
//
|
|
// If forCall is set, it is the call that initiated the stream, and the
|
|
// stream is complete when we receive its response. Otherwise, this is the
|
|
// standalone stream.
|
|
func (c *streamableClientConn) handleSSE(ctx context.Context, requestSummary string, resp *http.Response, forCall *jsonrpc2.Request) {
|
|
for {
|
|
// Connection was successful. Continue the loop with the new response.
|
|
//
|
|
// TODO(#679): we should set a reasonable limit on the number of times
|
|
// we'll try getting a response for a given request, or enforce that we
|
|
// actually make progress.
|
|
//
|
|
// Eventually, if we don't get the response, we should stop trying and
|
|
// fail the request.
|
|
lastEventID, reconnectDelay, clientClosed := c.processStream(ctx, requestSummary, resp, forCall)
|
|
|
|
// If the connection was closed by the client, we're done.
|
|
if clientClosed {
|
|
return
|
|
}
|
|
// If we don't have a last event ID, we can never get the call response, so
|
|
// there's nothing to resume. For the standalone stream, we can reconnect,
|
|
// but we may just miss messages.
|
|
if lastEventID == "" && forCall != nil {
|
|
return
|
|
}
|
|
|
|
// The stream was interrupted or ended by the server. Attempt to reconnect.
|
|
newResp, err := c.connectSSE(ctx, lastEventID, reconnectDelay, false)
|
|
if err != nil {
|
|
// If the client didn't cancel this request, any failure to execute it
|
|
// breaks the logical MCP session.
|
|
if ctx.Err() == nil {
|
|
// All reconnection attempts failed: fail the connection.
|
|
c.fail(fmt.Errorf("%s: failed to reconnect (session ID: %v): %v", requestSummary, c.sessionID, err))
|
|
}
|
|
return
|
|
}
|
|
|
|
resp = newResp
|
|
if err := c.checkResponse(requestSummary, resp); err != nil {
|
|
c.fail(err)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// checkResponse checks the status code of the provided response, and
|
|
// translates it into an error if the request was unsuccessful.
|
|
//
|
|
// The response body is close if a non-nil error is returned.
|
|
func (c *streamableClientConn) checkResponse(requestSummary string, resp *http.Response) (err error) {
|
|
defer func() {
|
|
if err != nil {
|
|
resp.Body.Close()
|
|
}
|
|
}()
|
|
// §2.5.3: "The server MAY terminate the session at any time, after
|
|
// which it MUST respond to requests containing that session ID with HTTP
|
|
// 404 Not Found."
|
|
if resp.StatusCode == http.StatusNotFound {
|
|
// Return an errSessionMissing to avoid sending a redundant DELETE when the
|
|
// session is already gone.
|
|
return fmt.Errorf("%s: failed to connect (session ID: %v): %w", requestSummary, c.sessionID, errSessionMissing)
|
|
}
|
|
// Transient server errors (502, 503, 504, 429) should not break the connection.
|
|
// Wrap them with ErrRejected so the jsonrpc2 layer doesn't set writeErr.
|
|
if isTransientHTTPStatus(resp.StatusCode) {
|
|
return fmt.Errorf("%w: %s: %v", jsonrpc2.ErrRejected, requestSummary, http.StatusText(resp.StatusCode))
|
|
}
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return fmt.Errorf("%s: %v", requestSummary, http.StatusText(resp.StatusCode))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// processStream reads from a single response body, sending events to the
|
|
// incoming channel. It returns the ID of the last processed event and a flag
|
|
// indicating if the connection was closed by the client. If resp is nil, it
|
|
// returns "", false.
|
|
func (c *streamableClientConn) processStream(ctx context.Context, requestSummary string, resp *http.Response, forCall *jsonrpc.Request) (lastEventID string, reconnectDelay time.Duration, clientClosed bool) {
|
|
defer func() {
|
|
// Drain any remaining unprocessed body. This allows the connection to be re-used after closing.
|
|
io.Copy(io.Discard, resp.Body)
|
|
resp.Body.Close()
|
|
}()
|
|
for evt, err := range scanEvents(resp.Body) {
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
return "", 0, true // don't reconnect: client cancelled
|
|
}
|
|
break
|
|
}
|
|
|
|
if evt.ID != "" {
|
|
lastEventID = evt.ID
|
|
}
|
|
|
|
if evt.Retry != "" {
|
|
if n, err := strconv.ParseInt(evt.Retry, 10, 64); err == nil {
|
|
reconnectDelay = time.Duration(n) * time.Millisecond
|
|
}
|
|
}
|
|
// According to SSE spec, events with no name default to "message"
|
|
if evt.Name != "" && evt.Name != "message" {
|
|
continue
|
|
}
|
|
|
|
msg, err := jsonrpc.DecodeMessage(evt.Data)
|
|
if err != nil {
|
|
c.fail(fmt.Errorf("%s: failed to decode event: %v", requestSummary, err))
|
|
return "", 0, true
|
|
}
|
|
|
|
select {
|
|
case c.incoming <- msg:
|
|
// Check if this is the response to our call, which terminates the request.
|
|
// (it could also be a server->client request or notification).
|
|
if jsonResp, ok := msg.(*jsonrpc.Response); ok && forCall != nil {
|
|
// TODO: we should never get a response when forReq is nil (the standalone SSE request).
|
|
// We should detect this case.
|
|
if jsonResp.ID == forCall.ID {
|
|
return "", 0, true
|
|
}
|
|
}
|
|
|
|
case <-c.done:
|
|
// The connection was closed by the client; exit gracefully.
|
|
return "", 0, true
|
|
}
|
|
}
|
|
// The loop finished without an error, indicating the server closed the stream.
|
|
//
|
|
// If the lastEventID is "", the stream is not retryable and we should
|
|
// report a synthetic error for the call.
|
|
//
|
|
// Note that this is different from the cancellation case above, since the
|
|
// caller is still waiting for a response that will never come.
|
|
if lastEventID == "" && forCall != nil {
|
|
errmsg := &jsonrpc2.Response{
|
|
ID: forCall.ID,
|
|
Error: fmt.Errorf("request terminated without response"),
|
|
}
|
|
select {
|
|
case c.incoming <- errmsg:
|
|
case <-c.done:
|
|
}
|
|
}
|
|
return lastEventID, reconnectDelay, false
|
|
}
|
|
|
|
// connectSSE handles the logic of connecting a text/event-stream connection.
|
|
//
|
|
// If lastEventID is set, it is the last-event ID of a stream being resumed.
|
|
//
|
|
// If connection fails, connectSSE retries with an exponential backoff
|
|
// strategy. It returns a new, valid HTTP response if successful, or an error
|
|
// if all retries are exhausted.
|
|
//
|
|
// reconnectDelay is the delay set by the server using the SSE retry field, or
|
|
// 0.
|
|
//
|
|
// If initial is set, this is the initial attempt.
|
|
//
|
|
// If connectSSE exits due to context cancellation, the result is (nil, ctx.Err()).
|
|
func (c *streamableClientConn) connectSSE(ctx context.Context, lastEventID string, reconnectDelay time.Duration, initial bool) (*http.Response, error) {
|
|
var finalErr error
|
|
attempt := 0
|
|
if !initial {
|
|
// We've already connected successfully once, so delay subsequent
|
|
// reconnections. Otherwise, if the server returns 200 but terminates the
|
|
// connection, we'll reconnect as fast as we can, ad infinitum.
|
|
//
|
|
// TODO: we should consider also setting a limit on total attempts for one
|
|
// logical request.
|
|
attempt = 1
|
|
}
|
|
delay := calculateReconnectDelay(attempt)
|
|
if reconnectDelay > 0 {
|
|
delay = reconnectDelay // honor the server's requested initial delay
|
|
}
|
|
for ; attempt <= c.maxRetries; attempt++ {
|
|
select {
|
|
case <-c.done:
|
|
return nil, fmt.Errorf("connection closed by client during reconnect")
|
|
|
|
case <-ctx.Done():
|
|
// If the connection context is canceled, the request below will not
|
|
// succeed anyway.
|
|
return nil, ctx.Err()
|
|
|
|
case <-time.After(delay):
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.url, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
c.setMCPHeaders(req)
|
|
if lastEventID != "" {
|
|
req.Header.Set(lastEventIDHeader, lastEventID)
|
|
}
|
|
req.Header.Set("Accept", "text/event-stream")
|
|
resp, err := c.client.Do(req)
|
|
if err != nil {
|
|
finalErr = err // Store the error and try again.
|
|
delay = calculateReconnectDelay(attempt + 1)
|
|
continue
|
|
}
|
|
return resp, nil
|
|
}
|
|
}
|
|
// If the loop completes, all retries have failed, or the client is closing.
|
|
if finalErr != nil {
|
|
return nil, fmt.Errorf("connection failed after %d attempts: %w", c.maxRetries, finalErr)
|
|
}
|
|
return nil, fmt.Errorf("connection aborted after %d attempts", c.maxRetries)
|
|
}
|
|
|
|
// Close implements the [Connection] interface.
|
|
func (c *streamableClientConn) Close() error {
|
|
c.closeOnce.Do(func() {
|
|
if errors.Is(c.failure(), errSessionMissing) {
|
|
// If the session is missing, no need to delete it.
|
|
} else {
|
|
req, err := http.NewRequestWithContext(c.ctx, http.MethodDelete, c.url, nil)
|
|
if err != nil {
|
|
c.closeErr = err
|
|
} else {
|
|
c.setMCPHeaders(req)
|
|
if _, err := c.client.Do(req); err != nil {
|
|
c.closeErr = err
|
|
}
|
|
}
|
|
}
|
|
|
|
// Cancel any hanging network requests after cleanup.
|
|
c.cancel()
|
|
close(c.done)
|
|
})
|
|
return c.closeErr
|
|
}
|
|
|
|
// calculateReconnectDelay calculates a delay using exponential backoff with full jitter.
|
|
func calculateReconnectDelay(attempt int) time.Duration {
|
|
if attempt == 0 {
|
|
return 0
|
|
}
|
|
// Calculate the exponential backoff using the grow factor.
|
|
backoffDuration := time.Duration(float64(reconnectInitialDelay) * math.Pow(reconnectGrowFactor, float64(attempt-1)))
|
|
// Cap the backoffDuration at maxDelay.
|
|
backoffDuration = min(backoffDuration, reconnectMaxDelay)
|
|
|
|
// Use a full jitter using backoffDuration
|
|
jitter := rand.N(backoffDuration)
|
|
|
|
return backoffDuration + jitter
|
|
}
|
|
|
|
// isTransientHTTPStatus reports whether the HTTP status code indicates a
|
|
// transient server error that should not permanently break the connection.
|
|
func isTransientHTTPStatus(statusCode int) bool {
|
|
switch statusCode {
|
|
case http.StatusInternalServerError, // 500
|
|
http.StatusBadGateway, // 502
|
|
http.StatusServiceUnavailable, // 503
|
|
http.StatusGatewayTimeout, // 504
|
|
http.StatusTooManyRequests: // 429
|
|
return true
|
|
}
|
|
return false
|
|
}
|