mcp: add StreamableHTTPOptions.EventStore (#593)

Allow users to enable streamable replay by setting a non-nil event
store.

Updates #587
This commit is contained in:
Robert Findley 2025-10-20 14:58:24 -04:00 • committed by GitHub
commit cfa7a515a9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 15 additions and 21 deletions

View file

@ -72,8 +72,11 @@ type StreamableHTTPOptions struct {
// If nil, do not log.
Logger *slog.Logger
// TODO(rfindley): file a proposal to export this option, or something equivalent.
configureTransport func(req *http.Request, transport *StreamableServerTransport)
// EventStore enables stream resumption.
//
// If set, EventStore will be used to persist stream events and replay them
// upon stream resumption.
EventStore EventStore
}
// NewStreamableHTTPHandler returns a new [StreamableHTTPHandler].
@ -237,12 +240,10 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
transport = &StreamableServerTransport{
SessionID: sessionID,
Stateless: h.opts.Stateless,
EventStore: h.opts.EventStore,
jsonResponse: h.opts.JSONResponse,
logger: h.opts.Logger,
}
if h.opts.configureTransport != nil {
h.opts.configureTransport(req, transport)
}
// To support stateless mode, we initialize the session with a default
// state, so that it doesn't reject subsequent requests.

View file

@ -125,14 +125,13 @@ func TestStreamableTransports(t *testing.T) {
// Start an httptest.Server with the StreamableHTTPHandler, wrapped in a
// cookie-checking middleware.
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{
opts := &StreamableHTTPOptions{
JSONResponse: test.useJSON,
configureTransport: func(_ *http.Request, transport *StreamableServerTransport) {
if test.replay {
transport.EventStore = NewMemoryEventStore(nil)
}
},
})
}
if test.replay {
opts.EventStore = NewMemoryEventStore(nil)
}
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts)
var (
headerMu sync.Mutex
@ -420,9 +419,7 @@ func testClientReplay(t *testing.T, test clientReplayTest) {
})
realServer := httptest.NewServer(mustNotPanic(t, NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
configureTransport: func(_ *http.Request, t *StreamableServerTransport) {
t.EventStore = NewMemoryEventStore(nil) // necessary for replay
},
EventStore: NewMemoryEventStore(nil), // necessary for replay
})))
t.Cleanup(func() {
t.Log("Closing real HTTP server")
@ -601,9 +598,7 @@ func TestServerInitiatedSSE(t *testing.T) {
// However, it shouldn't be necessary to use replay here, as we should be
// guaranteed that the standalone SSE stream is started by the time the
// client is connected.
configureTransport: func(_ *http.Request, transport *StreamableServerTransport) {
transport.EventStore = NewMemoryEventStore(nil)
},
EventStore: NewMemoryEventStore(nil),
}
httpServer := httptest.NewServer(mustNotPanic(t, NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, opts)))
defer httpServer.Close()
@ -976,9 +971,7 @@ func TestStreamableServerTransport(t *testing.T) {
opts := &StreamableHTTPOptions{}
if test.replay {
opts.configureTransport = func(_ *http.Request, t *StreamableServerTransport) {
t.EventStore = NewMemoryEventStore(nil)
}
opts.EventStore = NewMemoryEventStore(nil)
}
// Start the streamable handler.
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts)