fix(mcp): retry stdio transport initialization with exponential backoff

Heavy MCP servers (FastMCP with 80+ tools, OAuth servers) take 3-5s to
start their stdin read loop. Without retries, connectAndDiscover sends
the initialize request immediately after fork/exec, gets EOF, and
permanently fails with "transport closed".

Add a retry loop (4 attempts, 2s/4s/8s backoff) for stdio transports
only. SSE/HTTP transports fail immediately since connection-refused is
definitive. The retry respects context cancellation to avoid blocking
shutdown.

Closes #385
This commit is contained in:
viettranx committed 2026-04-02 18:43:41 +07:00
1 parent 1170a85eb4
commit 4c485d6fa1
2 files changed
+79 -2

No files matched your search

+49
View File
@@ -0,0 +1,49 @@
package mcp
import (
"context"
"testing"
"time"
)
func TestConnectAndDiscoverRetriesOnStdioInitFailure(t *testing.T) {
// "cat" starts instantly but doesn't speak JSON-RPC → Initialize fails every time.
// The retry loop should add observable backoff delay before exhausting attempts.
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
start := time.Now()
_, _, err := connectAndDiscover(ctx, "test-retry", "stdio",
"cat", nil, nil, "", nil, 2)
elapsed := time.Since(start)
if err == nil {
t.Fatal("expected error — cat doesn't speak MCP")
}
// First retry waits 2s → elapsed must be at least 2s proving retry happened.
if elapsed < 2*time.Second {
t.Errorf("elapsed %v — expected ≥2s proving retries happened", elapsed)
}
t.Logf("elapsed: %v (retries confirmed)", elapsed)
}
func TestConnectAndDiscoverNoRetryForSSE(t *testing.T) {
// SSE/HTTP transports should fail immediately without retrying.
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
start := time.Now()
_, _, err := connectAndDiscover(ctx, "test-no-retry", "sse",
"", nil, nil, "http://127.0.0.1:1", nil, 10)
elapsed := time.Since(start)
if err == nil {
t.Fatal("expected error for bad SSE URL")
}
// SSE should fail fast (no retries) — well under 2 seconds.
if elapsed > 2*time.Second {
t.Errorf("elapsed %v — SSE should fail fast without retries", elapsed)
}
}
+30 -2
View File
@@ -38,9 +38,37 @@ func connectAndDiscover(ctx context.Context, name, transportType, command string
Version: "1.0.0",
}
if _, err := client.Initialize(ctx, initReq); err != nil {
// Retry initialization with exponential backoff for slow-starting stdio servers.
// Heavy MCP servers (FastMCP with 80+ tools, OAuth servers) can take 3-5s to start
// their stdin read loop. Without retries, Initialize sends JSON-RPC before the
// server is ready, gets EOF, and permanently fails. SSE/HTTP transports don't need
// this because the HTTP server rejects connections until ready (connection refused).
const maxInitAttempts = 4 // backoff: 2s + 4s + 8s = ~14s total before giving up
var initErr error
for attempt := range maxInitAttempts {
if attempt > 0 {
backoff := time.Duration(1<<attempt) * time.Second // 2s, 4s, 8s
slog.Debug("mcp.init.retry", "server", name, "attempt", attempt+1, "backoff", backoff)
select {
case <-time.After(backoff):
case <-ctx.Done():
_ = client.Close()
return nil, nil, fmt.Errorf("initialize: context cancelled during retry: %w", ctx.Err())
}
}
if _, err := client.Initialize(ctx, initReq); err == nil {
break
} else {
initErr = err
// Non-stdio transports: connection errors are definitive, don't retry.
if transportType != "stdio" {
break
}
}
}
if initErr != nil {
_ = client.Close()
return nil, nil, fmt.Errorf("initialize: %w", err)
return nil, nil, fmt.Errorf("initialize: %w", initErr)
}
toolsResult, err := client.ListTools(ctx, mcpgo.ListToolsRequest{})