add Connect arg (#232)

Author: jba

mcp/example_middleware_test.go

@@ -114,7 +114,7 @@ func Example_loggingMiddleware() {
 	ctx := context.Background()
 
 	// Connect server and client
-	serverSession, _ := server.Connect(ctx, serverTransport)
+	serverSession, _ := server.Connect(ctx, serverTransport, nil)
 	defer serverSession.Close()
 
 	clientSession, _ := client.Connect(ctx, clientTransport)

mcp/logging.go

@@ -117,7 +117,7 @@ func (h *LoggingHandler) Enabled(ctx context.Context, level slog.Level) bool {
 	// This is also checked in ServerSession.LoggingMessage, so checking it here
 	// is just an optimization that skips building the JSON.
 	h.ss.mu.Lock()
-	mcpLevel := h.ss.state.LogLevel
+	mcpLevel := h.ss.opts.SessionState.LogLevel
 	h.ss.mu.Unlock()
 	return level >= mcpLevelToSlog(mcpLevel)
 }

mcp/mcp_test.go

@@ -104,7 +104,7 @@ func TestEndToEnd(t *testing.T) {
 	s.AddResource(resource2, readHandler)
 
 	// Connect the server.
-	ss, err := s.Connect(ctx, st)
+	ss, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -549,7 +549,7 @@ func basicConnection(t *testing.T, config func(*Server)) (*ServerSession, *Clien
 	if config != nil {
 		config(s)
 	}
-	ss, err := s.Connect(ctx, st)
+	ss, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -598,7 +598,7 @@ func TestBatching(t *testing.T) {
 	ct, st := NewInMemoryTransports()
 
 	s := NewServer(testImpl, nil)
-	_, err := s.Connect(ctx, st)
+	_, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -668,7 +668,7 @@ func TestMiddleware(t *testing.T) {
 	ct, st := NewInMemoryTransports()
 
 	s := NewServer(testImpl, nil)
-	ss, err := s.Connect(ctx, st)
+	ss, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -777,7 +777,7 @@ func TestNoJSONNull(t *testing.T) {
 	ct = NewLoggingTransport(ct, &logbuf)
 
 	s := NewServer(testImpl, nil)
-	ss, err := s.Connect(ctx, st)
+	ss, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -845,7 +845,7 @@ func TestKeepAlive(t *testing.T) {
 	s := NewServer(testImpl, serverOpts)
 	AddTool(s, greetTool(), sayHi)
 
-	ss, err := s.Connect(ctx, st)
+	ss, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -889,7 +889,7 @@ func TestKeepAliveFailure(t *testing.T) {
 	// Server without keepalive (to test one-sided keepalive)
 	s := NewServer(testImpl, nil)
 	AddTool(s, greetTool(), sayHi)
-	ss, err := s.Connect(ctx, st)
+	ss, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		t.Fatal(err)
 	}

mcp/server.go

@@ -10,6 +10,7 @@ import (
 	"encoding/base64"
 	"encoding/gob"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"iter"
 	"maps"
@@ -513,7 +514,8 @@ func (s *Server) unsubscribe(ctx context.Context, ss *ServerSession, params *Uns
 // If no tools have been added, the server will not have the tool capability.
 // The same goes for other features like prompts and resources.
 func (s *Server) Run(ctx context.Context, t Transport) error {
-	ss, err := s.Connect(ctx, t)
+	// TODO: provide a way to pass ServerSessionOptions?
+	ss, err := s.Connect(ctx, t, nil)
 	if err != nil {
 		return err
 	}
@@ -535,7 +537,7 @@ 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, state: &SessionState{}}
+	ss := &ServerSession{conn: conn, server: s}
 	s.mu.Lock()
 	s.sessions = append(s.sessions, ss)
 	s.mu.Unlock()
@@ -556,14 +558,33 @@ func (s *Server) disconnect(cc *ServerSession) {
 	}
 }
 
+type ServerSessionOptions struct {
+	SessionID    string
+	SessionState *SessionState
+	SessionStore SessionStore
+}
+
 // 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)
+func (s *Server) Connect(ctx context.Context, t Transport, opts *ServerSessionOptions) (*ServerSession, error) {
+	if opts != nil && opts.SessionState == nil && opts.SessionStore != nil {
+		return nil, errors.New("ServerSessionOptions has store but no state")
+	}
+	ss, err := connect(ctx, t, s)
+	if err != nil {
+		return nil, err
+	}
+	if opts != nil {
+		ss.opts = *opts
+	}
+	if ss.opts.SessionState == nil {
+		ss.opts.SessionState = &SessionState{}
+	}
+	return ss, nil
 }
 
 func (s *Server) callInitializedHandler(ctx context.Context, ss *ServerSession, params *InitializedParams) (Result, error) {
@@ -596,32 +617,20 @@ 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
+	server *Server
+	conn   *jsonrpc2.Connection
+	opts   ServerSessionOptions
+
 	mu              sync.Mutex
 	initialized     bool
 	keepaliveCancel context.CancelFunc
-
-	sessionID string
-	state     *SessionState
-	store     SessionStore
 }
 
 func (ss *ServerSession) setConn(c Connection) {
 }
 
-func (ss *ServerSession) ID() string { return ss.sessionID }
-
-// InitSession initializes the session with a session ID, state, and store.
-// If called, it must be called immediately after the session is connected.
-// If never called, the session will begin with a zero SessionState and no session
-// ID or store.
-func (ss *ServerSession) InitSession(sessionID string, state *SessionState, store SessionStore) {
-	ss.mu.Lock()
-	defer ss.mu.Unlock()
-	ss.sessionID = sessionID
-	ss.state = state
-	ss.store = store
+func (ss *ServerSession) ID() string {
+	return ss.opts.SessionID
 }
 
 // Ping pings the client.
@@ -645,7 +654,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.state.LogLevel
+	logLevel := ss.opts.SessionState.LogLevel
 	ss.mu.Unlock()
 	if logLevel == "" {
 		// The spec is unclear, but seems to imply that no log messages are sent until the client
@@ -759,10 +768,10 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam
 		return nil, fmt.Errorf("%w: \"params\" must be be provided", jsonrpc2.ErrInvalidParams)
 	}
 	ss.mu.Lock()
-	ss.state.InitializeParams = params
+	ss.opts.SessionState.InitializeParams = params
 	ss.mu.Unlock()
-	if ss.store != nil {
-		if err := ss.store.Store(ctx, ss.sessionID, ss.state); err != nil {
+	if store := ss.opts.SessionStore; store != nil {
+		if err := store.Store(ctx, ss.opts.SessionID, ss.opts.SessionState); err != nil {
 			return nil, fmt.Errorf("storing session state: %w", err)
 		}
 	}
@@ -802,9 +811,9 @@ func (ss *ServerSession) ping(context.Context, *PingParams) (*emptyResult, error
 func (ss *ServerSession) setLevel(ctx context.Context, params *SetLevelParams) (*emptyResult, error) {
 	ss.mu.Lock()
 	defer ss.mu.Unlock()
-	ss.state.LogLevel = params.Level
-	if ss.store != nil {
-		if err := ss.store.Store(ctx, ss.sessionID, ss.state); err != nil {
+	ss.opts.SessionState.LogLevel = params.Level
+	if store := ss.opts.SessionStore; store != nil {
+		if err := store.Store(ctx, ss.opts.SessionID, ss.opts.SessionState); err != nil {
 			return nil, err
 		}
 	}

mcp/server_example_test.go

@@ -31,7 +31,7 @@ func ExampleServer() {
 	server := mcp.NewServer(&mcp.Implementation{Name: "greeter", Version: "v0.0.1"}, nil)
 	mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi)
 
-	serverSession, err := server.Connect(ctx, serverTransport)
+	serverSession, err := server.Connect(ctx, serverTransport, nil)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -62,7 +62,7 @@ func createSessions(ctx context.Context) (*mcp.ClientSession, *mcp.ServerSession
 	server := mcp.NewServer(testImpl, nil)
 	client := mcp.NewClient(testImpl, nil)
 	serverTransport, clientTransport := mcp.NewInMemoryTransports()
-	serverSession, err := server.Connect(ctx, serverTransport)
+	serverSession, err := server.Connect(ctx, serverTransport, nil)
 	if err != nil {
 		log.Fatal(err)
 	}

mcp/sse.go

@@ -221,7 +221,7 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 		http.Error(w, "no server available", http.StatusBadRequest)
 		return
 	}
-	ss, err := server.Connect(req.Context(), transport)
+	ss, err := server.Connect(req.Context(), transport, nil)
 	if err != nil {
 		http.Error(w, "connection failed", http.StatusInternalServerError)
 		return

mcp/streamable.go

@@ -171,12 +171,15 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 		// 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)
+		_, err = server.Connect(req.Context(), transport, &ServerSessionOptions{
+			SessionID:    sessionID,
+			SessionState: state,
+			SessionStore: h.opts.SessionStore,
+		})
 		if err != nil {
 			http.Error(w, fmt.Sprintf("failed connection: %v", err), http.StatusInternalServerError)
 			return
 		}
-		session.InitSession(sessionID, state, h.opts.SessionStore)
 		h.transportMu.Lock()
 		h.transports[transport.sessionID] = transport
 		h.transportMu.Unlock()