fix(audit): detect JSON-RPC errors in responses larger than the detection buffer

pkg/audit/auditor.go

@@ -464,7 +464,13 @@ func (*Auditor) detectApplicationError(rw *responseWriter) *mcp.ParsedMCPRespons
 		prefix = rw.errorDetectionBody.Bytes()
 	}
 	if len(prefix) > 0 && prefix[0] == '{' {
-		return mcp.ParseMCPResponse(prefix)
+		parsed := mcp.ParseMCPResponse(prefix)
+		// A full buffer may be a cut-off response, which does not unmarshal.
+		// Fall back to a token walk that only needs the bytes before the cut.
+		if !parsed.HasError && len(prefix) >= errorDetectionBufferSize {
+			return mcp.ParseTruncatedMCPResponse(prefix)
+		}
+		return parsed
 	}
 	return nil
 }

pkg/audit/auditor_test.go

@@ -1253,6 +1253,63 @@ func TestMiddlewareDetectsJSONRPCErrors(t *testing.T) {
 		assert.Contains(t, logOutput, "jsonrpc_error_code")
 	})
 
+	t.Run("detects JSON-RPC error in a body larger than the detection buffer", func(t *testing.T) {
+		t.Parallel()
+		// Bodies of 511, 512 and 513+ bytes: the old prefix parse missed
+		// everything past the buffer size.
+		for _, size := range []int{errorDetectionBufferSize - 1, errorDetectionBufferSize, errorDetectionBufferSize + 1, 4 * errorDetectionBufferSize} {
+			var logBuf bytes.Buffer
+			detectErrors := true
+			config := &Config{DetectApplicationErrors: &detectErrors}
+			auditor, err := NewAuditorWithTransport(config, "streamable-http")
+			require.NoError(t, err)
+			auditor.auditLogger = NewAuditLogger(&logBuf)
+
+			const envelope = `{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":""}}`
+			errorResponse := `{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":"` +
+				strings.Repeat("a", size-len(envelope)) + `"}}`
+			require.Len(t, errorResponse, size)
+			handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+				w.WriteHeader(http.StatusOK)
+				_, err := w.Write([]byte(errorResponse))
+				require.NoError(t, err)
+			})
+
+			req := httptest.NewRequest("POST", "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":"1","method":"tools/call","params":{"name":"test"}}`))
+			req.Header.Set("Content-Type", "application/json")
+			rr := httptest.NewRecorder()
+			auditor.Middleware(handler).ServeHTTP(rr, req)
+
+			assert.Equal(t, errorResponse, rr.Body.String(), "size %d", size)
+			assert.Contains(t, logBuf.String(), OutcomeApplicationError, "size %d", size)
+			assert.Contains(t, logBuf.String(), "-32000", "size %d", size)
+		}
+	})
+
+	t.Run("keeps outcome=success for a large JSON-RPC result", func(t *testing.T) {
+		t.Parallel()
+		var logBuf bytes.Buffer
+		detectErrors := true
+		config := &Config{DetectApplicationErrors: &detectErrors}
+		auditor, err := NewAuditorWithTransport(config, "streamable-http")
+		require.NoError(t, err)
+		auditor.auditLogger = NewAuditLogger(&logBuf)
+
+		response := `{"jsonrpc":"2.0","id":1,"result":{"error":{"code":1,"message":"x"},"text":"` +
+			strings.Repeat("a", 4*errorDetectionBufferSize) + `"}}`
+		handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			w.WriteHeader(http.StatusOK)
+			_, err := w.Write([]byte(response))
+			require.NoError(t, err)
+		})
+
+		req := httptest.NewRequest("POST", "/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":"1","method":"tools/call","params":{"name":"test"}}`))
+		req.Header.Set("Content-Type", "application/json")
+		auditor.Middleware(handler).ServeHTTP(httptest.NewRecorder(), req)
+
+		assert.NotContains(t, logBuf.String(), OutcomeApplicationError)
+	})
+
 	t.Run("keeps outcome=success for valid JSON-RPC result", func(t *testing.T) {
 		t.Parallel()
 		var logBuf bytes.Buffer

pkg/mcp/response.go

@@ -4,6 +4,7 @@
 package mcp
 
 import (
+	"bytes"
 	"encoding/json"
 )
 
@@ -60,3 +61,72 @@
 		ErrorMessage: envelope.Error.Message,
 	}
 }
+
+// ParseTruncatedMCPResponse detects a top-level JSON-RPC "error" object in a
+// response body that may have been cut off, such as a fixed-size prefix of a
+// larger response. ParseMCPResponse cannot do this because a truncated JSON
+// document does not unmarshal.
+//
+// It walks the top-level object token by token and stops at the "error" key,
+// so only the bytes before the cut matter. "code" and "message" are filled in
+// when they were read completely; otherwise they keep their zero values and
+// HasError is still true. Values of other top-level keys are skipped, and
+// nested "error" keys are ignored. Like ParseMCPResponse it is lenient and
+// returns HasError=false when the body is not a JSON object or has no
+// top-level "error" object before the cut.
+func ParseTruncatedMCPResponse(body []byte) *ParsedMCPResponse {
+	dec := json.NewDecoder(bytes.NewReader(body))
+	if tok, err := dec.Token(); err != nil || tok != json.Delim('{') {
+		return &ParsedMCPResponse{}
+	}
+	for dec.More() {
+		keyTok, err := dec.Token()
+		if err != nil {
+			return &ParsedMCPResponse{}
+		}
+		key, ok := keyTok.(string)
+		if !ok {
+			return &ParsedMCPResponse{}
+		}
+		if key != "error" {
+			var skip json.RawMessage
+			if err := dec.Decode(&skip); err != nil {
+				return &ParsedMCPResponse{}
+			}
+			continue
+		}
+		if tok, err := dec.Token(); err != nil || tok != json.Delim('{') {
+			// "error": null or a non-object value is not a JSON-RPC error.
+			return &ParsedMCPResponse{}
+		}
+		res := &ParsedMCPResponse{HasError: true}
+		for dec.More() {
+			fieldTok, err := dec.Token()
+			if err != nil {
+				return res
+			}
+			field, _ := fieldTok.(string)
+			switch field {
+			case "code":
+				var code int
+				if err := dec.Decode(&code); err != nil {
+					return res
+				}
+				res.ErrorCode = code
+			case "message":
+				var msg string
+				if err := dec.Decode(&msg); err != nil {
+					return res
+				}
+				res.ErrorMessage = msg
+			default:
+				var skip json.RawMessage
+				if err := dec.Decode(&skip); err != nil {
+					return res
+				}
+			}
+		}
+		return res
+	}
+	return &ParsedMCPResponse{}
+}

pkg/mcp/response_test.go

@@ -90,3 +90,89 @@ func TestParseMCPResponse(t *testing.T) {
 		})
 	}
 }
+
+func TestParseTruncatedMCPResponse(t *testing.T) {
+	t.Parallel()
+
+	full := `{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"message":"` + strings.Repeat("a", 600) + `"}}`
+
+	tests := []struct {
+		name         string
+		body         string
+		wantHasError bool
+		wantCode     int
+		wantMessage  string
+	}{
+		{
+			name:         "complete error",
+			body:         `{"jsonrpc":"2.0","id":1,"error":{"code":-32603,"message":"boom"}}`,
+			wantHasError: true,
+			wantCode:     -32603,
+			wantMessage:  "boom",
+		},
+		{
+			name:         "cut inside message string keeps code",
+			body:         full[:200],
+			wantHasError: true,
+			wantCode:     -32000,
+		},
+		{
+			name:         "cut right after code",
+			body:         `{"jsonrpc":"2.0","id":1,"error":{"code":-32000,`,
+			wantHasError: true,
+			wantCode:     -32000,
+		},
+		{
+			name:         "cut inside the error key name",
+			body:         `{"jsonrpc":"2.0","id":1,"err`,
+			wantHasError: false,
+		},
+		{
+			name:         "cut right after the error key opens",
+			body:         `{"jsonrpc":"2.0","id":1,"error":{`,
+			wantHasError: true,
+		},
+		{
+			name:         "data before message is skipped",
+			body:         `{"jsonrpc":"2.0","id":1,"error":{"code":-32000,"data":{"a":[1,2]},"message":"late"}}`,
+			wantHasError: true,
+			wantCode:     -32000,
+			wantMessage:  "late",
+		},
+		{
+			name:         "truncated result is not an error",
+			body:         `{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"` + strings.Repeat("a", 600),
+			wantHasError: false,
+		},
+		{
+			name:         "nested error key is ignored",
+			body:         `{"jsonrpc":"2.0","id":1,"result":{"error":{"code":1,"message":"x"},"pad":"` + strings.Repeat("a", 600),
+			wantHasError: false,
+		},
+		{
+			name:         "null error",
+			body:         `{"jsonrpc":"2.0","id":1,"error":null}`,
+			wantHasError: false,
+		},
+		{
+			name:         "not an object",
+			body:         `[{"error":{"code":1}}`,
+			wantHasError: false,
+		},
+		{
+			name:         "empty",
+			body:         "",
+			wantHasError: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			got := ParseTruncatedMCPResponse([]byte(tt.body))
+			assert.Equal(t, tt.wantHasError, got.HasError)
+			assert.Equal(t, tt.wantCode, got.ErrorCode)
+			assert.Equal(t, tt.wantMessage, got.ErrorMessage)
+		})
+	}
+}