From 4b7ac63319356fe514f21c51a2519ef7cd735fa3 Mon Sep 17 00:00:00 2001 From: Deepam02 <116721751+Deepam02@users.noreply.github.com> Date: Sat, 4 Oct 2025 10:18:58 +0530 Subject: [PATCH 1/3] Add tool error detection to telemetry middleware Closes #2084 Signed-off-by: Deepam02 <116721751+Deepam02@users.noreply.github.com> --- pkg/telemetry/integration_test.go | 40 +++++++++++++++++++++++++++++++ pkg/telemetry/middleware.go | 33 ++++++++++++++++++++++--- pkg/telemetry/middleware_test.go | 25 +++++++++++++++++++ 3 files changed, 95 insertions(+), 3 deletions(-) diff --git a/pkg/telemetry/integration_test.go b/pkg/telemetry/integration_test.go index e65167f19..68bfef9cb 100644 --- a/pkg/telemetry/integration_test.go +++ b/pkg/telemetry/integration_test.go @@ -483,3 +483,43 @@ func TestTelemetryIntegration_MultipleRequests(t *testing.T) { assert.Contains(t, metricsBody, "toolhive_mcp_requests") assert.Contains(t, metricsBody, `server="multi-test"`) } + +func TestTelemetryIntegration_ToolErrorDetection(t *testing.T) { + // Setup test providers + exporter := tracetest.NewInMemoryExporter() + tracerProvider := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) + meterProvider := sdkmetric.NewMeterProvider() + + config := Config{ServiceName: "test", ServiceVersion: "1.0.0"} + middleware := NewHTTPMiddleware(config, tracerProvider, meterProvider, "test", "stdio") + + // Test tool call with error + testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`{"result":{"isError":true}}`)) + }) + + mcpRequest := &mcp.ParsedMCPRequest{Method: "tools/call", ID: "test", IsRequest: true} + req := httptest.NewRequest("POST", "/messages", nil) + ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, mcpRequest) + req = req.WithContext(ctx) + + rec := httptest.NewRecorder() + middleware(testHandler).ServeHTTP(rec, req) + + // Verify span has error attribute + tracerProvider.ForceFlush(ctx) + spans := exporter.GetSpans() + require.Len(t, spans, 1) + + span := spans[0] + assert.Equal(t, "mcp.tools/call", span.Name) + + // Check for tool error attribute + for _, attr := range span.Attributes { + if attr.Key == "mcp.tool.error" { + assert.True(t, attr.Value.AsBool()) + return + } + } + t.Error("Expected mcp.tool.error attribute not found") +} diff --git a/pkg/telemetry/middleware.go b/pkg/telemetry/middleware.go index 19e943bc3..66a27a383 100644 --- a/pkg/telemetry/middleware.go +++ b/pkg/telemetry/middleware.go @@ -1,6 +1,7 @@ package telemetry import ( + "bytes" "context" "encoding/json" "fmt" @@ -130,6 +131,7 @@ func (m *HTTPMiddleware) Handler(next http.Handler) http.Handler { ResponseWriter: w, statusCode: http.StatusOK, bytesWritten: 0, + isToolCall: mcpparser.GetMCPMethod(ctx) == string(mcp.MethodToolsCall), } // Add HTTP attributes @@ -390,19 +392,36 @@ func (*HTTPMiddleware) finalizeSpan(span trace.Span, rw *responseWriter, duratio attribute.Float64("http.duration_ms", float64(duration.Nanoseconds())/1e6), ) - // Set span status based on HTTP status code + // Add MCP tool error indicator if detected + if rw.isToolCall { + span.SetAttributes(attribute.Bool("mcp.tool.error", rw.hasToolError)) + } + + // Set span status based on HTTP status code AND MCP tool errors if rw.statusCode >= 400 { span.SetStatus(codes.Error, fmt.Sprintf("HTTP %d", rw.statusCode)) + } else if rw.hasToolError { + span.SetStatus(codes.Error, "MCP tool execution error") } else { span.SetStatus(codes.Ok, "") } } +// detectMCPToolError performs lightweight detection of MCP tool execution errors +// Returns true if the response likely contains a tool execution error +func detectMCPToolError(data []byte) bool { + // Quick byte search for common error patterns + return bytes.Contains(data, []byte(`"isError":true`)) || + bytes.Contains(data, []byte(`"isError": true`)) +} + // responseWriter wraps http.ResponseWriter to capture response details. type responseWriter struct { http.ResponseWriter statusCode int bytesWritten int64 + hasToolError bool // tracks if MCP tool execution error is detected + isToolCall bool // tracks if this is a tools/call request } // WriteHeader captures the status code. @@ -411,10 +430,16 @@ func (rw *responseWriter) WriteHeader(statusCode int) { rw.ResponseWriter.WriteHeader(statusCode) } -// Write captures the number of bytes written. +// Write captures the number of bytes written and detects tool errors. func (rw *responseWriter) Write(data []byte) (int, error) { n, err := rw.ResponseWriter.Write(data) rw.bytesWritten += int64(n) + + // Only check for tool errors if this is a tool call and we haven't found one yet + if rw.isToolCall && !rw.hasToolError { + rw.hasToolError = detectMCPToolError(data) + } + return n, err } @@ -426,10 +451,12 @@ func (m *HTTPMiddleware) recordMetrics(ctx context.Context, r *http.Request, rw mcpMethod = "unknown" } - // Determine status (success/error) + // Determine status (success/error/tool_error) status := "success" if rw.statusCode >= 400 { status = "error" + } else if rw.hasToolError { + status = "tool_error" } // Common attributes for all metrics diff --git a/pkg/telemetry/middleware_test.go b/pkg/telemetry/middleware_test.go index 5a7a36822..fc2e9da72 100644 --- a/pkg/telemetry/middleware_test.go +++ b/pkg/telemetry/middleware_test.go @@ -1491,3 +1491,28 @@ func TestFactoryMiddleware_Integration(t *testing.T) { assert.NoError(t, err) }) } + +func TestDetectMCPToolError(t *testing.T) { + assert.False(t, detectMCPToolError([]byte(`{"result":{"isError":false}}`))) + assert.True(t, detectMCPToolError([]byte(`{"result":{"isError":true}}`))) + assert.False(t, detectMCPToolError([]byte(`{"result":{"content":"test"}}`))) +} + +func TestResponseWriter_ToolErrorDetection(t *testing.T) { + rec := httptest.NewRecorder() + + // Tool call with error + rw := &responseWriter{ResponseWriter: rec, isToolCall: true} + rw.Write([]byte(`{"result":{"isError":true}}`)) + assert.True(t, rw.hasToolError) + + // Tool call without error + rw = &responseWriter{ResponseWriter: rec, isToolCall: true} + rw.Write([]byte(`{"result":{"isError":false}}`)) + assert.False(t, rw.hasToolError) + + // Non-tool call should not detect errors + rw = &responseWriter{ResponseWriter: rec, isToolCall: false} + rw.Write([]byte(`{"result":{"isError":true}}`)) + assert.False(t, rw.hasToolError) +} From 51ecc34ff1f4ec0cd59afcc76e5c2c5e85e84a61 Mon Sep 17 00:00:00 2001 From: Deepam02 <116721751+Deepam02@users.noreply.github.com> Date: Sat, 4 Oct 2025 10:59:48 +0530 Subject: [PATCH 2/3] telemetry: fix test parallelism and formatting Signed-off-by: Deepam02 <116721751+Deepam02@users.noreply.github.com> --- pkg/telemetry/integration_test.go | 15 ++++++++------- pkg/telemetry/middleware.go | 4 ++-- pkg/telemetry/middleware_test.go | 10 ++++++---- 3 files changed, 16 insertions(+), 13 deletions(-) diff --git a/pkg/telemetry/integration_test.go b/pkg/telemetry/integration_test.go index 68bfef9cb..4d4f47291 100644 --- a/pkg/telemetry/integration_test.go +++ b/pkg/telemetry/integration_test.go @@ -485,35 +485,36 @@ func TestTelemetryIntegration_MultipleRequests(t *testing.T) { } func TestTelemetryIntegration_ToolErrorDetection(t *testing.T) { + t.Parallel() // Setup test providers exporter := tracetest.NewInMemoryExporter() tracerProvider := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter)) meterProvider := sdkmetric.NewMeterProvider() - + config := Config{ServiceName: "test", ServiceVersion: "1.0.0"} middleware := NewHTTPMiddleware(config, tracerProvider, meterProvider, "test", "stdio") // Test tool call with error - testHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + testHandler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Write([]byte(`{"result":{"isError":true}}`)) }) - + mcpRequest := &mcp.ParsedMCPRequest{Method: "tools/call", ID: "test", IsRequest: true} req := httptest.NewRequest("POST", "/messages", nil) ctx := context.WithValue(req.Context(), mcp.MCPRequestContextKey, mcpRequest) req = req.WithContext(ctx) - + rec := httptest.NewRecorder() middleware(testHandler).ServeHTTP(rec, req) - + // Verify span has error attribute tracerProvider.ForceFlush(ctx) spans := exporter.GetSpans() require.Len(t, spans, 1) - + span := spans[0] assert.Equal(t, "mcp.tools/call", span.Name) - + // Check for tool error attribute for _, attr := range span.Attributes { if attr.Key == "mcp.tool.error" { diff --git a/pkg/telemetry/middleware.go b/pkg/telemetry/middleware.go index 66a27a383..7d44c4f62 100644 --- a/pkg/telemetry/middleware.go +++ b/pkg/telemetry/middleware.go @@ -434,12 +434,12 @@ func (rw *responseWriter) WriteHeader(statusCode int) { func (rw *responseWriter) Write(data []byte) (int, error) { n, err := rw.ResponseWriter.Write(data) rw.bytesWritten += int64(n) - + // Only check for tool errors if this is a tool call and we haven't found one yet if rw.isToolCall && !rw.hasToolError { rw.hasToolError = detectMCPToolError(data) } - + return n, err } diff --git a/pkg/telemetry/middleware_test.go b/pkg/telemetry/middleware_test.go index fc2e9da72..98920648f 100644 --- a/pkg/telemetry/middleware_test.go +++ b/pkg/telemetry/middleware_test.go @@ -1493,24 +1493,26 @@ func TestFactoryMiddleware_Integration(t *testing.T) { } func TestDetectMCPToolError(t *testing.T) { + t.Parallel() assert.False(t, detectMCPToolError([]byte(`{"result":{"isError":false}}`))) assert.True(t, detectMCPToolError([]byte(`{"result":{"isError":true}}`))) assert.False(t, detectMCPToolError([]byte(`{"result":{"content":"test"}}`))) } func TestResponseWriter_ToolErrorDetection(t *testing.T) { + t.Parallel() rec := httptest.NewRecorder() - + // Tool call with error rw := &responseWriter{ResponseWriter: rec, isToolCall: true} rw.Write([]byte(`{"result":{"isError":true}}`)) assert.True(t, rw.hasToolError) - - // Tool call without error + + // Tool call without error rw = &responseWriter{ResponseWriter: rec, isToolCall: true} rw.Write([]byte(`{"result":{"isError":false}}`)) assert.False(t, rw.hasToolError) - + // Non-tool call should not detect errors rw = &responseWriter{ResponseWriter: rec, isToolCall: false} rw.Write([]byte(`{"result":{"isError":true}}`)) From 4fcc067a807bfa476448243f90c1874d95c2e8ed Mon Sep 17 00:00:00 2001 From: Deepam02 <116721751+Deepam02@users.noreply.github.com> Date: Sun, 5 Oct 2025 00:44:48 +0530 Subject: [PATCH 3/3] Improve error detection with JSON parsing and response buffering Signed-off-by: Deepam02 <116721751+Deepam02@users.noreply.github.com> --- pkg/telemetry/middleware.go | 42 +++++++++++++++++++++++--------- pkg/telemetry/middleware_test.go | 14 +++++++++++ 2 files changed, 45 insertions(+), 11 deletions(-) diff --git a/pkg/telemetry/middleware.go b/pkg/telemetry/middleware.go index 7d44c4f62..8dae47085 100644 --- a/pkg/telemetry/middleware.go +++ b/pkg/telemetry/middleware.go @@ -1,7 +1,6 @@ package telemetry import ( - "bytes" "context" "encoding/json" "fmt" @@ -149,6 +148,9 @@ func (m *HTTPMiddleware) Handler(next http.Handler) http.Handler { // Call the next handler with the instrumented context next.ServeHTTP(rw, r.WithContext(ctx)) + // Finalize tool error detection now that response is complete + rw.finalizeToolErrorDetection() + // Record completion metrics and finalize span duration := time.Since(startTime) m.finalizeSpan(span, rw, duration) @@ -410,18 +412,26 @@ func (*HTTPMiddleware) finalizeSpan(span trace.Span, rw *responseWriter, duratio // detectMCPToolError performs lightweight detection of MCP tool execution errors // Returns true if the response likely contains a tool execution error func detectMCPToolError(data []byte) bool { - // Quick byte search for common error patterns - return bytes.Contains(data, []byte(`"isError":true`)) || - bytes.Contains(data, []byte(`"isError": true`)) + // Attempt to parse JSON and check for isError field + var resp struct { + Result struct { + IsError bool `json:"isError"` + } `json:"result"` + } + if err := json.Unmarshal(data, &resp); err != nil { + return false + } + return resp.Result.IsError } // responseWriter wraps http.ResponseWriter to capture response details. type responseWriter struct { http.ResponseWriter - statusCode int - bytesWritten int64 - hasToolError bool // tracks if MCP tool execution error is detected - isToolCall bool // tracks if this is a tools/call request + statusCode int + bytesWritten int64 + hasToolError bool // tracks if MCP tool execution error is detected + isToolCall bool // tracks if this is a tools/call request + responseBuffer []byte // buffer to collect response data for tool calls } // WriteHeader captures the status code. @@ -430,19 +440,29 @@ func (rw *responseWriter) WriteHeader(statusCode int) { rw.ResponseWriter.WriteHeader(statusCode) } -// Write captures the number of bytes written and detects tool errors. +// Write captures the number of bytes written and buffers data for tool calls. func (rw *responseWriter) Write(data []byte) (int, error) { n, err := rw.ResponseWriter.Write(data) rw.bytesWritten += int64(n) - // Only check for tool errors if this is a tool call and we haven't found one yet + // Buffer response data for tool calls to enable proper error detection if rw.isToolCall && !rw.hasToolError { - rw.hasToolError = detectMCPToolError(data) + rw.responseBuffer = append(rw.responseBuffer, data...) } return n, err } +// finalizeToolErrorDetection performs error detection on the complete buffered response. +// This should be called after the response is completely written. +func (rw *responseWriter) finalizeToolErrorDetection() { + if rw.isToolCall && !rw.hasToolError && len(rw.responseBuffer) > 0 { + rw.hasToolError = detectMCPToolError(rw.responseBuffer) + // Clear buffer to free memory + rw.responseBuffer = nil + } +} + // recordMetrics records request metrics. func (m *HTTPMiddleware) recordMetrics(ctx context.Context, r *http.Request, rw *responseWriter, duration time.Duration) { // Get MCP method from context if available diff --git a/pkg/telemetry/middleware_test.go b/pkg/telemetry/middleware_test.go index 98920648f..1da7dc745 100644 --- a/pkg/telemetry/middleware_test.go +++ b/pkg/telemetry/middleware_test.go @@ -1497,6 +1497,10 @@ func TestDetectMCPToolError(t *testing.T) { assert.False(t, detectMCPToolError([]byte(`{"result":{"isError":false}}`))) assert.True(t, detectMCPToolError([]byte(`{"result":{"isError":true}}`))) assert.False(t, detectMCPToolError([]byte(`{"result":{"content":"test"}}`))) + + // Test invalid JSON - should return false, not panic + assert.False(t, detectMCPToolError([]byte(`invalid json`))) + assert.False(t, detectMCPToolError([]byte(`{"malformed": json}`))) } func TestResponseWriter_ToolErrorDetection(t *testing.T) { @@ -1506,15 +1510,25 @@ func TestResponseWriter_ToolErrorDetection(t *testing.T) { // Tool call with error rw := &responseWriter{ResponseWriter: rec, isToolCall: true} rw.Write([]byte(`{"result":{"isError":true}}`)) + rw.finalizeToolErrorDetection() // Now we need to explicitly finalize assert.True(t, rw.hasToolError) // Tool call without error rw = &responseWriter{ResponseWriter: rec, isToolCall: true} rw.Write([]byte(`{"result":{"isError":false}}`)) + rw.finalizeToolErrorDetection() assert.False(t, rw.hasToolError) // Non-tool call should not detect errors rw = &responseWriter{ResponseWriter: rec, isToolCall: false} rw.Write([]byte(`{"result":{"isError":true}}`)) + rw.finalizeToolErrorDetection() assert.False(t, rw.hasToolError) + + // Test chunked writes (multiple Write calls) + rw = &responseWriter{ResponseWriter: rec, isToolCall: true} + rw.Write([]byte(`{"result":{"isError":`)) + rw.Write([]byte(`true}}`)) + rw.finalizeToolErrorDetection() + assert.True(t, rw.hasToolError, "Should detect error in chunked response") }