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

View file

@ -5,6 +5,8 @@ import (
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"os"
"path/filepath"
"strings"
@ -628,3 +630,79 @@ func (t *scriptedTransport) Close() error {
func (t *scriptedTransport) SessionID() string {
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)
}
}