gitea-mcp/operation/operation.go

package operation

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"strings"
	"syscall"
	"time"

	"gitea.com/gitea/gitea-mcp/operation/actions"
	"gitea.com/gitea/gitea-mcp/operation/issue"
	"gitea.com/gitea/gitea-mcp/operation/label"
	"gitea.com/gitea/gitea-mcp/operation/milestone"
	"gitea.com/gitea/gitea-mcp/operation/notification"
	"gitea.com/gitea/gitea-mcp/operation/org"
	"gitea.com/gitea/gitea-mcp/operation/pull"
	"gitea.com/gitea/gitea-mcp/operation/repo"
	"gitea.com/gitea/gitea-mcp/operation/search"
	"gitea.com/gitea/gitea-mcp/operation/timetracking"
	"gitea.com/gitea/gitea-mcp/operation/user"
	"gitea.com/gitea/gitea-mcp/operation/version"
	"gitea.com/gitea/gitea-mcp/operation/wiki"
	mcpContext "gitea.com/gitea/gitea-mcp/pkg/context"
	"gitea.com/gitea/gitea-mcp/pkg/flag"
	"gitea.com/gitea/gitea-mcp/pkg/log"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
)

var mcpServer *server.MCPServer

// serverNames returns a sorted list of configured server names for use in
// enum validation and the list_servers tool.
func serverNames() []string {
	names := make([]string, 0, len(flag.Servers))
	for name := range flag.Servers {
		names = append(names, name)
	}
	// Deterministic ordering for tool schemas.
	sort.Strings(names)
	return names
}

// injectServerParam adds an optional "server" parameter to a tool and wraps
// its handler so the chosen server's host/token/insecure are injected into
// the context before the original handler runs.
func injectServerParam(tools []server.ServerTool) []server.ServerTool {
	names := serverNames()
	for i, st := range tools {
		// Add "server" to the JSON Schema properties.
		if st.Tool.InputSchema.Properties == nil {
			st.Tool.InputSchema.Properties = make(map[string]any)
		}
		st.Tool.InputSchema.Properties["server"] = map[string]any{
			"type":        "string",
			"description": "Target Gitea server name. Available: " + strings.Join(names, ", ") + ". Omit to use the default server.",
			"enum":        names,
		}
		tools[i].Tool = st.Tool

		// Wrap the handler to resolve server config → context values.
		origHandler := st.Handler
		tools[i].Handler = func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
			if serverName, ok := req.GetArguments()["server"].(string); ok && serverName != "" {
				cfg, exists := flag.Servers[serverName]
				if !exists {
					return mcp.NewToolResultText("unknown server: " + serverName + ". Available: " + strings.Join(serverNames(), ", ")), nil
				}
				ctx = context.WithValue(ctx, mcpContext.HostContextKey, cfg.Host)
				ctx = context.WithValue(ctx, mcpContext.TokenContextKey, cfg.Token)
				ctx = context.WithValue(ctx, mcpContext.InsecureContextKey, cfg.Insecure)
			}
			return origHandler(ctx, req)
		}
	}
	return tools
}

// listServersFn handles the "list_servers" tool — returns the configured servers.
func listServersFn(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	if len(flag.Servers) == 0 {
		defaultInfo := fmt.Sprintf("default: %s (no additional servers configured)", flag.Host)
		return mcp.NewToolResultText(defaultInfo), nil
	}
	var b strings.Builder
	b.WriteString(fmt.Sprintf("default: %s\n", flag.Host))
	for _, name := range serverNames() {
		cfg := flag.Servers[name]
		b.WriteString(fmt.Sprintf("%s: %s", name, cfg.Host))
		if cfg.Insecure {
			b.WriteString(" (insecure)")
		}
		b.WriteString("\n")
	}
	return mcp.NewToolResultText(b.String()), nil
}

