change tokensources into resolved strings

Author: dishaprakash

core/transport/mcp/v20250326/mcp.go

@@ -27,7 +27,6 @@ import (
 	"github.com/google/uuid"
 	"github.com/googleapis/mcp-toolbox-sdk-go/core/transport"
 	"github.com/googleapis/mcp-toolbox-sdk-go/core/transport/mcp"
-	"golang.org/x/oauth2"
 )
 
 const (
@@ -63,16 +62,11 @@ func New(baseURL string, client *http.Client) (*McpTransport, error) {
 }
 
 // ListTools fetches available tools
-func (t *McpTransport) ListTools(ctx context.Context, toolsetName string, headers map[string]oauth2.TokenSource) (*transport.ManifestSchema, error) {
+func (t *McpTransport) ListTools(ctx context.Context, toolsetName string, headers map[string]string) (*transport.ManifestSchema, error) {
 	if err := t.EnsureInitialized(ctx); err != nil {
 		return nil, err
 	}
 
-	finalHeaders, err := t.resolveHeaders(headers)
-	if err != nil {
-		return nil, err
-	}
-
 	// Append toolset name to base URL if provided
 	requestURL := t.BaseURL()
 	if toolsetName != "" {
@@ -84,7 +78,7 @@ func (t *McpTransport) ListTools(ctx context.Context, toolsetName string, header
 	}
 
 	var result listToolsResult
-	if _, err := t.sendRequest(ctx, requestURL, "tools/list", map[string]any{}, finalHeaders, &result); err != nil {
+	if _, err := t.sendRequest(ctx, requestURL, "tools/list", map[string]any{}, headers, &result); err != nil {
 		return nil, fmt.Errorf("failed to list tools: %w", err)
 	}
 
@@ -117,7 +111,7 @@ func (t *McpTransport) ListTools(ctx context.Context, toolsetName string, header
 }
 
 // GetTool fetches a single tool
-func (t *McpTransport) GetTool(ctx context.Context, toolName string, headers map[string]oauth2.TokenSource) (*transport.ManifestSchema, error) {
+func (t *McpTransport) GetTool(ctx context.Context, toolName string, headers map[string]string) (*transport.ManifestSchema, error) {
 	manifest, err := t.ListTools(ctx, "", headers)
 	if err != nil {
 		return nil, err
@@ -135,22 +129,17 @@ func (t *McpTransport) GetTool(ctx context.Context, toolName string, headers map
 }
 
 // InvokeTool executes a tool
-func (t *McpTransport) InvokeTool(ctx context.Context, toolName string, payload map[string]any, headers map[string]oauth2.TokenSource) (any, error) {
+func (t *McpTransport) InvokeTool(ctx context.Context, toolName string, payload map[string]any, headers map[string]string) (any, error) {
 	if err := t.EnsureInitialized(ctx); err != nil {
 		return "", err
 	}
 
-	finalHeaders, err := t.resolveHeaders(headers)
-	if err != nil {
-		return "", err
-	}
-
 	params := callToolRequestParams{
 		Name:      toolName,
 		Arguments: payload,
 	}
 	var result callToolResult
-	if _, err := t.sendRequest(ctx, t.BaseURL(), "tools/call", params, finalHeaders, &result); err != nil {
+	if _, err := t.sendRequest(ctx, t.BaseURL(), "tools/call", params, headers, &result); err != nil {
 		return "", fmt.Errorf("failed to invoke tool '%s': %w", toolName, err)
 	}
 
@@ -223,27 +212,6 @@ func (t *McpTransport) initializeSession(ctx context.Context) error {
 	return err
 }
 
-// resolveHeaders converts a map of TokenSources into standard HTTP headers.
-func (t *McpTransport) resolveHeaders(sources map[string]oauth2.TokenSource) (map[string]string, error) {
-	if sources == nil {
-		return nil, nil
-	}
-
-	headers := make(map[string]string, len(sources))
-	for headerKey, source := range sources {
-		if source == nil {
-			continue
-		}
-
-		token, err := source.Token()
-		if err != nil {
-			return nil, fmt.Errorf("failed to get token for header %s: %w", headerKey, err)
-		}
-		headers[headerKey] = token.AccessToken
-	}
-	return headers, nil
-}
-
 // sendRequest sends a JSON-RPC request and injects the Session ID if active.
 func (t *McpTransport) sendRequest(ctx context.Context, url string, method string, params any, headers map[string]string, dest any) (http.Header, error) {
 

core/transport/mcp/v20250326/mcp_test.go

@@ -27,7 +27,6 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"golang.org/x/oauth2"
 )
 
 // mockMCPServer is a helper to mock MCP JSON-RPC responses
@@ -316,8 +315,7 @@ func TestListTools_WithAuthHeaders(t *testing.T) {
 	}
 
 	client, _ := New(server.URL, server.Client())
-	ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "secret"})
-	headers := map[string]oauth2.TokenSource{"Authorization": ts}
+	headers := map[string]string{"Authorization": "secret"}
 
 	_, err := client.ListTools(context.Background(), "", headers)
 	require.NoError(t, err)
@@ -442,31 +440,6 @@ func TestListTools_InitFailure(t *testing.T) {
 	assert.Contains(t, err.Error(), "http request failed")
 }
 
-type failingTokenSource struct{}
-
-func (f *failingTokenSource) Token() (*oauth2.Token, error) {
-	return nil, errors.New("token failure")
-}
-
-func TestHeaders_ResolutionError(t *testing.T) {
-	// Fix: Use mock server to pass initialization so we hit the header resolution logic
-	server := newMockMCPServer()
-	defer server.Close()
-
-	client, _ := New(server.URL, server.Client())
-	headers := map[string]oauth2.TokenSource{"auth": &failingTokenSource{}}
-
-	// ListTools: EnsureInitialized succeeds, then header resolution fails
-	_, err := client.ListTools(context.Background(), "", headers)
-	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "token failure")
-
-	// InvokeTool: EnsureInitialized succeeds, then header resolution fails
-	_, err = client.InvokeTool(context.Background(), "tool", nil, headers)
-	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "token failure")
-}
-
 func TestInit_NotificationFailure(t *testing.T) {
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		var req jsonRPCRequest