fix(mcp): sanitize non-finite floats before SDK marshal (#3866)

Signed-off-by: Tushar Verma <tusharmyself06@gmail.com>

pkg/mcp/sanitize_test.go

@@ -0,0 +1,88 @@
+package mcp
+
+import (
+	"encoding/json"
+	"math"
+	"testing"
+)
+
+// TestSanitizeNonFiniteFloatsAssetResponseMarshals reproduces the integration
+// failure (TestMCPAssetVsHTTP: "marshaling output: json: unsupported value:
+// NaN") and verifies the sanitizer fixes it: encoding/json must reject the
+// response before sanitization and accept it after, with non-finite floats
+// zeroed and finite ones preserved.
+func TestSanitizeNonFiniteFloatsAssetResponseMarshals(t *testing.T) {
+	usedBytes := math.NaN()
+	resp := &AssetResponse{
+		Assets: map[string]*AssetSet{
+			"assets": &AssetSet{
+				Name: "assets",
+				Assets: []*Asset{
+					&Asset{
+						Type:          "Node",
+						Minutes:       math.NaN(),
+						Adjustment:    math.Inf(1),
+						TotalCost:     math.Inf(-1),
+						CPUCost:       math.NaN(),
+						GPUCost:       5.0, // finite, must be preserved
+						ByteHoursUsed: &usedBytes,
+						Overhead:      &NodeOverhead{OverheadCostFraction: math.NaN()},
+						CPUBreakdown:  &AssetBreakdown{Idle: math.NaN()},
+					},
+				},
+			},
+		},
+	}
+
+	if _, err := json.Marshal(resp); err == nil {
+		t.Fatal("expected json.Marshal to fail before sanitization (NaN/Inf present)")
+	}
+
+	resp = sanitizeNonFiniteFloats(resp).(*AssetResponse)
+
+	if _, err := json.Marshal(resp); err != nil {
+		t.Fatalf("expected json.Marshal to succeed after sanitization, got %v", err)
+	}
+
+	a := resp.Assets["assets"].Assets[0]
+	if a.Minutes != 0 || a.Adjustment != 0 || a.TotalCost != 0 || a.CPUCost != 0 {
+		t.Fatalf("expected non-finite base floats zeroed, got %+v", a)
+	}
+	if a.GPUCost != 5.0 {
+		t.Fatalf("expected finite GPUCost preserved, got %v", a.GPUCost)
+	}
+	if a.ByteHoursUsed == nil || *a.ByteHoursUsed != 0 {
+		t.Fatalf("expected non-finite *float64 zeroed, got %v", a.ByteHoursUsed)
+	}
+	if a.Overhead.OverheadCostFraction != 0 {
+		t.Fatalf("expected nested overhead fraction zeroed, got %v", a.Overhead.OverheadCostFraction)
+	}
+	if a.CPUBreakdown.Idle != 0 {
+		t.Fatalf("expected nested breakdown value zeroed, got %v", a.CPUBreakdown.Idle)
+	}
+}
+
+// TestSanitizeNonFiniteFloatsValueType verifies a non-pointer (value) input is
+// sanitized via the returned copy, not just pointers.
+func TestSanitizeNonFiniteFloatsValueType(t *testing.T) {
+	in := Asset{TotalCost: math.NaN(), GPUCost: 3.0}
+	out, ok := sanitizeNonFiniteFloats(in).(Asset)
+	if !ok {
+		t.Fatalf("expected Asset back, got %T", sanitizeNonFiniteFloats(in))
+	}
+	if out.TotalCost != 0 {
+		t.Fatalf("expected NaN zeroed in returned value, got %v", out.TotalCost)
+	}
+	if out.GPUCost != 3.0 {
+		t.Fatalf("expected finite value preserved, got %v", out.GPUCost)
+	}
+}
+
+func TestSanitizeNonFiniteFloatsNilSafe(t *testing.T) {
+	if got := sanitizeNonFiniteFloats(nil); got != nil {
+		t.Fatalf("expected nil, got %v", got)
+	}
+	var p *AssetResponse
+	sanitizeNonFiniteFloats(p)
+	sanitizeNonFiniteFloats(&AssetResponse{})
+}

pkg/mcp/server.go

@@ -5,6 +5,8 @@ import (
 	"crypto/rand"
 	"encoding/hex"
 	"fmt"
+	"math"
+	"reflect"
 	"strings"
 	"sync"
 	"time"
@@ -410,6 +412,12 @@ func (s *MCPServer) ProcessMCPRequest(ctx context.Context, request *MCPRequest)
 		return nil, err
 	}
 
+	// The MCP SDK marshals tool output with encoding/json, which errors on
+	// non-finite floats. Upstream cost calculations can yield NaN or Inf (e.g.
+	// a 0/0 breakdown or overhead fraction), so scrub them before they reach the
+	// SDK and fail the whole tool call.
+	data = sanitizeNonFiniteFloats(data)
+
 	processingTime := time.Since(queryStart)
 
 	// 3. Construct Final Response
@@ -424,6 +432,68 @@ func (s *MCPServer) ProcessMCPRequest(ctx context.Context, request *MCPRequest)
 	return mcpResponse, nil
 }
 
