labmcp/internal/mcp/server.go

package mcp

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"

	"git.t-juice.club/torjus/labmcp/internal/database"
)

// ServerConfig contains configuration for the MCP server.
type ServerConfig struct {
	// Name is the server name reported in initialization.
	Name string
	// Version is the server version.
	Version string
	// Instructions are the server instructions sent to clients.
	Instructions string
	// DefaultChannel is the default channel to use when no revision is specified.
	DefaultChannel string
	// SourceName is the name of the source repository (e.g., "nixpkgs", "home-manager").
	SourceName string
}

// DefaultNixOSConfig returns the default configuration for NixOS options server.
func DefaultNixOSConfig() ServerConfig {
	return ServerConfig{
		Name:           "nixos-options",
		Version:        "0.1.1",
		DefaultChannel: "nixos-stable",
		SourceName:     "nixpkgs",
		Instructions: `NixOS Options MCP Server - Search and query NixOS configuration options.

If the current project contains a flake.lock file, you can index the exact nixpkgs revision used by the project:
1. Read the flake.lock file to find the nixpkgs "rev" field
2. Call index_revision with that git hash to index options for that specific version

Example: If flake.lock contains "rev": "abc123...", call index_revision with revision "abc123...".

This ensures option documentation matches the nixpkgs version the project actually uses.`,
	}
}

// DefaultHomeManagerConfig returns the default configuration for Home Manager options server.
func DefaultHomeManagerConfig() ServerConfig {
	return ServerConfig{
		Name:           "hm-options",
		Version:        "0.1.1",
		DefaultChannel: "hm-stable",
		SourceName:     "home-manager",
		Instructions: `Home Manager Options MCP Server - Search and query Home Manager configuration options.

If the current project contains a flake.lock file, you can index the exact home-manager revision used by the project:
1. Read the flake.lock file to find the home-manager "rev" field
2. Call index_revision with that git hash to index options for that specific version

Example: If flake.lock contains "rev": "abc123...", call index_revision with revision "abc123...".

This ensures option documentation matches the home-manager version the project actually uses.`,
	}
}

// Server is an MCP server that handles JSON-RPC requests.
type Server struct {
	store       database.Store
	config      ServerConfig
	tools       map[string]ToolHandler
	initialized bool
	logger      *log.Logger
}

// ToolHandler is a function that handles a tool call.
type ToolHandler func(ctx context.Context, args map[string]interface{}) (CallToolResult, error)

// NewServer creates a new MCP server with the given configuration.
func NewServer(store database.Store, logger *log.Logger, config ServerConfig) *Server {
	if logger == nil {
		logger = log.New(io.Discard, "", 0)
	}
	s := &Server{
		store:  store,
		config: config,
		tools:  make(map[string]ToolHandler),
		logger: logger,
	}
	s.registerTools()
	return s
}

// registerTools registers all available tools.
func (s *Server) registerTools() {
	// Tools will be implemented in handlers.go
}

// Run starts the server using STDIO transport (backward compatibility).
func (s *Server) Run(ctx context.Context, r io.Reader, w io.Writer) error {
	transport := NewStdioTransport(s, r, w)
	return transport.Run(ctx)
}

// HandleMessage parses a JSON-RPC message and returns the response.
// Returns (nil, nil) for notifications that don't require a response.
func (s *Server) HandleMessage(ctx context.Context, data []byte) (*Response, error) {
	var req Request
	if err := json.Unmarshal(data, &req); err != nil {
		return &Response{
			JSONRPC: "2.0",
			Error: &Error{
				Code:    ParseError,
				Message: "Parse error",
				Data:    err.Error(),
			},
		}, nil
	}

	return s.HandleRequest(ctx, &req), nil
}

// HandleRequest processes a single request and returns a response.
// Returns nil for notifications that don't require a response.
func (s *Server) HandleRequest(ctx context.Context, req *Request) *Response {
	return s.handleRequest(ctx, req)
}

// handleRequest processes a single request and returns a response.
func (s *Server) handleRequest(ctx context.Context, req *Request) *Response {
	s.logger.Printf("Received request: method=%s id=%v", req.Method, req.ID)

	switch req.Method {
	case MethodInitialize:
		return s.handleInitialize(req)
	case MethodInitialized:
		// This is a notification, no response needed
		s.initialized = true
		return nil
	case MethodToolsList:
		return s.handleToolsList(req)
	case MethodToolsCall:
		return s.handleToolsCall(ctx, req)
	default:
		return &Response{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error: &Error{
				Code:    MethodNotFound,
				Message: "Method not found",
				Data:    req.Method,
			},
		}
	}
}

// handleInitialize processes the initialize request.
func (s *Server) handleInitialize(req *Request) *Response {
	var params InitializeParams
	if req.Params != nil {
		if err := json.Unmarshal(req.Params, &params); err != nil {
			return &Response{
				JSONRPC: "2.0",
				ID:      req.ID,
				Error: &Error{
					Code:    InvalidParams,
					Message: "Invalid params",
					Data:    err.Error(),
				},
			}
		}
	}

	s.logger.Printf("Client: %s %s, protocol: %s",
		params.ClientInfo.Name, params.ClientInfo.Version, params.ProtocolVersion)

	result := InitializeResult{
		ProtocolVersion: ProtocolVersion,
		Capabilities: Capabilities{
			Tools: &ToolsCapability{
				ListChanged: false,
			},
		},
		ServerInfo: Implementation{
			Name:    s.config.Name,
			Version: s.config.Version,
		},
		Instructions: s.config.Instructions,
	}

	return &Response{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result:  result,
	}
}

