diff --git a/pkg/vmcp/client/modern.go b/pkg/vmcp/client/modern.go index d397e4cf5f..7f9f8edfd5 100644 --- a/pkg/vmcp/client/modern.go +++ b/pkg/vmcp/client/modern.go @@ -407,8 +407,9 @@ func readModernEnvelope( // with a Method) are not the response: when onNotification is non-nil their // method and params are handed to it (so a caller can relay the // notifications/message and notifications/progress the request's logLevel / -// progressToken elicited); when nil they are dropped as before. A stream that -// ends without a matching response yields errWrongEra. +// progressToken elicited); when nil they are dropped as before. Consecutive +// data: lines in one SSE event are joined with "\n" per the SSE spec before JSON +// decoding. A stream that ends without a matching response yields errWrongEra. func readModernSSE( body io.Reader, wantID int64, onNotification func(method string, params json.RawMessage), ) (json.RawMessage, *modernRPCError, error) { @@ -417,37 +418,67 @@ func readModernSSE( // data: event up to that size decodes; the outer io.LimitReader already bounds // the total, so this cannot over-allocate. sc.Buffer(make([]byte, 0, 64*1024), maxResponseSize) + + var data [][]byte + flush := func() (json.RawMessage, *modernRPCError, bool, error) { + result, rpcErr, matched, err := readModernSSEEvent(data, wantID, onNotification) + data = nil + return result, rpcErr, matched, err + } + for sc.Scan() { - data, ok := strings.CutPrefix(sc.Text(), "data:") - if !ok { - continue - } - var env modernRPCEnvelope - if json.Unmarshal([]byte(strings.TrimSpace(data)), &env) != nil { - continue - } - if env.Method != "" { - // server->client request/notification; not our response. Relay it when a - // listener is bound, otherwise drop it (historical behavior). - if onNotification != nil { - onNotification(env.Method, env.Params) + line := sc.Text() + if line == "" { + result, rpcErr, matched, err := flush() + if matched || err != nil { + return result, rpcErr, err } continue } - if !modernIDMatches(env.ID, wantID) { + value, ok := strings.CutPrefix(line, "data:") + if !ok { continue } - if env.Error == nil && len(env.Result) == 0 { - return nil, nil, errWrongEra - } - return env.Result, env.Error, nil + data = append(data, []byte(strings.TrimPrefix(value, " "))) } if err := sc.Err(); err != nil { return nil, nil, fmt.Errorf("%w: reading SSE stream: %w", errModernTransient, err) } + result, rpcErr, matched, err := flush() + if matched || err != nil { + return result, rpcErr, err + } return nil, nil, errWrongEra } +func readModernSSEEvent( + data [][]byte, wantID int64, onNotification func(method string, params json.RawMessage), +) (json.RawMessage, *modernRPCError, bool, error) { + if len(data) == 0 { + return nil, nil, false, nil + } + + var env modernRPCEnvelope + if json.Unmarshal(bytes.Join(data, []byte("\n")), &env) != nil { + return nil, nil, false, nil + } + if env.Method != "" { + // Server->client request/notification; not our response. Relay it when a + // listener is bound, otherwise drop it (historical behavior). + if onNotification != nil { + onNotification(env.Method, env.Params) + } + return nil, nil, false, nil + } + if !modernIDMatches(env.ID, wantID) { + return nil, nil, false, nil + } + if env.Error == nil && len(env.Result) == 0 { + return nil, nil, true, errWrongEra + } + return env.Result, env.Error, true, nil +} + // modernIDMatches reports whether the raw JSON id equals wantID. func modernIDMatches(raw json.RawMessage, wantID int64) bool { if len(raw) == 0 { diff --git a/pkg/vmcp/client/modern_test.go b/pkg/vmcp/client/modern_test.go index 5aef8b7d24..90fe91dea6 100644 --- a/pkg/vmcp/client/modern_test.go +++ b/pkg/vmcp/client/modern_test.go @@ -166,34 +166,70 @@ func TestModernCall_Decode(t *testing.T) { assert.Equal(t, []string{"2026-07-28"}, out.SupportedVersions) } -// TestModernCall_SSEResponse verifies the dual-body reader handles a -// text/event-stream response, ignoring interleaved notifications and returning +// TestModernCall_SSEResponse verifies the dual-body reader handles +// text/event-stream responses, ignoring interleaved notifications and returning // the final matching JSON-RPC response. func TestModernCall_SSEResponse(t *testing.T) { t.Parallel() - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Echo the request's JSON-RPC id: modernCall's SSE reader matches the - // response frame by id (unlike the JSON path), and the id comes from a - // shared counter, so it cannot be hardcoded. - var req struct { - ID json.RawMessage `json:"id"` - } - body, _ := io.ReadAll(r.Body) - require.NoError(t, json.Unmarshal(body, &req)) + tests := []struct { + name string + writeEvent func(t *testing.T, w http.ResponseWriter, id json.RawMessage) + }{ + { + name: "single-line response after notification", + writeEvent: func(_ *testing.T, w http.ResponseWriter, id json.RawMessage) { + _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{}}\n\n")) + _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"id\":" + string(id) + + ",\"result\":{\"resultType\":\"complete\",\"ok\":true}}\n\n")) + }, + }, + { + name: "multi-line data response", + writeEvent: func(_ *testing.T, w http.ResponseWriter, id json.RawMessage) { + _, _ = w.Write([]byte("event: message\n")) + _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"id\":" + string(id) + ",\n")) + _, _ = w.Write([]byte("data: \"result\":{\"resultType\":\"complete\",\"ok\":true}}\n\n")) + }, + }, + { + name: "CRLF response without trailing blank line", + writeEvent: func(_ *testing.T, w http.ResponseWriter, id json.RawMessage) { + _, _ = w.Write([]byte("event: message\r\n")) + _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\"}\r\n\r\n")) + _, _ = w.Write([]byte("event: message\r\n")) + _, _ = w.Write([]byte("data:{\"jsonrpc\":\"2.0\",\"id\":" + string(id) + + ",\"result\":{\"resultType\":\"complete\",\"ok\":true}}")) + }, + }, + } - w.Header().Set("Content-Type", "text/event-stream") - // A progress notification (has "method", no id match) then the response. - _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{}}\n\n")) - _, _ = w.Write([]byte("data: {\"jsonrpc\":\"2.0\",\"id\":" + string(req.ID) + - ",\"result\":{\"resultType\":\"complete\",\"ok\":true}}\n\n")) - })) - t.Cleanup(srv.Close) + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() - var out map[string]any - require.NoError(t, modernCall(context.Background(), srv.Client(), srv.URL, "server/discover", nil, "", nil, &out, "", nil)) - assert.Equal(t, "complete", out["resultType"]) - assert.Equal(t, true, out["ok"]) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Echo the request's JSON-RPC id: modernCall's SSE reader matches the + // response frame by id (unlike the JSON path), and the id comes from a + // shared counter, so it cannot be hardcoded. + var req struct { + ID json.RawMessage `json:"id"` + } + body, _ := io.ReadAll(r.Body) + require.NoError(t, json.Unmarshal(body, &req)) + + w.Header().Set("Content-Type", "text/event-stream") + tt.writeEvent(t, w, req.ID) + })) + t.Cleanup(srv.Close) + + var out map[string]any + require.NoError(t, modernCall(context.Background(), srv.Client(), srv.URL, "server/discover", nil, "", nil, &out, "", nil)) + assert.Equal(t, "complete", out["resultType"]) + assert.Equal(t, true, out["ok"]) + }) + } } // TestModernCall_LogLevelMeta verifies the logLevel argument is overlaid onto