fix(mcp): use transport-level timeouts and context timeouts for SSE support

Replace http.Client.Timeout with transport-level timeouts to prevent
terminating long-lived SSE streams while still protecting against
connection hangs during MCP server initialization.

Changes:
- Use http.Transport with DialContext, TLSHandshakeTimeout, and
  ResponseHeaderTimeout instead of http.Client.Timeout
- Add timeout-scoped contexts for connect and list operations
- SSE streams can now remain open indefinitely without being
  interrupted by a global request timeout
- Long-running tool calls are no longer affected by the 30s timeout

Fixes review comment: http.Client.Timeout applies to the full request
lifetime, including reading the response body, which would terminate
SSE streams after 30s.
This commit is contained in:
dtapps 2026-05-06 16:00:54 +08:00
parent bd18186187
commit d8b77086c2
2 changed files with 99 additions and 8 deletions

View file

@ -5,6 +5,7 @@ import (
"context" "context"
"errors" "errors"
"fmt" "fmt"
"net"
"net/http" "net/http"
"os" "os"
"os/exec" "os/exec"
@ -382,17 +383,24 @@ func connectServer(
DisableStandaloneSSE: disableStandaloneSSE, DisableStandaloneSSE: disableStandaloneSSE,
} }
// Set up HTTP client with a connection timeout to avoid hanging // Set up HTTP client with transport-level timeouts to avoid hanging
// indefinitely when the MCP server is unreachable. // indefinitely during connection establishment, while allowing SSE
// streams to remain open without a global request timeout.
baseTransport := &http.Transport{
DialContext: (&net.Dialer{Timeout: 10 * time.Second}).DialContext,
TLSHandshakeTimeout: 10 * time.Second,
ResponseHeaderTimeout: 30 * time.Second,
}
mcpHTTPClient := &http.Client{ mcpHTTPClient := &http.Client{
Timeout: 30 * time.Second, Transport: baseTransport,
} }
// Add custom headers if provided // Add custom headers if provided
if len(cfg.Headers) > 0 { if len(cfg.Headers) > 0 {
// Create a custom HTTP client with header-injecting transport // Create a custom HTTP client with header-injecting transport
mcpHTTPClient.Transport = &headerTransport{ mcpHTTPClient.Transport = &headerTransport{
base: http.DefaultTransport, base: baseTransport,
headers: cfg.Headers, headers: cfg.Headers,
} }
logger.DebugCF("mcp", "Added custom HTTP headers", logger.DebugCF("mcp", "Added custom HTTP headers",
@ -464,8 +472,11 @@ func connectServer(
) )
} }
// Connect to server // Connect to server with a timeout-scoped context to avoid hanging
session, err := client.Connect(ctx, transport, nil) // indefinitely when the MCP server is unreachable.
connectCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
session, err := client.Connect(connectCtx, transport, nil)
cancel()
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to connect: %w", err) return nil, fmt.Errorf("failed to connect: %w", err)
} }
@ -480,8 +491,10 @@ func connectServer(
"protocol": initResult.ProtocolVersion, "protocol": initResult.ProtocolVersion,
}) })
// List available tools if supported // List available tools if supported, with a timeout-scoped context
tools, err := listServerTools(ctx, name, session, initResult) listCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
tools, err := listServerTools(listCtx, name, session, initResult)
cancel()
if err != nil { if err != nil {
_ = session.Close() _ = session.Close()
return nil, err return nil, err

View file

@ -5,6 +5,8 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"io" "io"
"net"
"net/http"
"os" "os"
"path/filepath" "path/filepath"
"strings" "strings"
@ -628,3 +630,79 @@ func (t *scriptedTransport) Close() error {
func (t *scriptedTransport) SessionID() string { func (t *scriptedTransport) SessionID() string {
return t.sessionID return t.sessionID
} }
// TestHTTPClientTransportTimeouts verifies that the HTTP client is configured
// with transport-level timeouts instead of client-wide timeout.
// This addresses the PR review comment:
// "http.Client.Timeout applies to the full request lifetime, including reading
// the response body. On StreamableClientTransport in sse mode that will terminate
// the long-lived SSE stream after 30s"
func TestHTTPClientTransportTimeouts(t *testing.T) {
// Verify that transport-level timeouts are configured correctly
// instead of using http.Client.Timeout which would affect SSE streams
transport := &http.Transport{
DialContext: (&net.Dialer{Timeout: 10 * time.Second}).DialContext,
TLSHandshakeTimeout: 10 * time.Second,
ResponseHeaderTimeout: 30 * time.Second,
}
// Verify transport timeouts are set correctly
// This ensures SSE streams won't be terminated by client-wide timeout
if transport.TLSHandshakeTimeout != 10*time.Second {
t.Errorf("TLSHandshakeTimeout should be 10s, got %v", transport.TLSHandshakeTimeout)
}
if transport.ResponseHeaderTimeout != 30*time.Second {
t.Errorf("ResponseHeaderTimeout should be 30s, got %v", transport.ResponseHeaderTimeout)
}
// Verify DialContext is set (not DialTimeout which is deprecated)
if transport.DialContext == nil {
t.Error("DialContext should be set for connection timeout")
}
}
// TestContextTimeoutForConnect verifies that connect operations use timeout-scoped context.
func TestContextTimeoutForConnect(t *testing.T) {
// Test that a context with timeout properly cancels the operation
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
// Create a listener that delays acceptance
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("Failed to create listener: %v", err)
}
defer listener.Close()
// Start a goroutine that delays accepting connections
go func() {
time.Sleep(200 * time.Millisecond)
for {
conn, acceptErr := listener.Accept()
if acceptErr != nil {
return
}
conn.Close()
}
}()
mgr := NewManager()
defer mgr.Close()
start := time.Now()
err = mgr.ConnectServer(ctx, "timeout-test", config.MCPServerConfig{
Type: "sse",
URL: "http://" + listener.Addr().String() + "/mcp",
Enabled: true,
})
duration := time.Since(start)
if err == nil {
t.Fatal("Expected connection to fail due to context timeout")
}
// Should fail quickly due to context timeout, not hang
if duration > 500*time.Millisecond {
t.Fatalf("Operation took too long: %v, expected context timeout around 100ms", duration)
}
}