func RegisterTool(s *server.MCPServer) {
	// Collect all tools from every domain package.
	allTools := make([]server.ServerTool, 0, 200)
	allTools = append(allTools, user.Tool.Tools()...)
	allTools = append(allTools, actions.Tool.Tools()...)
	allTools = append(allTools, repo.Tool.Tools()...)
	allTools = append(allTools, notification.Tool.Tools()...)
	allTools = append(allTools, issue.Tool.Tools()...)
	allTools = append(allTools, label.Tool.Tools()...)
	allTools = append(allTools, milestone.Tool.Tools()...)
	allTools = append(allTools, pull.Tool.Tools()...)
	allTools = append(allTools, search.Tool.Tools()...)
	allTools = append(allTools, version.Tool.Tools()...)
	allTools = append(allTools, wiki.Tool.Tools()...)
	allTools = append(allTools, timetracking.Tool.Tools()...)
	allTools = append(allTools, org.Tool.Tools()...)

	// If multiple servers are configured, inject the "server" parameter into every tool.
	if len(flag.Servers) > 0 {
		allTools = injectServerParam(allTools)

		// Add a list_servers meta-tool so callers can discover available servers.
		allTools = append(allTools, server.ServerTool{
			Tool: mcp.NewTool("list_servers",
				mcp.WithDescription("List all configured Gitea server connections"),
			),
			Handler: listServersFn,
		})
	}

	s.AddTools(allTools...)
	s.DeleteTools("")
}

// parseAuthToken extracts the token from an Authorization header.
// Supports "Bearer <token>" (case-insensitive per RFC 7235) and
// Gitea-style "token <token>" formats.
// Returns the token and true if valid, empty string and false otherwise.
func parseAuthToken(authHeader string) (string, bool) {
	if len(authHeader) > 7 && strings.EqualFold(authHeader[:7], "Bearer ") {
		token := strings.TrimSpace(authHeader[7:])
		if token != "" {
			return token, true
		}
	}
	if len(authHeader) > 6 && strings.EqualFold(authHeader[:6], "token ") {
		token := strings.TrimSpace(authHeader[6:])
		if token != "" {
			return token, true
		}
	}
	return "", false
}

func getContextWithToken(ctx context.Context, r *http.Request) context.Context {
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" {
		return ctx
	}

	token, ok := parseAuthToken(authHeader)
	if !ok {
		return ctx
	}

	return context.WithValue(ctx, mcpContext.TokenContextKey, token)
}

func Run() error {
	mcpServer = newMCPServer(flag.Version)
	RegisterTool(mcpServer)
	switch flag.Mode {
	case "stdio":
		if err := server.ServeStdio(
			mcpServer,
		); err != nil {
			return err
		}
	case "http":
		httpServer := server.NewStreamableHTTPServer(
			mcpServer,
			server.WithLogger(log.New()),
			server.WithHeartbeatInterval(30*time.Second),
			server.WithHTTPContextFunc(getContextWithToken),
		)
		log.Infof("Gitea MCP HTTP server listening on :%d", flag.Port)

		// Graceful shutdown setup
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
		shutdownDone := make(chan struct{})

		go func() {
			<-sigCh
			log.Infof("Shutdown signal received, gracefully stopping HTTP server...")
			shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
			defer cancel()
			if err := httpServer.Shutdown(shutdownCtx); err != nil {
				log.Errorf("HTTP server shutdown error: %v", err)
			}
			close(shutdownDone)
		}()

		if err := httpServer.Start(fmt.Sprintf(":%d", flag.Port)); err != nil && !errors.Is(err, http.ErrServerClosed) {
			return err
		}
		<-shutdownDone // Wait for shutdown to finish
	default:
		return fmt.Errorf("invalid transport type: %s. Must be 'stdio' or 'http'", flag.Mode)
	}
	return nil
}

func newMCPServer(version string) *server.MCPServer {
	return server.NewMCPServer(
		"Gitea MCP Server",
		version,
		server.WithToolCapabilities(true),
		server.WithLogging(),
		server.WithRecovery(),
	)
}