From d8b77086c203e74c88c5faee5742883c928934ef Mon Sep 17 00:00:00 2001 From: dtapps Date: Wed, 6 May 2026 16:00:54 +0800 Subject: [PATCH] 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. --- pkg/mcp/manager.go | 29 ++++++++++----- pkg/mcp/manager_test.go | 78 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 8 deletions(-) diff --git a/pkg/mcp/manager.go b/pkg/mcp/manager.go index 25788b9c0..fee8e0766 100644 --- a/pkg/mcp/manager.go +++ b/pkg/mcp/manager.go @@ -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 diff --git a/pkg/mcp/manager_test.go b/pkg/mcp/manager_test.go index 5789a37a9..4002df322 100644 --- a/pkg/mcp/manager_test.go +++ b/pkg/mcp/manager_test.go @@ -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) + } +}