Proposal: Export callback to StreamableHTTPHandler for closed transports
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
I'm happy to provide a full PR including updates to unit tests.
I see you added the release blocker tag so I was so bold to open a PR :)
There's one (known) bug that we're closing the connection, not transport, when the server is deleted: https://github.com/modelcontextprotocol/go-sdk/blob/22f86c4dfdf440e9980a18d9d3a7a89618f6be3f/mcp/streamable.go#L46 It looks like your patch fixes that.
A couple quick comments, but as I was writing I see you added a PR, so we can discuss there.
- When we terminate the session, we should call
transport.session.close, nottransport.connection.close. The session owns the connection, so will close it when it is closed. - I don't think we should add an exported
StreamableServerTransport.Closeyet, since it feels misplaced and isn't necessary for your use case (unless I'm missing something) - Probably
OnConnectionCloseshould beOnSessionClose, because 'connection' is too ambigious (there are lots of HTTP connections). Furthermore, it accepts a 'SessionID', so we should use that naming.
Otherwise, I think this makes sense. Will look at the PR.
I'm not sure this needs to be a release blocker. Given that it's a new API, we'd like to let it soak as a proposal before committing, and we want to cut a release this week, so I advocate for merging after the release.
An alternative would be to add a callback
OnNewSession func(*ServerSession)
Then you could spin off a goroutine blocked in session.Wait(). Maybe the goroutine is not ideal, but I can see the OnNewSession callback being more generally useful. WDYT?
Just to make sure we are not blocked here. I continued the discussion with an answer to the last question at the PR here: https://github.com/modelcontextprotocol/go-sdk/pull/480#issuecomment-3303286390
@findleyr should we add the label to mark the proposal as accepted?