labmcp/internal/mcp/server.go

package mcp

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

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

// ServerMode indicates which type of tools the server should expose.
type ServerMode string

const (
	// ModeOptions exposes only option-related tools.
	ModeOptions ServerMode = "options"
	// ModePackages exposes only package-related tools.
	ModePackages ServerMode = "packages"
	// ModeCustom exposes externally registered tools (no database required).
	ModeCustom ServerMode = "custom"
)

// 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
	// InstructionsFunc, if set, is called during initialization to generate
	// dynamic instructions. Its return value is appended to Instructions.
	InstructionsFunc func() 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
	// Mode specifies which tools to expose (options or packages).
	Mode ServerMode
}

// DefaultNixOSConfig returns the default configuration for NixOS options server.
func DefaultNixOSConfig() ServerConfig {
	return ServerConfig{
		Name:           "nixos-options",
		Version:        "0.4.0",
		DefaultChannel: "nixos-stable",
		SourceName:     "nixpkgs",
		Mode:           ModeOptions,
		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.

Note: index_revision also indexes packages when available, so both options and packages become searchable.`,
	}
}

// DefaultNixpkgsPackagesConfig returns the default configuration for nixpkgs packages server.
func DefaultNixpkgsPackagesConfig() ServerConfig {
	return ServerConfig{
		Name:           "nixpkgs-packages",
		Version:        "0.4.0",
		DefaultChannel: "nixos-stable",
		SourceName:     "nixpkgs",
		Mode:           ModePackages,
		Instructions: `Nixpkgs Packages MCP Server - Search and query Nix packages from nixpkgs.

If the current project contains a flake.lock file, you can search packages from 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 packages for that specific version

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

This ensures package information matches the nixpkgs version the project actually uses.`,
	}
}

// DefaultMonitoringConfig returns the default configuration for the lab monitoring server.
func DefaultMonitoringConfig() ServerConfig {
	return ServerConfig{
		Name:    "lab-monitoring",
		Version: "0.3.1",
		Mode:    ModeCustom,
		Instructions: `Lab Monitoring MCP Server - Query Prometheus metrics and Alertmanager alerts.

Tools for querying your monitoring stack:
- Search and query Prometheus metrics with PromQL
- List and inspect alerts from Alertmanager
- View scrape target health status
- Manage alert silences
- Query logs via LogQL (when Loki is configured)

All queries are executed against live Prometheus, Alertmanager, and Loki HTTP APIs.`,
	}
}

// DefaultHomeManagerConfig returns the default configuration for Home Manager options server.
func DefaultHomeManagerConfig() ServerConfig {
	return ServerConfig{
		Name:           "hm-options",
		Version:        "0.3.0",
		DefaultChannel: "hm-stable",
		SourceName:     "home-manager",
		Mode:           ModeOptions,
		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.`,
	}
}

// DefaultGitExplorerConfig returns the default configuration for the git-explorer server.
func DefaultGitExplorerConfig() ServerConfig {
	return ServerConfig{
		Name:    "git-explorer",
		Version: "0.1.0",
		Mode:    ModeCustom,
		Instructions: `Git Explorer MCP Server - Read-only access to git repository information.

Tools for exploring git repositories:
- Resolve refs (branches, tags, commits) to commit hashes
- View commit logs with filtering by author, path, or range
- Get full commit details including file change statistics
- Compare commits to see which files changed
- Read file contents at any commit
- Check ancestry relationships between commits
- Search commit messages

All operations are read-only and will never modify the repository.`,
	}
}

// Server is an MCP server that handles JSON-RPC requests.
type Server struct {
	store       database.Store
	config      ServerConfig
	tools       map[string]ToolHandler
	toolDefs    []Tool
	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 a database store.
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
}

// NewGenericServer creates a new MCP server without a database store.
// Use RegisterTool to add tools externally.
func NewGenericServer(logger *log.Logger, config ServerConfig) *Server {
	if logger == nil {
		logger = log.New(io.Discard, "", 0)
	}
	return &Server{
		config: config,
		tools:  make(map[string]ToolHandler),
		logger: logger,
	}
}

// RegisterTool registers an externally defined tool with its handler.
func (s *Server) RegisterTool(tool Tool, handler ToolHandler) {
	s.toolDefs = append(s.toolDefs, tool)
	s.tools[tool.Name] = handler
}

