From 91c885b5b0a0eb5d4ca887a74730a0c87a4acd85 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Sat, 7 Jun 2025 20:41:57 +0800
Subject: [PATCH] chore: add more unit tests for mcp (#4928)

---
 mcp/server_test.go | 14 ++++++++++++++
 mcp/types.go       | 13 +++++--------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/mcp/server_test.go b/mcp/server_test.go
index 1d6731a0c..db718da26 100644
--- a/mcp/server_test.go
+++ b/mcp/server_test.go
@@ -175,6 +175,20 @@ func TestHandleRequest_badRequest(t *testing.T) {
 		mock.server.handleRequest(w, r)
 		assert.Equal(t, http.StatusBadRequest, w.Code)
 	})
+
+	t.Run("bad id", func(t *testing.T) {
+		mock := newMockMcpServer(t)
+		defer mock.shutdown()
+
+		addTestClient(mock.server, "test-session", true)
+
+		body := `{"jsonrpc": "2.0", "id": {}, "method": "tools.call", "params": {}}`
+		r := httptest.NewRequest(http.MethodPost, "/?session_id=test-session", bytes.NewReader([]byte(body)))
+		w := httptest.NewRecorder()
+		mock.server.handleRequest(w, r)
+		assert.Equal(t, http.StatusBadRequest, w.Code)
+		assert.Contains(t, w.Body.String(), "Invalid request.ID")
+	})
 }
 
 func TestRegisterTool(t *testing.T) {
diff --git a/mcp/types.go b/mcp/types.go
index 5a056313f..86119e1c6 100644
--- a/mcp/types.go
+++ b/mcp/types.go
@@ -22,23 +22,20 @@ type Request struct {
 }
 
 func (r Request) isNotification() (bool, error) {
-	var isNotification bool
 	switch val := r.ID.(type) {
 	case int:
-		isNotification = val == 0
+		return val == 0, nil
 	case int64:
-		isNotification = val == 0
+		return val == 0, nil
 	case float64:
-		isNotification = val == 0.0
+		return val == 0.0, nil
 	case string:
-		isNotification = len(val) == 0
+		return len(val) == 0, nil
 	case nil:
-		isNotification = true
+		return true, nil
 	default:
 		return false, fmt.Errorf("invalid type %T", val)
 	}
-
-	return isNotification, nil
 }
 
 type PaginatedParams struct {