coder/coderd/x/chatd/mcpclient/mcpclient.go

package mcpclient

import (
	"cmp"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"slices"
	"strings"
	"sync"
	"time"

	"charm.land/fantasy"
	"github.com/google/uuid"
	"github.com/modelcontextprotocol/go-sdk/mcp"
	"golang.org/x/oauth2"
	"golang.org/x/sync/errgroup"
	"golang.org/x/xerrors"

	"cdr.dev/slog/v3"
	aidmcp "github.com/coder/coder/v2/aibridge/mcp"
	"github.com/coder/coder/v2/buildinfo"
	"github.com/coder/coder/v2/coderd/database"
)

// toolNameSep separates the server slug from the original tool
// name in prefixed tool names. Double underscore avoids collisions
// with tool names that may contain single underscores.
//
// TODO: tool names that themselves contain "__" produce ambiguous
// prefixed names (e.g. "srv__my__tool" is indistinguishable from
// slug "srv" + tool "my__tool" vs slug "srv__my" + tool "tool").
// This doesn't affect tool invocation since originalName is used
// directly when calling the remote server.
const toolNameSep = "__"

// truncateToolName caps the assembled tool name at MaxToolNameLen so
// it fits within provider limits (e.g. OpenAI 64, Bedrock 128).
func truncateToolName(name string) string {
	if len(name) > aidmcp.MaxToolNameLen {
		return name[:aidmcp.MaxToolNameLen]
	}
	return name
}

// connectTimeout bounds how long we wait for a single MCP server
// to start its transport and complete initialization. Servers that
// take longer are skipped so one slow server cannot block the
// entire chat startup.
const connectTimeout = 10 * time.Second

// toolCallTimeout bounds how long a single tool invocation may
// take before being canceled.
const toolCallTimeout = 60 * time.Second

// UserOIDCTokenSource resolves the OIDC access token for the calling
// user. Implementations attempt to refresh tokens that are expired
// or close to expiring and MUST return ("", nil) when the user has
// no OIDC link or a refresh attempt failed for any reason. A
// non-nil error is reserved for unexpected infrastructure failures
// (e.g. database errors) and skips header construction entirely.
// The empty-token-on-refresh-failure behavior matches
// provisionerdserver.ObtainOIDCAccessToken.
type UserOIDCTokenSource interface {
	OIDCAccessToken(ctx context.Context, userID uuid.UUID) (string, error)
}

