mcp: support stateless streamable sessions (#277)

Support stateless streamable sessions by adding a GetSessionID function
to StreamableHTTPOptions. If GetSessionID returns "", the session is
stateless, and no validation is performed. This is implemented by
providing the session a trivial initialization state.

To implement this, some parts of #232 (distributed sessions) are copied
over, since they add an API for creating an already-initialized session.

In total, the following new API is added:
- StreamableHTTPOptions.GetSessionID
- ServerSessionOptions (a new parameter to Server.Connect)
- ServerSessionState
- ClientSessionOptions (a new parameter to Client.Connect, for symmetry)

For #10

mcp/server.go

@@ -301,7 +301,7 @@ func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPr
 	})
 }
 
-func (s *Server) getPrompt(ctx context.Context, cc *ServerSession, params *GetPromptParams) (*GetPromptResult, error) {
+func (s *Server) getPrompt(ctx context.Context, ss *ServerSession, params *GetPromptParams) (*GetPromptResult, error) {
 	s.mu.Lock()
 	prompt, ok := s.prompts.get(params.Name)
 	s.mu.Unlock()
@@ -309,7 +309,7 @@ func (s *Server) getPrompt(ctx context.Context, cc *ServerSession, params *GetPr
 		// TODO: surface the error code over the wire, instead of flattening it into the string.
 		return nil, fmt.Errorf("%s: unknown prompt %q", jsonrpc2.ErrInvalidParams, params.Name)
 	}
-	return prompt.handler(ctx, cc, params)
+	return prompt.handler(ctx, ss, params)
 }
 
 func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListToolsParams) (*ListToolsResult, error) {
@@ -518,7 +518,7 @@ func (s *Server) unsubscribe(ctx context.Context, ss *ServerSession, params *Uns
 // It need not be called on servers that are used for multiple concurrent connections,
 // as with [StreamableHTTPHandler].
 func (s *Server) Run(ctx context.Context, t Transport) error {
-	ss, err := s.Connect(ctx, t)
+	ss, err := s.Connect(ctx, t, nil)
 	if err != nil {
 		return err
 	}
@@ -539,8 +539,12 @@ func (s *Server) Run(ctx context.Context, t Transport) error {
 
 // bind implements the binder[*ServerSession] interface, so that Servers can
 // be connected using [connect].
-func (s *Server) bind(conn *jsonrpc2.Connection) *ServerSession {
-	ss := &ServerSession{conn: conn, server: s}
+func (s *Server) bind(mcpConn Connection, conn *jsonrpc2.Connection, state *ServerSessionState) *ServerSession {
+	assert(mcpConn != nil && conn != nil, "nil connection")
+	ss := &ServerSession{conn: conn, mcpConn: mcpConn, server: s}
+	if state != nil {
+		ss.state = *state
+	}
 	s.mu.Lock()
 	s.sessions = append(s.sessions, ss)
 	s.mu.Unlock()
@@ -561,32 +565,50 @@ func (s *Server) disconnect(cc *ServerSession) {
 	}
 }
 
+// ServerSessionOptions configures the server session.
+type ServerSessionOptions struct {
+	State *ServerSessionState
+}
+
 // Connect connects the MCP server over the given transport and starts handling
 // messages.
 //
 // It returns a connection object that may be used to terminate the connection
 // (with [Connection.Close]), or await client termination (with
 // [Connection.Wait]).
-func (s *Server) Connect(ctx context.Context, t Transport) (*ServerSession, error) {
-	return connect(ctx, t, s)
+//
+// If opts.State is non-nil, it is the initial state for the server.
+func (s *Server) Connect(ctx context.Context, t Transport, opts *ServerSessionOptions) (*ServerSession, error) {
+	var state *ServerSessionState
+	if opts != nil {
+		state = opts.State
+	}
+	return connect(ctx, t, s, state)
 }
 
+// TODO: (nit) move all ServerSession methods below the ServerSession declaration.
 func (ss *ServerSession) initialized(ctx context.Context, params *InitializedParams) (Result, error) {
+	if params == nil {
+		// Since we use nilness to signal 'initialized' state, we must ensure that
+		// params are non-nil.
+		params = new(InitializedParams)
+	}
 	if ss.server.opts.KeepAlive > 0 {
 		ss.startKeepalive(ss.server.opts.KeepAlive)
 	}
-	ss.mu.Lock()
-	hasParams := ss.initializeParams != nil
-	wasInitialized := ss._initialized
-	if hasParams {
-		ss._initialized = true
-	}
-	ss.mu.Unlock()
+	var wasInit, wasInitd bool
+	ss.updateState(func(state *ServerSessionState) {
+		wasInit = state.InitializeParams != nil
+		wasInitd = state.InitializedParams != nil
+		if wasInit && !wasInitd {
+			state.InitializedParams = params
+		}
+	})
 
-	if !hasParams {
+	if !wasInit {
 		return nil, fmt.Errorf("%q before %q", notificationInitialized, methodInitialize)
 	}
-	if wasInitialized {
+	if wasInitd {
 		return nil, fmt.Errorf("duplicate %q received", notificationInitialized)
 	}
 	return callNotificationHandler(ctx, ss.server.opts.InitializedHandler, ss, params)
@@ -615,25 +637,30 @@ func (ss *ServerSession) NotifyProgress(ctx context.Context, params *ProgressNot
 // Call [ServerSession.Close] to close the connection, or await client
 // termination with [ServerSession.Wait].
 type ServerSession struct {
-	server           *Server
-	conn             *jsonrpc2.Connection
-	mcpConn          Connection
-	mu               sync.Mutex
-	logLevel         LoggingLevel
-	initializeParams *InitializeParams
-	_initialized     bool
-	keepaliveCancel  context.CancelFunc
+	server          *Server
+	conn            *jsonrpc2.Connection
+	mcpConn         Connection
+	keepaliveCancel context.CancelFunc // TODO: theory around why keepaliveCancel need not be guarded
+
+	mu    sync.Mutex
+	state ServerSessionState
 }
 
-func (ss *ServerSession) setConn(c Connection) {
-	ss.mcpConn = c
+func (ss *ServerSession) updateState(mut func(*ServerSessionState)) {
+	ss.mu.Lock()
+	mut(&ss.state)
+	copy := ss.state
+	ss.mu.Unlock()
+	if c, ok := ss.mcpConn.(serverConnection); ok {
+		c.sessionUpdated(copy)
+	}
 }
 
 func (ss *ServerSession) ID() string {
-	if ss.mcpConn == nil {
-		return ""
+	if c, ok := ss.mcpConn.(hasSessionID); ok {
+		return c.SessionID()
 	}
-	return ss.mcpConn.SessionID()
+	return ""
 }
 
 // Ping pings the client.
@@ -657,7 +684,7 @@ func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessag
 // is below that of the last SetLevel.
 func (ss *ServerSession) Log(ctx context.Context, params *LoggingMessageParams) error {
 	ss.mu.Lock()
-	logLevel := ss.logLevel
+	logLevel := ss.state.LogLevel
 	ss.mu.Unlock()
 	if logLevel == "" {
 		// The spec is unclear, but seems to imply that no log messages are sent until the client
@@ -747,7 +774,7 @@ func (ss *ServerSession) getConn() *jsonrpc2.Connection { return ss.conn }
 // handle invokes the method described by the given JSON RPC request.
 func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any, error) {
 	ss.mu.Lock()
-	initialized := ss._initialized
+	initialized := ss.state.InitializedParams != nil
 	ss.mu.Unlock()
 	// From the spec:
 	// "The client SHOULD NOT send requests other than pings before the server
@@ -770,9 +797,9 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam
 	if params == nil {
 		return nil, fmt.Errorf("%w: \"params\" must be be provided", jsonrpc2.ErrInvalidParams)
 	}
-	ss.mu.Lock()
-	ss.initializeParams = params
-	ss.mu.Unlock()
+	ss.updateState(func(state *ServerSessionState) {
+		state.InitializeParams = params
+	})
 
 	// If we support the client's version, reply with it. Otherwise, reply with our
 	// latest version.
@@ -796,9 +823,9 @@ func (ss *ServerSession) ping(context.Context, *PingParams) (*emptyResult, error
 }
 
 func (ss *ServerSession) setLevel(_ context.Context, params *SetLevelParams) (*emptyResult, error) {
-	ss.mu.Lock()
-	defer ss.mu.Unlock()
-	ss.logLevel = params.Level
+	ss.updateState(func(state *ServerSessionState) {
+		state.LogLevel = params.Level
+	})
 	return &emptyResult{}, nil
 }