From b3399e67309abd3886c0b017f09eb6193469cf32 Mon Sep 17 00:00:00 2001 From: Robert Findley Date: Thu, 11 Dec 2025 15:13:38 -0500 Subject: [PATCH] mcp: use args struct for CloseSSEStream (#718) Change CloseSSEStream from func(time.Duration) to func(CloseSSEStreamArgs) for forward compatibility, allowing new fields to be added without breaking existing callers. --- examples/server/conformance/main.go | 4 ++-- mcp/shared.go | 18 ++++++++++-------- mcp/streamable.go | 8 ++++---- mcp/streamable_test.go | 10 +++++----- 4 files changed, 21 insertions(+), 19 deletions(-) diff --git a/examples/server/conformance/main.go b/examples/server/conformance/main.go index ffee291..70bd4bc 100644 --- a/examples/server/conformance/main.go +++ b/examples/server/conformance/main.go @@ -503,8 +503,8 @@ func jsonSchema202012Handler(ctx context.Context, req *mcp.CallToolRequest, inpu func testReconnectionHandler(ctx context.Context, req *mcp.CallToolRequest, _ any) (*mcp.CallToolResult, any, error) { // Close the SSE stream to trigger client reconnection (SEP-1699) - if req.Extra != nil && req.Extra.CloseStream != nil { - req.Extra.CloseStream(10 * time.Millisecond) + if req.Extra != nil && req.Extra.CloseSSEStream != nil { + req.Extra.CloseSSEStream(mcp.CloseSSEStreamArgs{RetryAfter: 10 * time.Millisecond}) } // Wait for client to reconnect diff --git a/mcp/shared.go b/mcp/shared.go index 288ed80..d83eae7 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -479,22 +479,24 @@ type RequestExtra struct { TokenInfo *auth.TokenInfo // bearer token info (e.g. from OAuth) if any Header http.Header // header from HTTP request, if any - // If set, CloseStream explicitly closes the current request stream. + // If set, CloseSSEStream explicitly closes the current SSE request stream. // // [SEP-1699] introduced server-side SSE stream disconnection: for // long-running requests, servers may opt to close the SSE stream and - // ask the client to retry at a later time. CloseStream implements this - // feature; if reconnectAfter is set, an event is sent with a `retry:` field + // ask the client to retry at a later time. CloseSSEStream implements this + // feature; if RetryAfter is set, an event is sent with a `retry:` field // to configure the reconnection delay. // // [SEP-1699]: https://github.com/modelcontextprotocol/modelcontextprotocol/issues/1699 - CloseStream func(reconnectAfter time.Duration) + CloseSSEStream func(CloseSSEStreamArgs) } -// TODO(cleanup): switch to an args struct here for forwards compatibility. -// type CloseStreamArgs struct { -// After time.Duration -// } +// CloseSSEStreamArgs are arguments for [RequestExtra.CloseSSEStream]. +type CloseSSEStreamArgs struct { + // RetryAfter configures the reconnection delay sent to the client via the + // SSE retry field. If zero, no retry field is sent. + RetryAfter time.Duration +} func (*ClientRequest[P]) isRequest() {} func (*ServerRequest[P]) isRequest() {} diff --git a/mcp/streamable.go b/mcp/streamable.go index d6452b2..dcb0bdf 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -1065,9 +1065,9 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques } if jreq.IsCall() { calls[jreq.ID] = struct{}{} - // See the doc for CloseStream: allow the request handler to explicitly - // close the ongoing stream. - jreq.Extra.(*RequestExtra).CloseStream = func(reconnectAfter time.Duration) { + // See the doc for CloseSSEStream: allow the request handler to + // explicitly close the ongoing stream. + jreq.Extra.(*RequestExtra).CloseSSEStream = func(args CloseSSEStreamArgs) { c.mu.Lock() streamID, ok := c.requestStreams[jreq.ID] var stream *stream @@ -1077,7 +1077,7 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques c.mu.Unlock() if stream != nil { - stream.close(reconnectAfter) + stream.close(args.RetryAfter) } } } diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 9eff74b..b1c3f07 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -528,10 +528,10 @@ func TestStreamableServerDisconnect(t *testing.T) { testStream := func(ctx context.Context, session *ServerSession, extra *RequestExtra) { // Close the stream before the first message. We should have sent an // initial priming message already, so the client will be able to replay - extra.CloseStream(10 * time.Millisecond) + extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: 10 * time.Millisecond}) session.NotifyProgress(ctx, &ProgressNotificationParams{Message: "msg1"}) time.Sleep(20 * time.Millisecond) - extra.CloseStream(10 * time.Millisecond) // Closing twice should still be supported. + extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: 10 * time.Millisecond}) // Closing twice should still be supported. session.NotifyProgress(ctx, &ProgressNotificationParams{Message: "msg2"}) } @@ -1069,7 +1069,7 @@ func TestStreamableServerTransport(t *testing.T) { name: "no close message on old protocol", replay: true, tool: func(t *testing.T, _ context.Context, req *CallToolRequest) { - req.Extra.CloseStream(time.Millisecond) + req.Extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: time.Millisecond}) }, requests: []streamableRequest{ initialize, @@ -1105,7 +1105,7 @@ func TestStreamableServerTransport(t *testing.T) { name: "close message on 2025-11-25", replay: true, tool: func(t *testing.T, _ context.Context, req *CallToolRequest) { - req.Extra.CloseStream(time.Millisecond) + req.Extra.CloseSSEStream(CloseSSEStreamArgs{RetryAfter: time.Millisecond}) }, requests: []streamableRequest{ initialize20251125, @@ -1126,7 +1126,7 @@ func TestStreamableServerTransport(t *testing.T) { name: "no close message", replay: true, tool: func(t *testing.T, _ context.Context, req *CallToolRequest) { - req.Extra.CloseStream(0) + req.Extra.CloseSSEStream(CloseSSEStreamArgs{}) }, requests: []streamableRequest{ initialize,