// ConnectAll connects to all configured MCP servers, discovers
// their tools, and returns them as fantasy.AgentTool values.
// Tools are sorted by their prefixed name so callers
// receive a deterministic order. It skips servers that fail to
// connect and logs warnings. The returned cleanup function
// must be called to close all connections.
func ConnectAll(
	ctx context.Context,
	logger slog.Logger,
	configs []database.MCPServerConfig,
	tokens []database.MCPServerUserToken,
	userID uuid.UUID,
	oidcSrc UserOIDCTokenSource,
	coderHeaders map[string]string,
) ([]fantasy.AgentTool, func()) {
	// Index tokens by server config ID so auth header
	// construction is O(1) per server.
	tokensByConfigID := make(
		map[uuid.UUID]database.MCPServerUserToken, len(tokens),
	)
	for _, tok := range tokens {
		tokensByConfigID[tok.MCPServerConfigID] = tok
	}

	var (
		mu       sync.Mutex
		sessions []*mcp.ClientSession
		tools    []fantasy.AgentTool
	)

	// Build cleanup eagerly so it always closes any sessions
	// that connected, even if a later connection fails.
	cleanup := func() {
		mu.Lock()
		defer mu.Unlock()
		for _, s := range sessions {
			_ = s.Close()
		}
		sessions = nil
	}

	var eg errgroup.Group
	for _, cfg := range configs {
		if !cfg.Enabled {
			continue
		}

		eg.Go(func() error {
			serverTools, session, connectErr := connectOne(
				ctx, logger, cfg, tokensByConfigID, userID, oidcSrc, coderHeaders,
			)
			if connectErr != nil {
				logger.Warn(ctx,
					"skipping MCP server due to connection failure",
					slog.F("server_slug", cfg.Slug),
					slog.F("server_url", RedactURL(cfg.Url)),
					slog.F("error", redactErrorURL(connectErr)),
				)
				// Connection failures are not propagated — the
				// LLM simply won't have this server's tools.
				return nil
			}

			mu.Lock()
			if session != nil {
				sessions = append(sessions, session)
			}
			tools = append(tools, serverTools...)
			mu.Unlock()
			return nil
		})
	}

	// All goroutines return nil; error is intentionally
	// discarded.
	_ = eg.Wait()

	// Sort tools by prefixed name for deterministic ordering
	// regardless of goroutine completion order. Ties, possible
	// when the __ separator produces ambiguous prefixed names,
	// are broken by config ID. Stable prompt construction
	// depends on consistent tool ordering.
	slices.SortFunc(tools, func(a, b fantasy.AgentTool) int {
		// All tools in this slice are mcpToolWrapper values
		// created by connectOne above, so these checked
		// assertions should always succeed. The config ID
		// tiebreaker resolves the __ separator ambiguity
		// documented at the top of this file.
		aTool, ok := a.(MCPToolIdentifier)
		if !ok {
			panic(fmt.Sprintf("unexpected tool type %T", a))
		}
		bTool, ok := b.(MCPToolIdentifier)
		if !ok {
			panic(fmt.Sprintf("unexpected tool type %T", b))
		}
		return cmp.Or(
			cmp.Compare(a.Info().Name, b.Info().Name),
			cmp.Compare(aTool.MCPServerConfigID().String(), bTool.MCPServerConfigID().String()),
		)
	})

	// Warn about name collisions that may result from sanitization
	// and truncation. When two tools resolve to the same name, the
	// LLM tool-call dispatch map keeps only one, so the other
	// becomes silently unreachable.
	for i := 1; i < len(tools); i++ {
		if tools[i-1].Info().Name == tools[i].Info().Name {
			prevTool, ok := tools[i-1].(MCPToolIdentifier)
			if !ok {
				continue
			}
			currTool, ok := tools[i].(MCPToolIdentifier)
			if !ok {
				continue
			}
			if prevTool.MCPServerConfigID() != currTool.MCPServerConfigID() {
				logger.Warn(ctx,
					"duplicate tool name after sanitization; one tool will be unreachable",
					slog.F("tool_name", tools[i].Info().Name),
					slog.F("prev_config_id", prevTool.MCPServerConfigID()),
					slog.F("curr_config_id", currTool.MCPServerConfigID()),
				)
			}
		}
	}

	return tools, cleanup
}

// connectOne establishes a connection to a single MCP server,
// discovers its tools, and wraps each one as an AgentTool with
// the server slug prefix applied.
func connectOne(
	ctx context.Context,
	logger slog.Logger,
	cfg database.MCPServerConfig,
	tokensByConfigID map[uuid.UUID]database.MCPServerUserToken,
	userID uuid.UUID,
	oidcSrc UserOIDCTokenSource,
	coderHeaders map[string]string,
) ([]fantasy.AgentTool, *mcp.ClientSession, error) {
	headers := buildAuthHeaders(ctx, logger, cfg, tokensByConfigID, userID, oidcSrc)

	// When opted-in, merge Coder identity headers BEFORE the
	// transport is created so any auth header already set above
	// wins on a conflict. Conflict detection uses
	// http.CanonicalHeaderKey because the upstream transport applies
	// http.Header.Set, which canonicalizes keys; without that, an
	// admin-configured header that differs only in case from a Coder
	// identity header would land in the request map twice and the
	// surviving value would be non-deterministic.
	if cfg.ForwardCoderHeaders {
		canonicalAuth := make(map[string]struct{}, len(headers))
		for k := range headers {
			canonicalAuth[http.CanonicalHeaderKey(k)] = struct{}{}
		}
		for k, v := range coderHeaders {
			if _, exists := canonicalAuth[http.CanonicalHeaderKey(k)]; exists {
				continue
			}
			headers[k] = v
		}
	}

	tr, err := createTransport(cfg, headers)
	if err != nil {
		return nil, nil, xerrors.Errorf(
			"create transport: %w", err,
		)
	}

	mcpClient := mcp.NewClient(&mcp.Implementation{
		Name:    "coder",
		Version: buildinfo.Version(),
	}, nil)

	// The timeout covers the entire connect+list sequence, not
	// each phase individually. The SDK negotiates the protocol
	// version during Connect; the session outlives connectCtx.
	connectCtx, cancel := context.WithTimeout(
		ctx, connectTimeout,
	)
	defer cancel()

	session, err := mcpClient.Connect(connectCtx, tr, nil)
	if err != nil {
		return nil, nil, xerrors.Errorf("connect: %w", err)
	}

	toolsResult, err := session.ListTools(connectCtx, nil)
	if err != nil {
		_ = session.Close()
		return nil, nil, xerrors.Errorf("list tools: %w", err)
	}

	var tools []fantasy.AgentTool
	for _, mcpTool := range toolsResult.Tools {
		if !isToolAllowed(
			mcpTool.Name,
			cfg.ToolAllowList,
			cfg.ToolDenyList,
		) {
			logger.Debug(ctx, "skipping denied MCP tool",
				slog.F("server_slug", cfg.Slug),
				slog.F("tool_name", mcpTool.Name),
			)
			continue
		}

		tools = append(
			tools, newMCPTool(cfg.ID, cfg.Slug, mcpTool, session, cfg.ModelIntent),
		)
	}

	if len(tools) == 0 {
		_ = session.Close()
		return nil, nil, nil
	}

	return tools, session, nil
}

