|
1 | 1 | package server
|
2 | 2 |
|
3 | 3 | import (
|
| 4 | + "bufio" |
4 | 5 | "bytes"
|
5 | 6 | "context"
|
6 | 7 | "encoding/json"
|
| 8 | + "io" |
7 | 9 | "fmt"
|
8 | 10 | "math/rand"
|
9 | 11 | "net/http"
|
@@ -739,4 +741,117 @@ func TestSSEServer(t *testing.T) {
|
739 | 741 | }
|
740 | 742 | }
|
741 | 743 | })
|
| 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 | + }) |
742 | 857 | }
|
0 commit comments