fix race (#232)

Author: jba

mcp/conformance_test.go

@@ -135,7 +135,7 @@ func runServerTest(t *testing.T, test *conformanceTest) {
 	// Connect the server, and connect the client stream,
 	// but don't connect an actual client.
 	cTransport, sTransport := NewInMemoryTransports()
-	ss, err := s.Connect(ctx, sTransport)
+	ss, err := s.Connect(ctx, sTransport, nil)
 	if err != nil {
 		t.Fatal(err)
 	}

mcp/server.go

@@ -559,7 +559,7 @@ func (s *Server) disconnect(cc *ServerSession) {
 	}
 }
 
-type ServerSessionOptions struct {
+type SessionOptions struct {
 	SessionID    string
 	SessionState *SessionState
 	SessionStore SessionStore
@@ -571,7 +571,7 @@ type ServerSessionOptions struct {
 // 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, opts *ServerSessionOptions) (*ServerSession, error) {
+func (s *Server) Connect(ctx context.Context, t Transport, opts *SessionOptions) (*ServerSession, error) {
 	if opts != nil && opts.SessionState == nil && opts.SessionStore != nil {
 		return nil, errors.New("ServerSessionOptions has store but no state")
 	}
@@ -634,7 +634,7 @@ func (ss *ServerSession) NotifyProgress(ctx context.Context, params *ProgressNot
 type ServerSession struct {
 	server          *Server
 	conn            *jsonrpc2.Connection
-	opts            ServerSessionOptions
+	opts            SessionOptions
 	mu              sync.Mutex
 	logLevel        LoggingLevel
 	_initialized    bool

mcp/streamable.go

@@ -117,26 +117,26 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 		return
 	}
 
-	var session *StreamableServerTransport
+	var transport *StreamableServerTransport
 	sessionID := req.Header.Get(sessionIDHeader)
 	if sessionID != "" {
 		h.transportMu.Lock()
-		session, _ = h.transports[sessionID]
+		transport, _ = h.transports[sessionID]
 		h.transportMu.Unlock()
 	}
 
 	// TODO(rfindley): simplify the locking so that each request has only one
 	// critical section.
 	if req.Method == http.MethodDelete {
-		if session == nil {
+		if transport == nil {
 			// => 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.transportMu.Lock()
-		delete(h.transports, session.sessionID)
+		delete(h.transports, transport.sessionID)
 		h.transportMu.Unlock()
-		session.Close()
+		transport.Close()
 		w.WriteHeader(http.StatusNoContent)
 		return
 	}
@@ -149,7 +149,7 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 		return
 	}
 
-	if session == nil {
+	if transport == nil {
 		var state *SessionState
 		var err error
 		if sessionID != "" {
@@ -162,16 +162,15 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 				http.Error(w, fmt.Sprintf("SessionStore.Load(%q): %v", sessionID, err), http.StatusInternalServerError)
 				return
 			}
-			session = NewStreamableServerTransport(sessionID, nil)
 		} else {
 			state = &SessionState{}
 			sessionID = randText()
 			if err := h.opts.SessionStore.Store(req.Context(), sessionID, state); err != nil {
 				http.Error(w, fmt.Sprintf("SessionStore.Store, new session: %v", err), http.StatusInternalServerError)
 				return
 			}
-			session = NewStreamableServerTransport(sessionID, nil)
 		}
+		transport = NewStreamableServerTransport(sessionID, nil)
 		server := h.getServer(req)
 		if server == nil {
 			http.Error(w, "no server available", http.StatusBadRequest)
@@ -180,7 +179,8 @@ 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.
-		_, err = server.Connect(req.Context(), session, &ServerSessionOptions{
+		// TODO: rename SessionOptions to ConnectOptions?
+		_, err = server.Connect(req.Context(), transport, &SessionOptions{
 			SessionID:    sessionID,
 			SessionState: state,
 			SessionStore: h.opts.SessionStore,
@@ -190,11 +190,14 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 			return
 		}
 		h.transportMu.Lock()
-		h.transports[session.sessionID] = session
+		// Check in case another request with the same stored session ID got here first.
+		if _, ok := h.transports[transport.sessionID]; !ok {
+			h.transports[transport.sessionID] = transport
+		}
 		h.transportMu.Unlock()
 	}
 
-	session.ServeHTTP(w, req)
+	transport.ServeHTTP(w, req)
 }
 
 type StreamableServerTransportOptions struct {