func createTransport(
	cfg database.MCPServerConfig,
	headers map[string]string,
) (mcp.Transport, error) {
	httpClient := httpClientWithHeaders(headers)

	switch cfg.Transport {
	case "sse":
		return &mcp.SSEClientTransport{
			Endpoint:   cfg.Url,
			HTTPClient: httpClient,
		}, nil
	case "", "streamable_http":
		// Default to streamable HTTP, the newer transport.
		return &mcp.StreamableClientTransport{
			Endpoint:   cfg.Url,
			HTTPClient: httpClient,
		}, nil
	default:
		return nil, xerrors.Errorf(
			"unsupported transport %q", cfg.Transport,
		)
	}
}

// buildAuthHeaders constructs HTTP headers for authenticating
// with the MCP server based on the configured auth type.
func buildAuthHeaders(
	ctx context.Context,
	logger slog.Logger,
	cfg database.MCPServerConfig,
	tokensByConfigID map[uuid.UUID]database.MCPServerUserToken,
	userID uuid.UUID,
	oidcSrc UserOIDCTokenSource,
) map[string]string {
	headers := make(map[string]string)

	switch cfg.AuthType {
	case "oauth2":
		tok, ok := tokensByConfigID[cfg.ID]
		if !ok {
			logger.Warn(ctx,
				"no oauth2 token found for MCP server",
				slog.F("server_slug", cfg.Slug),
			)
			break
		}
		if tok.OauthRefreshFailureReason != "" {
			// The grant is permanently unusable (e.g. revoked
			// upstream) and the user must reconnect. Do not attach
			// any leftover token material.
			logger.Warn(ctx,
				"oauth2 token for MCP server requires reconnect, skipping auth header",
				slog.F("server_slug", cfg.Slug),
			)
			break
		}
		if tok.Expiry.Valid && tok.Expiry.Time.Before(time.Now()) {
			logger.Warn(ctx,
				"oauth2 token for MCP server is expired",
				slog.F("server_slug", cfg.Slug),
				slog.F("expired_at", tok.Expiry.Time),
			)
		}
		if tok.AccessToken == "" {
			logger.Warn(ctx,
				"oauth2 token record has empty access token",
				slog.F("server_slug", cfg.Slug),
			)
			break
		}
		tokenType := tok.TokenType
		if tokenType == "" {
			tokenType = "Bearer"
		}
		// RFC 6750 says the scheme is case-insensitive, but
		// some servers (e.g. Linear) reject lowercase
		// "bearer". Normalize to the canonical form.
		if strings.EqualFold(tokenType, "bearer") {
			tokenType = "Bearer"
		}
		headers["Authorization"] = tokenType + " " + tok.AccessToken
	case "api_key":
		if cfg.APIKeyHeader != "" && cfg.APIKeyValue != "" {
			headers[cfg.APIKeyHeader] = cfg.APIKeyValue
		}
	case "custom_headers":
		if cfg.CustomHeaders != "" {
			var custom map[string]string
			if err := json.Unmarshal(
				[]byte(cfg.CustomHeaders), &custom,
			); err != nil {
				logger.Warn(ctx,
					"failed to parse custom headers JSON",
					slog.F("server_slug", cfg.Slug),
					slog.Error(err),
				)
			} else {
				for k, v := range custom {
					headers[k] = v
				}
			}
		}
	case "user_oidc":
		// Forward the calling user's OIDC access token from
		// user_links as Authorization: Bearer <token>. The token
		// source is responsible for refreshing tokens that are
		// expired or close to expiring before returning them.
		if oidcSrc == nil || userID == uuid.Nil {
			logger.Warn(ctx,
				"user_oidc auth requested but no token source available",
				slog.F("server_slug", cfg.Slug),
			)
			break
		}
		token, err := oidcSrc.OIDCAccessToken(ctx, userID)
		if err != nil {
			logger.Warn(ctx,
				"failed to obtain user OIDC token for MCP server",
				slog.F("server_slug", cfg.Slug),
				slog.Error(err),
			)
			break
		}
		if token == "" {
			// The user has no OIDC link, or a non-fatal refresh
			// failure occurred. Fall through with no header and let
			// the upstream MCP server decide how to respond
			// (typically 401). Logged at debug so password and
			// GitHub users don't generate noise for every chat turn.
			logger.Debug(ctx,
				"no user OIDC token available for MCP server",
				slog.F("server_slug", cfg.Slug),
			)
			break
		}
		headers["Authorization"] = "Bearer " + token
	case "none", "":
		// No auth headers needed.
	}

	return headers
}

