Fix CE mode OAuth and preserve raw tool arguments through gateway

Three fixes for running the MCP Gateway without Docker Desktop (CE mode):

1. **Forward tool arguments as raw JSON** (`handlers.go`):
   The gateway previously unmarshaled `CallToolParamsRaw.Arguments` into
   `any` before forwarding to tool handlers. This loses type fidelity for
   tools that rely on structured/typed inputs. Arguments are now forwarded
   as `json.RawMessage` unchanged, keeping the gateway schema-agnostic.

2. **Normalize argument types in clientpool** (`clientpool.go`):
   `runToolContainer` used a single type assertion (`map[string]any`) which
   silently dropped arguments arriving as `json.RawMessage` or `[]byte`.
   A `normalizeArguments` function now handles all expected argument
   representations safely via type switch.

3. **CE mode OAuth redirect URI and state validation** (`manager.go`):
   - When `DOCKER_MCP_USE_CE=true`, the redirect URI now defaults to the
     local callback (`http://localhost:5000/callback`) instead of the SaaS
     endpoint (`mcp.docker.com`), with override via
     `DOCKER_MCP_OAUTH_REDIRECT_URI`.
   - `ExchangeCode` now strips the `mcp-gateway:PORT:` prefix from the
     state parameter before validation, fixing a mismatch where
     `BuildAuthorizationURL` adds the prefix for proxy routing but
     `StateManager` only stores the base UUID.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: nathanielosullivanmolloy

pkg/gateway/clientpool.go

@@ -2,6 +2,7 @@ package gateway
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"os"
 	"os/exec"
@@ -211,6 +212,36 @@ func (cp *clientPool) InvalidateOAuthClients(provider string) {
 	}
 }
 
