-
Notifications
You must be signed in to change notification settings - Fork 316
fix: handle unframed JSON body with text/event-stream content type in MCP proxy #2229
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 3 commits
40f3b3d
5411b6e
4bbd643
be3f7b3
d56acce
36c194f
477145a
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -6,6 +6,7 @@ | |
| package mcpproxy | ||
|
|
||
| import ( | ||
| "bufio" | ||
| "bytes" | ||
| "cmp" | ||
| "context" | ||
|
|
@@ -782,13 +783,73 @@ func copyProxyHeaders(resp *http.Response, w http.ResponseWriter) { | |
| } | ||
| } | ||
|
|
||
| // writeSingleJSONRPCMessage applies request/response modifications to a single decoded | ||
| // JSON-RPC message and writes it to w as a plain JSON body. | ||
| func (m *mcpRequestContext) writeSingleJSONRPCMessage(ctx context.Context, w http.ResponseWriter, statusCode int, | ||
| _msg jsonrpc.Message, body []byte, req *jsonrpc.Request, backend filterapi.MCPBackend, | ||
| ) error { | ||
| var responseError error | ||
| switch msg := _msg.(type) { | ||
| case *jsonrpc.Request: | ||
| if err := m.maybeServerToClientRequestModify(ctx, msg, backend.Name); err != nil { | ||
| m.l.Error("failed to modify server->client request", slog.String("error", err.Error())) | ||
| return err | ||
| } | ||
| body, _ = jsonrpc.EncodeMessage(msg) | ||
| case *jsonrpc.Response: | ||
| if req != nil { | ||
| if err := m.maybeResponseModify(ctx, req, msg, backend.Name); err != nil { | ||
| m.l.Error("failed to modify response", slog.String("error", err.Error())) | ||
| return err | ||
| } | ||
| msg.ID = req.ID | ||
|
|
||
| // Check if this is a JSON-RPC error response | ||
| if msg.Error != nil { | ||
| responseError = msg.Error | ||
| } else if toolErr := checkToolCallError(req, msg, backend.Name); toolErr != nil { | ||
| // Check if this is a tools/call response with isError=true | ||
| responseError = toolErr | ||
| } | ||
|
|
||
| body, _ = jsonrpc.EncodeMessage(msg) | ||
| } | ||
| m.recordResponse(ctx, msg) | ||
| } | ||
|
|
||
| // We need to update the content length since we might have modified the ID. | ||
| w.Header().Set("Content-Length", fmt.Sprintf("%d", len(body))) | ||
| w.WriteHeader(statusCode) | ||
| _, _ = w.Write(body) | ||
|
|
||
| return responseError | ||
| } | ||
|
|
||
| // startsWithJSONObject reports whether the stream begins with '{', consuming only | ||
| // leading whitespace, which is insignificant to both SSE and JSON parsing. | ||
| func startsWithJSONObject(br *bufio.Reader) bool { | ||
| for { | ||
| b, err := br.ReadByte() | ||
| if err != nil { | ||
| return false | ||
| } | ||
| switch b { | ||
| case ' ', '\t', '\r', '\n': | ||
| continue | ||
| default: | ||
| _ = br.UnreadByte() | ||
| return b == '{' | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func (m *mcpRequestContext) proxyResponseBody(ctx context.Context, s *session, w http.ResponseWriter, resp *http.Response, | ||
| req *jsonrpc.Request, backend filterapi.MCPBackend, | ||
| ) error { | ||
| // Some backends (e.g. Slack MCP) send SSE data despite Content-Type: application/json. | ||
| // Try to decode as a single JSON-RPC message first; if that fails, fall through to the | ||
| // SSE parser using the already-read bytes. | ||
| var sseReader io.Reader = resp.Body | ||
| var sseReader io.Reader | ||
|
Comment on lines
-801
to
+865
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. can you explain this change?
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. sure. the |
||
| if resp.Header.Get("Content-Type") == "application/json" { | ||
| body, err := io.ReadAll(resp.Body) | ||
| if err != nil { | ||
|
|
@@ -797,48 +858,35 @@ func (m *mcpRequestContext) proxyResponseBody(ctx context.Context, s *session, w | |
| } | ||
| _msg, ok := tryDecodeJSONRPCMessage(body) | ||
| if ok { | ||
| var responseError error | ||
| switch msg := _msg.(type) { | ||
| case *jsonrpc.Request: | ||
| if err = m.maybeServerToClientRequestModify(ctx, msg, backend.Name); err != nil { | ||
| m.l.Error("failed to modify server->client request", slog.String("error", err.Error())) | ||
| return err | ||
| } | ||
| body, _ = jsonrpc.EncodeMessage(msg) | ||
| case *jsonrpc.Response: | ||
| if req != nil { | ||
| if err = m.maybeResponseModify(ctx, req, msg, backend.Name); err != nil { | ||
| m.l.Error("failed to modify response", slog.String("error", err.Error())) | ||
| return err | ||
| } | ||
| msg.ID = req.ID | ||
|
|
||
| // Check if this is a JSON-RPC error response | ||
| if msg.Error != nil { | ||
| responseError = msg.Error | ||
| } else if toolErr := checkToolCallError(req, msg, backend.Name); toolErr != nil { | ||
| // Check if this is a tools/call response with isError=true | ||
| responseError = toolErr | ||
| } | ||
|
|
||
| body, _ = jsonrpc.EncodeMessage(msg) | ||
| } | ||
| m.recordResponse(ctx, msg) | ||
| } | ||
|
|
||
| // We need to update the content length since we might have modified the ID. | ||
| w.Header().Set("Content-Length", fmt.Sprintf("%d", len(body))) | ||
| w.WriteHeader(resp.StatusCode) | ||
| _, _ = w.Write(body) | ||
|
|
||
| return responseError | ||
| return m.writeSingleJSONRPCMessage(ctx, w, resp.StatusCode, _msg, body, req, backend) | ||
| } | ||
| // Body claimed application/json but isn't valid JSON-RPC (e.g. some backends | ||
| // send SSE data despite the content type). Fall through to the SSE parser | ||
| // using the already-read body bytes since resp.Body is now drained. | ||
| m.l.Info("response Content-Type is application/json but body is not valid JSON-RPC, falling back to SSE parsing", | ||
| slog.String("backend", backend.Name)) | ||
| sseReader = bytes.NewReader(body) | ||
| } else { | ||
| // The mirror case: some backends send a plain JSON body despite Content-Type: | ||
| // text/event-stream. A valid SSE stream never starts with '{', so peek before | ||
| // handing the body to the SSE parser, which would silently drop the payload. | ||
| br := bufio.NewReader(resp.Body) | ||
| sseReader = br | ||
| if startsWithJSONObject(br) { | ||
| body, err := io.ReadAll(br) | ||
| if err != nil { | ||
| m.l.Error("failed to read response body", slog.String("error", err.Error())) | ||
| return err | ||
| } | ||
| if _msg, ok := tryDecodeJSONRPCMessage(body); ok { | ||
| m.l.Info("response Content-Type is text/event-stream but body is an unframed JSON-RPC message, handling as JSON", | ||
| slog.String("backend", backend.Name)) | ||
| w.Header().Set("Content-Type", "application/json") | ||
| w.Header().Del("Transfer-Encoding") | ||
| return m.writeSingleJSONRPCMessage(ctx, w, resp.StatusCode, _msg, body, req, backend) | ||
| } | ||
| sseReader = bytes.NewReader(body) | ||
| } | ||
| } | ||
|
|
||
| // io.Copy won't flush until the end, which doesn't happen for streaming responses. | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can use Peek here? Peek(n) returns a view of the upcoming bytes without advancing the stream.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
good point, switched the BOM check over to Peek. i left the whitespace skip
as a byte scan on purpose though, since this branch also handles real SSE
streaming. Peek blocks until it has the full count buffered (or hits EOF), so
peeking a big window could hold back the first event from flushing. just
peeking the 3 BOM bytes keeps that path cheap. can make the whole thing a
single Peek if you'd rather.