// isToolAllowed checks a tool name against the allow and deny
// lists. When the allow list is non-empty only tools in it are
// permitted and the deny list is ignored. When the allow list
// is empty and the deny list is non-empty, tools in the deny
// list are rejected. Both lists use exact string matching
// against the original (non-prefixed) tool name.
func isToolAllowed(
	toolName string,
	allowList []string,
	denyList []string,
) bool {
	if len(allowList) > 0 {
		for _, allowed := range allowList {
			if allowed == toolName {
				return true
			}
		}
		// Allow list is set but the tool isn't in it.
		return false
	}

	for _, denied := range denyList {
		if denied == toolName {
			return false
		}
	}

	return true
}

// RedactURL strips userinfo and query parameters from a URL
// to avoid logging embedded credentials. Query params are
// removed because API keys are sometimes passed as
// ?api_key=sk-... in server URLs.
func RedactURL(rawURL string) string {
	u, err := url.Parse(rawURL)
	if err != nil {
		return rawURL
	}
	u.User = nil
	u.RawQuery = ""
	u.Fragment = ""
	return u.String()
}

// redactErrorURL rewrites URLs in an error string to strip
// credentials. Go's net/http embeds the full request URL in
// *url.Error messages, which can leak userinfo.
func redactErrorURL(err error) string {
	if err == nil {
		return ""
	}
	var urlErr *url.Error
	if errors.As(err, &urlErr) {
		urlErr.URL = RedactURL(urlErr.URL)
		return urlErr.Error()
	}
	return err.Error()
}

// MCPToolIdentifier is implemented by tools that originate from
// an MCP server config and can report the config's database ID.
type MCPToolIdentifier interface {
	MCPServerConfigID() uuid.UUID
}

// mcpToolWrapper adapts a single MCP tool into a
// fantasy.AgentTool. It stores the prefixed name for Info() but
// strips the prefix when forwarding calls to the remote server.
type mcpToolWrapper struct {
	configID        uuid.UUID
	prefixedName    string
	originalName    string
	description     string
	parameters      map[string]any
	required        []string
	modelIntent     bool
	session         *mcp.ClientSession
	providerOptions fantasy.ProviderOptions
}

// MCPServerConfigID returns the database ID of the MCP server
// config that this tool originates from.
func (t *mcpToolWrapper) MCPServerConfigID() uuid.UUID {
	return t.configID
}

// newMCPTool creates an mcpToolWrapper from an mcp.Tool
// discovered on a remote server.
func newMCPTool(
	configID uuid.UUID,
	serverSlug string,
	tool *mcp.Tool,
	session *mcp.ClientSession,
	modelIntent bool,
) *mcpToolWrapper {
	properties, required := splitInputSchema(tool.InputSchema)
	return &mcpToolWrapper{
		configID:     configID,
		prefixedName: truncateToolName(aidmcp.SanitizeToolName(serverSlug) + toolNameSep + aidmcp.SanitizeToolName(tool.Name)),
		originalName: tool.Name,
		description:  tool.Description,
		parameters:   properties,
		required:     required,
		modelIntent:  modelIntent,
		session:      session,
	}
}

func splitInputSchema(schema any) (map[string]any, []string) {
	m, ok := schema.(map[string]any)
	if !ok {
		return nil, nil
	}
	properties, _ := m["properties"].(map[string]any)
	var required []string
	if rawRequired, ok := m["required"].([]any); ok {
		for _, r := range rawRequired {
			if str, ok := r.(string); ok {
				required = append(required, str)
			}
		}
	}
	return properties, required
}

