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:
parent
bd18186187
commit
d8b77086c2
2 changed files with 99 additions and 8 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue