WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

2026-09-22 16:19:06 +08:00
commit f7fc992f57
3929 changed files with 1006607 additions and 0 deletions
@@ -0,0 +1,567 @@
+package mcp
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	secutils "github.com/Tencent/WeKnora/internal/utils"
+	"github.com/mark3labs/mcp-go/client"
+	"github.com/mark3labs/mcp-go/client/transport"
+	"github.com/mark3labs/mcp-go/mcp"
+)
+
+// MCPClient defines the interface for MCP client implementations
+type MCPClient interface {
+	// Connect establishes connection to the MCP service
+	Connect(ctx context.Context) error
+
+	// Disconnect closes the connection to the MCP service
+	Disconnect() error
+
+	// Initialize performs the MCP initialize handshake
+	Initialize(ctx context.Context) (*InitializeResult, error)
+
+	// ListTools retrieves the list of available tools from the MCP service
+	ListTools(ctx context.Context) ([]*types.MCPTool, error)
+
+	// ListResources retrieves the list of available resources from the MCP service
+	ListResources(ctx context.Context) ([]*types.MCPResource, error)
+
+	// CallTool calls a tool on the MCP service
+	CallTool(ctx context.Context, name string, args map[string]interface{}) (*CallToolResult, error)
+
+	// ReadResource reads a resource from the MCP service
+	ReadResource(ctx context.Context, uri string) (*ReadResourceResult, error)
+
+	// IsConnected returns true if the client is connected
+	IsConnected() bool
+
+	// GetServiceID returns the service ID this client is connected to
+	GetServiceID() string
+}
+
+// ClientConfig represents configuration for creating an MCP client
+type ClientConfig struct {
+	Service *types.MCPService
+
+	// OAuth wiring (only used when Service.AuthConfig.AuthType == oauth).
+	// The token store is scoped to (TenantID, Principal, Service.ID) so each
+	// identity connects with its own access/refresh token.
+	TenantID  uint64
+	Principal types.Principal
+	// UserID is kept for compatibility with older call sites/tests. New code
+	// should pass Principal.
+	UserID    string
+	OAuthRepo interfaces.MCPOAuthRepository
+}
+
+// mcpGoClient wraps mark3labs/mcp-go client to implement our MCPClient interface
+type mcpGoClient struct {
+	service     *types.MCPService
+	client      *client.Client
+	oauth       *oauthRuntime
+	connected   bool
+	initialized bool
+}
+
+// applyAuthHeaders injects the auth header for the SELECTED strategy only —
+// driven by AuthType so static API key / bearer are mutually exclusive (the old
+// code emitted both whenever the fields happened to be set, which double-authed
+// after a strategy switch). OAuth is handled separately by the caller.
+// CustomHeaders are always layered on top regardless of strategy and may
+// override the strategy header.
+func applyAuthHeaders(headers map[string]string, ac *types.MCPAuthConfig) {
+	if ac == nil {
+		return
+	}
+	switch ac.AuthType {
+	case types.MCPAuthAPIKey:
+		if ac.APIKey != "" {
+			name := ac.APIKeyHeader
+			if name == "" {
+				name = "X-API-Key"
+			}
+			headers[name] = ac.APIKey
+		}
+	case types.MCPAuthBearer:
+		if ac.Token != "" {
+			headers["Authorization"] = "Bearer " + ac.Token
+		}
+	case types.MCPAuthNone:
+		// Backward compatibility for rows that predate AuthType: infer from
+		// whichever static credential is set, preserving the historical
+		// behavior so existing services keep authenticating after upgrade.
+		if ac.APIKey != "" {
+			headers["X-API-Key"] = ac.APIKey
+		}
+		if ac.Token != "" {
+			headers["Authorization"] = "Bearer " + ac.Token
+		}
+	}
+	for key, value := range ac.CustomHeaders {
+		headers[key] = value
+	}
+}
+
+// OAuthRequiredError signals that the target MCP server requires OAuth
+// authorization — it answered the connect/initialize handshake with a 401 that
+// advertised RFC 9728 protected-resource metadata — even though the service was
+// NOT configured to use OAuth. Callers use this to guide the user to switch the
+// auth strategy to OAuth instead of surfacing a generic "401" failure.
+type OAuthRequiredError struct {
+	// MetadataURL is the RFC 9728 protected-resource metadata URL advertised by
+	// the server via the WWW-Authenticate header. Non-empty by construction
+	// (asOAuthRequired only wraps when the server advertised it).
+	MetadataURL string
+	Err         error
+}
+
+func (e *OAuthRequiredError) Error() string {
+	return fmt.Sprintf("the MCP server requires OAuth authorization: %v", e.Err)
+}
+
+func (e *OAuthRequiredError) Unwrap() error { return e.Err }
+
+// asOAuthRequired inspects err for a transport-level authorization-required
+// signal that carries RFC 9728 protected-resource metadata. It returns a
+// non-nil *OAuthRequiredError ONLY when the server advertised a metadata URL —
+// a bare 401 without metadata is treated as an ordinary auth failure (e.g. a
+// wrong/missing API key) so we don't misdirect the user toward OAuth.
+func asOAuthRequired(err error) *OAuthRequiredError {
+	if err == nil {
+		return nil
+	}
+	var authErr *transport.AuthorizationRequiredError
+	if errors.As(err, &authErr) && authErr.ResourceMetadataURL != "" {
+		return &OAuthRequiredError{MetadataURL: authErr.ResourceMetadataURL, Err: err}
+	}
+	return nil
+}
+
+// NewMCPClient creates a new MCP client based on the transport type
+func NewMCPClient(config *ClientConfig) (MCPClient, error) {
+	if config == nil || config.Service == nil {
+		return nil, fmt.Errorf("MCP client config and service are required")
+	}
+	if err := ValidateServiceOutboundURLs(config.Service); err != nil {
+		return nil, err
+	}
+
+	// Create HTTP client with timeout
+	timeout := 30 * time.Second
+	if config.Service.AdvancedConfig != nil && config.Service.AdvancedConfig.Timeout > 0 {
+		timeout = time.Duration(config.Service.AdvancedConfig.Timeout) * time.Second
+	}
+
+	clientCfg := secutils.DefaultSSRFSafeHTTPClientConfig()
+	clientCfg.Timeout = timeout
+	httpClient := secutils.NewSSRFSafeHTTPClient(clientCfg)
+
+	// Build headers
+	headers := make(map[string]string)
+	for key, value := range config.Service.Headers {
+		headers[key] = value
+	}
+	applyAuthHeaders(headers, config.Service.AuthConfig)
+
+	// Build OAuth config when this service uses the OAuth strategy. The
+	// client_id comes from the dynamically-registered client persisted at
+	// authorization time; the token store loads the invoking user's token
+	// and transparently refreshes it.
+	oauthConfig, useOAuth, err := buildOAuthConfig(config, httpClient)
+	if err != nil {
+		return nil, err
+	}
+
+	// Create client based on transport type
+	var mcpClient *client.Client
+	switch config.Service.TransportType {
+	case types.MCPTransportSSE:
+		if config.Service.URL == nil || *config.Service.URL == "" {
+			return nil, fmt.Errorf("URL is required for SSE transport")
+		}
+		if useOAuth {
+			mcpClient, err = client.NewOAuthSSEClient(*config.Service.URL, oauthConfig,
+				transport.WithHTTPClient(httpClient),
+				transport.WithHeaders(headers),
+			)
+		} else {
+			mcpClient, err = client.NewSSEMCPClient(*config.Service.URL,
+				client.WithHTTPClient(httpClient),
+				client.WithHeaders(headers),
+			)
+		}
+		if err != nil {
+			return nil, fmt.Errorf("failed to create SSE client: %w", err)
+		}
+	case types.MCPTransportHTTPStreamable:
+		if config.Service.URL == nil || *config.Service.URL == "" {
+			return nil, fmt.Errorf("URL is required for HTTP Streamable transport")
+		}
+		if useOAuth {
+			mcpClient, err = client.NewOAuthStreamableHttpClient(*config.Service.URL, oauthConfig,
+				transport.WithHTTPBasicClient(httpClient),
+				transport.WithHTTPHeaders(headers),
+			)
+		} else {
+			// For HTTP streamable, we need to use transport options
+			mcpClient, err = client.NewStreamableHttpClient(*config.Service.URL,
+				transport.WithHTTPBasicClient(httpClient),
+				transport.WithHTTPHeaders(headers),
+			)
+		}
+		if err != nil {
+			return nil, fmt.Errorf("failed to create HTTP streamable client: %w", err)
+		}
+	case types.MCPTransportStdio:
+		// Stdio transport is disabled for security reasons (potential command injection vulnerabilities)
+		return nil, fmt.Errorf("stdio transport is disabled for security reasons; please use SSE or HTTP Streamable transport instead")
+	default:
+		return nil, ErrUnsupportedTransport
+	}
+
+	instance := &mcpGoClient{
+		service: config.Service,
+		client:  mcpClient,
+	}
+	if useOAuth {
+		instance.oauth = newOAuthRuntime(
+			config.OAuthRepo,
+			config.TenantID,
+			config.Principal,
+			config.Service.ID,
+			*config.Service.URL,
+			oauthConfig,
+		)
+	}
+	mcpClient.OnConnectionLost(instance.onConnectionLost)
+	return instance, nil
+}
+
+// buildOAuthConfig returns the OAuth configuration for an OAuth-enabled MCP
+// service, or (_, false, nil) when the service does not use OAuth. It loads
+// the dynamically-registered client_id and wires a per-user token store so
+// the transport injects the invoking user's bearer token and refreshes it.
+func buildOAuthConfig(config *ClientConfig, httpClient *http.Client) (transport.OAuthConfig, bool, error) {
+	svc := config.Service
+	if !svc.AuthConfig.IsOAuth() {
+		return transport.OAuthConfig{}, false, nil
+	}
+	if config.OAuthRepo == nil {
+		return transport.OAuthConfig{}, false, fmt.Errorf("OAuth repository is required for OAuth MCP services")
+	}
+	principal := config.Principal.Normalize()
+	if !principal.Valid() && config.UserID != "" {
+		principal = types.Principal{Type: types.PrincipalWebUser, ID: config.UserID}.Normalize()
+	}
+	if !principal.Valid() {
+		return transport.OAuthConfig{}, false, fmt.Errorf("principal context is required to connect to an OAuth MCP service")
+	}
+	config.Principal = principal
+
+	oauthCfg := transport.OAuthConfig{
+		Scopes:                svc.AuthConfig.Scopes,
+		TokenStore:            newManagedTokenStore(config.OAuthRepo, config.TenantID, principal, svc.ID),
+		PKCEEnabled:           true,
+		AuthServerMetadataURL: svc.AuthConfig.AuthServerMetadataURL,
+		HTTPClient:            httpClient,
+	}
+	if regClient, err := config.OAuthRepo.GetClient(context.Background(), config.TenantID, svc.ID); err == nil && regClient != nil {
+		oauthCfg.ClientID = regClient.ClientID
+		oauthCfg.ClientSecret = regClient.ClientSecret
+		oauthCfg.RedirectURI = regClient.RedirectURI
+	}
+	return oauthCfg, true, nil
+}
+
+// onConnectionLost callback when the connection is lost
+func (c *mcpGoClient) onConnectionLost(err error) {
+	_ = c.Disconnect()
+	logger.Warnf(context.Background(), "MCP server connection has been lost, URL:%s, error:%v", *c.service.URL, err)
+}
+
+// checkErrorAndDisconnectIfNeeded checks for transport errors that indicate the
+// session is no longer valid and proactively disconnects the client so that
+// subsequent GetOrCreateClient calls will establish a fresh connection.
+// Both SSE and HTTP Streamable transports use server-assigned sessions
+// (via Mcp-Session-Id header) that can expire or be invalidated.
+func (c *mcpGoClient) checkErrorAndDisconnectIfNeeded(err error) {
+	var transportErr *transport.Error
+	if !errors.As(err, &transportErr) || transportErr.Err == nil {
+		return
+	}
+	errMsg := transportErr.Err.Error()
+	// Known session invalidation errors from MCP servers:
+	//   - "Invalid session ID"  — server recognises the header but rejects the value
+	//   - "No active connection" — server has no record of the session at all
+	if strings.Contains(errMsg, "Invalid session ID") ||
+		strings.Contains(errMsg, "No active connection") {
+		_ = c.Disconnect()
+	}
+}
+
+// oauthCall runs one MCP operation with WeKnora-owned token lifecycle checks.
+// A resource-server 401 forces exactly one refresh and one retry. Other errors
+// are never retried, which avoids duplicating tool side effects after ambiguous
+// network failures.
+func oauthCall[T any](ctx context.Context, c *mcpGoClient, operation func() (T, error)) (T, error) {
+	var zero T
+	if c.oauth != nil {
+		if err := c.oauth.ensureFresh(ctx, false, nil); err != nil {
+			return zero, err
+		}
+	}
+	result, err := operation()
+	if err == nil || c.oauth == nil || !isOAuthAuthorizationFailure(err) {
+		return result, err
+	}
+	if refreshErr := c.oauth.ensureFresh(ctx, true, client.GetOAuthHandler(err)); refreshErr != nil {
+		return zero, refreshErr
+	}
+	return operation()
+}
+
+// Connect establishes connection to the MCP service
+func (c *mcpGoClient) Connect(ctx context.Context) error {
+	if c.connected {
+		return ErrAlreadyConnected
+	}
+
+	_, err := oauthCall(ctx, c, func() (struct{}, error) {
+		return struct{}{}, c.client.Start(ctx)
+	})
+	if err != nil {
+		if oerr := asOAuthRequired(err); oerr != nil {
+			return oerr
+		}
+		return fmt.Errorf("failed to start client: %w", err)
+	}
+	c.connected = true
+	if c.service.TransportType == types.MCPTransportStdio {
+		logger.GetLogger(ctx).Infof("MCP stdio client connected: %s %v",
+			c.service.StdioConfig.Command, c.service.StdioConfig.Args)
+	} else {
+		logger.GetLogger(ctx).Infof("MCP client connected to %s", *c.service.URL)
+	}
+	return nil
+}
+
+// Disconnect closes the connection
+func (c *mcpGoClient) Disconnect() error {
+	if !c.connected {
+		return nil
+	}
+
+	// Close the client
+	if c.client != nil {
+		c.client.Close()
+	}
+	c.connected = false
+	c.initialized = false
+	return nil
+}
+
+// Initialize performs the MCP initialize handshake
+func (c *mcpGoClient) Initialize(ctx context.Context) (*InitializeResult, error) {
+	if !c.connected {
+		return nil, ErrNotConnected
+	}
+
+	// Initialize the client
+	req := mcp.InitializeRequest{
+		Params: mcp.InitializeParams{
+			ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
+			Capabilities:    mcp.ClientCapabilities{},
+			ClientInfo: mcp.Implementation{
+				Name:    "WeKnora",
+				Version: "1.0.0",
+			},
+		},
+	}
+
+	result, err := oauthCall(ctx, c, func() (*mcp.InitializeResult, error) {
+		return c.client.Initialize(ctx, req)
+	})
+	if err != nil {
+		c.checkErrorAndDisconnectIfNeeded(err)
+		if oerr := asOAuthRequired(err); oerr != nil {
+			return nil, oerr
+		}
+		return nil, fmt.Errorf("failed to initialize: %w", err)
+	}
+
+	c.initialized = true
+
+	return &InitializeResult{
+		ProtocolVersion: result.ProtocolVersion,
+		ServerInfo: ServerInfo{
+			Name:        result.ServerInfo.Name,
+			Version:     result.ServerInfo.Version,
+			Title:       result.ServerInfo.Title,
+			Description: result.ServerInfo.Description,
+		},
+	}, nil
+}
+
+// ListTools retrieves the list of available tools
+func (c *mcpGoClient) ListTools(ctx context.Context) ([]*types.MCPTool, error) {
+	if !c.initialized {
+		return nil, ErrNotConnected
+	}
+
+	req := mcp.ListToolsRequest{}
+	result, err := oauthCall(ctx, c, func() (*mcp.ListToolsResult, error) {
+		return c.client.ListTools(ctx, req)
+	})
+	if err != nil {
+		c.checkErrorAndDisconnectIfNeeded(err)
+		return nil, fmt.Errorf("failed to list tools: %w", err)
+	}
+
+	// Convert to our types
+	tools := make([]*types.MCPTool, len(result.Tools))
+	for i, tool := range result.Tools {
+		data, _ := json.Marshal(tool.InputSchema)
+		tools[i] = &types.MCPTool{
+			Name:        tool.Name,
+			Description: tool.Description,
+			InputSchema: data,
+		}
+	}
+
+	return tools, nil
+}
+
+// ListResources retrieves the list of available resources
+func (c *mcpGoClient) ListResources(ctx context.Context) ([]*types.MCPResource, error) {
+	if !c.initialized {
+		return nil, ErrNotConnected
+	}
+
+	req := mcp.ListResourcesRequest{}
+	result, err := oauthCall(ctx, c, func() (*mcp.ListResourcesResult, error) {
+		return c.client.ListResources(ctx, req)
+	})
+	if err != nil {
+		c.checkErrorAndDisconnectIfNeeded(err)
+		return nil, fmt.Errorf("failed to list resources: %w", err)
+	}
+
+	// Convert to our types
+	resources := make([]*types.MCPResource, len(result.Resources))
+	for i, resource := range result.Resources {
+		resources[i] = &types.MCPResource{
+			URI:         resource.URI,
+			Name:        resource.Name,
+			Description: resource.Description,
+			MimeType:    resource.MIMEType,
+		}
+	}
+
+	return resources, nil
+}
+
+// CallTool calls a tool on the MCP service
+func (c *mcpGoClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*CallToolResult, error) {
+	if !c.initialized {
+		return nil, ErrNotConnected
+	}
+
+	req := mcp.CallToolRequest{
+		Params: mcp.CallToolParams{
+			Name:      name,
+			Arguments: args,
+		},
+	}
+
+	result, err := oauthCall(ctx, c, func() (*mcp.CallToolResult, error) {
+		return c.client.CallTool(ctx, req)
+	})
+	if err != nil {
+		c.checkErrorAndDisconnectIfNeeded(err)
+		return nil, fmt.Errorf("failed to call tool: %w", err)
+	}
+
+	// Convert to our types
+	content := make([]ContentItem, 0, len(result.Content))
+	for _, item := range result.Content {
+		if textContent, ok := mcp.AsTextContent(item); ok {
+			content = append(content, ContentItem{
+				Type: "text",
+				Text: textContent.Text,
+			})
+		} else if imageContent, ok := mcp.AsImageContent(item); ok {
+			content = append(content, ContentItem{
+				Type:     "image",
+				Data:     imageContent.Data,
+				MimeType: imageContent.MIMEType,
+			})
+		}
+	}
+
+	return &CallToolResult{
+		IsError: result.IsError,
+		Content: content,
+	}, nil
+}
+
+// ReadResource reads a resource from the MCP service
+func (c *mcpGoClient) ReadResource(ctx context.Context, uri string) (*ReadResourceResult, error) {
+	if !c.initialized {
+		return nil, ErrNotConnected
+	}
+
+	req := mcp.ReadResourceRequest{
+		Params: mcp.ReadResourceParams{
+			URI: uri,
+		},
+	}
+
+	result, err := oauthCall(ctx, c, func() (*mcp.ReadResourceResult, error) {
+		return c.client.ReadResource(ctx, req)
+	})
+	if err != nil {
+		c.checkErrorAndDisconnectIfNeeded(err)
+		return nil, fmt.Errorf("failed to read resource: %w", err)
+	}
+
+	// Convert to our types
+	contents := make([]ResourceContent, 0, len(result.Contents))
+	for _, item := range result.Contents {
+		if textContent, ok := mcp.AsTextResourceContents(item); ok {
+			contents = append(contents, ResourceContent{
+				URI:      textContent.URI,
+				MimeType: textContent.MIMEType,
+				Text:     textContent.Text,
+			})
+		} else if blobContent, ok := mcp.AsBlobResourceContents(item); ok {
+			contents = append(contents, ResourceContent{
+				URI:      blobContent.URI,
+				MimeType: blobContent.MIMEType,
+				Blob:     blobContent.Blob,
+			})
+		}
+	}
+
+	return &ReadResourceResult{
+		Contents: contents,
+	}, nil
+}
+
+// IsConnected returns true if the client is connected
+func (c *mcpGoClient) IsConnected() bool {
+	return c.connected
+}
+
+// GetServiceID returns the service ID
+func (c *mcpGoClient) GetServiceID() string {
+	return c.service.ID
+}
@@ -0,0 +1,123 @@
+package mcp
+
+import (
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/mark3labs/mcp-go/client/transport"
+)
+
+func TestAsOAuthRequired(t *testing.T) {
+	t.Run("nil error", func(t *testing.T) {
+		if got := asOAuthRequired(nil); got != nil {
+			t.Fatalf("got %v, want nil", got)
+		}
+	})
+
+	t.Run("401 with RFC 9728 metadata is treated as OAuth required", func(t *testing.T) {
+		meta := "https://example.com/.well-known/oauth-protected-resource"
+		err := fmt.Errorf("wrap: %w", &transport.AuthorizationRequiredError{ResourceMetadataURL: meta})
+		got := asOAuthRequired(err)
+		if got == nil {
+			t.Fatal("expected non-nil OAuthRequiredError")
+		}
+		if got.MetadataURL != meta {
+			t.Errorf("MetadataURL = %q, want %q", got.MetadataURL, meta)
+		}
+	})
+
+	t.Run("bare 401 without metadata is NOT OAuth required", func(t *testing.T) {
+		err := &transport.AuthorizationRequiredError{ResourceMetadataURL: ""}
+		if got := asOAuthRequired(err); got != nil {
+			t.Fatalf("got %v, want nil (bare 401 should not suggest OAuth)", got)
+		}
+	})
+
+	t.Run("unrelated error is ignored", func(t *testing.T) {
+		if got := asOAuthRequired(errors.New("connection refused")); got != nil {
+			t.Fatalf("got %v, want nil", got)
+		}
+	})
+}
+
+func TestApplyAuthHeaders(t *testing.T) {
+	tests := []struct {
+		name string
+		ac   *types.MCPAuthConfig
+		want map[string]string
+	}{
+		{
+			name: "nil config injects nothing",
+			ac:   nil,
+			want: map[string]string{},
+		},
+		{
+			name: "api_key uses default X-API-Key header",
+			ac:   &types.MCPAuthConfig{AuthType: types.MCPAuthAPIKey, APIKey: "k1"},
+			want: map[string]string{"X-API-Key": "k1"},
+		},
+		{
+			name: "api_key honors custom header name (e.g. raw token in Authorization)",
+			ac: &types.MCPAuthConfig{
+				AuthType:     types.MCPAuthAPIKey,
+				APIKey:       "f7bfde",
+				APIKeyHeader: "Authorization",
+			},
+			want: map[string]string{"Authorization": "f7bfde"},
+		},
+		{
+			name: "bearer adds Bearer prefix",
+			ac:   &types.MCPAuthConfig{AuthType: types.MCPAuthBearer, Token: "t1"},
+			want: map[string]string{"Authorization": "Bearer t1"},
+		},
+		{
+			name: "selected strategy is exclusive — stale token is not emitted",
+			ac: &types.MCPAuthConfig{
+				AuthType: types.MCPAuthAPIKey,
+				APIKey:   "k1",
+				Token:    "stale",
+			},
+			want: map[string]string{"X-API-Key": "k1"},
+		},
+		{
+			name: "empty AuthType keeps legacy behavior (infer from fields)",
+			ac: &types.MCPAuthConfig{
+				AuthType: types.MCPAuthNone,
+				APIKey:   "k1",
+				Token:    "t1",
+			},
+			want: map[string]string{"X-API-Key": "k1", "Authorization": "Bearer t1"},
+		},
+		{
+			name: "custom headers are always layered on top",
+			ac: &types.MCPAuthConfig{
+				AuthType:      types.MCPAuthBearer,
+				Token:         "t1",
+				CustomHeaders: map[string]string{"X-Trace": "abc"},
+			},
+			want: map[string]string{"Authorization": "Bearer t1", "X-Trace": "abc"},
+		},
+		{
+			name: "oauth strategy emits no static header (handled elsewhere)",
+			ac:   &types.MCPAuthConfig{AuthType: types.MCPAuthOAuth, Token: "ignored"},
+			want: map[string]string{},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			headers := map[string]string{}
+			applyAuthHeaders(headers, tt.ac)
+			if len(headers) != len(tt.want) {
+				t.Fatalf("header count = %d, want %d (%v)", len(headers), len(tt.want), headers)
+			}
+			for k, v := range tt.want {
+				if headers[k] != v {
+					t.Errorf("header[%q] = %q, want %q", k, headers[k], v)
+				}
+			}
+		})
+	}
+}
@@ -0,0 +1,32 @@
+package mcp
+
+import "errors"
+
+var (
+	// ErrUnsupportedTransport is returned when transport type is not supported
+	ErrUnsupportedTransport = errors.New("unsupported transport type")
+
+	// ErrNotConnected is returned when operation requires connection but client is not connected
+	ErrNotConnected = errors.New("client not connected")
+
+	// ErrAlreadyConnected is returned when trying to connect an already connected client
+	ErrAlreadyConnected = errors.New("client already connected")
+
+	// ErrInitializeFailed is returned when MCP initialize handshake fails
+	ErrInitializeFailed = errors.New("MCP initialize handshake failed")
+
+	// ErrToolNotFound is returned when requested tool is not found
+	ErrToolNotFound = errors.New("tool not found")
+
+	// ErrResourceNotFound is returned when requested resource is not found
+	ErrResourceNotFound = errors.New("resource not found")
+
+	// ErrInvalidResponse is returned when server response is invalid
+	ErrInvalidResponse = errors.New("invalid response from server")
+
+	// ErrTimeout is returned when operation times out
+	ErrTimeout = errors.New("operation timed out")
+
+	// ErrConnectionClosed is returned when connection is closed unexpectedly
+	ErrConnectionClosed = errors.New("connection closed")
+)
@@ -0,0 +1,261 @@
+package mcp
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+)
+
+// MCPManager manages MCP client connections
+type MCPManager struct {
+	clients   map[string]MCPClient // cacheKey -> client
+	clientsMu sync.RWMutex
+	oauthRepo interfaces.MCPOAuthRepository
+	ctx       context.Context
+	cancel    context.CancelFunc
+}
+
+// NewMCPManager creates a new MCP manager. oauthRepo is used to wire per-user
+// OAuth token stores for OAuth-enabled MCP services.
+func NewMCPManager(oauthRepo interfaces.MCPOAuthRepository) *MCPManager {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	manager := &MCPManager{
+		clients:   make(map[string]MCPClient),
+		oauthRepo: oauthRepo,
+		ctx:       ctx,
+		cancel:    cancel,
+	}
+
+	// Start cleanup goroutine
+	go manager.cleanupIdleConnections()
+
+	return manager
+}
+
+// cacheKey computes the connection-cache key for a service. OAuth services are
+// keyed per principal (each identity connects with its own token); all other
+// services share a single connection per service ID.
+func cacheKey(service *types.MCPService, principal types.Principal) string {
+	if service.AuthConfig.IsOAuth() {
+		return service.ID + "\x00" + principal.Normalize().StorageID()
+	}
+	return service.ID
+}
+
+// GetOrCreateClient gets an existing client or creates a new one
+// Caches and reuses existing connections for SSE/HTTP Streamable
+// Note: Stdio transport is disabled for security reasons
+//
+// For OAuth-enabled services the connection is keyed per principal (derived from
+// ctx) so each identity connects with its own token.
+func (m *MCPManager) GetOrCreateClient(ctx context.Context, service *types.MCPService) (MCPClient, error) {
+	// Check if service is enabled
+	if !service.Enabled {
+		return nil, fmt.Errorf("MCP service %s is not enabled", service.Name)
+	}
+
+	// Stdio transport is disabled for security reasons
+	if service.TransportType == types.MCPTransportStdio {
+		return nil, fmt.Errorf("stdio transport is disabled for security reasons; please use SSE or HTTP Streamable transport instead")
+	}
+
+	var tenantID uint64
+	var principal types.Principal
+	if service.AuthConfig.IsOAuth() {
+		tenantID, _ = types.TenantIDFromContext(ctx)
+		principal, _ = types.PrincipalFromContext(ctx)
+		principal = types.MCPOAuthPrincipalFromContext(ctx)
+		if !principal.Valid() {
+			return nil, fmt.Errorf("principal context is required to connect to OAuth MCP service %s", service.Name)
+		}
+	}
+	key := cacheKey(service, principal)
+
+	// For SSE/HTTP Streamable, check if client already exists and reuse
+	m.clientsMu.RLock()
+	client, exists := m.clients[key]
+	m.clientsMu.RUnlock()
+
+	if exists && client.IsConnected() {
+		return client, nil
+	}
+
+	// Create new client
+	m.clientsMu.Lock()
+	defer m.clientsMu.Unlock()
+
+	// Double check after acquiring write lock
+	client, exists = m.clients[key]
+	if exists && client.IsConnected() {
+		return client, nil
+	}
+
+	// Create new client
+	config := &ClientConfig{
+		Service:   service,
+		TenantID:  tenantID,
+		Principal: principal,
+		OAuthRepo: m.oauthRepo,
+	}
+
+	client, err := NewMCPClient(config)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create MCP client: %w", err)
+	}
+
+	// For SSE connections, Connect() starts a persistent connection that needs a long-lived context
+	// Use manager's context (m.ctx) which persists for the lifetime of the manager
+	// The HTTP client's timeout will handle connection timeouts, not context cancellation
+	if err := client.Connect(m.ctx); err != nil {
+		return nil, fmt.Errorf("failed to connect to MCP service: %w", err)
+	}
+
+	if err := m.initializeClient(service, client, "failed to initialize MCP client"); err != nil {
+		return nil, err
+	}
+
+	// Store client (only for non-stdio transports)
+	m.clients[key] = client
+
+	logger.GetLogger(m.ctx).Infof("MCP client created and initialized for service: %s", service.Name)
+	return client, nil
+}
+
+// initializeClient handles the shared initialization flow with timeout enforcement.
+func (m *MCPManager) initializeClient(service *types.MCPService, client MCPClient, errPrefix string) error {
+	initTimeout := 30 * time.Second
+	if service.AdvancedConfig != nil && service.AdvancedConfig.Timeout > 0 {
+		initTimeout = time.Duration(service.AdvancedConfig.Timeout) * time.Second
+		if initTimeout > 60*time.Second {
+			initTimeout = 60 * time.Second
+		}
+	}
+
+	initCtx, initCancel := context.WithTimeout(m.ctx, initTimeout)
+	defer initCancel()
+
+	if _, err := client.Initialize(initCtx); err != nil {
+		client.Disconnect()
+		if errPrefix == "" {
+			errPrefix = "failed to initialize MCP client"
+		}
+		return fmt.Errorf("%s: %w", errPrefix, err)
+	}
+
+	return nil
+}
+
+// GetClient gets an existing client
+func (m *MCPManager) GetClient(serviceID string) (MCPClient, bool) {
+	m.clientsMu.RLock()
+	defer m.clientsMu.RUnlock()
+
+	client, exists := m.clients[serviceID]
+	return client, exists
+}
+
+// CloseClient closes and removes all cached connections for a service. For
+// OAuth services this spans every per-principal connection (keys are prefixed with
+// the service ID).
+func (m *MCPManager) CloseClient(serviceID string) error {
+	m.clientsMu.Lock()
+	defer m.clientsMu.Unlock()
+
+	for key, client := range m.clients {
+		// Match the plain service-ID key as well as per-principal OAuth keys
+		// ("<serviceID>\x00<principal>").
+		if key != serviceID && !strings.HasPrefix(key, serviceID+"\x00") {
+			continue
+		}
+		if err := client.Disconnect(); err != nil {
+			logger.GetLogger(m.ctx).Errorf("Failed to disconnect MCP client %s: %v", key, err)
+		}
+		delete(m.clients, key)
+		logger.GetLogger(m.ctx).Infof("MCP client closed: %s", key)
+	}
+	return nil
+}
+
+// CloseAll closes all clients
+func (m *MCPManager) CloseAll() {
+	m.clientsMu.Lock()
+	defer m.clientsMu.Unlock()
+
+	for serviceID, client := range m.clients {
+		if err := client.Disconnect(); err != nil {
+			logger.GetLogger(m.ctx).Errorf("Failed to disconnect MCP client %s: %v", serviceID, err)
+		}
+	}
+
+	m.clients = make(map[string]MCPClient)
+	logger.GetLogger(m.ctx).Info("All MCP clients closed")
+}
+
+// Shutdown gracefully shuts down the manager
+func (m *MCPManager) Shutdown() {
+	m.cancel()
+	m.CloseAll()
+}
+
+// cleanupIdleConnections periodically cleans up disconnected clients
+func (m *MCPManager) cleanupIdleConnections() {
+	ticker := time.NewTicker(5 * time.Minute)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-m.ctx.Done():
+			return
+		case <-ticker.C:
+			m.removeDisconnectedClients()
+		}
+	}
+}
+
+// removeDisconnectedClients removes clients that are no longer connected
+func (m *MCPManager) removeDisconnectedClients() {
+	m.clientsMu.Lock()
+	defer m.clientsMu.Unlock()
+
+	for serviceID, client := range m.clients {
+		if !client.IsConnected() {
+			delete(m.clients, serviceID)
+			logger.GetLogger(m.ctx).Infof("Removed disconnected MCP client: %s", serviceID)
+		}
+	}
+}
+
+// GetActiveClients returns the number of active clients
+func (m *MCPManager) GetActiveClients() int {
+	m.clientsMu.RLock()
+	defer m.clientsMu.RUnlock()
+
+	count := 0
+	for _, client := range m.clients {
+		if client.IsConnected() {
+			count++
+		}
+	}
+	return count
+}
+
+// ListActiveServices returns IDs of services with active connections
+func (m *MCPManager) ListActiveServices() []string {
+	m.clientsMu.RLock()
+	defer m.clientsMu.RUnlock()
+
+	services := make([]string, 0, len(m.clients))
+	for serviceID, client := range m.clients {
+		if client.IsConnected() {
+			services = append(services, serviceID)
+		}
+	}
+	return services
+}
@@ -0,0 +1,280 @@
+package mcp
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	"github.com/google/uuid"
+	"github.com/mark3labs/mcp-go/client"
+	"github.com/mark3labs/mcp-go/client/transport"
+)
+
+const (
+	oauthRefreshSkew       = 30 * time.Second
+	oauthRefreshLease      = 45 * time.Second
+	oauthRefreshPoll       = 100 * time.Millisecond
+	oauthStateAuthorized   = "authorized"
+	oauthStateRefreshable  = "refreshable"
+	oauthStateReauthNeeded = "reauth_required"
+)
+
+// OAuthAuthorizationStatus distinguishes a currently usable access token from
+// an expired token that can still be refreshed. This prevents a stale database
+// row from being presented as an already successful authorization.
+type OAuthAuthorizationStatus struct {
+	Authorized       bool       `json:"authorized"`
+	State            string     `json:"state"`
+	RefreshAvailable bool       `json:"refresh_available"`
+	ExpiresAt        *time.Time `json:"expires_at,omitempty"`
+}
+
+// OAuthReauthorizationRequiredError means no usable access token can be
+// recovered without interactive user consent.
+type OAuthReauthorizationRequiredError struct {
+	Reason string
+}
+
+func (e *OAuthReauthorizationRequiredError) Error() string {
+	if e.Reason == "" {
+		return "MCP OAuth authorization required"
+	}
+	return "MCP OAuth authorization required: " + e.Reason
+}
+
+// OAuthRefreshTemporaryError preserves a token when refresh failed for a
+// transient reason. Callers must surface/retry this as an operational failure,
+// not open a new consent popup.
+type OAuthRefreshTemporaryError struct {
+	Err error
+}
+
+func (e *OAuthRefreshTemporaryError) Error() string {
+	return fmt.Sprintf("MCP OAuth token refresh temporarily failed: %v", e.Err)
+}
+
+func (e *OAuthRefreshTemporaryError) Unwrap() error { return e.Err }
+
+type oauthRuntime struct {
+	repo          interfaces.MCPOAuthRepository
+	tenantID      uint64
+	principal     types.Principal
+	serviceID     string
+	handler       *transport.OAuthHandler
+	leaseDuration time.Duration
+}
+
+func newOAuthRuntime(
+	repo interfaces.MCPOAuthRepository,
+	tenantID uint64,
+	principal types.Principal,
+	serviceID, baseURL string,
+	cfg transport.OAuthConfig,
+) *oauthRuntime {
+	h := transport.NewOAuthHandler(cfg)
+	h.SetBaseURL(baseURL)
+	leaseDuration := oauthRefreshLease
+	if cfg.HTTPClient != nil && cfg.HTTPClient.Timeout > 0 && cfg.HTTPClient.Timeout+15*time.Second > leaseDuration {
+		leaseDuration = cfg.HTTPClient.Timeout + 15*time.Second
+	}
+	return &oauthRuntime{
+		repo:          repo,
+		tenantID:      tenantID,
+		principal:     principal.Normalize(),
+		serviceID:     serviceID,
+		handler:       h,
+		leaseDuration: leaseDuration,
+	}
+}
+
+func tokenStatus(token *types.MCPOAuthToken, now time.Time) OAuthAuthorizationStatus {
+	status := OAuthAuthorizationStatus{State: oauthStateReauthNeeded}
+	if token == nil || token.AccessToken == "" {
+		return status
+	}
+	status.RefreshAvailable = token.RefreshToken != ""
+	if !token.ExpiresAt.IsZero() {
+		expiresAt := token.ExpiresAt
+		status.ExpiresAt = &expiresAt
+	}
+	if token.ExpiresAt.IsZero() || token.ExpiresAt.After(now) {
+		status.Authorized = true
+		status.State = oauthStateAuthorized
+		return status
+	}
+	if status.RefreshAvailable {
+		status.State = oauthStateRefreshable
+	}
+	return status
+}
+
+func (r *oauthRuntime) ensureFresh(ctx context.Context, force bool, override *transport.OAuthHandler) error {
+	row, err := r.repo.GetTokenForPrincipal(ctx, r.tenantID, r.principal, r.serviceID)
+	if err != nil {
+		return fmt.Errorf("load MCP OAuth token: %w", err)
+	}
+	if row == nil || row.AccessToken == "" {
+		return &OAuthReauthorizationRequiredError{Reason: "no token is stored"}
+	}
+	now := time.Now()
+	if !force {
+		if row.ExpiresAt.IsZero() || row.ExpiresAt.After(now.Add(oauthRefreshSkew)) {
+			return nil
+		}
+		// Tokens issued without refresh_token remain usable through their actual
+		// expiry; the refresh skew must not shorten their lifetime.
+		if row.RefreshToken == "" && row.ExpiresAt.After(now) {
+			return nil
+		}
+	}
+	if row.RefreshToken == "" {
+		_ = r.repo.DeleteTokenForPrincipal(ctx, r.tenantID, r.principal, r.serviceID)
+		return &OAuthReauthorizationRequiredError{Reason: "the access token expired and no refresh token is available"}
+	}
+	return r.refreshWithLease(ctx, row, override)
+}
+
+func (r *oauthRuntime) refreshWithLease(
+	ctx context.Context, observed *types.MCPOAuthToken, override *transport.OAuthHandler,
+) error {
+	for {
+		leaseID := uuid.NewString()
+		leaseDuration := r.leaseDuration
+		if leaseDuration <= 0 {
+			leaseDuration = oauthRefreshLease
+		}
+		leaseUntil := time.Now().Add(leaseDuration)
+		acquired, err := r.repo.TryAcquireTokenRefreshLease(
+			ctx, r.tenantID, r.principal, r.serviceID, leaseID, leaseUntil,
+		)
+		if err != nil {
+			return fmt.Errorf("claim MCP OAuth token refresh: %w", err)
+		}
+		if acquired {
+			return r.refreshAsLeaseOwner(ctx, observed, leaseID, override)
+		}
+
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(oauthRefreshPoll):
+		}
+		current, err := r.repo.GetTokenForPrincipal(ctx, r.tenantID, r.principal, r.serviceID)
+		if err != nil {
+			return fmt.Errorf("reload MCP OAuth token after concurrent refresh: %w", err)
+		}
+		if current == nil || current.AccessToken == "" {
+			return &OAuthReauthorizationRequiredError{Reason: "the refresh token is no longer valid"}
+		}
+		if oauthTokenMaterialChanged(current, observed) {
+			if current.ExpiresAt.IsZero() || current.ExpiresAt.After(time.Now().Add(oauthRefreshSkew)) {
+				return nil
+			}
+			observed = current
+		}
+	}
+}
+
+func (r *oauthRuntime) refreshAsLeaseOwner(
+	ctx context.Context, observed *types.MCPOAuthToken, leaseID string, override *transport.OAuthHandler,
+) error {
+	defer func() {
+		releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
+		defer cancel()
+		if err := r.repo.ReleaseTokenRefreshLease(
+			releaseCtx, r.tenantID, r.principal, r.serviceID, leaseID,
+		); err != nil {
+			logger.GetLogger(releaseCtx).Warnf("failed to release MCP OAuth refresh lease: %v", err)
+		}
+	}()
+
+	current, err := r.repo.GetTokenForPrincipal(ctx, r.tenantID, r.principal, r.serviceID)
+	if err != nil {
+		return fmt.Errorf("reload MCP OAuth token before refresh: %w", err)
+	}
+	if current == nil || current.AccessToken == "" {
+		return &OAuthReauthorizationRequiredError{Reason: "no token is stored"}
+	}
+	// Another owner may have completed a refresh immediately before this lease
+	// was acquired. Never consume its newly rotated refresh token unnecessarily.
+	if oauthTokenMaterialChanged(current, observed) &&
+		(current.ExpiresAt.IsZero() || current.ExpiresAt.After(time.Now().Add(oauthRefreshSkew))) {
+		return nil
+	}
+	if current.RefreshToken == "" {
+		return r.invalidateToken(ctx, false, "no refresh token is available")
+	}
+
+	handler := override
+	if handler == nil {
+		handler = r.handler
+	}
+	refreshed, refreshErr := handler.RefreshToken(ctx, current.RefreshToken)
+	if refreshErr == nil && refreshed != nil && refreshed.AccessToken != "" {
+		logger.GetLogger(ctx).Infof("MCP OAuth token refreshed: service=%s principal=%s", r.serviceID, r.principal.StorageID())
+		return nil
+	}
+	if refreshErr == nil {
+		refreshErr = errors.New("authorization server returned an empty access token")
+	}
+	permanent, resetClient := permanentRefreshFailure(refreshErr)
+	if permanent {
+		return r.invalidateToken(ctx, resetClient, "the refresh token or OAuth client is no longer valid")
+	}
+	return &OAuthRefreshTemporaryError{Err: refreshErr}
+}
+
+func oauthTokenMaterialChanged(current, observed *types.MCPOAuthToken) bool {
+	if current == nil || observed == nil {
+		return current != observed
+	}
+	return current.AccessToken != observed.AccessToken ||
+		current.RefreshToken != observed.RefreshToken ||
+		!current.ExpiresAt.Equal(observed.ExpiresAt)
+}
+
+func (r *oauthRuntime) invalidateToken(ctx context.Context, resetClient bool, reason string) error {
+	if err := r.repo.DeleteTokenForPrincipal(ctx, r.tenantID, r.principal, r.serviceID); err != nil {
+		return fmt.Errorf("delete invalid MCP OAuth token: %w", err)
+	}
+	if resetClient {
+		if err := r.repo.DeleteClient(ctx, r.tenantID, r.serviceID); err != nil {
+			return fmt.Errorf("delete invalid MCP OAuth client registration: %w", err)
+		}
+	}
+	return &OAuthReauthorizationRequiredError{Reason: reason}
+}
+
+func permanentRefreshFailure(err error) (permanent bool, resetClient bool) {
+	var oauthErr transport.OAuthError
+	if errors.As(err, &oauthErr) {
+		switch strings.ToLower(oauthErr.ErrorCode) {
+		case "invalid_grant", "invalid_token", "bad_refresh_token", "expired_token":
+			return true, false
+		case "invalid_client", "unauthorized_client":
+			return true, true
+		}
+	}
+	lower := strings.ToLower(err.Error())
+	if strings.Contains(lower, "status 400") {
+		return true, false
+	}
+	if strings.Contains(lower, "status 401") {
+		return true, true
+	}
+	return false, false
+}
+
+func isOAuthAuthorizationFailure(err error) bool {
+	if err == nil {
+		return false
+	}
+	var reauth *OAuthReauthorizationRequiredError
+	return errors.As(err, &reauth) || client.IsOAuthAuthorizationRequiredError(err) ||
+		client.IsAuthorizationRequiredError(err)
+}
@@ -0,0 +1,316 @@
+package mcp
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/mark3labs/mcp-go/client/transport"
+	"github.com/stretchr/testify/require"
+)
+
+type lockedOAuthRepo struct {
+	*fakeOAuthRepo
+	mu sync.Mutex
+}
+
+func newLockedOAuthRepo() *lockedOAuthRepo {
+	return &lockedOAuthRepo{fakeOAuthRepo: newFakeOAuthRepo()}
+}
+
+func cloneOAuthToken(token *types.MCPOAuthToken) *types.MCPOAuthToken {
+	if token == nil {
+		return nil
+	}
+	clone := *token
+	if token.RefreshLeaseUntil != nil {
+		leaseUntil := *token.RefreshLeaseUntil
+		clone.RefreshLeaseUntil = &leaseUntil
+	}
+	return &clone
+}
+
+func (r *lockedOAuthRepo) GetTokenForPrincipal(
+	_ context.Context, tenantID uint64, principal types.Principal, serviceID string,
+) (*types.MCPOAuthToken, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return cloneOAuthToken(r.tokens[fakeOAuthKey(tenantID, principal, serviceID)]), nil
+}
+
+func (r *lockedOAuthRepo) SaveTokenForPrincipal(_ context.Context, token *types.MCPOAuthToken) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	principal := types.Principal{Type: token.PrincipalType, ID: token.PrincipalID}.Normalize()
+	token = cloneOAuthToken(token)
+	token.UpdatedAt = time.Now()
+	r.tokens[fakeOAuthKey(token.TenantID, principal, token.ServiceID)] = token
+	return nil
+}
+
+func (r *lockedOAuthRepo) DeleteTokenForPrincipal(
+	_ context.Context, tenantID uint64, principal types.Principal, serviceID string,
+) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	delete(r.tokens, fakeOAuthKey(tenantID, principal, serviceID))
+	return nil
+}
+
+func (r *lockedOAuthRepo) TryAcquireTokenRefreshLease(
+	_ context.Context,
+	tenantID uint64,
+	principal types.Principal,
+	serviceID, leaseID string,
+	leaseUntil time.Time,
+) (bool, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	row := r.tokens[fakeOAuthKey(tenantID, principal, serviceID)]
+	if row == nil || (row.RefreshLeaseUntil != nil && row.RefreshLeaseUntil.After(time.Now())) {
+		return false, nil
+	}
+	row.RefreshLeaseID = leaseID
+	row.RefreshLeaseUntil = &leaseUntil
+	row.UpdatedAt = time.Now()
+	return true, nil
+}
+
+func (r *lockedOAuthRepo) ReleaseTokenRefreshLease(
+	_ context.Context,
+	tenantID uint64,
+	principal types.Principal,
+	serviceID, leaseID string,
+) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	row := r.tokens[fakeOAuthKey(tenantID, principal, serviceID)]
+	if row != nil && row.RefreshLeaseID == leaseID {
+		row.RefreshLeaseID = ""
+		row.RefreshLeaseUntil = nil
+		row.UpdatedAt = time.Now()
+	}
+	return nil
+}
+
+func newOAuthLifecycleFixture(
+	t *testing.T, tokenStatus int, tokenBody map[string]any,
+) (*oauthRuntime, *lockedOAuthRepo, *atomic.Int32, func()) {
+	t.Helper()
+	var requests atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		switch req.URL.Path {
+		case "/metadata":
+			_ = json.NewEncoder(w).Encode(map[string]any{
+				"issuer":                                "http://" + req.Host,
+				"authorization_endpoint":                "http://" + req.Host + "/authorize",
+				"token_endpoint":                        "http://" + req.Host + "/token",
+				"response_types_supported":              []string{"code"},
+				"token_endpoint_auth_methods_supported": []string{"none"},
+			})
+		case "/token":
+			requests.Add(1)
+			require.NoError(t, req.ParseForm())
+			require.Equal(t, "refresh_token", req.Form.Get("grant_type"))
+			require.Equal(t, "old-refresh", req.Form.Get("refresh_token"))
+			w.WriteHeader(tokenStatus)
+			_ = json.NewEncoder(w).Encode(tokenBody)
+		default:
+			http.NotFound(w, req)
+		}
+	}))
+
+	repo := newLockedOAuthRepo()
+	principal := types.Principal{Type: types.PrincipalWebUser, ID: "user-1"}
+	row := &types.MCPOAuthToken{
+		TenantID:      7,
+		PrincipalType: principal.Type,
+		PrincipalID:   principal.ID,
+		UserID:        principal.StorageID(),
+		ServiceID:     "svc-1",
+		AccessToken:   "old-access",
+		RefreshToken:  "old-refresh",
+		TokenType:     "Bearer",
+		ExpiresAt:     time.Now().Add(-time.Minute),
+		UpdatedAt:     time.Now().Add(-time.Hour),
+	}
+	repo.tokens[fakeOAuthKey(7, principal, "svc-1")] = row
+	store := newDBTokenStore(repo, 7, principal, "svc-1")
+	runtime := newOAuthRuntime(repo, 7, principal, "svc-1", server.URL, transport.OAuthConfig{
+		ClientID:              "client-1",
+		AuthServerMetadataURL: server.URL + "/metadata",
+		TokenStore:            store,
+		HTTPClient:            server.Client(),
+	})
+	return runtime, repo, &requests, server.Close
+}
+
+func TestOAuthRuntimeRefreshesExpiredToken(t *testing.T) {
+	runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusOK, map[string]any{
+		"access_token":  "new-access",
+		"refresh_token": "rotated-refresh",
+		"token_type":    "Bearer",
+		"expires_in":    3600,
+	})
+	defer closeServer()
+
+	require.NoError(t, runtime.ensureFresh(context.Background(), false, nil))
+	require.EqualValues(t, 1, requests.Load())
+	row, err := repo.GetTokenForPrincipal(context.Background(), 7, runtime.principal, "svc-1")
+	require.NoError(t, err)
+	require.Equal(t, "new-access", row.AccessToken)
+	require.Equal(t, "rotated-refresh", row.RefreshToken)
+	require.True(t, row.ExpiresAt.After(time.Now()))
+}
+
+func TestOAuthRuntimeDeletesPermanentlyInvalidRefreshToken(t *testing.T) {
+	runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusBadRequest, map[string]any{
+		"error":             "invalid_grant",
+		"error_description": "refresh token expired",
+	})
+	defer closeServer()
+
+	err := runtime.ensureFresh(context.Background(), false, nil)
+	var reauth *OAuthReauthorizationRequiredError
+	require.ErrorAs(t, err, &reauth)
+	require.EqualValues(t, 1, requests.Load())
+	row, getErr := repo.GetTokenForPrincipal(context.Background(), 7, runtime.principal, "svc-1")
+	require.NoError(t, getErr)
+	require.Nil(t, row)
+}
+
+func TestOAuthRuntimePreservesTokenOnTemporaryRefreshFailure(t *testing.T) {
+	runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusServiceUnavailable, map[string]any{
+		"error": "temporarily_unavailable",
+	})
+	defer closeServer()
+
+	err := runtime.ensureFresh(context.Background(), false, nil)
+	var temporary *OAuthRefreshTemporaryError
+	require.ErrorAs(t, err, &temporary)
+	require.EqualValues(t, 1, requests.Load())
+	row, getErr := repo.GetTokenForPrincipal(context.Background(), 7, runtime.principal, "svc-1")
+	require.NoError(t, getErr)
+	require.NotNil(t, row)
+	require.Equal(t, "old-refresh", row.RefreshToken)
+}
+
+func TestOAuthRuntimeSerializesRotatingRefreshToken(t *testing.T) {
+	runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusOK, map[string]any{
+		"access_token":  "new-access",
+		"refresh_token": "rotated-refresh",
+		"token_type":    "Bearer",
+		"expires_in":    3600,
+	})
+	defer closeServer()
+
+	const callers = 12
+	start := make(chan struct{})
+	errs := make(chan error, callers)
+	var wg sync.WaitGroup
+	for range callers {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			errs <- runtime.ensureFresh(context.Background(), false, nil)
+		}()
+	}
+	close(start)
+	wg.Wait()
+	close(errs)
+	for err := range errs {
+		require.NoError(t, err)
+	}
+	require.EqualValues(t, 1, requests.Load(), "a rotating refresh token must be consumed once")
+	row, err := repo.GetTokenForPrincipal(context.Background(), 7, runtime.principal, "svc-1")
+	require.NoError(t, err)
+	require.Equal(t, "rotated-refresh", row.RefreshToken)
+}
+
+func TestOAuthCallRefreshesAndRetriesResource401Once(t *testing.T) {
+	runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusOK, map[string]any{
+		"access_token":  "new-access",
+		"refresh_token": "rotated-refresh",
+		"token_type":    "Bearer",
+		"expires_in":    3600,
+	})
+	defer closeServer()
+	row := repo.tokens[fakeOAuthKey(7, runtime.principal, "svc-1")]
+	row.ExpiresAt = time.Now().Add(time.Hour)
+
+	calls := 0
+	result, err := oauthCall(context.Background(), &mcpGoClient{oauth: runtime}, func() (string, error) {
+		calls++
+		if calls == 1 {
+			return "", &transport.OAuthAuthorizationRequiredError{Handler: runtime.handler}
+		}
+		return "ok", nil
+	})
+	require.NoError(t, err)
+	require.Equal(t, "ok", result)
+	require.Equal(t, 2, calls)
+	require.EqualValues(t, 1, requests.Load())
+}
+
+func TestOAuthCallDoesNotRetryMoreThanOnce(t *testing.T) {
+	runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusOK, map[string]any{
+		"access_token":  "new-access",
+		"refresh_token": "rotated-refresh",
+		"token_type":    "Bearer",
+		"expires_in":    3600,
+	})
+	defer closeServer()
+	row := repo.tokens[fakeOAuthKey(7, runtime.principal, "svc-1")]
+	row.ExpiresAt = time.Now().Add(time.Hour)
+
+	calls := 0
+	_, err := oauthCall(context.Background(), &mcpGoClient{oauth: runtime}, func() (string, error) {
+		calls++
+		return "", &transport.OAuthAuthorizationRequiredError{Handler: runtime.handler}
+	})
+	require.Error(t, err)
+	require.Equal(t, 2, calls)
+	require.EqualValues(t, 1, requests.Load())
+}
+
+func TestTokenStatusDoesNotTreatExpiredRowAsAuthorized(t *testing.T) {
+	expired := &types.MCPOAuthToken{
+		AccessToken:  "stale-access",
+		RefreshToken: "refresh",
+		ExpiresAt:    time.Now().Add(-time.Minute),
+	}
+	status := tokenStatus(expired, time.Now())
+	require.False(t, status.Authorized)
+	require.Equal(t, oauthStateRefreshable, status.State)
+	require.True(t, status.RefreshAvailable)
+
+	expired.RefreshToken = ""
+	status = tokenStatus(expired, time.Now())
+	require.False(t, status.Authorized)
+	require.Equal(t, oauthStateReauthNeeded, status.State)
+}
+
+func TestOAuthRuntimeDoesNotExpireNonRefreshableTokenEarly(t *testing.T) {
+	repo := newLockedOAuthRepo()
+	principal := types.Principal{Type: types.PrincipalWebUser, ID: "user-1"}
+	repo.tokens[fakeOAuthKey(7, principal, "svc-1")] = &types.MCPOAuthToken{
+		TenantID:      7,
+		PrincipalType: principal.Type,
+		PrincipalID:   principal.ID,
+		ServiceID:     "svc-1",
+		AccessToken:   "access",
+		ExpiresAt:     time.Now().Add(10 * time.Second),
+	}
+	runtime := &oauthRuntime{repo: repo, tenantID: 7, principal: principal, serviceID: "svc-1"}
+	require.NoError(t, runtime.ensureFresh(context.Background(), false, nil))
+	row, err := repo.GetTokenForPrincipal(context.Background(), 7, principal, "svc-1")
+	require.NoError(t, err)
+	require.NotNil(t, row)
+}
@@ -0,0 +1,280 @@
+package mcp
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	secutils "github.com/Tencent/WeKnora/internal/utils"
+	"github.com/mark3labs/mcp-go/client/transport"
+	"github.com/redis/go-redis/v9"
+)
+
+// clientRegistrationName is sent as client_name during dynamic client
+// registration (RFC 7591).
+const clientRegistrationName = "WeKnora"
+
+// oauthCallbackTimeout bounds token exchange after the browser lands on the
+// public callback route. The Gin request context is canceled once the client
+// receives the redirect, so CompleteAuthorization must detach from it.
+const oauthCallbackTimeout = 60 * time.Second
+
+// OAuthManager orchestrates the MCP OAuth2 authorization-code flow:
+// discovery, dynamic client registration, the authorize redirect, and the
+// callback code exchange. Tokens are persisted per (tenant, principal, service);
+// the registered client is persisted per (tenant, service) and reused.
+type OAuthManager struct {
+	repo        interfaces.MCPOAuthRepository
+	serviceRepo interfaces.MCPServiceRepository
+	states      *oauthStateStore
+}
+
+// NewOAuthManager constructs the OAuth manager. rdb may be nil (Lite mode),
+// in which case in-flight authorization states are kept in memory.
+func NewOAuthManager(
+	repo interfaces.MCPOAuthRepository,
+	serviceRepo interfaces.MCPServiceRepository,
+	rdb *redis.Client,
+) *OAuthManager {
+	return &OAuthManager{
+		repo:        repo,
+		serviceRepo: serviceRepo,
+		states:      newOAuthStateStore(rdb),
+	}
+}
+
+// newHandler builds an OAuth handler bound to a service + per-principal token store.
+func (m *OAuthManager) newHandler(
+	ctx context.Context, service *types.MCPService, tenantID uint64, principal types.Principal, redirectURI string,
+) (*transport.OAuthHandler, error) {
+	if service.URL == nil || *service.URL == "" {
+		return nil, fmt.Errorf("MCP service URL is required for OAuth")
+	}
+	if err := ValidateServiceOutboundURLs(service); err != nil {
+		return nil, err
+	}
+	httpCfg := secutils.DefaultSSRFSafeHTTPClientConfig()
+	httpCfg.Timeout = 30 * time.Second
+	cfg := transport.OAuthConfig{
+		RedirectURI:           redirectURI,
+		Scopes:                service.AuthConfig.Scopes,
+		TokenStore:            newDBTokenStore(m.repo, tenantID, principal, service.ID),
+		PKCEEnabled:           true,
+		AuthServerMetadataURL: service.AuthConfig.AuthServerMetadataURL,
+		HTTPClient:            secutils.NewSSRFSafeHTTPClient(httpCfg),
+	}
+	if existing, err := m.repo.GetClient(ctx, tenantID, service.ID); err == nil && existing != nil {
+		cfg.ClientID = existing.ClientID
+		cfg.ClientSecret = existing.ClientSecret
+	}
+	h := transport.NewOAuthHandler(cfg)
+	h.SetBaseURL(*service.URL)
+	return h, nil
+}
+
+// StartAuthorization performs discovery + (one-time) dynamic client
+// registration, then returns the authorization URL and an opaque attempt ID.
+// redirectURI is the backend callback URL registered with the auth server;
+// frontendRedirect is where the callback bounces the browser when finished.
+func (m *OAuthManager) StartAuthorization(
+	ctx context.Context,
+	service *types.MCPService,
+	tenantID uint64,
+	principal types.Principal,
+	redirectURI, frontendRedirect string,
+) (authorizationURL, attemptID string, err error) {
+	if !service.AuthConfig.IsOAuth() {
+		return "", "", fmt.Errorf("MCP service %s does not use OAuth", service.ID)
+	}
+	principal = principal.Normalize()
+	if !principal.Valid() {
+		return "", "", fmt.Errorf("principal context is required to authorize OAuth MCP service %s", service.ID)
+	}
+
+	h, err := m.newHandler(ctx, service, tenantID, principal, redirectURI)
+	if err != nil {
+		return "", "", err
+	}
+
+	// Register a client dynamically if we don't have one yet for this service.
+	existing, _ := m.repo.GetClient(ctx, tenantID, service.ID)
+	if existing == nil {
+		if err := h.RegisterClient(ctx, clientRegistrationName); err != nil {
+			return "", "", fmt.Errorf("dynamic client registration failed: %w", err)
+		}
+		clientID := h.GetClientID()
+		if clientID == "" {
+			return "", "", fmt.Errorf("dynamic client registration returned an empty client_id")
+		}
+		if err := m.repo.SaveClient(ctx, &types.MCPOAuthClient{
+			TenantID:    tenantID,
+			ServiceID:   service.ID,
+			ClientID:    clientID,
+			RedirectURI: redirectURI,
+		}); err != nil {
+			logger.GetLogger(ctx).Warnf("failed to persist MCP oauth client: %v", err)
+		}
+	}
+
+	verifier, err := transport.GenerateCodeVerifier()
+	if err != nil {
+		return "", "", fmt.Errorf("failed to generate PKCE verifier: %w", err)
+	}
+	challenge := transport.GenerateCodeChallenge(verifier)
+	state, err := transport.GenerateState()
+	if err != nil {
+		return "", "", fmt.Errorf("failed to generate state: %w", err)
+	}
+
+	authURL, err := h.GetAuthorizationURL(ctx, state, challenge)
+	if err != nil {
+		return "", "", fmt.Errorf("failed to build authorization URL: %w", err)
+	}
+
+	if err := m.states.Put(ctx, state, OAuthState{
+		TenantID:         tenantID,
+		UserID:           principal.StorageID(),
+		Principal:        principal,
+		ServiceID:        service.ID,
+		CodeVerifier:     verifier,
+		ClientID:         h.GetClientID(),
+		RedirectURI:      redirectURI,
+		FrontendRedirect: frontendRedirect,
+	}); err != nil {
+		return "", "", fmt.Errorf("failed to persist authorization state: %w", err)
+	}
+
+	return authURL, state, nil
+}
+
+// StartAuthorizationForService loads the MCP service by ID and starts the
+// authorization-code flow, returning the URL the user must open. It is a
+// convenience for callers (e.g. IM channels) that only hold a service ID and
+// cannot reach the MCP service lookup directly.
+func (m *OAuthManager) StartAuthorizationForService(
+	ctx context.Context,
+	tenantID uint64,
+	principal types.Principal,
+	serviceID, redirectURI, frontendRedirect string,
+) (string, error) {
+	service, err := m.serviceRepo.GetByID(ctx, tenantID, serviceID)
+	if err != nil {
+		return "", fmt.Errorf("failed to load MCP service: %w", err)
+	}
+	if service == nil {
+		return "", fmt.Errorf("MCP service not found")
+	}
+	authURL, _, err := m.StartAuthorization(ctx, service, tenantID, principal, redirectURI, frontendRedirect)
+	return authURL, err
+}
+
+// CompleteAuthorization handles the provider callback: it validates state,
+// exchanges the code for tokens (PKCE), and persists the per-user token.
+// Returns the frontend redirect URL and service ID recorded at
+// StartAuthorization time so the caller can recycle any cached transport that
+// still carries the previous OAuth client registration.
+func (m *OAuthManager) CompleteAuthorization(
+	ctx context.Context, state, code string,
+) (frontendRedirect, serviceID string, err error) {
+	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), oauthCallbackTimeout)
+	defer cancel()
+
+	st, err := m.states.Take(ctx, state)
+	if err != nil {
+		return "", "", err
+	}
+	frontendRedirect = st.FrontendRedirect
+	serviceID = st.ServiceID
+	principal := st.Principal.Normalize()
+	if !principal.Valid() && st.UserID != "" {
+		principal = types.Principal{Type: types.PrincipalWebUser, ID: st.UserID}.Normalize()
+	}
+	if !principal.Valid() {
+		return frontendRedirect, serviceID, fmt.Errorf("principal context is missing from OAuth state")
+	}
+
+	service, err := m.serviceRepo.GetByID(ctx, st.TenantID, st.ServiceID)
+	if err != nil {
+		return frontendRedirect, serviceID, fmt.Errorf("failed to load MCP service: %w", err)
+	}
+	if service == nil {
+		return frontendRedirect, serviceID, fmt.Errorf("MCP service not found")
+	}
+
+	h, err := m.newHandler(ctx, service, st.TenantID, principal, st.RedirectURI)
+	if err != nil {
+		return frontendRedirect, serviceID, err
+	}
+	// Re-prime the expected state so the library's CSRF check passes after
+	// reconstructing the handler in this separate request.
+	h.SetExpectedState(state)
+
+	if err := h.ProcessAuthorizationResponse(ctx, code, state, st.CodeVerifier); err != nil {
+		return frontendRedirect, serviceID, fmt.Errorf("token exchange failed: %w", err)
+	}
+	if err := m.states.CompleteAttempt(ctx, state); err != nil {
+		return frontendRedirect, serviceID, fmt.Errorf("failed to record authorization completion: %w", err)
+	}
+	// ProcessAuthorizationResponse persists the token via the TokenStore.
+	logger.GetLogger(ctx).Infof(
+		"MCP OAuth authorized: service=%s principal=%s", st.ServiceID, principal.StorageID(),
+	)
+	return frontendRedirect, serviceID, nil
+}
+
+// IsAuthorizationAttemptComplete reports whether this exact authorization
+// attempt completed for the requested principal and service. A pre-existing
+// token must never satisfy a newly opened OAuth popup.
+func (m *OAuthManager) IsAuthorizationAttemptComplete(
+	ctx context.Context,
+	tenantID uint64,
+	principal types.Principal,
+	serviceID, attemptID string,
+) (bool, error) {
+	attempt, err := m.states.Attempt(ctx, attemptID)
+	if err != nil {
+		return false, err
+	}
+	principal = principal.Normalize()
+	attemptPrincipal := attempt.Principal.Normalize()
+	if attempt.TenantID != tenantID || attempt.ServiceID != serviceID ||
+		attemptPrincipal.Type != principal.Type || attemptPrincipal.ID != principal.ID {
+		return false, fmt.Errorf("oauth authorization attempt does not match the current principal or service")
+	}
+	return attempt.Completed, nil
+}
+
+// AuthorizationStatus reports whether the stored access token is usable now,
+// or is expired but still has a refresh token that runtime use can rotate.
+func (m *OAuthManager) AuthorizationStatus(
+	ctx context.Context, tenantID uint64, principal types.Principal, serviceID string,
+) (OAuthAuthorizationStatus, error) {
+	tok, err := m.repo.GetTokenForPrincipal(ctx, tenantID, principal, serviceID)
+	if err != nil {
+		return OAuthAuthorizationStatus{}, err
+	}
+	return tokenStatus(tok, time.Now()), nil
+}
+
+// IsAuthorized reports whether the given principal has an access token that is
+// usable now. An expired row is not authorization success merely because its
+// encrypted token columns remain non-empty.
+func (m *OAuthManager) IsAuthorized(
+	ctx context.Context, tenantID uint64, principal types.Principal, serviceID string,
+) (bool, error) {
+	status, err := m.AuthorizationStatus(ctx, tenantID, principal, serviceID)
+	if err != nil {
+		return false, err
+	}
+	return status.Authorized, nil
+}
+
+// Revoke removes the principal's stored token for the service.
+func (m *OAuthManager) Revoke(
+	ctx context.Context, tenantID uint64, principal types.Principal, serviceID string,
+) error {
+	return m.repo.DeleteTokenForPrincipal(ctx, tenantID, principal, serviceID)
+}
@@ -0,0 +1,173 @@
+package mcp
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/mark3labs/mcp-go/client/transport"
+	"github.com/stretchr/testify/require"
+)
+
+type fakeOAuthRepo struct {
+	clients map[string]*types.MCPOAuthClient
+	tokens  map[string]*types.MCPOAuthToken
+}
+
+func newFakeOAuthRepo() *fakeOAuthRepo {
+	return &fakeOAuthRepo{
+		clients: map[string]*types.MCPOAuthClient{},
+		tokens:  map[string]*types.MCPOAuthToken{},
+	}
+}
+
+func fakeOAuthKey(tenantID uint64, principal types.Principal, serviceID string) string {
+	return fmt.Sprintf("%d|%s|%s", tenantID, principal.Normalize().StorageID(), serviceID)
+}
+
+func (r *fakeOAuthRepo) GetClient(
+	_ context.Context, tenantID uint64, serviceID string,
+) (*types.MCPOAuthClient, error) {
+	return r.clients[fmt.Sprintf("%d|%s", tenantID, serviceID)], nil
+}
+
+func (r *fakeOAuthRepo) SaveClient(_ context.Context, client *types.MCPOAuthClient) error {
+	r.clients[fmt.Sprintf("%d|%s", client.TenantID, client.ServiceID)] = client
+	return nil
+}
+
+func (r *fakeOAuthRepo) DeleteClient(_ context.Context, tenantID uint64, serviceID string) error {
+	delete(r.clients, fmt.Sprintf("%d|%s", tenantID, serviceID))
+	return nil
+}
+
+func (r *fakeOAuthRepo) GetToken(
+	ctx context.Context, tenantID uint64, userID, serviceID string,
+) (*types.MCPOAuthToken, error) {
+	return r.GetTokenForPrincipal(ctx, tenantID, types.Principal{Type: types.PrincipalWebUser, ID: userID}, serviceID)
+}
+
+func (r *fakeOAuthRepo) GetTokenForPrincipal(
+	_ context.Context, tenantID uint64, principal types.Principal, serviceID string,
+) (*types.MCPOAuthToken, error) {
+	return r.tokens[fakeOAuthKey(tenantID, principal, serviceID)], nil
+}
+
+func (r *fakeOAuthRepo) SaveToken(_ context.Context, token *types.MCPOAuthToken) error {
+	return r.SaveTokenForPrincipal(context.Background(), token)
+}
+
+func (r *fakeOAuthRepo) SaveTokenForPrincipal(_ context.Context, token *types.MCPOAuthToken) error {
+	principal := types.Principal{Type: token.PrincipalType, ID: token.PrincipalID}.Normalize()
+	if !principal.Valid() {
+		principal = types.Principal{Type: types.PrincipalWebUser, ID: token.UserID}.Normalize()
+	}
+	r.tokens[fakeOAuthKey(token.TenantID, principal, token.ServiceID)] = token
+	return nil
+}
+
+func (r *fakeOAuthRepo) DeleteToken(
+	ctx context.Context, tenantID uint64, userID, serviceID string,
+) error {
+	return r.DeleteTokenForPrincipal(ctx, tenantID, types.Principal{Type: types.PrincipalWebUser, ID: userID}, serviceID)
+}
+
+func (r *fakeOAuthRepo) DeleteTokenForPrincipal(
+	_ context.Context, tenantID uint64, principal types.Principal, serviceID string,
+) error {
+	delete(r.tokens, fakeOAuthKey(tenantID, principal, serviceID))
+	return nil
+}
+
+func (r *fakeOAuthRepo) TryAcquireTokenRefreshLease(
+	_ context.Context,
+	tenantID uint64,
+	principal types.Principal,
+	serviceID, leaseID string,
+	leaseUntil time.Time,
+) (bool, error) {
+	row := r.tokens[fakeOAuthKey(tenantID, principal, serviceID)]
+	if row == nil || (row.RefreshLeaseUntil != nil && row.RefreshLeaseUntil.After(time.Now())) {
+		return false, nil
+	}
+	row.RefreshLeaseID = leaseID
+	row.RefreshLeaseUntil = &leaseUntil
+	return true, nil
+}
+
+func (r *fakeOAuthRepo) ReleaseTokenRefreshLease(
+	_ context.Context,
+	tenantID uint64,
+	principal types.Principal,
+	serviceID, leaseID string,
+) error {
+	row := r.tokens[fakeOAuthKey(tenantID, principal, serviceID)]
+	if row != nil && row.RefreshLeaseID == leaseID {
+		row.RefreshLeaseID = ""
+		row.RefreshLeaseUntil = nil
+	}
+	return nil
+}
+
+func TestDBTokenStoreUsesPrincipal(t *testing.T) {
+	repo := newFakeOAuthRepo()
+	principal := types.Principal{Type: types.PrincipalAPIExternalUser, ID: "7:external-42"}
+	store := newDBTokenStore(repo, 7, principal, "svc-1")
+
+	expiresAt := time.Now().Add(time.Hour).UTC()
+	require.NoError(t, store.SaveToken(context.Background(), &transport.Token{
+		AccessToken:  "access",
+		RefreshToken: "refresh",
+		TokenType:    "Bearer",
+		ExpiresAt:    expiresAt,
+	}))
+
+	row, err := repo.GetTokenForPrincipal(context.Background(), 7, principal, "svc-1")
+	require.NoError(t, err)
+	require.NotNil(t, row)
+	require.Equal(t, types.PrincipalAPIExternalUser, row.PrincipalType)
+	require.Equal(t, "7:external-42", row.PrincipalID)
+	require.Equal(t, principal.StorageID(), row.UserID)
+
+	token, err := store.GetToken(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, "access", token.AccessToken)
+	require.Equal(t, "refresh", token.RefreshToken)
+	require.Equal(t, expiresAt, token.ExpiresAt)
+}
+
+func TestManagedTokenStoreLeavesRefreshToOAuthRuntime(t *testing.T) {
+	repo := newFakeOAuthRepo()
+	principal := types.Principal{Type: types.PrincipalWebUser, ID: "user-1"}
+	store := newManagedTokenStore(repo, 7, principal, "svc-1")
+	require.NoError(t, store.SaveToken(context.Background(), &transport.Token{
+		AccessToken:  "access",
+		RefreshToken: "refresh",
+		ExpiresAt:    time.Now().Add(-time.Minute),
+	}))
+
+	token, err := store.GetToken(context.Background())
+	require.NoError(t, err)
+	require.True(t, token.ExpiresAt.IsZero(), "mcp-go must not race WeKnora's coordinated refresh")
+	row, err := repo.GetTokenForPrincipal(context.Background(), 7, principal, "svc-1")
+	require.NoError(t, err)
+	require.False(t, row.ExpiresAt.IsZero(), "the database must retain the real expiry for preflight checks")
+}
+
+func TestOAuthCacheKeyUsesPrincipalForOAuthServices(t *testing.T) {
+	service := &types.MCPService{
+		ID:         "svc-1",
+		AuthConfig: &types.MCPAuthConfig{AuthType: types.MCPAuthOAuth},
+	}
+	alice := types.Principal{Type: types.PrincipalAPIExternalUser, ID: "7:alice"}
+	bob := types.Principal{Type: types.PrincipalAPIExternalUser, ID: "7:bob"}
+
+	require.NotEqual(t, cacheKey(service, alice), cacheKey(service, bob))
+	require.Contains(t, cacheKey(service, alice), alice.StorageID())
+
+	service.AuthConfig.AuthType = types.MCPAuthAPIKey
+	require.Equal(t, "svc-1", cacheKey(service, alice))
+	require.Equal(t, cacheKey(service, alice), cacheKey(service, bob))
+}
@@ -0,0 +1,233 @@
+package mcp
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"os"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/redis/go-redis/v9"
+)
+
+// oauthStateTTL bounds how long an in-flight authorization may take from
+// "authorize-url issued" to "callback received".
+const oauthStateTTL = 10 * time.Minute
+
+// OAuthState is the transient data needed to complete an authorization-code
+// exchange. It is keyed by the opaque OAuth `state` parameter and MUST hold
+// the PKCE code_verifier, which is a secret that must never reach the
+// authorization server — hence server-side storage rather than encoding it
+// into the state parameter.
+type OAuthState struct {
+	TenantID     uint64          `json:"tenant_id"`
+	UserID       string          `json:"user_id"`
+	Principal    types.Principal `json:"principal"`
+	ServiceID    string          `json:"service_id"`
+	CodeVerifier string          `json:"code_verifier"`
+	ClientID     string          `json:"client_id"`
+	RedirectURI  string          `json:"redirect_uri"`
+	// FrontendRedirect is where the backend callback redirects the browser
+	// after completing (or failing) the exchange.
+	FrontendRedirect string `json:"frontend_redirect"`
+}
+
+// oauthStateStore persists in-flight OAuth states. Backed by Redis when
+// available (so the callback can land on any backend replica); falls back to
+// a TTL in-memory map for single-instance / Lite deployments.
+type oauthStateStore struct {
+	rdb *redis.Client
+
+	mu       sync.Mutex
+	mem      map[string]memStateEntry
+	attempts map[string]memAttemptEntry
+}
+
+type memStateEntry struct {
+	value     OAuthState
+	expiresAt time.Time
+}
+
+// OAuthAttempt is the non-secret, authenticated status of one authorization
+// flow. It is kept separately from OAuthState because Take consumes the PKCE
+// state before the token exchange finishes, while the opener still needs to
+// distinguish this callback from an older token stored for the same service.
+type OAuthAttempt struct {
+	TenantID  uint64          `json:"tenant_id"`
+	Principal types.Principal `json:"principal"`
+	ServiceID string          `json:"service_id"`
+	Completed bool            `json:"completed"`
+}
+
+type memAttemptEntry struct {
+	value     OAuthAttempt
+	expiresAt time.Time
+}
+
+func newOAuthStateStore(rdb *redis.Client) *oauthStateStore {
+	s := &oauthStateStore{
+		rdb:      rdb,
+		mem:      make(map[string]memStateEntry),
+		attempts: make(map[string]memAttemptEntry),
+	}
+	if rdb == nil {
+		go s.gcLoop()
+	}
+	return s
+}
+
+func (s *oauthStateStore) key(state string) string {
+	ns := strings.TrimSpace(os.Getenv("WEKNORA_REDIS_NAMESPACE"))
+	if ns != "" {
+		return "weknora:mcp_oauth_state:" + ns + ":" + state
+	}
+	return "weknora:mcp_oauth_state:" + state
+}
+
+func (s *oauthStateStore) attemptKey(state string) string {
+	return s.key(state) + ":attempt"
+}
+
+// Put stores a state with a fixed TTL.
+func (s *oauthStateStore) Put(ctx context.Context, state string, value OAuthState) error {
+	attempt := OAuthAttempt{
+		TenantID:  value.TenantID,
+		Principal: value.Principal.Normalize(),
+		ServiceID: value.ServiceID,
+	}
+	if s.rdb != nil {
+		data, err := json.Marshal(value)
+		if err != nil {
+			return err
+		}
+		attemptData, err := json.Marshal(attempt)
+		if err != nil {
+			return err
+		}
+		pipe := s.rdb.TxPipeline()
+		pipe.Set(ctx, s.key(state), data, oauthStateTTL)
+		pipe.Set(ctx, s.attemptKey(state), attemptData, oauthStateTTL)
+		_, err = pipe.Exec(ctx)
+		return err
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	expiresAt := time.Now().Add(oauthStateTTL)
+	s.mem[state] = memStateEntry{value: value, expiresAt: expiresAt}
+	s.attempts[state] = memAttemptEntry{value: attempt, expiresAt: expiresAt}
+	return nil
+}
+
+// CompleteAttempt marks an authorization attempt complete only after the code
+// exchange has successfully persisted a token.
+func (s *oauthStateStore) CompleteAttempt(ctx context.Context, state string) error {
+	if s.rdb != nil {
+		data, err := s.rdb.Get(ctx, s.attemptKey(state)).Bytes()
+		if err != nil {
+			if err == redis.Nil {
+				return fmt.Errorf("oauth attempt not found or expired")
+			}
+			return err
+		}
+		var attempt OAuthAttempt
+		if err := json.Unmarshal(data, &attempt); err != nil {
+			return err
+		}
+		attempt.Completed = true
+		data, err = json.Marshal(attempt)
+		if err != nil {
+			return err
+		}
+		return s.rdb.Set(ctx, s.attemptKey(state), data, oauthStateTTL).Err()
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	entry, ok := s.attempts[state]
+	if !ok || time.Now().After(entry.expiresAt) {
+		delete(s.attempts, state)
+		return fmt.Errorf("oauth attempt not found or expired")
+	}
+	entry.value.Completed = true
+	entry.expiresAt = time.Now().Add(oauthStateTTL)
+	s.attempts[state] = entry
+	return nil
+}
+
+// Attempt returns the status record for one authorization flow.
+func (s *oauthStateStore) Attempt(ctx context.Context, state string) (OAuthAttempt, error) {
+	if s.rdb != nil {
+		data, err := s.rdb.Get(ctx, s.attemptKey(state)).Bytes()
+		if err != nil {
+			if err == redis.Nil {
+				return OAuthAttempt{}, fmt.Errorf("oauth attempt not found or expired")
+			}
+			return OAuthAttempt{}, err
+		}
+		var attempt OAuthAttempt
+		if err := json.Unmarshal(data, &attempt); err != nil {
+			return OAuthAttempt{}, err
+		}
+		return attempt, nil
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	entry, ok := s.attempts[state]
+	if !ok || time.Now().After(entry.expiresAt) {
+		delete(s.attempts, state)
+		return OAuthAttempt{}, fmt.Errorf("oauth attempt not found or expired")
+	}
+	return entry.value, nil
+}
+
+// Take retrieves and deletes a state (single-use). Returns an error if the
+// state is unknown or expired.
+func (s *oauthStateStore) Take(ctx context.Context, state string) (OAuthState, error) {
+	if s.rdb != nil {
+		data, err := s.rdb.GetDel(ctx, s.key(state)).Bytes()
+		if err != nil {
+			if err == redis.Nil {
+				return OAuthState{}, fmt.Errorf("oauth state not found or expired")
+			}
+			return OAuthState{}, err
+		}
+		var v OAuthState
+		if err := json.Unmarshal(data, &v); err != nil {
+			return OAuthState{}, err
+		}
+		return v, nil
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	entry, ok := s.mem[state]
+	if !ok {
+		return OAuthState{}, fmt.Errorf("oauth state not found or expired")
+	}
+	delete(s.mem, state)
+	if time.Now().After(entry.expiresAt) {
+		return OAuthState{}, fmt.Errorf("oauth state not found or expired")
+	}
+	return entry.value, nil
+}
+
+func (s *oauthStateStore) gcLoop() {
+	ticker := time.NewTicker(time.Minute)
+	defer ticker.Stop()
+	for range ticker.C {
+		now := time.Now()
+		s.mu.Lock()
+		for k, v := range s.mem {
+			if now.After(v.expiresAt) {
+				delete(s.mem, k)
+			}
+		}
+		for k, v := range s.attempts {
+			if now.After(v.expiresAt) {
+				delete(s.attempts, k)
+			}
+		}
+		s.mu.Unlock()
+	}
+}
@@ -0,0 +1,71 @@
+package mcp
+
+import (
+	"context"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestOAuthAttemptCompletesOnlyAfterCallbackExchange(t *testing.T) {
+	ctx := context.Background()
+	store := newOAuthStateStore(nil)
+	principal := types.Principal{Type: types.PrincipalWebUser, ID: "user-1"}
+	state := OAuthState{
+		TenantID:  7,
+		Principal: principal,
+		ServiceID: "service-1",
+	}
+
+	require.NoError(t, store.Put(ctx, "attempt-1", state))
+
+	attempt, err := store.Attempt(ctx, "attempt-1")
+	require.NoError(t, err)
+	require.False(t, attempt.Completed)
+
+	// Consuming the OAuth state means the provider callback started. It must
+	// not look successful until the code exchange and token persistence finish.
+	_, err = store.Take(ctx, "attempt-1")
+	require.NoError(t, err)
+	attempt, err = store.Attempt(ctx, "attempt-1")
+	require.NoError(t, err)
+	require.False(t, attempt.Completed)
+
+	require.NoError(t, store.CompleteAttempt(ctx, "attempt-1"))
+	attempt, err = store.Attempt(ctx, "attempt-1")
+	require.NoError(t, err)
+	require.True(t, attempt.Completed)
+	require.Equal(t, uint64(7), attempt.TenantID)
+	require.Equal(t, principal.Normalize(), attempt.Principal)
+	require.Equal(t, "service-1", attempt.ServiceID)
+}
+
+func TestAuthorizationAttemptStatusIsScopedToPrincipalAndService(t *testing.T) {
+	ctx := context.Background()
+	manager := &OAuthManager{states: newOAuthStateStore(nil)}
+	principal := types.Principal{Type: types.PrincipalWebUser, ID: "user-1"}
+
+	require.NoError(t, manager.states.Put(ctx, "attempt-1", OAuthState{
+		TenantID:  7,
+		Principal: principal,
+		ServiceID: "service-1",
+	}))
+	require.NoError(t, manager.states.CompleteAttempt(ctx, "attempt-1"))
+
+	completed, err := manager.IsAuthorizationAttemptComplete(
+		ctx, 7, principal, "service-1", "attempt-1",
+	)
+	require.NoError(t, err)
+	require.True(t, completed)
+
+	_, err = manager.IsAuthorizationAttemptComplete(
+		ctx, 7, types.Principal{Type: types.PrincipalWebUser, ID: "user-2"}, "service-1", "attempt-1",
+	)
+	require.ErrorContains(t, err, "does not match")
+
+	_, err = manager.IsAuthorizationAttemptComplete(
+		ctx, 7, principal, "service-2", "attempt-1",
+	)
+	require.ErrorContains(t, err, "does not match")
+}
@@ -0,0 +1,109 @@
+package mcp
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	"github.com/mark3labs/mcp-go/client/transport"
+)
+
+// dbTokenStore is a transport.TokenStore backed by the MCPOAuthRepository,
+// scoped to a single (tenant, principal, service) tuple. The mcp-go OAuth handler
+// calls SaveToken after a successful authorization or refresh. Runtime MCP
+// transports receive the managedTokenStore wrapper below so refresh decisions
+// stay in WeKnora's coordinated lifecycle instead of the dependency.
+type dbTokenStore struct {
+	repo      interfaces.MCPOAuthRepository
+	tenantID  uint64
+	principal types.Principal
+	serviceID string
+}
+
+// managedTokenStore hides local expiry from mcp-go transports. WeKnora checks
+// the persisted ExpiresAt before every operation and performs the coordinated
+// refresh itself; allowing the dependency to also auto-refresh would bypass
+// the cross-instance lease and collapse refresh failures into a generic
+// authorization-required error.
+type managedTokenStore struct {
+	*dbTokenStore
+}
+
+func newManagedTokenStore(
+	repo interfaces.MCPOAuthRepository, tenantID uint64, principal types.Principal, serviceID string,
+) *managedTokenStore {
+	return &managedTokenStore{dbTokenStore: newDBTokenStore(repo, tenantID, principal, serviceID)}
+}
+
+func (s *managedTokenStore) GetToken(ctx context.Context) (*transport.Token, error) {
+	token, err := s.dbTokenStore.GetToken(ctx)
+	if err != nil {
+		return nil, err
+	}
+	token.ExpiresAt = time.Time{}
+	return token, nil
+}
+
+// newDBTokenStore creates a per-principal, per-service token store.
+func newDBTokenStore(
+	repo interfaces.MCPOAuthRepository, tenantID uint64, principal types.Principal, serviceID string,
+) *dbTokenStore {
+	return &dbTokenStore{
+		repo:      repo,
+		tenantID:  tenantID,
+		principal: principal.Normalize(),
+		serviceID: serviceID,
+	}
+}
+
+// GetToken returns the persisted token, or transport.ErrNoToken when the user
+// has not authorized this service yet.
+func (s *dbTokenStore) GetToken(ctx context.Context) (*transport.Token, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	row, err := s.repo.GetTokenForPrincipal(ctx, s.tenantID, s.principal, s.serviceID)
+	if err != nil {
+		return nil, err
+	}
+	if row == nil || row.AccessToken == "" {
+		return nil, transport.ErrNoToken
+	}
+	return &transport.Token{
+		AccessToken:  row.AccessToken,
+		RefreshToken: row.RefreshToken,
+		TokenType:    row.TokenType,
+		ExpiresAt:    row.ExpiresAt,
+	}, nil
+}
+
+// SaveToken persists a freshly issued or refreshed token.
+func (s *dbTokenStore) SaveToken(ctx context.Context, token *transport.Token) error {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	if token == nil || token.AccessToken == "" {
+		return fmt.Errorf("OAuth token response did not contain an access_token")
+	}
+	if token.TokenType == "" {
+		token.TokenType = "Bearer"
+	}
+	expiresAt := token.ExpiresAt
+	if expiresAt.IsZero() && token.ExpiresIn > 0 {
+		expiresAt = time.Now().Add(time.Duration(token.ExpiresIn) * time.Second)
+	}
+	principal := s.principal.Normalize()
+	return s.repo.SaveTokenForPrincipal(ctx, &types.MCPOAuthToken{
+		TenantID:      s.tenantID,
+		PrincipalType: principal.Type,
+		PrincipalID:   principal.ID,
+		UserID:        principal.StorageID(),
+		ServiceID:     s.serviceID,
+		AccessToken:   token.AccessToken,
+		RefreshToken:  token.RefreshToken,
+		TokenType:     token.TokenType,
+		ExpiresAt:     expiresAt,
+	})
+}
@@ -0,0 +1,36 @@
+package mcp
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	secutils "github.com/Tencent/WeKnora/internal/utils"
+)
+
+// ValidateServiceOutboundURLs validates every URL that the MCP transport or
+// OAuth discovery flow may contact. It is intentionally called both at
+// persistence boundaries and immediately before client construction so stale
+// or imported rows cannot bypass the current SSRF policy.
+func ValidateServiceOutboundURLs(service *types.MCPService) error {
+	if service == nil {
+		return fmt.Errorf("MCP service is required")
+	}
+	if service.URL != nil {
+		serviceURL := strings.TrimSpace(*service.URL)
+		if serviceURL != "" {
+			if err := secutils.ValidateURLForSSRF(serviceURL); err != nil {
+				return fmt.Errorf("MCP service URL failed SSRF validation: %w", err)
+			}
+		}
+	}
+	if service.AuthConfig != nil {
+		metadataURL := strings.TrimSpace(service.AuthConfig.AuthServerMetadataURL)
+		if metadataURL != "" {
+			if err := secutils.ValidateURLForSSRF(metadataURL); err != nil {
+				return fmt.Errorf("MCP OAuth metadata URL failed SSRF validation: %w", err)
+			}
+		}
+	}
+	return nil
+}
@@ -0,0 +1,31 @@
+package mcp
+
+import (
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/types"
+)
+
+func TestValidateServiceOutboundURLsRejectsOAuthMetadataSSRF(t *testing.T) {
+	serviceURL := "https://example.com/mcp"
+	err := ValidateServiceOutboundURLs(&types.MCPService{
+		URL: &serviceURL,
+		AuthConfig: &types.MCPAuthConfig{
+			AuthServerMetadataURL: "http://169.254.169.254/latest/meta-data",
+		},
+	})
+	if err == nil {
+		t.Fatal("expected OAuth metadata URL to be rejected")
+	}
+}
+
+func TestNewMCPClientRejectsStoredUnsafeURL(t *testing.T) {
+	serviceURL := "http://127.0.0.1:8080/mcp"
+	_, err := NewMCPClient(&ClientConfig{Service: &types.MCPService{
+		URL:           &serviceURL,
+		TransportType: types.MCPTransportHTTPStreamable,
+	}})
+	if err == nil {
+		t.Fatal("expected stored unsafe MCP URL to be rejected at client construction")
+	}
+}
@@ -0,0 +1,68 @@
+package mcp
+
+// InitializeResult represents the result of initialize request
+type InitializeResult struct {
+	ProtocolVersion string             `json:"protocolVersion"`
+	Capabilities    ServerCapabilities `json:"capabilities"`
+	ServerInfo      ServerInfo         `json:"serverInfo"`
+}
+
+// ServerCapabilities represents server capabilities
+type ServerCapabilities struct {
+	Tools        *ToolsCapability       `json:"tools,omitempty"`
+	Resources    *ResourcesCapability   `json:"resources,omitempty"`
+	Prompts      *PromptsCapability     `json:"prompts,omitempty"`
+	Logging      map[string]interface{} `json:"logging,omitempty"`
+	Experimental map[string]interface{} `json:"experimental,omitempty"`
+}
+
+// ToolsCapability represents tools capability
+type ToolsCapability struct {
+	ListChanged bool `json:"listChanged,omitempty"`
+}
+
+// ResourcesCapability represents resources capability
+type ResourcesCapability struct {
+	Subscribe   bool `json:"subscribe,omitempty"`
+	ListChanged bool `json:"listChanged,omitempty"`
+}
+
+// PromptsCapability represents prompts capability
+type PromptsCapability struct {
+	ListChanged bool `json:"listChanged,omitempty"`
+}
+
+// ServerInfo represents information about the server
+type ServerInfo struct {
+	Name        string `json:"name"`
+	Version     string `json:"version"`
+	Title       string `json:"title,omitempty"`
+	Description string `json:"description,omitempty"`
+}
+
+// CallToolResult represents the result of tools/call request
+type CallToolResult struct {
+	Content []ContentItem `json:"content"`
+	IsError bool          `json:"isError,omitempty"`
+}
+
+// ContentItem represents a content item in tool result
+type ContentItem struct {
+	Type     string `json:"type"` // "text", "image", "resource"
+	Text     string `json:"text,omitempty"`
+	Data     string `json:"data,omitempty"`
+	MimeType string `json:"mimeType,omitempty"`
+}
+
+// ReadResourceResult represents the result of resources/read request
+type ReadResourceResult struct {
+	Contents []ResourceContent `json:"contents"`
+}
+
+// ResourceContent represents resource content
+type ResourceContent struct {
+	URI      string `json:"uri"`
+	MimeType string `json:"mimeType,omitempty"`
+	Text     string `json:"text,omitempty"`
+	Blob     string `json:"blob,omitempty"` // Base64 encoded
+}