mcp/streamable: add persistent SSE GET listener (#206)

This CL adds the optional persistent SSE GET listener as specified in
section 2.2.
https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#listening-for-messages-from-the-server

This enables server initiated SSE streams.

For #10
This commit is contained in:
Sam Thanawalla 2025-07-30 17:00:50 -04:00 • committed by GitHub
commit a9a503f74a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 57 additions and 15 deletions

View file

@ -743,6 +743,12 @@ func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, er
ctx: connCtx, ctx: connCtx,
cancel: cancel, cancel: cancel,
} }
// Start the persistent SSE listener right away.
// Section 2.2: The client MAY issue an HTTP GET to the MCP endpoint.
// This can be used to open an SSE stream, allowing the server to
// communicate to the client, without the client first sending data via HTTP POST.
go conn.handleSSE(nil, true)
return conn, nil return conn, nil
} }
@ -867,7 +873,7 @@ func (s *streamableClientConn) postMessage(ctx context.Context, sessionID string
switch ct := resp.Header.Get("Content-Type"); ct { switch ct := resp.Header.Get("Content-Type"); ct {
case "text/event-stream": case "text/event-stream":
// Section 2.1: The SSE stream is initiated after a POST. // Section 2.1: The SSE stream is initiated after a POST.
go s.handleSSE(resp) go s.handleSSE(resp, false)
case "application/json": case "application/json":
body, err := io.ReadAll(resp.Body) body, err := io.ReadAll(resp.Body)
resp.Body.Close() resp.Body.Close()
@ -887,13 +893,11 @@ func (s *streamableClientConn) postMessage(ctx context.Context, sessionID string
return sessionID, nil return sessionID, nil
} }
// handleSSE manages the entire lifecycle of an SSE connection. It processes // handleSSE manages the lifecycle of an SSE connection. It can be either
// an incoming Server-Sent Events stream and automatically handles reconnection // persistent (for the main GET listener) or temporary (for a POST response).
// logic if the stream breaks. func (s *streamableClientConn) handleSSE(initialResp *http.Response, persistent bool) {
func (s *streamableClientConn) handleSSE(initialResp *http.Response) {
resp := initialResp resp := initialResp
var lastEventID string var lastEventID string
for { for {
eventID, clientClosed := s.processStream(resp) eventID, clientClosed := s.processStream(resp)
lastEventID = eventID lastEventID = eventID
@ -902,6 +906,11 @@ func (s *streamableClientConn) handleSSE(initialResp *http.Response) {
if clientClosed { if clientClosed {
return return
} }
// If the stream has ended, then do not reconnect if the stream is
// temporary (POST initiated SSE).
if lastEventID == "" && !persistent {
return
}
// The stream was interrupted or ended by the server. Attempt to reconnect. // The stream was interrupted or ended by the server. Attempt to reconnect.
newResp, err := s.reconnect(lastEventID) newResp, err := s.reconnect(lastEventID)
@ -923,9 +932,13 @@ func (s *streamableClientConn) handleSSE(initialResp *http.Response) {
// processStream reads from a single response body, sending events to the // processStream reads from a single response body, sending events to the
// incoming channel. It returns the ID of the last processed event, any error // incoming channel. It returns the ID of the last processed event, any error
// that occurred, and a flag indicating if the connection was closed by the client. // that occurred, and a flag indicating if the connection was closed by the client.
// If resp is nil, it returns "", false.
func (s *streamableClientConn) processStream(resp *http.Response) (lastEventID string, clientClosed bool) { func (s *streamableClientConn) processStream(resp *http.Response) (lastEventID string, clientClosed bool) {
defer resp.Body.Close() if resp == nil {
return "", false
}
defer resp.Body.Close()
for evt, err := range scanEvents(resp.Body) { for evt, err := range scanEvents(resp.Body) {
if err != nil { if err != nil {
return lastEventID, false return lastEventID, false
@ -939,13 +952,11 @@ func (s *streamableClientConn) processStream(resp *http.Response) (lastEventID s
case s.incoming <- evt.Data: case s.incoming <- evt.Data:
case <-s.done: case <-s.done:
// The connection was closed by the client; exit gracefully. // The connection was closed by the client; exit gracefully.
return lastEventID, true return "", true
} }
} }
// The loop finished without an error, indicating the server closed the stream. // The loop finished without an error, indicating the server closed the stream.
// We'll attempt to reconnect, so this is not a client-side close. return "", false
return lastEventID, false
} }
// reconnect handles the logic of retrying a connection with an exponential // reconnect handles the logic of retrying a connection with an exponential

View file

@ -168,7 +168,7 @@ func TestClientReplay(t *testing.T) {
clientSession.CallTool(ctx, &CallToolParams{Name: "multiMessageTool"}) clientSession.CallTool(ctx, &CallToolParams{Name: "multiMessageTool"})
// 4. Read and verify messages until the server signals it's ready for the proxy kill. // 4. Read and verify messages until the server signals it's ready for the proxy kill.
receivedNotifications := readProgressNotifications(t, ctx, notifications, 2) receivedNotifications := readNotifications(t, ctx, notifications, 2)
wantReceived := []string{"msg1", "msg2"} wantReceived := []string{"msg1", "msg2"}
if diff := cmp.Diff(wantReceived, receivedNotifications); diff != "" { if diff := cmp.Diff(wantReceived, receivedNotifications); diff != "" {
@ -201,7 +201,7 @@ func TestClientReplay(t *testing.T) {
// 7. Continue reading from the same connection object. // 7. Continue reading from the same connection object.
// Its internal logic should successfully retry, reconnect to the new proxy, // Its internal logic should successfully retry, reconnect to the new proxy,
// and receive the replayed messages. // and receive the replayed messages.
recoveredNotifications := readProgressNotifications(t, ctx, notifications, 2) recoveredNotifications := readNotifications(t, ctx, notifications, 2)
// 8. Verify the correct messages were received on the recovered connection. // 8. Verify the correct messages were received on the recovered connection.
wantRecovered := []string{"msg3", "msg4"} wantRecovered := []string{"msg3", "msg4"}
@ -211,8 +211,39 @@ func TestClientReplay(t *testing.T) {
} }
} }
// Helper to read a specific number of progress notifications. // TestServerInitiatedSSE verifies that the persistent SSE connection remains
func readProgressNotifications(t *testing.T, ctx context.Context, notifications chan string, count int) []string { // open and can receive server-initiated events.
func TestServerInitiatedSSE(t *testing.T) {
notifications := make(chan string)
server := NewServer(testImpl, nil)
httpServer := httptest.NewServer(NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil))
defer httpServer.Close()
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
client := NewClient(testImpl, &ClientOptions{ToolListChangedHandler: func(ctx context.Context, cc *ClientSession, params *ToolListChangedParams) {
notifications <- "toolListChanged"
},
})
clientSession, err := client.Connect(ctx, NewStreamableClientTransport(httpServer.URL, nil))
if err != nil {
t.Fatalf("client.Connect() failed: %v", err)
}
defer clientSession.Close()
server.AddTool(&Tool{Name: "testTool", InputSchema: &jsonschema.Schema{}},
func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
return &CallToolResult{}, nil
})
receivedNotifications := readNotifications(t, ctx, notifications, 1)
wantReceived := []string{"toolListChanged"}
if diff := cmp.Diff(wantReceived, receivedNotifications); diff != "" {
t.Errorf("Received notifications mismatch (-want +got):\n%s", diff)
}
}
// Helper to read a specific number of notifications.
func readNotifications(t *testing.T, ctx context.Context, notifications chan string, count int) []string {
t.Helper() t.Helper()
var collectedNotifications []string var collectedNotifications []string
for { for {