Proposal: Export callback to StreamableHTTPHandler for closed transports #479
Description
Is your feature request related to a problem? Please describe
I'm using the StreamableHTTPHandler to implement an MCP server. Just as the underlying mcp.Server, my server is stateful, and I need to clean up resources when a session ends (e.g., deleting entries from a map). I am currently performing this cleanup by wrapping the StreamableHTTPHandler in my own HTTP handler implementation, which allows me to directly handle DELETE requests from the client. While this works fine, I'm still not able to clean up sessions when the connection breaks due to timeouts or other network errors.
Describe the solution you'd like
Browsing around the SDK internals, I found there is already such a cleanup, but it's not yet exported: StreamableHTTPHandler.onTransportDeletion(...). I experimented with this locally to see if it fits my use case and was able to make it work. However, I found out that this cleanup isn't actually called yet when the client closes the session with a DELETE call explicitly. Ideally, I would like to replace my current HTTP handler for DELETE requests with a single cleanup function, so it would be beneficial if the StreamableHTTPHandler would actually close the transport connection and session when handling the DELETE call.
I've attached a full working diff in the details below:
diff --git a/mcp/streamable.go b/mcp/streamable.go
index f359b7c..9c563cf 100644
--- a/mcp/streamable.go
+++ b/mcp/streamable.go
@@ -40,8 +40,6 @@ type StreamableHTTPHandler struct {
getServer func(*http.Request) *Server
opts StreamableHTTPOptions
- onTransportDeletion func(sessionID string) // for testing only
-
mu sync.Mutex
// TODO: we should store the ServerSession along with the transport, because
// we need to cancel keepalive requests when closing the transport.
@@ -67,6 +65,11 @@ type StreamableHTTPOptions struct {
//
// [§2.1.5]: https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#sending-messages-to-the-server
JSONResponse bool
+
+ // OnConnectionClose is a callback function that is invoked when a [Connection]
+ // is closed. A connection is closed when the session is ended explicitly by
+ // the client or when it is interrupted due to a timeout or other errors.
+ OnConnectionClose func(sessionID string)
}
// NewStreamableHTTPHandler returns a new [StreamableHTTPHandler].
@@ -153,7 +156,7 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
h.mu.Lock()
delete(h.transports, transport.SessionID)
h.mu.Unlock()
- transport.connection.Close()
+ _ = transport.Close()
}
w.WriteHeader(http.StatusNoContent)
return
@@ -286,8 +289,8 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
h.mu.Lock()
delete(h.transports, transport.SessionID)
h.mu.Unlock()
- if h.onTransportDeletion != nil {
- h.onTransportDeletion(transport.SessionID)
+ if h.opts.OnConnectionClose != nil {
+ h.opts.OnConnectionClose(transport.SessionID)
}
},
}
@@ -307,6 +310,7 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
} else {
// Otherwise, save the transport so that it can be reused
h.mu.Lock()
+ transport.session = ss
h.transports[transport.SessionID] = transport
h.mu.Unlock()
}
@@ -369,6 +373,9 @@ type StreamableServerTransport struct {
// connection is non-nil if and only if the transport has been connected.
connection *streamableServerConn
+
+ // the server session associated with this transport
+ session *ServerSession
}
// Connect implements the [Transport] interface.
@@ -550,6 +557,19 @@ func (t *StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.R
}
}
+// Close releases resources related to this transport if it has already been connected.
+func (t *StreamableServerTransport) Close() error {
+ var sessionErr, connErr error
+ if t.session != nil {
+ sessionErr = t.session.Close()
+ }
+ if t.connection != nil {
+ connErr = t.connection.Close()
+ }
+
+ return errors.Join(sessionErr, connErr)
+}
+
// serveGET streams messages to a hanging http GET, with stream ID and last
// message parsed from the Last-Event-ID header.
//Describe alternatives you've considered
Cleaning up resources just with a DELETE handler, but this alternative doesn't allow cleanup in error cases, which means my server will leak memory over time.
Additional context
none