mcp: use http.ResponseController to ensure writes are flushed (#870)

If the ResponseWriter is a wrapped one (which is quite common when using
middleware to log things like response sizes), the direct cast to
`http.Flusher` doesn't find the flush capability. Instead, use
`http.ResponseController` as introduced in go1.20 which handles
unwrapping to find the base response writer.
This commit is contained in:
Dan McGee 2026-03-31 06:13:16 -05:00 • committed by GitHub
commit d3fd25b68a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 10 additions and 9 deletions

View file

@ -41,7 +41,7 @@ func (e Event) Empty() bool {
} }
// writeEvent writes the event to w, and flushes. // writeEvent writes the event to w, and flushes.
func writeEvent(w io.Writer, evt Event) (int, error) { func writeEvent(w http.ResponseWriter, evt Event) (int, error) {
var b bytes.Buffer var b bytes.Buffer
if evt.Name != "" { if evt.Name != "" {
fmt.Fprintf(&b, "event: %s\n", evt.Name) fmt.Fprintf(&b, "event: %s\n", evt.Name)
@ -54,9 +54,9 @@ func writeEvent(w io.Writer, evt Event) (int, error) {
} }
fmt.Fprintf(&b, "data: %s\n\n", string(evt.Data)) fmt.Fprintf(&b, "data: %s\n\n", string(evt.Data))
n, err := w.Write(b.Bytes()) n, err := w.Write(b.Bytes())
if f, ok := w.(http.Flusher); ok { rc := http.NewResponseController(w)
f.Flush() // Ignore returned error as flushing is best-effort.
} _ = rc.Flush()
return n, err return n, err
} }

View file

@ -1053,9 +1053,9 @@ func (c *streamableServerConn) acquireStream(ctx context.Context, w http.Respons
// Issue #410: the standalone SSE stream is likely not to receive messages // Issue #410: the standalone SSE stream is likely not to receive messages
// for a long time. Ensure that headers are flushed. // for a long time. Ensure that headers are flushed.
w.WriteHeader(http.StatusOK) w.WriteHeader(http.StatusOK)
if f, ok := w.(http.Flusher); ok { rc := http.NewResponseController(w)
f.Flush() // Ignore returned error as flushing is best-effort.
} _ = rc.Flush()
} }
for _, data := range toReplay { for _, data := range toReplay {

View file

@ -119,14 +119,15 @@ func (s *fakeStreamableServer) ServeHTTP(w http.ResponseWriter, req *http.Reques
for k, v := range resp.header { for k, v := range resp.header {
w.Header().Set(k, v) w.Header().Set(k, v)
} }
rc := http.NewResponseController(w)
w.WriteHeader(status) w.WriteHeader(status)
w.(http.Flusher).Flush() // flush response headers rc.Flush() // flush response headers
if v := req.Header.Get(protocolVersionHeader); v != resp.wantProtocolVersion && resp.wantProtocolVersion != "" { if v := req.Header.Get(protocolVersionHeader); v != resp.wantProtocolVersion && resp.wantProtocolVersion != "" {
s.t.Errorf("%v: bad protocol version header: got %q, want %q", key, v, resp.wantProtocolVersion) s.t.Errorf("%v: bad protocol version header: got %q, want %q", key, v, resp.wantProtocolVersion)
} }
w.Write([]byte(body)) w.Write([]byte(body))
w.(http.Flusher).Flush() // flush response rc.Flush() // flush response
if resp.done != nil { if resp.done != nil {
<-resp.done <-resp.done