Skip to content

Commit 3aa0aea

Browse files
committed
feat(server): convert ping messages to be spec compliant
Replace string ping messages with MCP spec compliant JSON-RPC requests. - Update request handler to ignore ping responses (add `Result` field to `baseMessage` to disambiguate empty responses from other incoming messages) - Implement MCP spec compliant [ping request](https://modelcontextprotocol.io/specification/2024-11-05/basic/utilities/ping)
1 parent 71b910b commit 3aa0aea

File tree

4 files changed

+140
-2
lines changed

4 files changed

+140
-2
lines changed

server/internal/gen/request_handler.go.tmpl

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ func (s *MCPServer) HandleMessage(
2424
JSONRPC string `json:"jsonrpc"`
2525
Method mcp.MCPMethod `json:"method"`
2626
ID any `json:"id,omitempty"`
27+
Result any `json:"result,omitempty"`
2728
}
2829

2930
if err := json.Unmarshal(message, &baseMessage); err != nil {
@@ -56,6 +57,12 @@ func (s *MCPServer) HandleMessage(
5657
return nil // Return nil for notifications
5758
}
5859

60+
if baseMessage.Result != nil {
61+
// this is a response to a request sent by the server (e.g. from a ping
62+
// sent due to WithKeepAlive option)
63+
return nil
64+
}
65+
5966
switch baseMessage.Method {
6067
{{- range .}}
6168
case mcp.{{.MethodName}}:

server/request_handler.go

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

server/sse.go

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ type sseSession struct {
2323
done chan struct{}
2424
eventQueue chan string // Channel for queuing events
2525
sessionID string
26+
requestID atomic.Int64
2627
notificationChannel chan mcp.JSONRPCNotification
2728
initialized atomic.Bool
2829
}
@@ -282,8 +283,16 @@ func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) {
282283
for {
283284
select {
284285
case <-ticker.C:
285-
//: ping - 2025-03-27 07:44:38.682659+00:00
286-
session.eventQueue <- fmt.Sprintf(":ping - %s\n\n", time.Now().Format(time.RFC3339))
286+
message := mcp.JSONRPCRequest{
287+
JSONRPC: "2.0",
288+
ID: session.requestID.Add(1),
289+
Request: mcp.Request{
290+
Method: "ping",
291+
},
292+
}
293+
messageBytes, _ := json.Marshal(message)
294+
pingMsg := fmt.Sprintf("event: message\ndata:%s\n\n", messageBytes)
295+
session.eventQueue <- pingMsg
287296
case <-session.done:
288297
return
289298
case <-r.Context().Done():

server/sse_test.go

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
package server
22

33
import (
4+
"bufio"
45
"bytes"
56
"context"
67
"encoding/json"
8+
"io"
79
"fmt"
810
"math/rand"
911
"net/http"
@@ -739,4 +741,117 @@ func TestSSEServer(t *testing.T) {
739741
}
740742
}
741743
})
744+
745+
t.Run("Client receives and can respond to ping messages", func(t *testing.T) {
746+
mcpServer := NewMCPServer("test", "1.0.0")
747+
testServer := NewTestServer(mcpServer,
748+
WithKeepAlive(true),
749+
WithKeepAliveInterval(50*time.Millisecond),
750+
)
751+
defer testServer.Close()
752+
753+
sseResp, err := http.Get(fmt.Sprintf("%s/sse", testServer.URL))
754+
if err != nil {
755+
t.Fatalf("Failed to connect to SSE endpoint: %v", err)
756+
}
757+
defer sseResp.Body.Close()
758+
759+
reader := bufio.NewReader(sseResp.Body)
760+
761+
var messageURL string
762+
var pingID float64
763+
764+
for {
765+
line, err := reader.ReadString('\n')
766+
if err != nil {
767+
t.Fatalf("Failed to read SSE event: %v", err)
768+
}
769+
770+
if strings.HasPrefix(line, "event: endpoint") {
771+
dataLine, err := reader.ReadString('\n')
772+
if err != nil {
773+
t.Fatalf("Failed to read endpoint data: %v", err)
774+
}
775+
messageURL = strings.TrimSpace(strings.TrimPrefix(dataLine, "data: "))
776+
777+
_, err = reader.ReadString('\n')
778+
if err != nil {
779+
t.Fatalf("Failed to read blank line: %v", err)
780+
}
781+
}
782+
783+
if strings.HasPrefix(line, "event: message") {
784+
dataLine, err := reader.ReadString('\n')
785+
if err != nil {
786+
t.Fatalf("Failed to read message data: %v", err)
787+
}
788+
789+
pingData := strings.TrimSpace(strings.TrimPrefix(dataLine, "data:"))
790+
var pingMsg mcp.JSONRPCRequest
791+
if err := json.Unmarshal([]byte(pingData), &pingMsg); err != nil {
792+
t.Fatalf("Failed to parse ping message: %v", err)
793+
}
794+
795+
if pingMsg.Method == "ping" {
796+
pingID = pingMsg.ID.(float64)
797+
t.Logf("Received ping with ID: %f", pingID)
798+
break // We got the ping, exit the loop
799+
}
800+
801+
_, err = reader.ReadString('\n')
802+
if err != nil {
803+
t.Fatalf("Failed to read blank line: %v", err)
804+
}
805+
}
806+
807+
if messageURL != "" && pingID != 0 {
808+
break
809+
}
810+
}
811+
812+
if messageURL == "" {
813+
t.Fatal("Did not receive message endpoint URL")
814+
}
815+
816+
pingResponse := map[string]any{
817+
"jsonrpc": "2.0",
818+
"id": pingID,
819+
"result": map[string]any{},
820+
}
821+
822+
requestBody, err := json.Marshal(pingResponse)
823+
if err != nil {
824+
t.Fatalf("Failed to marshal ping response: %v", err)
825+
}
826+
827+
resp, err := http.Post(
828+
messageURL,
829+
"application/json",
830+
bytes.NewBuffer(requestBody),
831+
)
832+
if err != nil {
833+
t.Fatalf("Failed to send ping response: %v", err)
834+
}
835+
defer resp.Body.Close()
836+
837+
if resp.StatusCode != http.StatusAccepted {
838+
t.Errorf("Expected status 202 for ping response, got %d", resp.StatusCode)
839+
}
840+
841+
body, err := io.ReadAll(resp.Body)
842+
if err != nil {
843+
t.Fatalf("Failed to read response body: %v", err)
844+
}
845+
846+
if len(body) > 0 {
847+
var response map[string]any
848+
if err := json.Unmarshal(body, &response); err != nil {
849+
t.Fatalf("Failed to parse response body: %v", err)
850+
}
851+
852+
if response["error"] != nil {
853+
t.Errorf("Expected no error in response, got %v", response["error"])
854+
}
855+
}
856+
})
742857
}

0 commit comments

Comments
 (0)