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.

Author: findleyr

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

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() {}

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)
 					}
 				}
 			}

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,