+// normalizeArguments converts tool arguments from various representations
+// (json.RawMessage, []byte, map[string]any, nil) into a consistent
+// map[string]any for template evaluation. MCP transports may deliver
+// arguments in any of these forms depending on the caller.
+func normalizeArguments(args any) map[string]any {
+	switch v := args.(type) {
+	case map[string]any:
+		return v
+	case json.RawMessage:
+		var m map[string]any
+		if err := json.Unmarshal(v, &m); err != nil {
+			log.Logf("Warning: failed to decode tool arguments RawMessage: %v", err)
+			return make(map[string]any)
+		}
+		return m
+	case []byte:
+		var m map[string]any
+		if err := json.Unmarshal(v, &m); err != nil {
+			log.Logf("Warning: failed to decode tool arguments JSON: %v", err)
+			return make(map[string]any)
+		}
+		return m
+	case nil:
+		return make(map[string]any)
+	default:
+		log.Logf("Warning: unsupported tool arguments type: %T", args)
+		return make(map[string]any)
+	}
+}
+
 func (cp *clientPool) runToolContainer(ctx context.Context, tool catalog.Tool, params *mcp.CallToolParams) (*mcp.CallToolResult, error) {
 	args := cp.baseArgs(tool.Name)
 
@@ -219,11 +250,7 @@ func (cp *clientPool) runToolContainer(ctx context.Context, tool catalog.Tool, p
 		args = append(args, "--network", network)
 	}
 
-	// Convert params.Arguments to map[string]any
-	arguments, ok := params.Arguments.(map[string]any)
-	if !ok {
-		arguments = make(map[string]any)
-	}
+	arguments := normalizeArguments(params.Arguments)
 
 	// Volumes
 	for _, mount := range eval.EvaluateList(tool.Container.Volumes, arguments) {

pkg/gateway/clientpool_test.go

@@ -2,6 +2,7 @@ package gateway
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"os"
 	"testing"
@@ -278,6 +279,59 @@ func parseConfig(t *testing.T, contentYAML string) map[string]any {
 	return config
 }
 
+func TestNormalizeArguments_MapStringAny(t *testing.T) {
+	input := map[string]any{"key": "value", "num": float64(42)}
+	result := normalizeArguments(input)
+	assert.Equal(t, input, result)
+}
+
+func TestNormalizeArguments_JSONRawMessage(t *testing.T) {
+	raw := json.RawMessage(`{"url":"https://example.com","count":3}`)
+	result := normalizeArguments(raw)
+	assert.Equal(t, "https://example.com", result["url"])
+	assert.Equal(t, float64(3), result["count"])
+}
+
+func TestNormalizeArguments_ByteSlice(t *testing.T) {
+	raw := []byte(`{"path":"/tmp/data","verbose":true}`)
+	result := normalizeArguments(raw)
+	assert.Equal(t, "/tmp/data", result["path"])
+	assert.Equal(t, true, result["verbose"])
+}
+
+func TestNormalizeArguments_Nil(t *testing.T) {
+	result := normalizeArguments(nil)
+	assert.NotNil(t, result)
+	assert.Empty(t, result)
+}
+
+func TestNormalizeArguments_UnexpectedType(t *testing.T) {
+	result := normalizeArguments("unexpected string")
+	assert.NotNil(t, result)
+	assert.Empty(t, result)
+}
+
+func TestNormalizeArguments_InvalidJSON(t *testing.T) {
+	raw := json.RawMessage(`{not valid json}`)
+	result := normalizeArguments(raw)
+	assert.NotNil(t, result)
+	assert.Empty(t, result)
+}
+
+func TestNormalizeArguments_InvalidByteSlice(t *testing.T) {
+	raw := []byte(`{not valid json}`)
+	result := normalizeArguments(raw)
+	assert.NotNil(t, result)
+	assert.Empty(t, result)
+}
+
+func TestNormalizeArguments_EmptyJSONObject(t *testing.T) {
+	raw := json.RawMessage(`{}`)
+	result := normalizeArguments(raw)
+	assert.NotNil(t, result)
+	assert.Empty(t, result)
+}
+
 func TestInvalidateOAuthClients_MatchesCommunityServer(t *testing.T) {
 	// Community server: remote URL set, but no Spec.OAuth metadata.
 	// This verifies Gap 3: InvalidateOAuthClients matches community servers

pkg/gateway/handlers.go

@@ -2,7 +2,6 @@ package gateway
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 	"os"
 	"time"
@@ -43,18 +42,22 @@ func inferServerTransportType(serverConfig *catalog.ServerConfig) string {
 
 func (g *Gateway) mcpToolHandler(tool catalog.Tool) mcp.ToolHandler {
 	return func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-		// Convert CallToolParamsRaw to CallToolParams
-		var args any
-		if len(req.Params.Arguments) > 0 {
-			if err := json.Unmarshal(req.Params.Arguments, &args); err != nil {
-				return nil, fmt.Errorf("failed to unmarshal arguments: %w", err)
-			}
-		}
+		// Convert CallToolParamsRaw to CallToolParams.
+		//
+		// Arguments are forwarded as raw JSON (json.RawMessage) and intentionally
+		// not unmarshaled here. The gateway must remain schema-agnostic and avoid
+		// coercing tool inputs, preserving full argument fidelity for tools that
+		// rely on structured or typed inputs.
 		params := &mcp.CallToolParams{
-			Meta:      req.Params.Meta,
-			Name:      req.Params.Name,
-			Arguments: args,
+			Meta: req.Params.Meta,
+			Name: req.Params.Name,
 		}
+
+		// Forward raw arguments unchanged, if present.
+		if len(req.Params.Arguments) > 0 {
+			params.Arguments = req.Params.Arguments
+		}
+
 		return g.clientPool.runToolContainer(ctx, tool, params)
 	}
 }
@@ -132,19 +135,21 @@ func (g *Gateway) mcpServerToolHandler(serverName string, server *mcp.Server, _
 		}
 		defer g.clientPool.ReleaseClient(client)
 
-		// Convert CallToolParamsRaw to CallToolParams
-		var args any
-		if len(req.Params.Arguments) > 0 {
-			if jsonErr := json.Unmarshal(req.Params.Arguments, &args); jsonErr != nil {
-				telemetry.RecordToolError(ctx, span, serverConfig.Name, serverTransportType, req.Params.Name)
-				span.SetStatus(codes.Error, "Failed to unmarshal arguments")
-				return nil, fmt.Errorf("failed to unmarshal arguments: %w", jsonErr)
-			}
-		}
+		// Convert CallToolParamsRaw to CallToolParams.
+		//
+		// NOTE: Arguments are forwarded as raw JSON (json.RawMessage) instead of being
+		// unmarshaled here. The gateway must not interpret or coerce tool arguments,
+		// as it does not own the tool schema. Preserving the raw payload ensures full
+		// fidelity for schema-based and typed tools and matches the MCP Go SDK
+		// expectations.
 		params := &mcp.CallToolParams{
-			Meta:      req.Params.Meta,
-			Name:      originalToolName,
-			Arguments: args,
+			Meta: req.Params.Meta,
+			Name: originalToolName,
+		}
+
+		// Forward raw arguments unchanged, if present.
+		if len(req.Params.Arguments) > 0 {
+			params.Arguments = req.Params.Arguments
 		}
 
 		// Execute the tool call

pkg/oauth/manager.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 	"net/url"
+	"os"
+	"strings"
 
 	"github.com/docker/docker-credential-helpers/credentials"
 	"golang.org/x/oauth2"
@@ -23,13 +25,31 @@ type Manager struct {
 	redirectURI  string
 }
 
-// NewManager creates a new OAuth manager for CE mode
+// NewManager creates a new OAuth manager.
+// In CE mode, the redirect URI must point to the local gateway callback
+// instead of the SaaS endpoint (mcp.docker.com).
 func NewManager(credHelper credentials.Helper) *Manager {
+	redirectURI := DefaultRedirectURI
+
+	// CE mode requires a local redirect URI because the OAuth callback
+	// is handled by the local mcp-gateway process, not by Docker SaaS.
+	//
+	// Example:
+	//   http://localhost:5000/callback
+	if os.Getenv("DOCKER_MCP_USE_CE") == "true" {
+		if v := os.Getenv("DOCKER_MCP_OAUTH_REDIRECT_URI"); v != "" {
+			redirectURI = v
+		} else {
+			// Default CE callback used by the local OAuth proxy
+			redirectURI = "http://localhost:5000/callback"
+		}
+	}
+
 	return &Manager{
-		dcrManager:   dcr.NewManager(credHelper, DefaultRedirectURI),
+		dcrManager:   dcr.NewManager(credHelper, redirectURI),
 		tokenStore:   NewTokenStore(credHelper),
 		stateManager: NewStateManager(),
-		redirectURI:  DefaultRedirectURI,
+		redirectURI:  redirectURI,
 	}
 }
 
@@ -125,6 +145,16 @@ func (m *Manager) BuildAuthorizationURL(_ context.Context, serverName string, sc
 
 // ExchangeCode exchanges an authorization code for an access token
 func (m *Manager) ExchangeCode(ctx context.Context, code string, state string) error {
+	// Strip the mcp-gateway:PORT: prefix if present.
+	// BuildAuthorizationURL formats state as "mcp-gateway:PORT:UUID" for proxy routing,
+	// but the StateManager only stores the base UUID.
+	if strings.HasPrefix(state, "mcp-gateway:") {
+		parts := strings.SplitN(state, ":", 3)
+		if len(parts) == 3 {
+			state = parts[2]
+		}
+	}
+
 	// Validate state and retrieve verifier
 	serverName, verifier, err := m.stateManager.Validate(state)
 	if err != nil {

pkg/oauth/manager_test.go

@@ -2,6 +2,7 @@ package oauth
 
 import (
 	"context"
+	"os"
 	"strings"
 	"testing"
 	"time"
@@ -344,6 +345,101 @@ func TestManager_CallbackURLParsing(t *testing.T) {
 	}
 }
 
+func TestManager_ExchangeCode_StripsStatePrefix(t *testing.T) {
+	manager := setupTestManager(t)
+	serverName := "test-server"
+
+	setupTestDCRClient(t, manager, serverName)
+
+	// Generate a state via BuildAuthorizationURL with a callback URL,
+	// which produces "mcp-gateway:PORT:UUID" format.
+	_, baseState, _, err := manager.BuildAuthorizationURL(
+		context.Background(),
+		serverName,
+		[]string{"read"},
+		"http://localhost:8080/callback",
+	)
+	require.NoError(t, err)
+
+	// Simulate the prefixed state that would come back from the OAuth callback
+	prefixedState := "mcp-gateway:8080:" + baseState
+
+	// ExchangeCode will fail at token exchange (no real server), but it should
+	// get past state validation — meaning the prefix was correctly stripped.
+	err = manager.ExchangeCode(context.Background(), "test-code", prefixedState)
+	require.Error(t, err)
+	// If prefix stripping failed, we'd get "invalid state parameter".
+	// If it succeeded, we get a token exchange error instead.
+	assert.NotContains(t, err.Error(), "invalid state parameter")
+}
+
+func TestManager_ExchangeCode_PlainStateStillWorks(t *testing.T) {
+	manager := setupTestManager(t)
+	serverName := "test-server"
+
+	setupTestDCRClient(t, manager, serverName)
+
+	// Generate a state without callback URL (no prefix)
+	_, baseState, _, err := manager.BuildAuthorizationURL(
+		context.Background(),
+		serverName,
+		[]string{"read"},
+		"",
+	)
+	require.NoError(t, err)
+
+	// ExchangeCode should still validate plain UUIDs (no prefix to strip)
+	err = manager.ExchangeCode(context.Background(), "test-code", baseState)
+	require.Error(t, err)
+	assert.NotContains(t, err.Error(), "invalid state parameter")
+}
+
+func TestManager_NewManager_CEModeRedirectURI(t *testing.T) {
+	// Save and restore env vars
+	origCE := os.Getenv("DOCKER_MCP_USE_CE")
+	origURI := os.Getenv("DOCKER_MCP_OAUTH_REDIRECT_URI")
+	defer func() {
+		os.Setenv("DOCKER_MCP_USE_CE", origCE)
+		os.Setenv("DOCKER_MCP_OAUTH_REDIRECT_URI", origURI)
+	}()
+
+	t.Run("default mode uses SaaS redirect", func(t *testing.T) {
+		os.Setenv("DOCKER_MCP_USE_CE", "")
+		os.Setenv("DOCKER_MCP_OAUTH_REDIRECT_URI", "")
+
+		helper := newFakeCredentialHelper()
+		manager := NewManager(helper)
+		assert.Equal(t, DefaultRedirectURI, manager.redirectURI)
+	})
+
+	t.Run("CE mode uses localhost redirect", func(t *testing.T) {
+		os.Setenv("DOCKER_MCP_USE_CE", "true")
+		os.Setenv("DOCKER_MCP_OAUTH_REDIRECT_URI", "")
+
+		helper := newFakeCredentialHelper()
+		manager := NewManager(helper)
+		assert.Equal(t, "http://localhost:5000/callback", manager.redirectURI)
+	})
+
+	t.Run("CE mode with custom redirect URI", func(t *testing.T) {
+		os.Setenv("DOCKER_MCP_USE_CE", "true")
+		os.Setenv("DOCKER_MCP_OAUTH_REDIRECT_URI", "http://localhost:9999/custom")
+
+		helper := newFakeCredentialHelper()
+		manager := NewManager(helper)
+		assert.Equal(t, "http://localhost:9999/custom", manager.redirectURI)
+	})
+
+	t.Run("non-CE mode ignores custom redirect URI", func(t *testing.T) {
+		os.Setenv("DOCKER_MCP_USE_CE", "false")
+		os.Setenv("DOCKER_MCP_OAUTH_REDIRECT_URI", "http://localhost:9999/custom")
+
+		helper := newFakeCredentialHelper()
+		manager := NewManager(helper)
+		assert.Equal(t, DefaultRedirectURI, manager.redirectURI)
+	})
+}
+
 func TestManager_StateFormatWithPort(t *testing.T) {
 	manager := setupTestManager(t)
 	serverName := "test-server"