Merge pull request #412 from rumpl/oauth

OAuth minor tweaks

Author: rumpl

pkg/teamloader/teamloader.go

@@ -349,7 +349,7 @@ func createTool(ctx context.Context, toolset latest.Toolset, a *latest.AgentConf
 			headers[k] = expanded
 		}
 
-		return mcp.NewToolsetRemote(toolset.Remote.URL, toolset.Remote.TransportType, headers, toolset.Tools, runtimeConfig.RedirectURI)
+		return mcp.NewRemoteToolset(toolset.Remote.URL, toolset.Remote.TransportType, headers, toolset.Tools, runtimeConfig.RedirectURI)
 
 	default:
 		return nil, fmt.Errorf("unknown toolset type: %s", toolset.Type)

pkg/tools/mcp/remote.go

@@ -78,40 +78,28 @@ func NewRemoteClient(url, transportType string, headers map[string]string, redir
 	var c *client.Client
 	var err error
 
-	if requiresOAuth {
-		oauthConfig := client.OAuthConfig{
-			RedirectURI: redirectURI,
-			TokenStore:  tokenStore,
-			PKCEEnabled: true,
-		}
+	oauthConfig := client.OAuthConfig{
+		RedirectURI: redirectURI,
+		TokenStore:  tokenStore,
+		PKCEEnabled: true,
+	}
 
-		if transportType == "sse" {
-			c, err = client.NewOAuthSSEClient(url, oauthConfig)
-			if err != nil {
-				slog.Error("Failed to create OAuth SSE remote MCP client", "error", err)
-				return nil, fmt.Errorf("failed to create OAuth SSE remote MCP client: %w", err)
-			}
-		} else {
-			c, err = client.NewOAuthStreamableHttpClient(url, oauthConfig)
-			if err != nil {
-				slog.Error("Failed to create OAuth streamable remote MCP client", "error", err)
-				return nil, fmt.Errorf("failed to create OAuth streamable remote MCP client: %w", err)
-			}
+	if transportType == "sse" {
+		options := []transport.ClientOption{transport.WithHeaders(headers)}
+		if requiresOAuth {
+			options = append(options, transport.WithOAuth(oauthConfig))
 		}
+		c, err = client.NewSSEMCPClient(url, options...)
 	} else {
-		if transportType == "sse" {
-			c, err = client.NewSSEMCPClient(url, client.WithHeaders(headers))
-			if err != nil {
-				slog.Error("Failed to create sse remote MCP client", "error", err)
-				return nil, fmt.Errorf("failed to create sse remote MCP client: %w", err)
-			}
-		} else {
-			c, err = client.NewStreamableHttpClient(url, transport.WithHTTPHeaders(headers))
-			if err != nil {
-				slog.Error("Failed to create streamable remote MCP client", "error", err)
-				return nil, fmt.Errorf("failed to create streamable remote MCP client: %w", err)
-			}
+		options := []transport.StreamableHTTPCOption{transport.WithHTTPHeaders(headers)}
+		if requiresOAuth {
+			options = append(options, transport.WithHTTPOAuth(oauthConfig))
 		}
+		c, err = client.NewStreamableHttpClient(url, options...)
+	}
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to create MCP client: %w", err)
 	}
 
 	slog.Debug("Created remote MCP client successfully", "url", url, "transport", transportType, "requiresOAuth", requiresOAuth)

pkg/tools/mcp/toolset.go

@@ -32,8 +32,8 @@ func NewToolsetCommand(command string, args, env, toolFilter []string) *Toolset
 	}
 }
 
-// NewToolsetRemote creates a new MCP toolset from a remote MCP Server.
-func NewToolsetRemote(url, transport string, headers map[string]string, toolFilter []string, redirectURI string) (*Toolset, error) {
+// NewRemoteToolset creates a new MCP toolset from a remote MCP Server.
+func NewRemoteToolset(url, transport string, headers map[string]string, toolFilter []string, redirectURI string) (*Toolset, error) {
 	slog.Debug("Creating Remote MCP toolset", "url", url, "transport", transport, "headers", headers, "toolFilter", toolFilter, "redirectURI", redirectURI)
 
 	tokenStore := GetTokenStore(url)