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: use args struct for CloseSSEStream (#718)

Browse source 
Change CloseSSEStream from func(time.Duration) to
func(CloseSSEStreamArgs) for forward compatibility, allowing new fields
to be added without breaking existing callers.
This commit is contained in:
Robert Findley 2025-12-11 15:13:38 -05:00 • committed by GitHub
commit b3399e6730
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 21 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

4
examples/server/conformance/main.go
View file

@ -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

18
mcp/shared.go
View file

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

8
mcp/streamable.go
View file

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

10
mcp/streamable_test.go
View file

@ -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,