diff --git a/cmd/yardstick-client/main.go b/cmd/yardstick-client/main.go index a707e43..73f08f3 100644 --- a/cmd/yardstick-client/main.go +++ b/cmd/yardstick-client/main.go @@ -59,8 +59,11 @@ func (c *Client) Connect(ctx context.Context) error { return fmt.Errorf("failed to create transport: %w", err) } - c.client = mcp.NewClient("yardstick-client", "1.0.0", nil) - session, err := c.client.Connect(ctx, transport) + c.client = mcp.NewClient(&mcp.Implementation{ + Name: "yardstick-client", + Version: "1.0.0", + }, nil) + session, err := c.client.Connect(ctx, transport, nil) if err != nil { return err } @@ -76,7 +79,7 @@ func (c *Client) connectStdio() (mcp.Transport, error) { // #nosec G204 - Command and args are from user configuration, this is intentional cmd := exec.Command(c.config.Command, c.config.Args...) - return mcp.NewCommandTransport(cmd), nil + return &mcp.CommandTransport{Command: cmd}, nil } // connectSSE creates an SSE transport connection @@ -84,7 +87,7 @@ func (c *Client) connectStdio() (mcp.Transport, error) { //nolint:unparam func (c *Client) connectSSE() (mcp.Transport, error) { url := fmt.Sprintf("http://%s:%d/sse", c.config.Address, c.config.Port) - return mcp.NewSSEClientTransport(url, nil), nil + return &mcp.SSEClientTransport{Endpoint: url}, nil } // connectStreamableHTTP creates a streamable HTTP transport connection @@ -92,7 +95,7 @@ func (c *Client) connectSSE() (mcp.Transport, error) { //nolint:unparam func (c *Client) connectStreamableHTTP() (mcp.Transport, error) { url := fmt.Sprintf("http://%s:%d/mcp", c.config.Address, c.config.Port) - return mcp.NewStreamableClientTransport(url, nil), nil + return &mcp.StreamableClientTransport{Endpoint: url}, nil } // Close closes the client connection diff --git a/cmd/yardstick-client/main_test.go b/cmd/yardstick-client/main_test.go index 9b38634..bbcf1a1 100644 --- a/cmd/yardstick-client/main_test.go +++ b/cmd/yardstick-client/main_test.go @@ -168,26 +168,24 @@ func createMockMCPServer(_ *testing.T, transport string) *httptest.Server { mux := http.NewServeMux() // Create a mock MCP server - server := mcp.NewServer("test-server", "1.0.0", nil) + server := mcp.NewServer(&mcp.Implementation{ + Name: "test-server", + Version: "1.0.0", + }, nil) // Add a simple echo tool for testing - echoTool := mcp.NewServerTool("echo", "Echo tool for testing", - func(_ context.Context, _ *mcp.ServerSession, _ *mcp.CallToolParamsFor[map[string]interface{}]) (*mcp.CallToolResultFor[map[string]interface{}], error) { - return &mcp.CallToolResultFor[map[string]interface{}]{ - Content: []mcp.Content{ - &mcp.TextContent{Text: "test response"}, - }, - }, nil - }, - mcp.Input(), - ) - server.AddTools(echoTool) + mcp.AddTool(server, &mcp.Tool{ + Name: "echo", + Description: "Echo tool for testing", + }, func(_ context.Context, _ *mcp.CallToolRequest, _ map[string]interface{}) (*mcp.CallToolResult, map[string]interface{}, error) { + return nil, map[string]interface{}{"message": "test response"}, nil + }) switch transport { case "sse": handler := mcp.NewSSEHandler(func(_ *http.Request) *mcp.Server { return server - }) + }, nil) mux.Handle("/sse", handler) case "streamable-http": handler := mcp.NewStreamableHTTPHandler(func(_ *http.Request) *mcp.Server { diff --git a/cmd/yardstick-server/integration_test.go b/cmd/yardstick-server/integration_test.go index 8a294b6..879abd2 100644 --- a/cmd/yardstick-server/integration_test.go +++ b/cmd/yardstick-server/integration_test.go @@ -162,24 +162,22 @@ func TestEndToEndEchoFunctionality(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - params := &mcp.CallToolParamsFor[EchoRequest]{ - Arguments: EchoRequest{Input: tc.input}, - } + // Create a CallToolRequest for testing + req := &mcp.CallToolRequest{} + params := EchoRequest{Input: tc.input} - result, err := echoHandler(ctx, nil, params) + result, response, err := echoHandler(ctx, req, params) if tc.expectedError { - assert.Error(t, err) - assert.Nil(t, result) - } else { + // For invalid input, we expect an error result, not an error assert.NoError(t, err) assert.NotNil(t, result) - assert.Len(t, result.Content, 1) - - // Check that the content is TextContent - textContent, ok := result.Content[0].(*mcp.TextContent) - assert.True(t, ok) - assert.Equal(t, tc.input, textContent.Text) + assert.True(t, result.IsError) + } else { + assert.NoError(t, err) + assert.Nil(t, result) // When successful, result is nil and response contains the data + assert.NotNil(t, response) + assert.Equal(t, tc.input, response.Output) } }) } @@ -203,13 +201,13 @@ func TestServerStartup(t *testing.T) { assert.False(t, validateAlphanumeric("test@123")) // Verify echo handler works - params := &mcp.CallToolParamsFor[EchoRequest]{ - Arguments: EchoRequest{Input: "test123"}, - } + req := &mcp.CallToolRequest{} + params := EchoRequest{Input: "test123"} - result, err := echoHandler(context.Background(), nil, params) + result, response, err := echoHandler(context.Background(), req, params) assert.NoError(t, err) - assert.NotNil(t, result) + assert.Nil(t, result) + assert.NotNil(t, response) }) } } @@ -225,30 +223,23 @@ func TestConcurrentEchoRequests(t *testing.T) { for i := 0; i < numConcurrentRequests; i++ { go func(id int) { - params := &mcp.CallToolParamsFor[EchoRequest]{ - Arguments: EchoRequest{Input: fmt.Sprintf("test%d", id)}, - } + req := &mcp.CallToolRequest{} + params := EchoRequest{Input: fmt.Sprintf("test%d", id)} - result, err := echoHandler(ctx, nil, params) + result, response, err := echoHandler(ctx, req, params) if err != nil { results <- err return } - if result == nil || len(result.Content) == 0 { - results <- fmt.Errorf("invalid result for request %d", id) - return - } - - textContent, ok := result.Content[0].(*mcp.TextContent) - if !ok { - results <- fmt.Errorf("expected TextContent for request %d", id) + if result != nil { + results <- fmt.Errorf("expected nil result for request %d", id) return } expectedOutput := fmt.Sprintf("test%d", id) - if textContent.Text != expectedOutput { - results <- fmt.Errorf("expected %s, got %s", expectedOutput, textContent.Text) + if response.Output != expectedOutput { + results <- fmt.Errorf("expected %s, got %s", expectedOutput, response.Output) return } diff --git a/cmd/yardstick-server/main.go b/cmd/yardstick-server/main.go index 6d268d4..5722db3 100644 --- a/cmd/yardstick-server/main.go +++ b/cmd/yardstick-server/main.go @@ -11,7 +11,7 @@ import ( "strconv" "time" - "github.com/modelcontextprotocol/go-sdk/jsonschema" + "github.com/google/jsonschema-go/jsonschema" "github.com/modelcontextprotocol/go-sdk/mcp" ) @@ -33,22 +33,19 @@ func validateAlphanumeric(input string) bool { return alphanumericRegex.MatchString(input) } -func echoHandler(_ context.Context, _ *mcp.ServerSession, params *mcp.CallToolParamsFor[EchoRequest]) ( - *mcp.CallToolResultFor[EchoResponse], error, -) { - if !validateAlphanumeric(params.Arguments.Input) { - return nil, fmt.Errorf("input must be alphanumeric only") +func echoHandler(_ context.Context, _ *mcp.CallToolRequest, params EchoRequest) (*mcp.CallToolResult, EchoResponse, error) { + if !validateAlphanumeric(params.Input) { + return &mcp.CallToolResult{ + Content: []mcp.Content{&mcp.TextContent{Text: "input must be alphanumeric only"}}, + IsError: true, + }, EchoResponse{}, nil } response := EchoResponse{ - Output: params.Arguments.Input, + Output: params.Input, } - return &mcp.CallToolResultFor[EchoResponse]{ - Content: []mcp.Content{ - &mcp.TextContent{Text: response.Output}, - }, - }, nil + return nil, response, nil } func main() { @@ -56,30 +53,38 @@ func main() { parseConfig() // Create MCP server - server := mcp.NewServer("echo-server", "1.0.0", nil) - - // Add echo tool to server - echoTool := mcp.NewServerTool("echo", "Echo back an alphanumeric string for deterministic testing", echoHandler, - mcp.Input( - mcp.Property("input", - mcp.Description("Alphanumeric string to echo back"), - mcp.Schema(&jsonschema.Schema{ - Type: "string", - Pattern: "^[a-zA-Z0-9]+$", - }), - ), - ), - ) - - server.AddTools(echoTool) + server := mcp.NewServer(&mcp.Implementation{ + Name: "echo-server", + Version: "1.0.0", + }, nil) + + // Create custom schema for input validation + inputSchema := &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "input": { + Type: "string", + Pattern: "^[a-zA-Z0-9]+$", + Description: "Alphanumeric string to echo back", + }, + }, + Required: []string{"input"}, + } + + // Add echo tool to server using the new API + mcp.AddTool(server, &mcp.Tool{ + Name: "echo", + Description: "Echo back an alphanumeric string for deterministic testing", + InputSchema: inputSchema, + }, echoHandler) ctx := context.Background() switch transport { case "stdio": log.Println("Starting MCP server with stdio transport") - transport := mcp.NewStdioTransport() - if err := server.Run(ctx, transport); err != nil { + stdioTransport := &mcp.StdioTransport{} + if err := server.Run(ctx, stdioTransport); err != nil { log.Fatal("Failed to run server:", err) } @@ -90,7 +95,7 @@ func main() { handler := mcp.NewSSEHandler(func(_ *http.Request) *mcp.Server { return server - }) + }, nil) // Mount the SSE handler at /sse - it will handle both GET (SSE stream) and POST (messages) requests http.Handle("/sse", handler) diff --git a/cmd/yardstick-server/main_test.go b/cmd/yardstick-server/main_test.go index 0c25e1e..26fb6d8 100644 --- a/cmd/yardstick-server/main_test.go +++ b/cmd/yardstick-server/main_test.go @@ -79,26 +79,22 @@ func TestEchoHandler(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { // Create request with the new API - echoReq := EchoRequest{Input: tt.input} - params := &mcp.CallToolParamsFor[EchoRequest]{ - Arguments: echoReq, - } + req := &mcp.CallToolRequest{} + params := EchoRequest{Input: tt.input} // Call handler - result, err := echoHandler(context.Background(), nil, params) + result, response, err := echoHandler(context.Background(), req, params) if tt.expectError { - assert.Error(t, err) - assert.Nil(t, result) - } else { + // For invalid input, we expect an error result, not an error assert.NoError(t, err) assert.NotNil(t, result) - assert.Len(t, result.Content, 1) - - // Check that the content is TextContent - textContent, ok := result.Content[0].(*mcp.TextContent) - assert.True(t, ok) - assert.Equal(t, tt.expectedOutput, textContent.Text) + assert.True(t, result.IsError) + } else { + assert.NoError(t, err) + assert.Nil(t, result) + assert.NotNil(t, response) + assert.Equal(t, tt.expectedOutput, response.Output) } }) } diff --git a/go.mod b/go.mod index cc395c6..0b6181c 100644 --- a/go.mod +++ b/go.mod @@ -1,14 +1,16 @@ module github.com/stackloklabs/yardstick -go 1.24.0 +go 1.25.0 require ( - github.com/modelcontextprotocol/go-sdk v0.1.0 + github.com/google/jsonschema-go v0.3.0 + github.com/modelcontextprotocol/go-sdk v1.0.0 github.com/stretchr/testify v1.11.1 ) require ( github.com/davecgh/go-spew v1.1.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 3aabb8a..0870f0b 100644 --- a/go.sum +++ b/go.sum @@ -2,12 +2,16 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= -github.com/modelcontextprotocol/go-sdk v0.1.0 h1:ItzbFWYNt4EHcUrScX7P8JPASn1FVYb29G773Xkl+IU= -github.com/modelcontextprotocol/go-sdk v0.1.0/go.mod h1:DcXfbr7yl7e35oMpzHfKw2nUYRjhIGS2uou/6tdsTB0= +github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q= +github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= +github.com/modelcontextprotocol/go-sdk v1.0.0 h1:Z4MSjLi38bTgLrd/LjSmofqRqyBiVKRyQSJgw8q8V74= +github.com/modelcontextprotocol/go-sdk v1.0.0/go.mod h1:nYtYQroQ2KQiM0/SbyEPUWQ6xs4B95gJjEalc9AQyOs= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo= golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=