123 lines
3.2 KiB
Go
123 lines
3.2 KiB
Go
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
|
// Use of this source code is governed by an MIT-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
// The load command load tests a streamable MCP server
|
|
//
|
|
// Usage: loadtest <URL>
|
|
//
|
|
// For example:
|
|
//
|
|
// loadtest -tool=greet -args='{"name": "foo"}' http://localhost:8080
|
|
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"os/signal"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
)
|
|
|
|
var (
|
|
duration = flag.Duration("duration", 1*time.Minute, "duration of the load test")
|
|
tool = flag.String("tool", "", "tool to call")
|
|
jsonArgs = flag.String("args", "", "JSON arguments to pass")
|
|
workers = flag.Int("workers", 10, "number of concurrent workers")
|
|
timeout = flag.Duration("timeout", 1*time.Second, "request timeout")
|
|
qps = flag.Int("qps", 100, "tool calls per second, per worker")
|
|
verbose = flag.Bool("v", false, "if set, enable verbose logging")
|
|
cleanup = flag.Bool("cleanup", true, "whether to clean up sessions at the end of the test")
|
|
)
|
|
|
|
func main() {
|
|
flag.Usage = func() {
|
|
out := flag.CommandLine.Output()
|
|
fmt.Fprintf(out, "Usage: loadtest [flags] <URL>")
|
|
fmt.Fprintf(out, "Load test a streamable HTTP server (CTRL-C to end early)")
|
|
fmt.Fprintln(out)
|
|
fmt.Fprintf(out, "Example: loadtest -tool=greet -args='{\"name\": \"foo\"}' http://localhost:8080\n")
|
|
fmt.Fprintln(out)
|
|
fmt.Fprintln(out, "Flags:")
|
|
flag.PrintDefaults()
|
|
}
|
|
flag.Parse()
|
|
args := flag.Args()
|
|
if len(args) != 1 || *tool == "" {
|
|
flag.Usage()
|
|
os.Exit(2)
|
|
}
|
|
|
|
parentCtx, cancel := context.WithTimeout(context.Background(), *duration)
|
|
defer cancel()
|
|
parentCtx, stop := signal.NotifyContext(parentCtx, os.Interrupt)
|
|
defer stop()
|
|
|
|
var (
|
|
start = time.Now()
|
|
success atomic.Int64
|
|
failure atomic.Int64
|
|
)
|
|
|
|
// Run the test.
|
|
var wg sync.WaitGroup
|
|
for range *workers {
|
|
wg.Go(func() {
|
|
client := mcp.NewClient(&mcp.Implementation{Name: "mcp-client", Version: "v1.0.0"}, nil)
|
|
cs, err := client.Connect(parentCtx, &mcp.StreamableClientTransport{Endpoint: args[0]}, nil)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
if *cleanup {
|
|
defer cs.Close()
|
|
}
|
|
|
|
ticker := time.NewTicker(1 * time.Second / time.Duration(*qps))
|
|
defer ticker.Stop()
|
|
|
|
for range ticker.C {
|
|
ctx, cancel := context.WithTimeout(parentCtx, *timeout)
|
|
defer cancel()
|
|
|
|
res, err := cs.CallTool(ctx, &mcp.CallToolParams{Name: *tool, Arguments: json.RawMessage(*jsonArgs)})
|
|
if err != nil {
|
|
if parentCtx.Err() != nil {
|
|
return // test ended
|
|
}
|
|
failure.Add(1)
|
|
if *verbose {
|
|
log.Printf("FAILURE: %v", err)
|
|
}
|
|
} else {
|
|
success.Add(1)
|
|
if *verbose {
|
|
data, err := json.Marshal(res)
|
|
if err != nil {
|
|
log.Fatalf("marshalling result: %v", err)
|
|
}
|
|
log.Printf("SUCCESS: %s", string(data))
|
|
}
|
|
}
|
|
}
|
|
})
|
|
}
|
|
wg.Wait()
|
|
stop() // restore the interrupt signal
|
|
|
|
// Print stats.
|
|
var (
|
|
dur = time.Since(start)
|
|
succ = success.Load()
|
|
fail = failure.Load()
|
|
)
|
|
fmt.Printf("Results (in %s):\n", dur)
|
|
fmt.Printf("\tsuccess: %d (%g QPS)\n", succ, float64(succ)/dur.Seconds())
|
|
fmt.Printf("\tfailure: %d (%g QPS)\n", fail, float64(fail)/dur.Seconds())
|
|
}
|