func (t *mcpToolWrapper) Info() fantasy.ToolInfo {
	required := t.required
	if required == nil {
		required = []string{}
	}

	if !t.modelIntent {
		return fantasy.ToolInfo{
			Name:        t.prefixedName,
			Description: t.description,
			Parameters:  t.parameters,
			Required:    required,
			Parallel:    true,
		}
	}

	// Wrap original parameters under "properties" and add
	// "model_intent" so the LLM provides a human-readable
	// description of each tool call.
	wrapped := map[string]any{
		"model_intent": map[string]any{
			"type": "string",
			"description": "A short, natural-language, present-participle " +
				"phrase describing why you are calling this tool. " +
				"This is shown to the user as a status label while " +
				"the tool runs. Use plain English with no underscores " +
				"or technical jargon. Keep it under 100 characters. " +
				"Good examples: \"Reading the authentication module\", " +
				"\"Searching for configuration files\", " +
				"\"Creating a new workspace\".",
		},
		"properties": map[string]any{
			"type":       "object",
			"properties": t.parameters,
			"required":   required,
		},
	}
	return fantasy.ToolInfo{
		Name:        t.prefixedName,
		Description: t.description,
		Parameters:  wrapped,
		Required:    []string{"model_intent", "properties"},
		Parallel:    true,
	}
}

func (t *mcpToolWrapper) Run(
	ctx context.Context,
	params fantasy.ToolCall,
) (fantasy.ToolResponse, error) {
	input := params.Input
	if t.modelIntent {
		input = unwrapModelIntent(input)
	}

	var args map[string]any
	if input != "" {
		if err := json.Unmarshal(
			[]byte(input), &args,
		); err != nil {
			return fantasy.NewTextErrorResponse(
				"invalid JSON input: " + err.Error(),
			), nil
		}
	}

	callCtx, cancel := context.WithTimeout(ctx, toolCallTimeout)
	defer cancel()

	result, err := t.session.CallTool(
		callCtx,
		&mcp.CallToolParams{
			Name:      t.originalName,
			Arguments: args,
		},
	)
	if err != nil {
		return fantasy.NewTextErrorResponse(err.Error()), nil
	}

	return convertCallResult(result), nil
}

func (t *mcpToolWrapper) ProviderOptions() fantasy.ProviderOptions {
	return t.providerOptions
}

func (t *mcpToolWrapper) SetProviderOptions(
	opts fantasy.ProviderOptions,
) {
	t.providerOptions = opts
}

// unwrapModelIntent strips the model_intent wrapper from tool
// call input so the remote MCP server receives only the original
// arguments. It handles three shapes the model may produce:
//
//  1. { model_intent, properties: {...} } — correct format
//  2. { model_intent, key: val, ... } — flat, no properties wrapper
//  3. Anything else — returned as-is
func unwrapModelIntent(input string) string {
	var parsed map[string]any
	if err := json.Unmarshal([]byte(input), &parsed); err != nil {
		return input
	}

	delete(parsed, "model_intent")

	// Case 1: correct { model_intent, properties: {...} } format.
	if props, ok := parsed["properties"]; ok {
		if b, err := json.Marshal(props); err == nil {
			return string(b)
		}
	}

	// Case 2: flat { model_intent, key: val, ... } without wrapper.
	if b, err := json.Marshal(parsed); err == nil {
		return string(b)
	}

	return input
}

