mcp: add StreamableHTTPOptions.SessionTimeout

Add a timeout option for the streamable HTTP handler that automatically
cleans up idle sessions.

Also, fix a bug in the streamable client, where we hang on a
request even though the client can never get a response (because the
HTTP request terminated without a response or Last-Event-Id).

Fixes #499

Author: findleyr

mcp/server.go

@@ -825,7 +825,7 @@ func (s *Server) disconnect(cc *ServerSession) {
 type ServerSessionOptions struct {
 	State *ServerSessionState
 
-	onClose func()
+	onClose func() // used to clean up associated resources
 }
 
 // Connect connects the MCP server over the given transport and starts handling

mcp/streamable.go

@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"io"
 	"log/slog"
+	"maps"
 	"math"
 	"math/rand/v2"
 	"net/http"
@@ -40,12 +41,46 @@ type StreamableHTTPHandler struct {
 	getServer func(*http.Request) *Server
 	opts      StreamableHTTPOptions
 
-	onTransportDeletion func(sessionID string) // for testing only
+	onTransportDeletion func(sessionID string) // for testing
 
-	mu sync.Mutex
-	// TODO: we should store the ServerSession along with the transport, because
-	// we need to cancel keepalive requests when closing the transport.
-	transports map[string]*StreamableServerTransport // keyed by IDs (from Mcp-Session-Id header)
+	mu       sync.Mutex
+	sessions map[string]*sessionInfo // keyed by session ID
+}
+
+type sessionInfo struct {
+	session   *ServerSession
+	transport *StreamableServerTransport
+
+	// If timeout is set, automatically close the session after an idle period.
+	timeout time.Duration
+	timerMu sync.Mutex
+	timer   *time.Timer
+}
+
+// resetTimeout resets the inactivity timer.
+func (i *sessionInfo) resetTimeout() {
+	if i.timeout <= 0 {
+		return
+	}
+
+	i.timerMu.Lock()
+	defer i.timerMu.Unlock()
+
+	if i.timer == nil {
+		return
+	}
+	// Reset the timer if we successfully stopped it.
+	i.timer.Reset(i.timeout)
+}
+
+// stopTimer stops the inactivity timer.
+func (i *sessionInfo) stopTimer() {
+	i.timerMu.Lock()
+	defer i.timerMu.Unlock()
+	if i.timer != nil {
+		i.timer.Stop()
+		i.timer = nil
+	}
 }
 
 // StreamableHTTPOptions configures the StreamableHTTPHandler.
@@ -72,6 +107,14 @@ type StreamableHTTPOptions struct {
 	// If nil, do not log.
 	Logger *slog.Logger
 
+	// SessionTimeout configures a timeout for idle sessions.
+	//
+	// When sessions receive no new HTTP requests from the client for this
+	// duration, they are automatically closed.
+	//
+	// If SessionTimeout is the zero value, idle sessions are never closed.
+	SessionTimeout time.Duration
+
 	// TODO(rfindley): file a proposal to export this option, or something equivalent.
 	configureTransport func(req *http.Request, transport *StreamableServerTransport)
 }
@@ -83,8 +126,8 @@ type StreamableHTTPOptions struct {
 // If getServer returns nil, a 400 Bad Request will be served.
 func NewStreamableHTTPHandler(getServer func(*http.Request) *Server, opts *StreamableHTTPOptions) *StreamableHTTPHandler {
 	h := &StreamableHTTPHandler{
-		getServer:  getServer,
-		transports: make(map[string]*StreamableServerTransport),
+		getServer: getServer,
+		sessions:  make(map[string]*sessionInfo),
 	}
 	if opts != nil {
 		h.opts = *opts
@@ -97,20 +140,27 @@ func NewStreamableHTTPHandler(getServer func(*http.Request) *Server, opts *Strea
 	return h
 }
 
-// closeAll closes all ongoing sessions.
+// closeAll closes all ongoing sessions, for tests.
 //
 // TODO(rfindley): investigate the best API for callers to configure their
 // session lifecycle. (?)
 //
 // Should we allow passing in a session store? That would allow the handler to
 // be stateless.
 func (h *StreamableHTTPHandler) closeAll() {
+	// TODO: if we ever expose this outside of tests, we'll need to do better
+	// than simply collecting sessions while holding the lock: we need to prevent
+	// new sessions from being added.
+	//
+	// Currently, sessions remove themselves from h.sessions when closed, so we
+	// can't call Close while holding the lock.
 	h.mu.Lock()
-	defer h.mu.Unlock()
-	for _, s := range h.transports {
-		s.connection.Close()
+	sessionInfos := slices.Collect(maps.Values(h.sessions))
+	h.sessions = nil
+	h.mu.Unlock()
+	for _, s := range sessionInfos {
+		s.session.Close()
 	}
-	h.transports = nil
 }
 
 func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
@@ -141,12 +191,12 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 	}
 
 	sessionID := req.Header.Get(sessionIDHeader)
-	var transport *StreamableServerTransport
+	var sessInfo *sessionInfo
 	if sessionID != "" {
 		h.mu.Lock()
-		transport = h.transports[sessionID]
+		sessInfo = h.sessions[sessionID]
 		h.mu.Unlock()
-		if transport == nil && !h.opts.Stateless {
+		if sessInfo == nil && !h.opts.Stateless {
 			// Unless we're in 'stateless' mode, which doesn't perform any Session-ID
 			// validation, we require that the session ID matches a known session.
 			//
@@ -161,11 +211,10 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 			http.Error(w, "Bad Request: DELETE requires an Mcp-Session-Id header", http.StatusBadRequest)
 			return
 		}
-		if transport != nil { // transport may be nil in stateless mode
-			h.mu.Lock()
-			delete(h.transports, transport.SessionID)
-			h.mu.Unlock()
-			transport.connection.Close()
+		if sessInfo != nil { // sessInfo may be nil in stateless mode
+			// Closing the session also removes it from h.sessions, due to the
+			// onClose callback.
+			sessInfo.session.Close()
 		}
 		w.WriteHeader(http.StatusNoContent)
 		return
@@ -222,7 +271,7 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 		return
 	}
 
-	if transport == nil {
+	if sessInfo == nil {
 		server := h.getServer(req)
 		if server == nil {
 			// The getServer argument to NewStreamableHTTPHandler returned nil.
@@ -234,7 +283,7 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 			// existing transport.
 			sessionID = server.opts.GetSessionID()
 		}
-		transport = &StreamableServerTransport{
+		transport := &StreamableServerTransport{
 			SessionID:    sessionID,
 			Stateless:    h.opts.Stateless,
 			jsonResponse: h.opts.JSONResponse,
@@ -300,10 +349,13 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
 			connectOpts = &ServerSessionOptions{
 				onClose: func() {
 					h.mu.Lock()
-					delete(h.transports, transport.SessionID)
-					h.mu.Unlock()
-					if h.onTransportDeletion != nil {
-						h.onTransportDeletion(transport.SessionID)
+					defer h.mu.Unlock()
+					if info, ok := h.sessions[transport.SessionID]; ok {
+						info.stopTimer()
+						delete(h.sessions, transport.SessionID)
+						if h.onTransportDeletion != nil {
+							h.onTransportDeletion(transport.SessionID)
+						}
 					}
 				},
 			}
@@ -312,23 +364,36 @@ 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.
-		ss, err := server.Connect(req.Context(), transport, connectOpts)
+		session, err := server.Connect(req.Context(), transport, connectOpts)
 		if err != nil {
 			http.Error(w, "failed connection", http.StatusInternalServerError)
 			return
 		}
+		sessInfo = &sessionInfo{
+			session:   session,
+			transport: transport,
+		}
 		if h.opts.Stateless {
 			// Stateless mode: close the session when the request exits.
-			defer ss.Close() // close the fake session after handling the request
+			defer session.Close() // close the fake session after handling the request
 		} else {
 			// Otherwise, save the transport so that it can be reused
+
+			// Clean up the session when it times out.
+			if h.opts.SessionTimeout > 0 {
+				sessInfo.timeout = h.opts.SessionTimeout
+				sessInfo.timer = time.AfterFunc(sessInfo.timeout, func() {
+					sessInfo.session.Close()
+				})
+			}
 			h.mu.Lock()
-			h.transports[transport.SessionID] = transport
+			h.sessions[transport.SessionID] = sessInfo
 			h.mu.Unlock()
 		}
 	}
 
-	transport.ServeHTTP(w, req)
+	sessInfo.resetTimeout()
+	sessInfo.transport.ServeHTTP(w, req)
 }
 
 // A StreamableServerTransport implements the server side of the MCP streamable
@@ -1340,9 +1405,12 @@ func (c *streamableClientConn) Write(ctx context.Context, msg jsonrpc.Message) e
 		go c.handleJSON(requestSummary, resp)
 
 	case "text/event-stream":
-		jsonReq, _ := msg.(*jsonrpc.Request)
+		var forCall *jsonrpc.Request
+		if jsonReq, ok := msg.(*jsonrpc.Request); ok && jsonReq.IsCall() {
+			forCall = jsonReq
+		}
 		// TODO: should we cancel this logical SSE request if/when jsonReq is canceled?
-		go c.handleSSE(requestSummary, resp, false, jsonReq)
+		go c.handleSSE(requestSummary, resp, false, forCall)
 
 	default:
 		resp.Body.Close()
@@ -1392,9 +1460,9 @@ func (c *streamableClientConn) handleJSON(requestSummary string, resp *http.Resp
 // handleSSE manages the lifecycle of an SSE connection. It can be either
 // persistent (for the main GET listener) or temporary (for a POST response).
 //
-// If forReq is set, it is the request that initiated the stream, and the
+// If forCall is set, it is the call that initiated the stream, and the
 // stream is complete when we receive its response.
-func (c *streamableClientConn) handleSSE(requestSummary string, initialResp *http.Response, persistent bool, forReq *jsonrpc2.Request) {
+func (c *streamableClientConn) handleSSE(requestSummary string, initialResp *http.Response, persistent bool, forCall *jsonrpc2.Request) {
 	resp := initialResp
 	var lastEventID string
 	for {
@@ -1404,7 +1472,7 @@ func (c *streamableClientConn) handleSSE(requestSummary string, initialResp *htt
 		// Eventually, if we don't get the response, we should stop trying and
 		// fail the request.
 		if resp != nil {
-			eventID, clientClosed := c.processStream(requestSummary, resp, forReq)
+			eventID, clientClosed := c.processStream(requestSummary, resp, forCall)
 			lastEventID = eventID
 
 			// If the connection was closed by the client, we're done.
@@ -1467,11 +1535,11 @@ func (c *streamableClientConn) handleSSE(requestSummary string, initialResp *htt
 // incoming channel. It returns the ID of the last processed event and a flag
 // indicating if the connection was closed by the client. If resp is nil, it
 // returns "", false.
-func (c *streamableClientConn) processStream(requestSummary string, resp *http.Response, forReq *jsonrpc.Request) (lastEventID string, clientClosed bool) {
+func (c *streamableClientConn) processStream(requestSummary string, resp *http.Response, forCall *jsonrpc.Request) (lastEventID string, clientClosed bool) {
 	defer resp.Body.Close()
 	for evt, err := range scanEvents(resp.Body) {
 		if err != nil {
-			return lastEventID, false
+			break
 		}
 
 		if evt.ID != "" {
@@ -1486,10 +1554,10 @@ func (c *streamableClientConn) processStream(requestSummary string, resp *http.R
 
 		select {
 		case c.incoming <- msg:
-			if jsonResp, ok := msg.(*jsonrpc.Response); ok && forReq != nil {
+			if jsonResp, ok := msg.(*jsonrpc.Response); ok && forCall != nil {
 				// TODO: we should never get a response when forReq is nil (the standalone SSE request).
 				// We should detect this case.
-				if jsonResp.ID == forReq.ID {
+				if jsonResp.ID == forCall.ID {
 					return "", true
 				}
 			}
@@ -1499,7 +1567,20 @@ func (c *streamableClientConn) processStream(requestSummary string, resp *http.R
 		}
 	}
 	// The loop finished without an error, indicating the server closed the stream.
-	return "", false
+	//
+	// If the lastEventID is "", the stream is not retryable and we should
+	// report a synthetic error for the call.
+	if lastEventID == "" && forCall != nil {
+		errmsg := &jsonrpc2.Response{
+			ID:    forCall.ID,
+			Error: fmt.Errorf("request terminated without response"),
+		}
+		select {
+		case c.incoming <- errmsg:
+		case <-c.done:
+		}
+	}
+	return lastEventID, false
 }
 
 // reconnect handles the logic of retrying a connection with an exponential