gitea-search/internal/mcp/server.go

package mcp

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"strings"

	"gitea.rspworks.tech/rpert/gitea-search/internal/meili"
)

// JSON-RPC message types
type jsonRPCRequest struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id,omitempty"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}

type jsonRPCResponse struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id,omitempty"`
	Result  interface{}     `json:"result,omitempty"`
	Error   *jsonRPCError   `json:"error,omitempty"`
}

type jsonRPCError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

// MCP protocol types
type serverInfo struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

type serverCapabilities struct {
	Tools *toolsCapability `json:"tools,omitempty"`
}

type toolsCapability struct {
	ListChanged bool `json:"listChanged,omitempty"`
}

type initializeResult struct {
	ProtocolVersion string             `json:"protocolVersion"`
	ServerInfo      serverInfo         `json:"serverInfo"`
	Capabilities    serverCapabilities `json:"capabilities"`
}

type toolDefinition struct {
	Name        string      `json:"name"`
	Description string      `json:"description"`
	InputSchema inputSchema `json:"inputSchema"`
}

type inputSchema struct {
	Type       string              `json:"type"`
	Properties map[string]property `json:"properties"`
	Required   []string            `json:"required,omitempty"`
}

type property struct {
	Type        string `json:"type"`
	Description string `json:"description"`
	Default     any    `json:"default,omitempty"`
}

type toolsListResult struct {
	Tools []toolDefinition `json:"tools"`
}

type toolCallParams struct {
	Name      string          `json:"name"`
	Arguments json.RawMessage `json:"arguments"`
}

type searchArgs struct {
	Query    string `json:"query"`
	Repo     string `json:"repo"`
	Filetype string `json:"filetype"`
	Limit    int64  `json:"limit"`
}

type contentItem struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

type toolCallResult struct {
	Content []contentItem `json:"content"`
	IsError bool          `json:"isError,omitempty"`
}

// Server is the MCP server that handles stdio JSON-RPC.
type Server struct {
	meiliClient *meili.Client
	version     string
}

// NewServer creates a new MCP server.
func NewServer(meiliClient *meili.Client, version string) *Server {
	return &Server{
		meiliClient: meiliClient,
		version:     version,
	}
}

// Run starts the MCP server, reading from stdin and writing to stdout.
func (s *Server) Run() error {
	// Log to stderr so it doesn't interfere with JSON-RPC on stdout
	log.SetOutput(os.Stderr)
	log.SetPrefix("[mcp-server] ")

	reader := bufio.NewReader(os.Stdin)
	writer := os.Stdout

	log.Println("MCP server started, waiting for requests on stdin")

	for {
		line, err := reader.ReadBytes('\n')
		if err != nil {
			if err == io.EOF {
				log.Println("stdin closed, shutting down")
				return nil
			}
			return fmt.Errorf("reading stdin: %w", err)
		}

		line = []byte(strings.TrimSpace(string(line)))
		if len(line) == 0 {
			continue
		}

		var req jsonRPCRequest
		if err := json.Unmarshal(line, &req); err != nil {
			log.Printf("invalid JSON-RPC request: %s", string(line))
			continue
		}

		resp := s.handleRequest(req)
		if resp == nil {
			// Notification, no response needed
			continue
		}

		respBytes, err := json.Marshal(resp)
		if err != nil {
			log.Printf("error marshaling response: %v", err)
			continue
		}

		respBytes = append(respBytes, '\n')
		if _, err := writer.Write(respBytes); err != nil {
			return fmt.Errorf("writing response: %w", err)
		}
	}
}

func (s *Server) handleRequest(req jsonRPCRequest) *jsonRPCResponse {
	switch req.Method {
	case "initialize":
		return s.handleInitialize(req)
	case "notifications/initialized":
		log.Println("Client initialized")
		return nil // notification, no response
	case "tools/list":
		return s.handleToolsList(req)
	case "tools/call":
		return s.handleToolsCall(req)
	case "ping":
		return &jsonRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result:  map[string]interface{}{},
		}
	default:
		log.Printf("Unknown method: %s", req.Method)
		return &jsonRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error: &jsonRPCError{
				Code:    -32601,
				Message: fmt.Sprintf("method not found: %s", req.Method),
			},
		}
	}
}

