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

WIP: needs self review and more tests.

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

@@ -46,10 +46,14 @@ type StreamableHTTPHandler struct {
 type StreamableHTTPOptions struct {
 	// GetSessionID provides the next session ID to use for an incoming request.
 	//
-	// If GetSessionID returns an empty string, the session is 'stateless',
-	// meaning it is not persisted and no session validation is performed.
+	// FIXME: update doc.
 	GetSessionID func() string
 
+	// Stateless controls whether the session is 'stateless'.
+	//
+	// FIXME: update doc.
+	Stateless bool
+
 	// TODO: support session retention (?)
 
 	// jsonResponse is forwarded to StreamableServerTransport.jsonResponse.
@@ -119,11 +123,12 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 	}
 
 	var transport *StreamableServerTransport
-	if id := req.Header.Get(sessionIDHeader); id != "" {
+	sessionID := req.Header.Get(sessionIDHeader)
+	if id := sessionID; id != "" {
 		h.mu.Lock()
 		transport, _ = h.transports[id]
 		h.mu.Unlock()
-		if transport == nil {
+		if transport == nil && !h.opts.Stateless {
 			http.Error(w, "session not found", http.StatusNotFound)
 			return
 		}
@@ -132,22 +137,24 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 	// TODO(rfindley): simplify the locking so that each request has only one
 	// critical section.
 	if req.Method == http.MethodDelete {
-		if transport == nil {
+		if sessionID == "" {
 			// => Mcp-Session-Id was not set; else we'd have returned NotFound above.
 			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 +171,76 @@ 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 an initialize request.
+			// We want that 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
+			}
+			// Reset the body to 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
+						}
+					}
+				}
+			}
+			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 {
 			h.mu.Lock()
-			h.transports[s.SessionID] = s
+			h.transports[transport.SessionID] = transport
 			h.mu.Unlock()
 		}
-		transport = s
 	}
 
 	transport.ServeHTTP(w, req)
@@ -225,6 +271,9 @@ type StreamableServerTransport struct {
 	// generator to produce one, as with [crypto/rand.Text].)
 	SessionID string
 
+	// FIXME: doc
+	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 +314,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 +335,7 @@ func (t *StreamableServerTransport) Connect(context.Context) (Connection, error)
 
 type streamableServerConn struct {
 	sessionID    string
+	stateless    bool
 	jsonResponse bool
 	eventStore   EventStore
 
@@ -759,6 +810,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.
+		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)
@@ -955,19 +955,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",
@@ -987,7 +986,43 @@ 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)
+	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)
+	})
+
+	// 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)
+	})
 }