// handleToolsList returns the list of available tools.
func (s *Server) handleToolsList(req *Request) *Response {
	tools := s.getToolDefinitions()
	return &Response{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result:  ListToolsResult{Tools: tools},
	}
}

// getToolDefinitions returns the tool definitions.
func (s *Server) getToolDefinitions() []Tool {
	// Determine naming based on source
	optionType := "NixOS"
	sourceRepo := "nixpkgs"
	exampleOption := "services.nginx.enable"
	exampleNamespace := "services.nginx"
	exampleFilePath := "nixos/modules/services/web-servers/nginx/default.nix"
	exampleChannels := "'nixos-unstable', 'nixos-24.05'"

	if s.config.SourceName == "home-manager" {
		optionType = "Home Manager"
		sourceRepo = "home-manager"
		exampleOption = "programs.git.enable"
		exampleNamespace = "programs.git"
		exampleFilePath = "modules/programs/git.nix"
		exampleChannels = "'hm-unstable', 'release-24.11'"
	}

	return []Tool{
		{
			Name:        "search_options",
			Description: fmt.Sprintf("Search for %s configuration options by name or description", optionType),
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"query": {
						Type:        "string",
						Description: "Search query (matches option names and descriptions)",
					},
					"revision": {
						Type:        "string",
						Description: fmt.Sprintf("Git hash or channel name (e.g., %s). Uses default if not specified.", exampleChannels),
					},
					"type": {
						Type:        "string",
						Description: "Filter by option type (e.g., 'boolean', 'string', 'list')",
					},
					"namespace": {
						Type:        "string",
						Description: fmt.Sprintf("Filter by namespace prefix (e.g., '%s')", exampleNamespace),
					},
					"limit": {
						Type:        "integer",
						Description: "Maximum number of results (default: 50)",
						Default:     50,
					},
				},
				Required: []string{"query"},
			},
		},
		{
			Name:        "get_option",
			Description: fmt.Sprintf("Get full details for a specific %s option including its children", optionType),
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"name": {
						Type:        "string",
						Description: fmt.Sprintf("Full option path (e.g., '%s')", exampleOption),
					},
					"revision": {
						Type:        "string",
						Description: "Git hash or channel name. Uses default if not specified.",
					},
					"include_children": {
						Type:        "boolean",
						Description: "Include direct children of this option (default: true)",
						Default:     true,
					},
				},
				Required: []string{"name"},
			},
		},
		{
			Name:        "get_file",
			Description: fmt.Sprintf("Fetch the contents of a file from %s", sourceRepo),
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"path": {
						Type:        "string",
						Description: fmt.Sprintf("File path relative to %s root (e.g., '%s')", sourceRepo, exampleFilePath),
					},
					"revision": {
						Type:        "string",
						Description: "Git hash or channel name. Uses default if not specified.",
					},
				},
				Required: []string{"path"},
			},
		},
		{
			Name:        "index_revision",
			Description: fmt.Sprintf("Index a %s revision to make its options searchable", sourceRepo),
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"revision": {
						Type:        "string",
						Description: fmt.Sprintf("Git hash (full or short) or channel name (e.g., %s)", exampleChannels),
					},
				},
				Required: []string{"revision"},
			},
		},
		{
			Name:        "list_revisions",
			Description: fmt.Sprintf("List all indexed %s revisions", sourceRepo),
			InputSchema: InputSchema{
				Type:       "object",
				Properties: map[string]Property{},
			},
		},
		{
			Name:        "delete_revision",
			Description: "Delete an indexed revision and all its data",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"revision": {
						Type:        "string",
						Description: "Git hash or channel name of the revision to delete",
					},
				},
				Required: []string{"revision"},
			},
		},
	}
}

// handleToolsCall handles a tool invocation.
func (s *Server) handleToolsCall(ctx context.Context, req *Request) *Response {
	var params CallToolParams
	if err := json.Unmarshal(req.Params, &params); err != nil {
		return &Response{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error: &Error{
				Code:    InvalidParams,
				Message: "Invalid params",
				Data:    err.Error(),
			},
		}
	}

	s.logger.Printf("Tool call: %s with args %v", params.Name, params.Arguments)

	handler, ok := s.tools[params.Name]
	if !ok {
		return &Response{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result: CallToolResult{
				Content: []Content{TextContent(fmt.Sprintf("Unknown tool: %s", params.Name))},
				IsError: true,
			},
		}
	}

	result, err := handler(ctx, params.Arguments)
	if err != nil {
		s.logger.Printf("Tool error: %v", err)
		return &Response{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result:  ErrorContent(err),
		}
	}

	return &Response{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result:  result,
	}
}