Skip to content

Commit a1dd4ef

Browse files
authored
fix: replace bufio.Scanner with bufio.Reader to support large messages in stdio transport (#603)
The bufio.Scanner has a default 64KB token limit which causes 'token too long' errors when MCP servers send large messages (e.g., large tool responses, resource contents, or prompts). This change replaces Scanner with Reader.ReadString('\n') which can handle arbitrarily large lines. Changes: - client/transport/stdio.go: Changed stdout from *bufio.Scanner to *bufio.Reader - testdata/mockstdio_server.go: Applied same fix to mock server - client/transport/stdio_test.go: Added TestStdio_LargeMessages with tests for messages ranging from 1KB to 5MB to ensure the fix works correctly The original code (pre-commit 4e353ac) used bufio.Reader, but was incorrectly changed to Scanner claiming it would avoid panics with long lines. This fix reverts to the Reader approach which actually handles large messages correctly. Fixes issue where stdio clients fail with 'bufio.Scanner: token too long' error when communicating with servers that send large responses.
1 parent b8297f5 commit a1dd4ef

File tree

3 files changed

+139
-13
lines changed

3 files changed

+139
-13
lines changed

client/transport/stdio.go

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"io"
1010
"os"
1111
"os/exec"
12+
"strings"
1213
"sync"
1314

1415
"github.com/mark3labs/mcp-go/mcp"
@@ -27,7 +28,7 @@ type Stdio struct {
2728
cmd *exec.Cmd
2829
cmdFunc CommandFunc
2930
stdin io.WriteCloser
30-
stdout *bufio.Scanner
31+
stdout *bufio.Reader
3132
stderr io.ReadCloser
3233
responses map[string]chan *JSONRPCResponse
3334
mu sync.RWMutex
@@ -72,7 +73,7 @@ func WithCommandLogger(logger util.Logger) StdioOption {
7273
func NewIO(input io.Reader, output io.WriteCloser, logging io.ReadCloser) *Stdio {
7374
return &Stdio{
7475
stdin: output,
75-
stdout: bufio.NewScanner(input),
76+
stdout: bufio.NewReader(input),
7677
stderr: logging,
7778

7879
responses: make(map[string]chan *JSONRPCResponse),
@@ -180,7 +181,7 @@ func (c *Stdio) spawnCommand(ctx context.Context) error {
180181
c.cmd = cmd
181182
c.stdin = stdin
182183
c.stderr = stderr
183-
c.stdout = bufio.NewScanner(stdout)
184+
c.stdout = bufio.NewReader(stdout)
184185

185186
if err := cmd.Start(); err != nil {
186187
return fmt.Errorf("failed to start command: %w", err)
@@ -251,15 +252,15 @@ func (c *Stdio) readResponses() {
251252
case <-c.done:
252253
return
253254
default:
254-
if !c.stdout.Scan() {
255-
err := c.stdout.Err()
256-
if err != nil && !errors.Is(err, context.Canceled) {
255+
line, err := c.stdout.ReadString('\n')
256+
if err != nil {
257+
if err != io.EOF && !errors.Is(err, context.Canceled) {
257258
c.logger.Errorf("Error reading from stdout: %v", err)
258259
}
259260
return
260261
}
261262

262-
line := c.stdout.Text()
263+
line = strings.TrimRight(line, "\r\n")
263264
// First try to parse as a generic message to check for ID field
264265
var baseMessage struct {
265266
JSONRPC string `json:"jsonrpc"`

client/transport/stdio_test.go

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -703,3 +703,120 @@ func TestStdio_NewStdioWithOptions_AppliesOptions(t *testing.T) {
703703
require.NotNil(t, stdio)
704704
require.True(t, configured, "option was not applied")
705705
}
706+
707+
func TestStdio_LargeMessages(t *testing.T) {
708+
tempFile, err := os.CreateTemp("", "mockstdio_server")
709+
if err != nil {
710+
t.Fatalf("Failed to create temp file: %v", err)
711+
}
712+
tempFile.Close()
713+
mockServerPath := tempFile.Name()
714+
715+
if runtime.GOOS == "windows" {
716+
os.Remove(mockServerPath)
717+
mockServerPath += ".exe"
718+
}
719+
720+
if compileErr := compileTestServer(mockServerPath); compileErr != nil {
721+
t.Fatalf("Failed to compile mock server: %v", compileErr)
722+
}
723+
defer os.Remove(mockServerPath)
724+
725+
stdio := NewStdio(mockServerPath, nil)
726+
727+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
728+
defer cancel()
729+
730+
startErr := stdio.Start(ctx)
731+
if startErr != nil {
732+
t.Fatalf("Failed to start Stdio transport: %v", startErr)
733+
}
734+
defer stdio.Close()
735+
736+
testCases := []struct {
737+
name string
738+
dataSize int
739+
description string
740+
}{
741+
{"SmallMessage_1KB", 1024, "Small message under scanner default limit"},
742+
{"MediumMessage_32KB", 32 * 1024, "Medium message under scanner default limit"},
743+
{"AtLimit_64KB", 64 * 1024, "Message at default scanner limit"},
744+
{"OverLimit_128KB", 128 * 1024, "Message over default scanner limit - would fail with Scanner"},
745+
{"Large_256KB", 256 * 1024, "Large message well over scanner limit"},
746+
{"VeryLarge_1MB", 1024 * 1024, "Very large message"},
747+
{"Huge_5MB", 5 * 1024 * 1024, "Huge message to stress test"},
748+
}
749+
750+
for _, tc := range testCases {
751+
t.Run(tc.name, func(t *testing.T) {
752+
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
753+
defer cancel()
754+
755+
largeString := generateRandomString(tc.dataSize)
756+
757+
params := map[string]any{
758+
"data": largeString,
759+
"size": len(largeString),
760+
}
761+
762+
request := JSONRPCRequest{
763+
JSONRPC: "2.0",
764+
ID: mcp.NewRequestId(int64(1)),
765+
Method: "debug/echo",
766+
Params: params,
767+
}
768+
769+
response, err := stdio.SendRequest(ctx, request)
770+
if err != nil {
771+
t.Fatalf("SendRequest failed for %s: %v", tc.description, err)
772+
}
773+
774+
var result struct {
775+
JSONRPC string `json:"jsonrpc"`
776+
ID mcp.RequestId `json:"id"`
777+
Method string `json:"method"`
778+
Params map[string]any `json:"params"`
779+
}
780+
781+
if err := json.Unmarshal(response.Result, &result); err != nil {
782+
t.Fatalf("Failed to unmarshal result for %s: %v", tc.description, err)
783+
}
784+
785+
if result.JSONRPC != "2.0" {
786+
t.Errorf("Expected JSONRPC value '2.0', got '%s'", result.JSONRPC)
787+
}
788+
789+
returnedData, ok := result.Params["data"].(string)
790+
if !ok {
791+
t.Fatalf("Expected data to be string, got %T", result.Params["data"])
792+
}
793+
794+
if returnedData != largeString {
795+
t.Errorf("Data mismatch for %s: expected length %d, got length %d",
796+
tc.description, len(largeString), len(returnedData))
797+
}
798+
799+
returnedSize, ok := result.Params["size"].(float64)
800+
if !ok {
801+
t.Fatalf("Expected size to be number, got %T", result.Params["size"])
802+
}
803+
804+
if int(returnedSize) != tc.dataSize {
805+
t.Errorf("Size mismatch for %s: expected %d, got %d",
806+
tc.description, tc.dataSize, int(returnedSize))
807+
}
808+
809+
t.Logf("Successfully handled %s message of size %d bytes", tc.name, tc.dataSize)
810+
})
811+
}
812+
}
813+
814+
func generateRandomString(size int) string {
815+
const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 "
816+
817+
b := make([]byte, size)
818+
for i := range b {
819+
b[i] = charset[i%len(charset)]
820+
}
821+
return string(b)
822+
}

testdata/mockstdio_server.go

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"fmt"
77
"log/slog"
88
"os"
9+
"strings"
910

1011
"github.com/mark3labs/mcp-go/mcp"
1112
)
@@ -18,19 +19,26 @@ type JSONRPCRequest struct {
1819
}
1920

2021
type JSONRPCResponse struct {
21-
JSONRPC string `json:"jsonrpc"`
22-
ID *mcp.RequestId `json:"id,omitempty"`
23-
Result any `json:"result,omitempty"`
22+
JSONRPC string `json:"jsonrpc"`
23+
ID *mcp.RequestId `json:"id,omitempty"`
24+
Result any `json:"result,omitempty"`
2425
Error *mcp.JSONRPCErrorDetails `json:"error,omitempty"`
2526
}
2627

2728
func main() {
2829
logger := slog.New(slog.NewJSONHandler(os.Stderr, &slog.HandlerOptions{}))
2930
logger.Info("launch successful")
30-
scanner := bufio.NewScanner(os.Stdin)
31-
for scanner.Scan() {
31+
reader := bufio.NewReader(os.Stdin)
32+
for {
33+
line, err := reader.ReadString('\n')
34+
if err != nil {
35+
break
36+
}
37+
38+
line = strings.TrimRight(line, "\r\n")
39+
3240
var request JSONRPCRequest
33-
if err := json.Unmarshal(scanner.Bytes(), &request); err != nil {
41+
if err := json.Unmarshal([]byte(line), &request); err != nil {
3442
continue
3543
}
3644

0 commit comments

Comments
 (0)