// convertCallResult translates an MCP CallToolResult into a
// fantasy.ToolResponse. The fantasy response model supports a
// single content type per response, so we prioritize text. All
// text items are collected first. Binary items (image, audio,
// or embedded blob) are only returned when no text content is
// available.
func convertCallResult(
	result *mcp.CallToolResult,
) fantasy.ToolResponse {
	if result == nil {
		return fantasy.NewTextResponse("")
	}

	var (
		textParts    []string
		binaryResult *fantasy.ToolResponse
	)
	for _, item := range result.Content {
		switch c := item.(type) {
		case *mcp.TextContent:
			textParts = append(textParts, strings.ToValidUTF8(c.Text, "\uFFFD"))
		case *mcp.ImageContent:
			// The SDK decodes base64 payloads during unmarshal, so
			// Data is raw bytes.
			if binaryResult == nil {
				r := fantasy.ToolResponse{
					Type:      "image",
					Data:      c.Data,
					MediaType: c.MIMEType,
					IsError:   result.IsError,
				}
				binaryResult = &r
			}
		case *mcp.AudioContent:
			if binaryResult == nil {
				r := fantasy.ToolResponse{
					Type:      "media",
					Data:      c.Data,
					MediaType: c.MIMEType,
					IsError:   result.IsError,
				}
				binaryResult = &r
			}
		case *mcp.EmbeddedResource:
			// Embedded resources wrap either text or blob content
			// from an MCP resource. Exactly one of Text or Blob is
			// set per the spec; a nil Blob means text content.
			switch {
			case c.Resource == nil:
				textParts = append(textParts,
					"[embedded resource with no contents]",
				)
			case c.Resource.Blob != nil:
				if binaryResult == nil {
					blobType := "media"
					if strings.HasPrefix(c.Resource.MIMEType, "image/") {
						blobType = "image"
					}
					res := fantasy.ToolResponse{
						Type:      blobType,
						Data:      c.Resource.Blob,
						MediaType: c.Resource.MIMEType,
						IsError:   result.IsError,
					}
					binaryResult = &res
				}
			default:
				textParts = append(textParts, strings.ToValidUTF8(c.Resource.Text, "\uFFFD"))
			}
		case *mcp.ResourceLink:
			// Resource links point to content the LLM can
			// reference by URI. Surface the URI so the model
			// can use it in follow-ups.
			label := c.URI
			if c.Name != "" {
				label = fmt.Sprintf("%s (%s)", c.Name, c.URI)
			}
			if c.Description != "" {
				label += ": " + c.Description
			}
			textParts = append(textParts,
				fmt.Sprintf("[resource: %s]", label),
			)
		default:
			textParts = append(textParts,
				fmt.Sprintf("[unsupported content type: %T]", c),
			)
		}
	}

	// If structured content is present, marshal it to JSON and
	// append as a text part so the data is preserved for the LLM.
	if result.StructuredContent != nil {
		data, err := json.Marshal(result.StructuredContent)
		if err != nil {
			textParts = append(textParts,
				"[structured content marshal error: "+
					err.Error()+"]",
			)
		} else {
			textParts = append(textParts, string(data))
		}
	}

	// Prefer text content. Only fall back to binary when no
	// text was collected.
	if len(textParts) > 0 {
		resp := fantasy.NewTextResponse(
			strings.Join(textParts, "\n"),
		)
		resp.IsError = result.IsError
		return resp
	}
	if binaryResult != nil {
		return *binaryResult
	}
	return fantasy.NewTextResponse("")
}

// RefreshResult contains the outcome of an OAuth2 token refresh
// attempt.
type RefreshResult struct {
	// AccessToken is the new (or unchanged) access token.
	AccessToken string
	// RefreshToken is the new (or preserved original) refresh
	// token. Providers that don't rotate refresh tokens return
	// an empty value; in that case the original is kept.
	RefreshToken string
	// TokenType is the token type (usually "Bearer").
	TokenType string
	// Expiry is the new token expiry. Zero value means no expiry
	// was provided by the provider.
	Expiry time.Time
	// Refreshed is true when the access token actually changed,
	// meaning a refresh occurred. When false the token was still
	// valid and no network call was made.
	Refreshed bool
}

// refreshFailureReasonLimit caps the error text persisted to
// mcp_server_user_tokens.oauth_refresh_failure_reason, matching the
// external auth failure reason limit.
const refreshFailureReasonLimit = 400

// IsPermanentRefreshError reports whether an OAuth2 token refresh
// error means the user's grant is permanently unusable (for example
// the upstream grant was revoked) rather than a transient provider
// failure. Only error codes tied to the grant itself count: client
// or config problems (invalid_client, unauthorized_client, ...)
// affect every user of the server and cannot be fixed by the user
// reconnecting, so they are treated as transient here. See RFC 6749
// section 5.2.
func IsPermanentRefreshError(err error) bool {
	var oauthErr *oauth2.RetrieveError
	if !xerrors.As(err, &oauthErr) {
		return false
	}
	switch oauthErr.ErrorCode {
	case "invalid_grant", // RFC 6749: grant invalid, expired, or revoked
		"bad_refresh_token": // GitHub's equivalent, returned with HTTP 200
		return true
	}
	return false
}

// RefreshFailureReason converts a refresh error into a bounded string
// safe to persist as oauth_refresh_failure_reason. It is stored for
// operator debugging only and is never returned through the API.
func RefreshFailureReason(err error) string {
	reason := err.Error()
	if len(reason) > refreshFailureReasonLimit {
		reason = reason[:refreshFailureReasonLimit]
	}
	return reason
}

