modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

mcp: add StreamableHTTPOptions.EventStore (#593)

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

11
mcp/streamable.go
View file

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

25
mcp/streamable_test.go
View file

@ -125,14 +125,13 @@ func TestStreamableTransports(t *testing.T) {
// Start an httptest.Server with the StreamableHTTPHandler, wrapped in a // Start an httptest.Server with the StreamableHTTPHandler, wrapped in a
// cookie-checking middleware. // cookie-checking middleware.
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, &StreamableHTTPOptions{ opts := &StreamableHTTPOptions{
JSONResponse: test.useJSON, JSONResponse: test.useJSON,
configureTransport: func(_ *http.Request, transport *StreamableServerTransport) { }
if test.replay { if test.replay {
transport.EventStore = NewMemoryEventStore(nil) opts.EventStore = NewMemoryEventStore(nil)
} }
}, handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts)
})
var ( var (
headerMu sync.Mutex 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{ realServer := httptest.NewServer(mustNotPanic(t, NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{
configureTransport: func(_ *http.Request, t *StreamableServerTransport) { EventStore: NewMemoryEventStore(nil), // necessary for replay
t.EventStore = NewMemoryEventStore(nil) // necessary for replay
},
}))) })))
t.Cleanup(func() { t.Cleanup(func() {
t.Log("Closing real HTTP server") 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 // 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 // guaranteed that the standalone SSE stream is started by the time the
// client is connected. // client is connected.
configureTransport: func(_ *http.Request, transport *StreamableServerTransport) { EventStore: NewMemoryEventStore(nil),
transport.EventStore = NewMemoryEventStore(nil)
},
} }
httpServer := httptest.NewServer(mustNotPanic(t, NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, opts))) httpServer := httptest.NewServer(mustNotPanic(t, NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, opts)))
defer httpServer.Close() defer httpServer.Close()
@ -976,9 +971,7 @@ func TestStreamableServerTransport(t *testing.T) {
opts := &StreamableHTTPOptions{} opts := &StreamableHTTPOptions{}
if test.replay { if test.replay {
opts.configureTransport = func(_ *http.Request, t *StreamableServerTransport) { opts.EventStore = NewMemoryEventStore(nil)
t.EventStore = NewMemoryEventStore(nil)
}
} }
// Start the streamable handler. // Start the streamable handler.
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts) handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, opts)