+// sanitizeNonFiniteFloats returns v with every non-finite float (NaN, +Inf,
+// -Inf) replaced by 0 so the value can be marshaled by encoding/json, which the
+// MCP SDK uses and which rejects non-finite floats. Callers must use the return
+// value, since value-type inputs are sanitized on a copy. It is best-effort:
+// any reflection panic is recovered and the original value returned unchanged.
+func sanitizeNonFiniteFloats(v any) (out any) {
+	out = v
+	defer func() {
+		if r := recover(); r != nil {
+			log.Warnf("mcp: sanitizeNonFiniteFloats recovered: %v", r)
+			out = v
+		}
+	}()
+	if v == nil {
+		return nil
+	}
+	// Work on an addressable copy so value-type inputs are sanitized too, not
+	// only pointers. For a pointer input this copies the pointer and mutates the
+	// pointed-to value in place; for a value input it yields a sanitized copy.
+	box := reflect.New(reflect.TypeOf(v))
+	box.Elem().Set(reflect.ValueOf(v))
+	sanitizeFloatsValue(box.Elem())
+	return box.Elem().Interface()
+}
+
+func sanitizeFloatsValue(v reflect.Value) {
+	switch v.Kind() {
+	case reflect.Ptr, reflect.Interface:
+		if !v.IsNil() {
+			sanitizeFloatsValue(v.Elem())
+		}
+	case reflect.Struct:
+		for i := 0; i < v.NumField(); i++ {
+			sanitizeFloatsValue(v.Field(i))
+		}
+	case reflect.Slice, reflect.Array:
+		for i := 0; i < v.Len(); i++ {
+			sanitizeFloatsValue(v.Index(i))
+		}
+	case reflect.Map:
+		for _, key := range v.MapKeys() {
+			elem := v.MapIndex(key)
+			// Map elements aren't addressable. Pointer/interface/slice/map
+			// values are mutated in place by recursing; value-type elements
+			// (e.g. a float or struct) must be rebuilt and reassigned.
+			switch elem.Kind() {
+			case reflect.Ptr, reflect.Interface, reflect.Slice, reflect.Map:
+				sanitizeFloatsValue(elem)
+			default:
+				tmp := reflect.New(elem.Type()).Elem()
+				tmp.Set(elem)
+				sanitizeFloatsValue(tmp)
+				v.SetMapIndex(key, tmp)
+			}
+		}
+	case reflect.Float32, reflect.Float64:
+		if v.CanSet() && (math.IsNaN(v.Float()) || math.IsInf(v.Float(), 0)) {
+			v.SetFloat(0)
+		}
+	}
+}
+
 // validate is the singleton validator instance.
 var validate = validator.New()