fix(mcp): preserve valid numeric request IDs

Author: gnanam1990

internal/mcp/client.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"os"
 	"os/exec"
 	"strconv"
@@ -554,7 +555,7 @@ func rpcIDMatches(value any, id int) bool {
 }
 
 // jsonRPCIDEchoable reports whether id is a valid JSON-RPC 2.0 identifier type
-// (string, integer, or float with no fractional part) that is safe to echo back.
+// (string or finite number) that is safe to echo back.
 func jsonRPCIDEchoable(id any) bool {
 	if id == nil {
 		return false
@@ -565,9 +566,13 @@ func jsonRPCIDEchoable(id any) bool {
 	case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
 		return true
 	case float64:
-		return v == float64(int64(v))
+		return !math.IsNaN(v) && !math.IsInf(v, 0)
 	case json.Number:
-		_, err := v.Int64()
+		parsed, err := v.Float64()
+		if err != nil || math.IsNaN(parsed) || math.IsInf(parsed, 0) {
+			return false
+		}
+		_, err = json.Marshal(v)
 		return err == nil
 	default:
 		return false

internal/mcp/client_test.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -1005,6 +1006,72 @@ func TestStdioClientDropsInvalidServerRequestIDs(t *testing.T) {
 	}
 }
 
+func TestStdioClientRepliesToValidNonIntegerServerRequestIDs(t *testing.T) {
+	requests := []struct {
+		wire string
+		want float64
+	}{
+		{wire: `{"jsonrpc":"2.0","id":1.5,"method":"roots/list","params":{}}` + "\n", want: 1.5},
+		{wire: `{"jsonrpc":"2.0","id":2e2,"method":"roots/list","params":{}}` + "\n", want: 200},
+	}
+	for _, request := range requests {
+		t.Run(fmt.Sprint(request.want), func(t *testing.T) {
+			inReader, inWriter := io.Pipe()
+			outReader, outWriter := io.Pipe()
+			client := &Client{
+				reader:  newMessageReader(inReader),
+				writer:  newMessageWriter(outWriter),
+				pending: make(map[int]chan dispatchResult),
+			}
+			t.Cleanup(func() {
+				_ = inWriter.Close()
+				_ = outReader.Close()
+			})
+			client.ensureReader()
+			if _, err := inWriter.Write([]byte(request.wire)); err != nil {
+				t.Fatalf("write server request: %v", err)
+			}
+			result := make(chan dispatchResult, 1)
+			go func() {
+				response, err := newMessageReader(outReader).read()
+				result <- dispatchResult{message: response, err: err}
+			}()
+			select {
+			case response := <-result:
+				if response.err != nil {
+					t.Fatalf("read method-not-found response: %v", response.err)
+				}
+				id, ok := response.message.ID.(float64)
+				if !ok || id != request.want || response.message.Error == nil || response.message.Error.Code != -32601 {
+					t.Fatalf("response = %#v, want id %v and error -32601", response.message, request.want)
+				}
+			case <-time.After(time.Second):
+				t.Fatal("timed out waiting for method-not-found response")
+			}
+		})
+	}
+}
+
+func TestJSONRPCIDEchoableAcceptsFiniteJSONNumbers(t *testing.T) {
+	for _, test := range []struct {
+		name string
+		id   any
+		want bool
+	}{
+		{name: "fractional float", id: 1.5, want: true},
+		{name: "exponent json number", id: json.Number("2e2"), want: true},
+		{name: "not a number", id: math.NaN(), want: false},
+		{name: "positive infinity", id: math.Inf(1), want: false},
+		{name: "invalid json number", id: json.Number("not-a-number"), want: false},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			if got := jsonRPCIDEchoable(test.id); got != test.want {
+				t.Fatalf("jsonRPCIDEchoable(%v) = %v, want %v", test.id, got, test.want)
+			}
+		})
+	}
+}
+
 type gatedCaptureWriter struct {
 	started     chan struct{}
 	release     chan struct{}

internal/mcp/hang_test.go

@@ -62,6 +62,15 @@ type blockingReader struct {
 	release chan struct{}
 }
 
+type signalingWriter struct {
+	writes chan struct{}
+}
+
+func (writer *signalingWriter) Write(p []byte) (int, error) {
+	writer.writes <- struct{}{}
+	return len(p), nil
+}
+
 func newBlockingReader() *blockingReader {
 	return &blockingReader{release: make(chan struct{})}
 }
@@ -86,10 +95,11 @@ func (reader *blockingReader) Close() error {
 func TestClientRequestUnblocksOnContextCancel(t *testing.T) {
 	reader := newBlockingReader()
 	defer reader.Close()
+	writes := make(chan struct{}, 2)
 
 	client := &Client{
 		reader: newMessageReader(reader),
-		writer: newMessageWriter(io.Discard),
+		writer: newMessageWriter(&signalingWriter{writes: writes}),
 		nextID: 1,
 	}
 
@@ -99,6 +109,11 @@ func TestClientRequestUnblocksOnContextCancel(t *testing.T) {
 		done <- client.request(ctx, "tools/list", map[string]any{}, nil)
 	}()
 
+	select {
+	case <-writes:
+	case <-time.After(time.Second):
+		t.Fatal("first request did not reach the transport")
+	}
 	// The request is now parked waiting for a response that never comes.
 	cancel()
 
@@ -111,15 +126,21 @@ func TestClientRequestUnblocksOnContextCancel(t *testing.T) {
 		t.Fatal("request() hung on a non-responsive server")
 	}
 
-	// The lock must be free: a second request under an already-cancelled
-	// context must return immediately rather than block.
-	cancelled, cancel2 := context.WithCancel(context.Background())
-	cancel2()
+	// The shared state must be free: prove a second live request reaches the
+	// transport, then cancel it and verify cancellation releases the caller.
+	secondCtx, cancel2 := context.WithCancel(context.Background())
+	defer cancel2()
 	second := make(chan error, 1)
 	go func() {
-		second <- client.request(cancelled, "tools/list", map[string]any{}, nil)
+		second <- client.request(secondCtx, "tools/list", map[string]any{}, nil)
 	}()
 	select {
+	case <-writes:
+	case <-time.After(time.Second):
+		t.Fatal("second request did not reach the transport")
+	}
+	cancel2()
+	select {
 	case err := <-second:
 		if !errors.Is(err, context.Canceled) {
 			t.Fatalf("second request() error = %v, want context.Canceled", err)