mcp: gracefully handle a nil server in handlers (#164)

If the getServer function passed to NewSSEHandler or
NewStreamableHTTPHandler returns nil, serve a 400 instead of panicking.

Fixes #161.

Author: jba

examples/sse/main.go

@@ -53,5 +53,5 @@ func main() {
 			return nil
 		}
 	})
-	http.ListenAndServe(*httpAddr, handler)
+	log.Fatal(http.ListenAndServe(*httpAddr, handler))
 }

mcp/sse.go

@@ -55,13 +55,14 @@ type SSEHandler struct {
 // Sessions are created when the client issues a GET request to the server,
 // which must accept text/event-stream responses (server-sent events).
 // For each such request, a new [SSEServerTransport] is created with a distinct
-// messages endpoint, and connected to the server returned by getServer. It is
-// up to the user whether getServer returns a distinct [Server] for each new
-// request, or reuses an existing server.
-//
+// messages endpoint, and connected to the server returned by getServer.
 // The SSEHandler also handles requests to the message endpoints, by
 // delegating them to the relevant server transport.
 //
+// The getServer function may return a distinct [Server] for each new
+// request, or reuse an existing server. If it returns nil, the handler
+// will return a 400 Bad Request.
+//
 // TODO(rfindley): add options.
 func NewSSEHandler(getServer func(request *http.Request) *Server) *SSEHandler {
 	return &SSEHandler{
@@ -208,8 +209,12 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 		h.mu.Unlock()
 	}()
 
-	// TODO(hxjiang): getServer returns nil will panic.
 	server := h.getServer(req)
+	if server == nil {
+		// The getServer argument to NewSSEHandler returned nil.
+		http.Error(w, "no server available", http.StatusBadRequest)
+		return
+	}
 	ss, err := server.Connect(req.Context(), transport)
 	if err != nil {
 		http.Error(w, "connection failed", http.StatusInternalServerError)

mcp/streamable.go

@@ -50,6 +50,7 @@ type StreamableHTTPOptions struct {
 //
 // The getServer function is used to create or look up servers for new
 // sessions. It is OK for getServer to return the same server multiple times.
+// If getServer returns nil, a 400 Bad Request will be served.
 func NewStreamableHTTPHandler(getServer func(*http.Request) *Server, opts *StreamableHTTPOptions) *StreamableHTTPHandler {
 	return &StreamableHTTPHandler{
 		getServer: getServer,
@@ -135,6 +136,11 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 	if session == nil {
 		s := NewStreamableServerTransport(randText(), nil)
 		server := h.getServer(req)
+		if server == nil {
+			// The getServer argument to NewStreamableHTTPHandler returned nil.
+			http.Error(w, "no server available", http.StatusBadRequest)
+			return
+		}
 		// 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.