// 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)

	instructions := s.config.Instructions
	if s.config.InstructionsFunc != nil {
		if extra := s.config.InstructionsFunc(); extra != "" {
			instructions += "\n\n" + extra
		}
	}

	result := InitializeResult{
		ProtocolVersion: ProtocolVersion,
		Capabilities: Capabilities{
			Tools: &ToolsCapability{
				ListChanged: false,
			},
		},
		ServerInfo: Implementation{
			Name:    s.config.Name,
			Version: s.config.Version,
		},
		Instructions: 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 {
	// For custom mode, return externally registered tools
	if s.config.Mode == ModeCustom {
		return s.toolDefs
	}

	// For packages mode, return package tools
	if s.config.Mode == ModePackages {
		return s.getPackageToolDefinitions()
	}

	// Default: options mode
	return s.getOptionToolDefinitions()
}

// getOptionToolDefinitions returns the tool definitions for options mode.
func (s *Server) getOptionToolDefinitions() []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.",
					},
					"offset": {
						Type:        "integer",
						Description: "Line offset (0-based). Default: 0",
						Default:     0,
					},
					"limit": {
						Type:        "integer",
						Description: "Maximum lines to return. Default: 250, use 0 for all lines",
						Default:     250,
					},
				},
				Required: []string{"path"},
			},
		},
		{
			Name:        "index_revision",
			Description: s.indexRevisionDescription(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"},
			},
		},
	}
}

// indexRevisionDescription returns the description for the index_revision tool,
// adjusted based on whether packages are also indexed.
func (s *Server) indexRevisionDescription(sourceRepo string) string {
	if s.config.SourceName == "nixpkgs" {
		return fmt.Sprintf("Index a %s revision to make its options and packages searchable", sourceRepo)
	}
	return fmt.Sprintf("Index a %s revision to make its options searchable", sourceRepo)
}

// getPackageToolDefinitions returns the tool definitions for packages mode.
func (s *Server) getPackageToolDefinitions() []Tool {
	exampleChannels := "'nixos-unstable', 'nixos-24.05'"
	exampleFilePath := "pkgs/applications/networking/browsers/firefox/default.nix"

	return []Tool{
		{
			Name:        "search_packages",
			Description: "Search for Nix packages by name or description",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"query": {
						Type:        "string",
						Description: "Search query (matches package name, attr path, and description)",
					},
					"revision": {
						Type:        "string",
						Description: fmt.Sprintf("Git hash or channel name (e.g., %s). Uses default if not specified.", exampleChannels),
					},
					"broken": {
						Type:        "boolean",
						Description: "Filter by broken status (true = only broken, false = only working)",
					},
					"unfree": {
						Type:        "boolean",
						Description: "Filter by license (true = only unfree, false = only free)",
					},
					"limit": {
						Type:        "integer",
						Description: "Maximum number of results (default: 50)",
						Default:     50,
					},
				},
				Required: []string{"query"},
			},
		},
		{
			Name:        "get_package",
			Description: "Get full details for a specific Nix package by attribute path",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"attr_path": {
						Type:        "string",
						Description: "Package attribute path (e.g., 'firefox', 'python312Packages.requests')",
					},
					"revision": {
						Type:        "string",
						Description: "Git hash or channel name. Uses default if not specified.",
					},
				},
				Required: []string{"attr_path"},
			},
		},
		{
			Name:        "get_file",
			Description: "Fetch the contents of a file from nixpkgs",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"path": {
						Type:        "string",
						Description: fmt.Sprintf("File path relative to nixpkgs root (e.g., '%s')", exampleFilePath),
					},
					"revision": {
						Type:        "string",
						Description: "Git hash or channel name. Uses default if not specified.",
					},
					"offset": {
						Type:        "integer",
						Description: "Line offset (0-based). Default: 0",
						Default:     0,
					},
					"limit": {
						Type:        "integer",
						Description: "Maximum lines to return. Default: 250, use 0 for all lines",
						Default:     250,
					},
				},
				Required: []string{"path"},
			},
		},
		{
			Name:        "index_revision",
			Description: "Index a nixpkgs revision to make its packages searchable",
			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: "List all indexed nixpkgs revisions",
			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,
	}
}