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.3.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.`,
	}
}

// DefaultNixpkgsPackagesConfig returns the default configuration for nixpkgs packages server.
func DefaultNixpkgsPackagesConfig() ServerConfig {
	return ServerConfig{
		Name:           "nixpkgs-packages",
		Version:        "0.3.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. Ensure the revision is indexed (packages are indexed separately from options)

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.0",
		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: 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"},
			},
		},
	}
}

// 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:        "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,
	}
}