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,615 @@
+package service
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"maps"
+	"slices"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/Tencent/WeKnora/internal/mcp"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	secutils "github.com/Tencent/WeKnora/internal/utils"
+)
+
+// mcpServiceService implements MCPServiceService interface
+type mcpServiceService struct {
+	mcpServiceRepo interfaces.MCPServiceRepository
+	mcpManager     *mcp.MCPManager
+	oauthRepo      interfaces.MCPOAuthRepository
+}
+
+// NewMCPServiceService creates a new MCP service service
+func NewMCPServiceService(
+	mcpServiceRepo interfaces.MCPServiceRepository,
+	mcpManager *mcp.MCPManager,
+	oauthRepo interfaces.MCPOAuthRepository,
+) interfaces.MCPServiceService {
+	return &mcpServiceService{
+		mcpServiceRepo: mcpServiceRepo,
+		mcpManager:     mcpManager,
+		oauthRepo:      oauthRepo,
+	}
+}
+
+// CreateMCPService creates a new MCP service
+func (s *mcpServiceService) CreateMCPService(ctx context.Context, service *types.MCPService) error {
+	// Stdio transport is disabled for security reasons
+	if service.TransportType == types.MCPTransportStdio {
+		return fmt.Errorf("stdio transport is disabled for security reasons; please use SSE or HTTP Streamable transport instead")
+	}
+	if err := mcp.ValidateServiceOutboundURLs(service); err != nil {
+		return err
+	}
+
+	// Set default advanced config if not provided
+	if service.AdvancedConfig == nil {
+		service.AdvancedConfig = types.GetDefaultAdvancedConfig()
+	}
+
+	// Set timestamps
+	service.CreatedAt = time.Now()
+	service.UpdatedAt = time.Now()
+
+	if err := s.mcpServiceRepo.Create(ctx, service); err != nil {
+		logger.GetLogger(ctx).Errorf("Failed to create MCP service: %v", err)
+		return fmt.Errorf("failed to create MCP service: %w", err)
+	}
+
+	return nil
+}
+
+// GetMCPServiceByID retrieves an MCP service by ID.
+//
+// Returns the raw stored entity including any AuthConfig credentials in plain
+// form. Callers MUST convert to dto.MCPServiceResponse (which strips secret
+// fields by construction) before serializing to a response body. Internal
+// callers (e.g. MCP client construction, credential metadata derivation) need
+// the unredacted form to function correctly.
+func (s *mcpServiceService) GetMCPServiceByID(
+	ctx context.Context,
+	tenantID uint64,
+	id string,
+) (*types.MCPService, error) {
+	service, err := s.mcpServiceRepo.GetByID(ctx, tenantID, id)
+	if err != nil {
+		logger.GetLogger(ctx).Errorf("Failed to get MCP service: %v", err)
+		return nil, fmt.Errorf("failed to get MCP service: %w", err)
+	}
+
+	if service == nil {
+		return nil, fmt.Errorf("MCP service not found")
+	}
+	return service, nil
+}
+
+// ListMCPServices lists all MCP services for a tenant.
+//
+// Same contract as GetMCPServiceByID — returns raw entities; handlers MUST
+// convert to dto.MCPServiceResponse before responding.
+func (s *mcpServiceService) ListMCPServices(ctx context.Context, tenantID uint64) ([]*types.MCPService, error) {
+	services, err := s.mcpServiceRepo.List(ctx, tenantID)
+	if err != nil {
+		logger.GetLogger(ctx).Errorf("Failed to list MCP services: %v", err)
+		return nil, fmt.Errorf("failed to list MCP services: %w", err)
+	}
+	return services, nil
+}
+
+// ListMCPServicesByIDs retrieves multiple MCP services by IDs
+func (s *mcpServiceService) ListMCPServicesByIDs(
+	ctx context.Context,
+	tenantID uint64,
+	ids []string,
+) ([]*types.MCPService, error) {
+	if len(ids) == 0 {
+		return []*types.MCPService{}, nil
+	}
+
+	services, err := s.mcpServiceRepo.ListByIDs(ctx, tenantID, ids)
+	if err != nil {
+		logger.GetLogger(ctx).Errorf("Failed to list MCP services by IDs: %v", err)
+		return nil, fmt.Errorf("failed to list MCP services by IDs: %w", err)
+	}
+
+	return services, nil
+}
+
+// UpdateMCPService updates an MCP service
+func (s *mcpServiceService) UpdateMCPService(
+	ctx context.Context,
+	service *types.MCPService,
+	updateFields map[string]bool,
+) error {
+	// Check if service exists
+	existing, err := s.mcpServiceRepo.GetByID(ctx, service.TenantID, service.ID)
+	if err != nil {
+		return fmt.Errorf("failed to get MCP service: %w", err)
+	}
+	if existing == nil {
+		return fmt.Errorf("MCP service not found")
+	}
+
+	// Builtin MCP services cannot be updated
+	if existing.IsBuiltin {
+		return fmt.Errorf("builtin MCP services cannot be updated")
+	}
+
+	// Determine the final transport type after merge
+	finalTransportType := existing.TransportType
+	if service.TransportType != "" {
+		finalTransportType = service.TransportType
+	}
+
+	// Stdio transport is disabled for security reasons
+	if finalTransportType == types.MCPTransportStdio {
+		return fmt.Errorf("stdio transport is disabled for security reasons; please use SSE or HTTP Streamable transport instead")
+	}
+
+	// Store old enabled state BEFORE any updates
+	oldEnabled := existing.Enabled
+
+	// Snapshot pre-merge values of fields that drive configChanged. We need
+	// this because the in-place merge below reassigns pointer fields such as
+	// existing.URL = service.URL, after which any post-merge comparison
+	// between service.URL and existing.URL would trivially match.
+	//
+	// AuthConfig is intentionally NOT snapshotted/compared here — credential
+	// changes now flow through the dedicated /credentials subresource which
+	// handles its own CloseClient call. Main PUT does not accept secret
+	// fields (see handler comment on UpdateMCPService).
+	preURL := ""
+	preURLSet := existing.URL != nil
+	if preURLSet {
+		preURL = *existing.URL
+	}
+	var preStdioCommand string
+	var preStdioArgs []string
+	preStdioSet := existing.StdioConfig != nil
+	if preStdioSet {
+		preStdioCommand = existing.StdioConfig.Command
+		preStdioArgs = append([]string(nil), existing.StdioConfig.Args...)
+	}
+	preTransportType := existing.TransportType
+	preHeaders := map[string]string{}
+	if existing.AuthConfig != nil && existing.AuthConfig.CustomHeaders != nil {
+		maps.Copy(preHeaders, existing.AuthConfig.CustomHeaders)
+	}
+
+	preAuthType := types.MCPAuthNone
+	preAPIKeyHeader := ""
+	if existing.AuthConfig != nil {
+		preAuthType = existing.AuthConfig.AuthType
+		preAPIKeyHeader = existing.AuthConfig.APIKeyHeader
+	}
+
+	// CustomHeaders flows through main PUT (it's structural, not a secret) —
+	// nil preserves, non-nil replaces. Other AuthConfig fields (APIKey/Token)
+	// are never accepted via main PUT; the handler strips them up front.
+	//
+	// auth_type / scopes / auth_server_metadata_url are non-secret OAuth
+	// configuration and also flow through here.
+	if service.AuthConfig != nil {
+		if existing.AuthConfig == nil {
+			existing.AuthConfig = &types.MCPAuthConfig{}
+		}
+		if service.AuthConfig.CustomHeaders != nil {
+			existing.AuthConfig.CustomHeaders = service.AuthConfig.CustomHeaders
+		}
+		// Only overwrite OAuth config when explicitly provided, so a partial
+		// PUT that carries only custom_headers does not wipe an existing
+		// auth_type / scopes. (Empty/absent is treated as "no change".)
+		if service.AuthConfig.AuthType != types.MCPAuthNone {
+			existing.AuthConfig.AuthType = service.AuthConfig.AuthType
+		}
+		// APIKeyHeader is non-secret; empty means "use default X-API-Key", so a
+		// partial PUT that omits it is treated as no-change (mirrors scopes).
+		if service.AuthConfig.APIKeyHeader != "" {
+			existing.AuthConfig.APIKeyHeader = service.AuthConfig.APIKeyHeader
+		}
+		if service.AuthConfig.Scopes != nil {
+			existing.AuthConfig.Scopes = service.AuthConfig.Scopes
+		}
+		if service.AuthConfig.AuthServerMetadataURL != "" {
+			existing.AuthConfig.AuthServerMetadataURL = service.AuthConfig.AuthServerMetadataURL
+		}
+	}
+
+	// Scalar zero values cannot distinguish omission from an explicit update,
+	// so use the handler's field-presence map for name, description, and enabled.
+	if updateFields["name"] {
+		existing.Name = service.Name
+	}
+	if updateFields["description"] {
+		existing.Description = service.Description
+	}
+	if updateFields["enabled"] {
+		existing.Enabled = service.Enabled
+	}
+	if service.TransportType != "" {
+		existing.TransportType = service.TransportType
+	}
+	if service.URL != nil {
+		existing.URL = service.URL
+	}
+	if service.StdioConfig != nil {
+		existing.StdioConfig = service.StdioConfig
+	}
+	if service.EnvVars != nil {
+		existing.EnvVars = service.EnvVars
+	}
+	if service.Headers != nil {
+		existing.Headers = service.Headers
+	}
+	if service.AdvancedConfig != nil {
+		existing.AdvancedConfig = service.AdvancedConfig
+	}
+
+	// Update timestamp
+	if err := mcp.ValidateServiceOutboundURLs(existing); err != nil {
+		return err
+	}
+	existing.UpdatedAt = time.Now()
+
+	if err := s.mcpServiceRepo.Update(ctx, existing); err != nil {
+		logger.GetLogger(ctx).Errorf("Failed to update MCP service: %v", err)
+		return fmt.Errorf("failed to update MCP service: %w", err)
+	}
+
+	// Check if critical configuration changed (URL / StdioConfig / transport
+	// type / custom headers). Comparisons MUST be against the pre-merge
+	// snapshots captured above — after the in-place merge, service.URL and
+	// existing.URL point to the same memory, making any post-merge compare
+	// vacuously equal.
+	//
+	// AuthConfig API key / token changes do NOT go through this path; they
+	// are handled by the /credentials subresource which triggers CloseClient
+	// inline.
+	configChanged := false
+	currURLSet := existing.URL != nil
+	switch {
+	case currURLSet != preURLSet:
+		configChanged = true
+	case currURLSet && *existing.URL != preURL:
+		configChanged = true
+	}
+	currStdioSet := existing.StdioConfig != nil
+	switch {
+	case currStdioSet != preStdioSet:
+		configChanged = true
+	case currStdioSet && (existing.StdioConfig.Command != preStdioCommand ||
+		!slices.Equal(existing.StdioConfig.Args, preStdioArgs)):
+		configChanged = true
+	}
+	if existing.TransportType != preTransportType {
+		configChanged = true
+	}
+	currHeaders := map[string]string{}
+	if existing.AuthConfig != nil && existing.AuthConfig.CustomHeaders != nil {
+		currHeaders = existing.AuthConfig.CustomHeaders
+	}
+	if !maps.Equal(currHeaders, preHeaders) {
+		configChanged = true
+	}
+	if existing.AuthConfig != nil &&
+		(existing.AuthConfig.AuthType != preAuthType ||
+			existing.AuthConfig.APIKeyHeader != preAPIKeyHeader) {
+		configChanged = true
+	}
+	name := secutils.SanitizeForLog(existing.Name)
+	// Close existing client connection if:
+	// 1. Service is now disabled (need to close connection)
+	// 2. Critical configuration changed (need to reconnect with new config)
+	if !existing.Enabled {
+		s.mcpManager.CloseClient(service.ID)
+		logger.GetLogger(ctx).Infof("MCP service disabled, connection closed: %s (ID: %s)", name, service.ID)
+	} else if configChanged {
+		s.mcpManager.CloseClient(service.ID)
+		logger.GetLogger(ctx).Infof("MCP service config changed, connection closed: %s (ID: %s)", name, service.ID)
+	} else if oldEnabled != existing.Enabled && existing.Enabled {
+		// Service was just enabled (was disabled, now enabled)
+		// Close any existing connection to ensure clean state
+		s.mcpManager.CloseClient(service.ID)
+		logger.GetLogger(ctx).Infof("MCP service enabled, existing connection closed: %s (ID: %s)", name, service.ID)
+	}
+
+	logger.GetLogger(ctx).Infof("MCP service updated: %s (ID: %s), enabled: %v", name, service.ID, existing.Enabled)
+	return nil
+}
+
+// DeleteMCPService deletes an MCP service
+func (s *mcpServiceService) DeleteMCPService(ctx context.Context, tenantID uint64, id string) error {
+	// Check if service exists
+	existing, err := s.mcpServiceRepo.GetByID(ctx, tenantID, id)
+	if err != nil {
+		return fmt.Errorf("failed to get MCP service: %w", err)
+	}
+	if existing == nil {
+		return fmt.Errorf("MCP service not found")
+	}
+
+	// Builtin MCP services cannot be deleted
+	if existing.IsBuiltin {
+		return fmt.Errorf("builtin MCP services cannot be deleted")
+	}
+
+	// Close client connection
+	s.mcpManager.CloseClient(id)
+
+	if err := s.mcpServiceRepo.Delete(ctx, tenantID, id); err != nil {
+		logger.GetLogger(ctx).Errorf("Failed to delete MCP service: %v", err)
+		return fmt.Errorf("failed to delete MCP service: %w", err)
+	}
+
+	logger.GetLogger(ctx).Infof("MCP service deleted: %s (ID: %s)", secutils.SanitizeForLog(existing.Name), id)
+	return nil
+}
+
+// TestMCPService tests the connection to an MCP service and returns available tools/resources
+// mcpTestFailure builds a failed MCPTestResult, upgrading a generic connection
+// error into an explicit "OAuth required" signal when the server answered with
+// an RFC 9728 OAuth challenge. This lets the UI guide the user to switch the
+// auth strategy to OAuth instead of showing a bare 401.
+func mcpTestFailure(err error, prefix string) *types.MCPTestResult {
+	var oerr *mcp.OAuthRequiredError
+	if errors.As(err, &oerr) {
+		return &types.MCPTestResult{
+			Success:       false,
+			OAuthRequired: true,
+			Message: "This MCP server requires OAuth authorization. " +
+				"Switch the auth method to OAuth 2.0 and authorize.",
+		}
+	}
+	return &types.MCPTestResult{
+		Success: false,
+		Message: fmt.Sprintf("%s: %v", prefix, err),
+	}
+}
+
+func (s *mcpServiceService) TestMCPService(
+	ctx context.Context,
+	tenantID uint64,
+	id string,
+) (*types.MCPTestResult, error) {
+	// Get service
+	service, err := s.mcpServiceRepo.GetByID(ctx, tenantID, id)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get MCP service: %w", err)
+	}
+	if service == nil {
+		return nil, fmt.Errorf("MCP service not found")
+	}
+
+	// Create temporary client for testing. For OAuth services, wire the
+	// per-user token store so the test connects with the current user's
+	// authorization (and surfaces an authorization-required message when the
+	// user has not authorized yet).
+	config := &mcp.ClientConfig{
+		Service: service,
+	}
+	if service.AuthConfig.IsOAuth() {
+		config.OAuthRepo = s.oauthRepo
+		config.TenantID, _ = types.TenantIDFromContext(ctx)
+		config.Principal, _ = types.PrincipalFromContext(ctx)
+	}
+
+	client, err := mcp.NewMCPClient(config)
+	if err != nil {
+		return &types.MCPTestResult{
+			Success: false,
+			Message: fmt.Sprintf("Failed to create client: %v", err),
+		}, nil
+	}
+
+	// Connect
+	testCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+
+	if err := client.Connect(testCtx); err != nil {
+		return mcpTestFailure(err, "Connection failed"), nil
+	}
+	defer client.Disconnect()
+
+	// Initialize
+	initResult, err := client.Initialize(testCtx)
+	if err != nil {
+		return mcpTestFailure(err, "Initialization failed"), nil
+	}
+
+	// List tools
+	tools, err := client.ListTools(testCtx)
+	if err != nil {
+		logger.GetLogger(ctx).Warnf("Failed to list tools: %v", err)
+		tools = []*types.MCPTool{}
+	}
+
+	// List resources
+	resources, err := client.ListResources(testCtx)
+	if err != nil {
+		logger.GetLogger(ctx).Warnf("Failed to list resources: %v", err)
+		resources = []*types.MCPResource{}
+	}
+
+	return &types.MCPTestResult{
+		Success: true,
+		Message: fmt.Sprintf(
+			"Connected successfully to %s v%s",
+			initResult.ServerInfo.Name,
+			initResult.ServerInfo.Version,
+		),
+		Description: initResult.ServerInfo.Description,
+		Tools:       tools,
+		Resources:   resources,
+	}, nil
+}
+
+// GetMCPServiceTools retrieves the list of tools from an MCP service
+func (s *mcpServiceService) GetMCPServiceTools(
+	ctx context.Context,
+	tenantID uint64,
+	id string,
+) ([]*types.MCPTool, error) {
+	// Get service
+	service, err := s.mcpServiceRepo.GetByID(ctx, tenantID, id)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get MCP service: %w", err)
+	}
+	if service == nil {
+		return nil, fmt.Errorf("MCP service not found")
+	}
+
+	// Get or create client
+	client, err := s.mcpManager.GetOrCreateClient(ctx, service)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get MCP client: %w", err)
+	}
+
+	// List tools
+	tools, err := client.ListTools(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to list tools: %w", err)
+	}
+
+	return tools, nil
+}
+
+// UpdateMCPCredentials writes one or more credential fields and recycles any
+// active client connection so the next upstream call uses the new credential.
+//
+// Implementation notes:
+//   - apiKey == nil and token == nil → no-op, returns current state.
+//   - apiKey == &"" → explicit empty string; treated as no-op because clearing
+//     is the dedicated ClearMCPCredential path. The handler enforces this
+//     contract by accepting empty as no-op too; this is defense-in-depth.
+//   - apiKey == &"sk-..." → replaces stored value.
+//   - Builtin services cannot have credentials updated (mirrors the
+//     UpdateMCPService restriction).
+//   - Always re-fetches existing AuthConfig before merge to avoid clobbering
+//     CustomHeaders or the unrelated credential field.
+func (s *mcpServiceService) UpdateMCPCredentials(
+	ctx context.Context, tenantID uint64, id string, apiKey *string, token *string,
+) (*types.MCPService, error) {
+	existing, err := s.mcpServiceRepo.GetByID(ctx, tenantID, id)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get MCP service: %w", err)
+	}
+	if existing == nil {
+		return nil, fmt.Errorf("MCP service not found")
+	}
+	if existing.IsBuiltin {
+		return nil, fmt.Errorf("builtin MCP services cannot have credentials modified")
+	}
+
+	if existing.AuthConfig == nil {
+		existing.AuthConfig = &types.MCPAuthConfig{}
+	}
+	changed := false
+	if apiKey != nil && *apiKey != "" && *apiKey != existing.AuthConfig.APIKey {
+		existing.AuthConfig.APIKey = *apiKey
+		changed = true
+	}
+	if token != nil && *token != "" && *token != existing.AuthConfig.Token {
+		existing.AuthConfig.Token = *token
+		changed = true
+	}
+	if !changed {
+		return existing, nil
+	}
+
+	existing.UpdatedAt = time.Now()
+	if err := s.mcpServiceRepo.Update(ctx, existing); err != nil {
+		return nil, fmt.Errorf("failed to update MCP service: %w", err)
+	}
+
+	// Credential changed → recycle client so the next call reconnects.
+	s.mcpManager.CloseClient(id)
+	logger.GetLogger(ctx).Infof(
+		"MCP credentials updated, connection closed: %s (ID: %s)",
+		secutils.SanitizeForLog(existing.Name), id,
+	)
+	return existing, nil
+}
+
+// ClearMCPCredential removes a single credential field. Idempotent: clearing
+// an already-empty field returns nil without writing or reconnecting.
+func (s *mcpServiceService) ClearMCPCredential(
+	ctx context.Context, tenantID uint64, id, field string,
+) error {
+	existing, err := s.mcpServiceRepo.GetByID(ctx, tenantID, id)
+	if err != nil {
+		return fmt.Errorf("failed to get MCP service: %w", err)
+	}
+	if existing == nil {
+		return fmt.Errorf("MCP service not found")
+	}
+	if existing.IsBuiltin {
+		return fmt.Errorf("builtin MCP services cannot have credentials modified")
+	}
+	if existing.AuthConfig == nil {
+		return nil // nothing to clear
+	}
+
+	changed := false
+	switch field {
+	case "api_key":
+		if existing.AuthConfig.APIKey != "" {
+			existing.AuthConfig.APIKey = ""
+			changed = true
+		}
+	case "token":
+		if existing.AuthConfig.Token != "" {
+			existing.AuthConfig.Token = ""
+			changed = true
+		}
+	default:
+		return fmt.Errorf("unknown credential field: %s", field)
+	}
+	if !changed {
+		return nil
+	}
+
+	existing.UpdatedAt = time.Now()
+	if err := s.mcpServiceRepo.Update(ctx, existing); err != nil {
+		return fmt.Errorf("failed to update MCP service: %w", err)
+	}
+
+	s.mcpManager.CloseClient(id)
+	logger.GetLogger(ctx).Infof(
+		"MCP credential cleared by user: id=%s field=%s, connection closed",
+		secutils.SanitizeForLog(id), field,
+	)
+	return nil
+}
+
+// GetMCPServiceResources retrieves the list of resources from an MCP service
+func (s *mcpServiceService) GetMCPServiceResources(
+	ctx context.Context,
+	tenantID uint64,
+	id string,
+) ([]*types.MCPResource, error) {
+	// Get service
+	service, err := s.mcpServiceRepo.GetByID(ctx, tenantID, id)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get MCP service: %w", err)
+	}
+	if service == nil {
+		return nil, fmt.Errorf("MCP service not found")
+	}
+
+	// Get or create client
+	client, err := s.mcpManager.GetOrCreateClient(ctx, service)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get MCP client: %w", err)
+	}
+
+	// List resources
+	resources, err := client.ListResources(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to list resources: %w", err)
+	}
+
+	return resources, nil
+}