// RefreshOAuth2Token checks whether the given MCP user token is
// expired (or within 10 seconds of expiry) and refreshes it using
// the OAuth2 credentials from the server config. If the token is
// still valid, no network call is made and Refreshed is false.
//
// The caller is responsible for persisting the result when
// Refreshed is true.
func RefreshOAuth2Token(
	ctx context.Context,
	cfg database.MCPServerConfig,
	tok database.MCPServerUserToken,
) (RefreshResult, error) {
	oauth2Cfg := &oauth2.Config{
		ClientID:     cfg.OAuth2ClientID,
		ClientSecret: cfg.OAuth2ClientSecret,
		Endpoint: oauth2.Endpoint{
			TokenURL: cfg.OAuth2TokenURL,
		},
	}

	oldToken := &oauth2.Token{
		AccessToken:  tok.AccessToken,
		RefreshToken: tok.RefreshToken,
		TokenType:    tok.TokenType,
	}
	if tok.Expiry.Valid {
		oldToken.Expiry = tok.Expiry.Time
	}

	// Cap the refresh HTTP call so a stalled token endpoint
	// cannot block the entire MCP connection phase. The timeout
	// matches connectTimeout used for MCP server connections.
	refreshCtx, cancel := context.WithTimeout(ctx, connectTimeout)
	defer cancel()

	// TokenSource automatically refreshes expired tokens. It
	// uses a 10-second expiry window, so tokens about to expire
	// are also refreshed proactively.
	newToken, err := oauth2Cfg.TokenSource(refreshCtx, oldToken).Token()
	if err != nil {
		return RefreshResult{}, xerrors.Errorf("refresh oauth2 token: %w", err)
	}

	refreshed := newToken.AccessToken != tok.AccessToken

	// Preserve the old refresh token when the provider doesn't
	// rotate (returns empty).
	refreshToken := cmp.Or(newToken.RefreshToken, tok.RefreshToken)

	return RefreshResult{
		AccessToken:  newToken.AccessToken,
		RefreshToken: refreshToken,
		TokenType:    newToken.TokenType,
		Expiry:       newToken.Expiry,
		Refreshed:    refreshed,
	}, nil
}

// RevokeOAuth2Token revokes the user's token at the provider's RFC 7009
// endpoint. It prefers the refresh token, retrying with the access token
// only on unsupported_token_type; other failures do not fall back, since
// an access-token success would hide a possibly live refresh token.
// Returns false without error when there is no revocation endpoint or no
// stored token. Errors carry only the HTTP status because provider
// bodies may echo secrets.
func RevokeOAuth2Token(
	ctx context.Context,
	httpClient *http.Client,
	cfg database.MCPServerConfig,
	tok database.MCPServerUserToken,
) (bool, error) {
	if cfg.OAuth2RevocationURL == "" {
		return false, nil
	}
	if tok.RefreshToken == "" && tok.AccessToken == "" {
		return false, nil
	}
	if err := ValidateRevocationEndpoint(cfg.OAuth2RevocationURL); err != nil {
		return false, err
	}

	if httpClient == nil {
		httpClient = mcpHTTPClient()
	}
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	// Copy so CheckRedirect does not leak into the shared client.
	redirectSafe := *httpClient
	redirectSafe.CheckRedirect = checkRevocationRedirect
	httpClient = &redirectSafe

	token, hint := tok.AccessToken, "access_token"
	if tok.RefreshToken != "" {
		token, hint = tok.RefreshToken, "refresh_token"
	}
	status, errorCode, err := postTokenRevocation(ctx, httpClient, cfg, token, hint)
	if err != nil {
		return false, err
	}
	if isRevocationSuccessStatus(status) {
		return true, nil
	}

	if hint == "refresh_token" && tok.AccessToken != "" && errorCode == "unsupported_token_type" {
		fbStatus, _, fbErr := postTokenRevocation(ctx, httpClient, cfg, tok.AccessToken, "access_token")
		if fbErr != nil {
			return false, fbErr
		}
		if isRevocationSuccessStatus(fbStatus) {
			return true, nil
		}
		return false, xerrors.Errorf(
			"revocation endpoint returned HTTP %d for the refresh token and HTTP %d for the access token",
			status, fbStatus,
		)
	}
	return false, xerrors.Errorf(
		"revocation endpoint returned HTTP %d", status,
	)
}

func isLoopbackHost(host string) bool {
	if strings.EqualFold(host, "localhost") {
		return true
	}
	ip := net.ParseIP(host)
	return ip != nil && ip.IsLoopback()
}