func (s *Server) handleInitialize(req jsonRPCRequest) *jsonRPCResponse {
	log.Println("Handling initialize")
	return &jsonRPCResponse{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result: initializeResult{
			ProtocolVersion: "2024-11-05",
			ServerInfo: serverInfo{
				Name:    "gitea-search",
				Version: s.version,
			},
			Capabilities: serverCapabilities{
				Tools: &toolsCapability{},
			},
		},
	}
}

func (s *Server) handleToolsList(req jsonRPCRequest) *jsonRPCResponse {
	log.Println("Handling tools/list")
	return &jsonRPCResponse{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result: toolsListResult{
			Tools: []toolDefinition{
				{
					Name:        "gitea_search",
					Description: "Search across all Gitea repositories. Returns matching files with code snippets. Use this to find code, configuration, documentation, or any file content across the codebase.",
					InputSchema: inputSchema{
						Type: "object",
						Properties: map[string]property{
							"query": {
								Type:        "string",
								Description: "Search terms. Supports natural language queries and exact phrases.",
							},
							"repo": {
								Type:        "string",
								Description: "Filter to a specific repo by full name (e.g., 'rpert/infra-ssh'). Omit to search all repos.",
							},
							"filetype": {
								Type:        "string",
								Description: "Filter by file extension without dot (e.g., 'go', 'md', 'yaml', 'py').",
							},
							"limit": {
								Type:        "integer",
								Description: "Maximum number of results to return.",
								Default:     10,
							},
						},
						Required: []string{"query"},
					},
				},
			},
		},
	}
}

func (s *Server) handleToolsCall(req jsonRPCRequest) *jsonRPCResponse {
	var params toolCallParams
	if err := json.Unmarshal(req.Params, &params); err != nil {
		return &jsonRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error: &jsonRPCError{
				Code:    -32602,
				Message: fmt.Sprintf("invalid params: %v", err),
			},
		}
	}

	if params.Name != "gitea_search" {
		return &jsonRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error: &jsonRPCError{
				Code:    -32602,
				Message: fmt.Sprintf("unknown tool: %s", params.Name),
			},
		}
	}

	var args searchArgs
	if err := json.Unmarshal(params.Arguments, &args); err != nil {
		return &jsonRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result: toolCallResult{
				Content: []contentItem{{Type: "text", Text: fmt.Sprintf("Error parsing arguments: %v", err)}},
				IsError: true,
			},
		}
	}

	if args.Query == "" {
		return &jsonRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result: toolCallResult{
				Content: []contentItem{{Type: "text", Text: "Error: query parameter is required"}},
				IsError: true,
			},
		}
	}

	if args.Limit <= 0 {
		args.Limit = 10
	}

	log.Printf("Searching: query=%q repo=%q filetype=%q limit=%d", args.Query, args.Repo, args.Filetype, args.Limit)

	results, err := s.meiliClient.Search(args.Query, args.Repo, args.Filetype, args.Limit)
	if err != nil {
		return &jsonRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result: toolCallResult{
				Content: []contentItem{{Type: "text", Text: fmt.Sprintf("Search error: %v", err)}},
				IsError: true,
			},
		}
	}

	text := formatResults(args.Query, results)

	return &jsonRPCResponse{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result: toolCallResult{
			Content: []contentItem{{Type: "text", Text: text}},
		},
	}
}

func formatResults(query string, results []meili.SearchResult) string {
	if len(results) == 0 {
		return fmt.Sprintf("No results found for %q.", query)
	}

	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("Found %d results for %q:\n\n", len(results), query))

	for i, r := range results {
		sb.WriteString(fmt.Sprintf("### %d. %s — `%s`\n", i+1, r.Repo, r.Path))
		sb.WriteString(fmt.Sprintf("- Branch: `%s`\n", r.Branch))
		if r.Extension != "" {
			sb.WriteString(fmt.Sprintf("- Type: `%s`\n", r.Extension))
		}
		if r.Snippet != "" {
			sb.WriteString(fmt.Sprintf("- Snippet: ...%s...\n", r.Snippet))
		}
		sb.WriteString("\n")
	}

	return sb.String()
}