Fix OAuth flow breaking MCP session initialization

When remote MCP servers require OAuth authentication, the OAuth flow
was being triggered during the MCP session initialization handshake.
This caused the session state to become corrupted because:

1. client.Connect() starts the MCP protocol handshake
2. Initial HTTP request returns 401 Unauthorized
3. oauthTransport intercepts and runs the full OAuth flow
4. After OAuth succeeds, the retry finds the MCP session in a broken state
5. Subsequent requests fail with "400 Bad Request: broken session"

This fix adds retry logic that detects "broken session" errors during
initialization and automatically retries once after OAuth completes,
ensuring:

- OAuth completes at the HTTP transport layer first
- MCP session initialization happens with authentication already in place
- No permanent session corruption from the OAuth interruption

The retry is limited to OAuth-related "broken session" errors to avoid
masking other legitimate connection failures.

Fixes the issue where OAuth-protected MCP servers (like mcp.prisma.io)
would fail to initialize despite successful user authentication.

Author: trungutt

pkg/tools/mcp/remote.go

@@ -6,7 +6,9 @@ import (
 	"iter"
 	"log/slog"
 	"net/http"
+	"strings"
 	"sync"
+	"time"
 
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 
@@ -84,47 +86,76 @@ func (c *remoteMCPClient) Initialize(ctx context.Context, request *mcp.Initializ
 	// Create HTTP client with OAuth support
 	httpClient := c.createHTTPClient()
 
-	var transport mcp.Transport
+	// Attempt MCP initialization with retry logic for OAuth-related failures.
+	// When a server requires OAuth, the first connection attempt may fail with a "broken session"
+	// error because OAuth flow (even successful flow) interrupts the MCP handshake. Let's retry once after OAuth completes.
+	// Example of such MCP Server that broke the session with OAuth flow: https://mcp.prisma.io/mcp
+	const maxAttempts = 2
+	var lastErr error
 
-	switch c.transportType {
-	case "sse":
-		transport = &mcp.SSEClientTransport{
-			Endpoint:   c.url,
-			HTTPClient: httpClient,
+	for attempt := 1; attempt <= maxAttempts; attempt++ {
+		if attempt > 1 {
+			slog.Debug("Retrying MCP initialization after OAuth flow", "attempt", attempt)
 		}
-	case "streamable", "streamable-http":
-		transport = &mcp.StreamableClientTransport{
-			Endpoint:   c.url,
-			HTTPClient: httpClient,
+
+		var transport mcp.Transport
+
+		switch c.transportType {
+		case "sse":
+			transport = &mcp.SSEClientTransport{
+				Endpoint:   c.url,
+				HTTPClient: httpClient,
+			}
+		case "streamable", "streamable-http":
+			transport = &mcp.StreamableClientTransport{
+				Endpoint:   c.url,
+				HTTPClient: httpClient,
+			}
+		default:
+			return nil, fmt.Errorf("unsupported transport type: %s", c.transportType)
 		}
-	default:
-		return nil, fmt.Errorf("unsupported transport type: %s", c.transportType)
-	}
 
-	// Create an MCP client with elicitation support
-	impl := &mcp.Implementation{
-		Name:    "cagent",
-		Version: "1.0.0",
-	}
+		// Create an MCP client with elicitation support
+		impl := &mcp.Implementation{
+			Name:    "cagent",
+			Version: "1.0.0",
+		}
 
-	opts := &mcp.ClientOptions{
-		ElicitationHandler: c.handleElicitationRequest,
-	}
+		opts := &mcp.ClientOptions{
+			ElicitationHandler: c.handleElicitationRequest,
+		}
 
-	client := mcp.NewClient(impl, opts)
+		client := mcp.NewClient(impl, opts)
+
+		// Connect to the MCP server
+		session, err := client.Connect(ctx, transport, nil)
+		if err != nil {
+			lastErr = err
+
+			// Check if this is a "broken session" error that might be OAuth-related
+			if attempt < maxAttempts && isBrokenSessionError(err) {
+				slog.Debug("MCP connection failed with broken session error, retrying after OAuth", "error", err)
+				// Brief pause before retry to allow OAuth state to settle
+				select {
+				case <-ctx.Done():
+					return nil, fmt.Errorf("failed to connect to MCP server: %w", ctx.Err())
+				case <-time.After(100 * time.Millisecond):
+				}
+				continue
+			}
+
+			return nil, fmt.Errorf("failed to connect to MCP server: %w", err)
+		}
 
-	// Connect to the MCP server
-	session, err := client.Connect(ctx, transport, nil)
-	if err != nil {
-		return nil, fmt.Errorf("failed to connect to MCP server: %w", err)
-	}
+		c.mu.Lock()
+		c.session = session
+		c.mu.Unlock()
 
-	c.mu.Lock()
-	c.session = session
-	c.mu.Unlock()
+		slog.Debug("Remote MCP client connected successfully", "attempt", attempt)
+		return session.InitializeResult(), nil
+	}
 
-	slog.Debug("Remote MCP client connected successfully")
-	return session.InitializeResult(), nil
+	return nil, fmt.Errorf("failed to connect to MCP server after %d attempts: %w", maxAttempts, lastErr)
 }
 
 // createHTTPClient creates an HTTP client with OAuth support
@@ -194,3 +225,15 @@ func (c *remoteMCPClient) requestUserConsent(ctx context.Context) (bool, error)
 
 	return result.Action == "accept", nil
 }
+
+// isBrokenSessionError checks if an error is a "broken session" error from the MCP SDK
+// This error typically occurs when OAuth interrupts the MCP session handshake
+func isBrokenSessionError(err error) bool {
+	if err == nil {
+		return false
+	}
+	errMsg := strings.ToLower(err.Error())
+	// The error message comes from mcp-go/mcp/streamable.go:1211
+	// "broken session: 400 Bad Request"
+	return strings.Contains(errMsg, "broken session")
+}