// ValidateRevocationEndpoint enforces the RFC 7009 HTTPS requirement;
// the request carries token material and the client secret. Plain HTTP
// is allowed only for loopback hosts.
func ValidateRevocationEndpoint(rawURL string) error {
	parsed, err := url.Parse(rawURL)
	if err != nil {
		return xerrors.Errorf("parse revocation URL: %w", err)
	}
	// url.Parse accepts hostless forms like "https:/revoke" that can
	// never be POSTed to.
	if parsed.Hostname() == "" {
		return xerrors.Errorf(
			"revocation endpoint %q has no host", parsed.Redacted(),
		)
	}
	if !isAllowedRevocationScheme(parsed) {
		return xerrors.Errorf(
			"revocation endpoint %q must use https", parsed.Redacted(),
		)
	}
	return nil
}

func isAllowedRevocationScheme(u *url.URL) bool {
	if u.Scheme == "https" {
		return true
	}
	return u.Scheme == "http" && isLoopbackHost(u.Hostname())
}

// checkRevocationRedirect stops the revocation POST, which carries
// token material and client credentials, from following redirects off
// the provider's origin. Loopback to loopback is exempt.
func checkRevocationRedirect(req *http.Request, via []*http.Request) error {
	if len(via) >= 10 {
		return xerrors.New("stopped after 10 redirects")
	}
	// net/http follows 301/302/303 with a bodyless GET; the token never
	// reaches the endpoint and a trailing 200 would be a false success.
	if req.Method != http.MethodPost {
		return xerrors.New(
			"revocation redirect dropped the POST body",
		)
	}
	if !isAllowedRevocationScheme(req.URL) {
		return xerrors.New("revocation redirect target must use https")
	}
	origin := via[0].URL
	if isLoopbackHost(req.URL.Hostname()) && isLoopbackHost(origin.Hostname()) {
		return nil
	}
	if req.URL.Scheme != origin.Scheme ||
		!strings.EqualFold(req.URL.Hostname(), origin.Hostname()) ||
		normalizedPort(req.URL) != normalizedPort(origin) {
		return xerrors.Errorf(
			"revocation redirect must stay on origin %q",
			origin.Scheme+"://"+origin.Host,
		)
	}
	return nil
}

func normalizedPort(u *url.URL) string {
	if p := u.Port(); p != "" {
		return p
	}
	switch u.Scheme {
	case "https":
		return "443"
	case "http":
		return "80"
	default:
		return ""
	}
}

func isRevocationSuccessStatus(status int) bool {
	return status == http.StatusOK || status == http.StatusNoContent
}

// postTokenRevocation returns the HTTP status and the RFC 6749 error
// code from the body; the raw body never propagates.
func postTokenRevocation(
	ctx context.Context,
	httpClient *http.Client,
	cfg database.MCPServerConfig,
	token, tokenTypeHint string,
) (int, string, error) {
	form := url.Values{}
	form.Set("token", token)
	form.Set("token_type_hint", tokenTypeHint)
	// Only public clients send client_id in the body; mixing it with
	// Basic auth is malformed per RFC 6749 section 2.3.1.
	if cfg.OAuth2ClientSecret == "" {
		form.Set("client_id", cfg.OAuth2ClientID)
	}

	revokeCtx, cancel := context.WithTimeout(ctx, connectTimeout)
	defer cancel()

	req, err := http.NewRequestWithContext(
		revokeCtx, http.MethodPost,
		cfg.OAuth2RevocationURL, strings.NewReader(form.Encode()),
	)
	if err != nil {
		return 0, "", xerrors.Errorf("create revocation request: %w", err)
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	// Credentials are form-encoded per RFC 6749 section 2.3.1
	// (mirrors x/oauth2).
	if cfg.OAuth2ClientSecret != "" {
		req.SetBasicAuth(url.QueryEscape(cfg.OAuth2ClientID), url.QueryEscape(cfg.OAuth2ClientSecret))
	}

	resp, err := httpClient.Do(req)
	if err != nil {
		var urlErr *url.Error
		if errors.As(err, &urlErr) {
			err = urlErr.Err
		}
		return 0, "", xerrors.Errorf("revoke oauth2 token: %w", err)
	}
	defer resp.Body.Close()

	if isRevocationSuccessStatus(resp.StatusCode) {
		_, _ = io.Copy(io.Discard, resp.Body)
		return resp.StatusCode, "", nil
	}
	var errBody struct {
		Error string `json:"error"`
	}
	_ = json.NewDecoder(io.LimitReader(resp.Body, 4096)).Decode(&errBody)
	_, _ = io.Copy(io.Discard, resp.Body)
	return resp.StatusCode, errBody.Error, nil
}