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

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

Author: findleyr

internal/jsonrpc2/conn.go

@@ -725,7 +725,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

internal/jsonrpc2/wire.go

@@ -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"

mcp/streamable.go

@@ -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
 
@@ -756,6 +825,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

mcp/streamable_test.go

@@ -706,7 +706,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)
@@ -953,19 +953,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",
@@ -985,7 +984,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,
+		},
+	}
+
+	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)
+		}
 	}
 
-	testStreamableHandler(t, handler, requests)
+	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
 }

mcp/transport.go

@@ -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.