From 07f62a06c13923d0b3291064b6f6bd8a9717386f Mon Sep 17 00:00:00 2001
From: Rob Findley <rfindley@google.com>
Date: Fri, 29 Aug 2025 18:44:39 +0000
Subject: [PATCH] examples/client: add a loadtest command

Add a loadtest client example, to help confirm performance of our
streamable transport implementation.

For golang/go#190
---
 examples/client/listfeatures/main.go |   6 +-
 examples/client/loadtest/main.go     | 122 +++++++++++++++++++++++++++
 2 files changed, 125 insertions(+), 3 deletions(-)
 create mode 100644 examples/client/loadtest/main.go

diff --git a/examples/client/listfeatures/main.go b/examples/client/listfeatures/main.go
index 755a4f98..9d473f0b 100644
--- a/examples/client/listfeatures/main.go
+++ b/examples/client/listfeatures/main.go
@@ -31,10 +31,10 @@ func main() {
 	flag.Parse()
 	args := flag.Args()
 	if len(args) == 0 {
-		fmt.Fprintf(os.Stderr, "Usage: listfeatures <command> [<args>]")
-		fmt.Fprintf(os.Stderr, "List all features for a stdio MCP server")
+		fmt.Fprintln(os.Stderr, "Usage: listfeatures <command> [<args>]")
+		fmt.Fprintln(os.Stderr, "List all features for a stdio MCP server")
 		fmt.Fprintln(os.Stderr)
-		fmt.Fprintf(os.Stderr, "Example: listfeatures npx @modelcontextprotocol/server-everything")
+		fmt.Fprintln(os.Stderr, "Example:\n\tlistfeatures npx @modelcontextprotocol/server-everything")
 		os.Exit(2)
 	}
 
diff --git a/examples/client/loadtest/main.go b/examples/client/loadtest/main.go
new file mode 100644
index 00000000..2c6a5c03
--- /dev/null
+++ b/examples/client/loadtest/main.go
@@ -0,0 +1,122 @@
+// 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")
+	v        = flag.Bool("v", false, "if set, enable verbose logging of results")
+)
+
+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, restoreSignal := signal.NotifyContext(parentCtx, os.Interrupt)
+	defer restoreSignal()
+
+	var (
+		start   = time.Now()
+		success atomic.Int64
+		failure atomic.Int64
+	)
+
+	// Run the test.
+	var wg sync.WaitGroup
+	for range *workers {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			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)
+			}
+			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 *v {
+						log.Printf("FAILURE: %v", err)
+					}
+				} else {
+					success.Add(1)
+					if *v {
+						data, err := json.Marshal(res)
+						if err != nil {
+							log.Fatalf("marshalling result: %v", err)
+						}
+						log.Printf("SUCCESS: %s", string(data))
+					}
+				}
+			}
+		}()
+	}
+	wg.Wait()
+	restoreSignal() // call restore signal (redundantly) here to allow ctrl-c to work again
+
+	// 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())
+}