modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

mcp: improvements for 'stateless' streamable servers; 'distributed' mode

Browse source 
Several improvements for the stateless streamable mode, plus support for
a 'distributed' (or rather, distributable) version of the stateless
server.

- Add a 'Stateless' option to StreamableHTTPOptions and
  StreamableServerTransport, which controls stateless behavior.
  GetSessionID may still return a non-empty session ID.
- Audit validation of stateless mode to allow requests with a session
  id. Propagate this session ID to the temporary connection.
- Peek at requests to allow 'initialize' requests to go through to the
  session, so that version negotiation can occur (FIXME: add tests).

Fixes #284
For #148
This commit is contained in:
Rob Findley 2025-08-15 03:13:53 +00:00 • committed by Robert Findley
commit 98e67fdf62
5 changed files with 203 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

18
internal/jsonrpc2/conn.go
View file

@ -739,7 +739,23 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e
// write is used by all things that write outgoing messages, including replies.
// it makes sure that writes are atomic
func (c *Connection) write(ctx context.Context, msg Message) error {
err := c.writer.Write(ctx, msg)
var err error
// Fail writes immediately if the connection is shutting down.
//
// TODO(rfindley): should we allow cancellation notifations through? It could
// be the case that writes can still succeed.
c.updateInFlight(func(s *inFlightState) {
err = s.shuttingDown(ErrServerClosing)
})
if err == nil {
err = c.writer.Write(ctx, msg)
}
// For rejected requests, we don't set the writeErr (which would break the
// connection). They can just be returned to the caller.
if errors.Is(err, ErrRejected) {
return err
}
if err != nil && ctx.Err() == nil {
// The call to Write failed, and since ctx.Err() is nil we can't attribute

11
internal/jsonrpc2/wire.go
View file

@ -37,6 +37,17 @@ var (
ErrServerClosing = NewError(-32004, "server is closing")
// ErrClientClosing is a dummy error returned for calls initiated while the client is closing.
ErrClientClosing = NewError(-32003, "client is closing")
// The following errors have special semantics for MCP transports
// ErrRejected may be wrapped to return errors from calls to Writer.Write
// that signal that the request was rejected by the transport layer as
// invalid.
//
// Such failures do not indicate that the connection is broken, but rather
// should be returned to the caller to indicate that the specific request is
// invalid in the current context.
ErrRejected = NewError(-32004, "rejected by transport")
)
const wireVersion = "2.0"

125
mcp/streamable.go
View file

@ -38,18 +38,29 @@ type StreamableHTTPHandler struct {
getServer func(*http.Request) *Server
opts StreamableHTTPOptions
mu sync.Mutex
mu sync.Mutex
// TODO: we should store the ServerSession along with the transport, because
// we need to cancel keepalive requests when closing the transport.
transports map[string]*StreamableServerTransport // keyed by IDs (from Mcp-Session-Id header)
}
// StreamableHTTPOptions configures the StreamableHTTPHandler.
type StreamableHTTPOptions struct {
// GetSessionID provides the next session ID to use for an incoming request.
// If nil, a default randomly generated ID will be used.
//
// If GetSessionID returns an empty string, the session is 'stateless',
// meaning it is not persisted and no session validation is performed.
// As a special case, if GetSessionID returns the empty string, the
// Mcp-Session-Id header will not be set.
GetSessionID func() string
// 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.
Stateless bool
// TODO: support session retention (?)
// jsonResponse is forwarded to StreamableServerTransport.jsonResponse.
@ -118,36 +129,39 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
return
}
sessionID := req.Header.Get(sessionIDHeader)
var transport *StreamableServerTransport
if id := req.Header.Get(sessionIDHeader); id != "" {
if sessionID != "" {
h.mu.Lock()
transport = h.transports[id]
transport, _ = h.transports[sessionID]
h.mu.Unlock()
if transport == nil {
if transport == nil && !h.opts.Stateless {
// In stateless mode we allow a missing transport.
//
// A synthetic transport will be created below for the transient session.
http.Error(w, "session not found", http.StatusNotFound)
return
}
}
// TODO(rfindley): simplify the locking so that each request has only one
// critical section.
if req.Method == http.MethodDelete {
if transport == nil {
// => Mcp-Session-Id was not set; else we'd have returned NotFound above.
if sessionID == "" {
http.Error(w, "DELETE requires an Mcp-Session-Id header", http.StatusBadRequest)
return
}
h.mu.Lock()
delete(h.transports, transport.SessionID)
h.mu.Unlock()
transport.connection.Close()
if transport != nil { // transport may be nil in stateless mode
h.mu.Lock()
delete(h.transports, transport.SessionID)
h.mu.Unlock()
transport.connection.Close()
}
w.WriteHeader(http.StatusNoContent)
return
}
switch req.Method {
case http.MethodPost, http.MethodGet:
if req.Method == http.MethodGet && transport == nil {
if req.Method == http.MethodGet && sessionID == "" {
http.Error(w, "GET requires an active session", http.StatusMethodNotAllowed)
return
}
@ -164,37 +178,83 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
http.Error(w, "no server available", http.StatusBadRequest)
return
}
sessionID := h.opts.GetSessionID()
s := &StreamableServerTransport{SessionID: sessionID, jsonResponse: h.opts.jsonResponse}
if sessionID == "" {
// In stateless mode, sessionID may be nonempty even if there's no
// existing transport.
sessionID = h.opts.GetSessionID()
}
transport = &StreamableServerTransport{
SessionID: sessionID,
Stateless: h.opts.Stateless,
jsonResponse: h.opts.jsonResponse,
}
// To support stateless mode, we initialize the session with a default
// state, so that it doesn't reject subsequent requests.
var connectOpts *ServerSessionOptions
if sessionID == "" {
if h.opts.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.StatusBadRequest)
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 = new(InitializeParams)
}
if !hasInitialized {
state.InitializedParams = new(InitializedParams)
}
connectOpts = &ServerSessionOptions{
State: &ServerSessionState{
InitializeParams: new(InitializeParams),
InitializedParams: new(InitializedParams),
},
State: state,
}
}
// 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.
ss, err := server.Connect(req.Context(), s, connectOpts)
ss, err := server.Connect(req.Context(), transport, connectOpts)
if err != nil {
http.Error(w, "failed connection", http.StatusInternalServerError)
return
}
if sessionID == "" {
if h.opts.Stateless {
// Stateless mode: close the session when the request exits.
defer ss.Close() // close the fake session after handling the request
} else {
// Otherwise, save the transport so that it can be reused
h.mu.Lock()
h.transports[s.SessionID] = s
h.transports[transport.SessionID] = transport
h.mu.Unlock()
}
transport = s
}
transport.ServeHTTP(w, req)
@ -225,6 +285,13 @@ type StreamableServerTransport struct {
// generator to produce one, as with [crypto/rand.Text].)
SessionID string
// Stateless controls whether the eventstore is 'Stateless'. Servers sessions
// connected to a stateless transport are disallowed from making outgoing
// requests.
//
// See also [StreamableHTTPOptions.Stateless].
Stateless bool
// Storage for events, to enable stream resumption.
// If nil, a [MemoryEventStore] with the default maximum size will be used.
EventStore EventStore
@ -265,6 +332,7 @@ func (t *StreamableServerTransport) Connect(context.Context) (Connection, error)
}
t.connection = &streamableServerConn{
sessionID: t.SessionID,
stateless: t.Stateless,
eventStore: t.EventStore,
jsonResponse: t.jsonResponse,
incoming: make(chan jsonrpc.Message, 10),
@ -285,6 +353,7 @@ func (t *StreamableServerTransport) Connect(context.Context) (Connection, error)
type streamableServerConn struct {
sessionID string
stateless bool
jsonResponse bool
eventStore EventStore
@ -755,6 +824,10 @@ func (c *streamableServerConn) Read(ctx context.Context) (jsonrpc.Message, error
// Write implements the [Connection] interface.
func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) error {
if req, ok := msg.(*jsonrpc.Request); ok && req.ID.IsValid() && (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 forRequest jsonrpc.ID
isResponse := false

82
mcp/streamable_test.go
View file

@ -748,7 +748,7 @@ func testStreamableHandler(t *testing.T, handler http.Handler, requests []stream
if !request.ignoreResponse {
transform := cmpopts.AcyclicTransformer("jsonrpcid", func(id jsonrpc.ID) any { return id.Raw() })
if diff := cmp.Diff(request.wantMessages, got, transform); diff != "" {
t.Errorf("received unexpected messages (-want +got):\n%s", diff)
t.Errorf("request #%d: received unexpected messages (-want +got):\n%s", i, diff)
}
}
sessionID.CompareAndSwap("", gotSessionID)
@ -996,19 +996,18 @@ func TestEventID(t *testing.T) {
}
func TestStreamableStateless(t *testing.T) {
// This version of sayHi doesn't make a ping request (we can't respond to
// This version of sayHi expects
// that request from our client).
sayHi := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[hiParams]]) (*CallToolResult, error) {
if err := req.Session.Ping(ctx, nil); err == nil {
// ping should fail, but not break the connection
t.Errorf("ping succeeded unexpectedly")
}
return &CallToolResult{Content: []Content{&TextContent{Text: "hi " + req.Params.Arguments.Name}}}, nil
}
server := NewServer(testImpl, nil)
AddTool(server, &Tool{Name: "greet", Description: "say hi"}, sayHi)
// Test stateless mode.
handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
GetSessionID: func() string { return "" },
})
requests := []streamableRequest{
{
method: "POST",
@ -1028,7 +1027,74 @@ func TestStreamableStateless(t *testing.T) {
},
wantSessionID: false,
},
{
method: "POST",
wantStatusCode: http.StatusOK,
messages: []jsonrpc.Message{
req(2, "tools/call", &CallToolParams{Name: "greet", Arguments: hiParams{Name: "foo"}}),
},
wantMessages: []jsonrpc.Message{
resp(2, &CallToolResult{Content: []Content{&TextContent{Text: "hi foo"}}}, nil),
},
wantSessionID: false,
},
}
testStreamableHandler(t, handler, requests)
testClientCompatibility := func(t *testing.T, handler http.Handler) {
ctx := context.Background()
httpServer := httptest.NewServer(handler)
defer httpServer.Close()
cs, err := NewClient(testImpl, nil).Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil)
if err != nil {
t.Fatal(err)
}
res, err := cs.CallTool(ctx, &CallToolParams{Name: "greet", Arguments: hiParams{Name: "bar"}})
if err != nil {
t.Fatal(err)
}
if got, want := textContent(t, res), "hi bar"; got != want {
t.Errorf("Result = %q, want %q", got, want)
}
}
handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
GetSessionID: func() string { return "" },
Stateless: true,
})
// Test the default stateless mode.
t.Run("stateless", func(t *testing.T) {
testStreamableHandler(t, handler, requests)
testClientCompatibility(t, handler)
})
// Test a "distributed" variant of stateless mode, where it has non-empty
// session IDs, but is otherwise stateless.
//
// This can be used by tools to look up application state preserved across
// subsequent requests.
for i, req := range requests {
// Now, we want a session for all requests.
req.wantSessionID = true
requests[i] = req
}
distributableHandler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
Stateless: true,
})
t.Run("distributed", func(t *testing.T) {
testStreamableHandler(t, distributableHandler, requests)
testClientCompatibility(t, handler)
})
}
func textContent(t *testing.T, res *CallToolResult) string {
t.Helper()
if len(res.Content) != 1 {
t.Fatalf("len(Content) = %d, want 1", len(res.Content))
}
text, ok := res.Content[0].(*TextContent)
if !ok {
t.Fatalf("Content[0] is %T, want *TextContent", res.Content[0])
}
return text.Text
}

4
mcp/transport.go
View file

@ -40,8 +40,8 @@ type Transport interface {
type Connection interface {
// Read reads the next message to process off the connection.
//
// Read need not be safe for concurrent use: Read is called in a
// concurrency-safe manner by the JSON-RPC library.
// Connections must allow Read to be called concurrently with Close. In
// particular, calling Close should unblock a Read waiting for input.
Read(context.Context) (jsonrpc.Message, error)
// Write writes a new message to the connection.