labmcp/internal/mcp/server.go

package mcp

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

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

// Server is an MCP server that handles JSON-RPC requests.
type Server struct {
	store       database.Store
	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.
func NewServer(store database.Store, logger *log.Logger) *Server {
	if logger == nil {
		logger = log.New(io.Discard, "", 0)
	}
	s := &Server{
		store:  store,
		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:    "nixos-options",
			Version: "0.1.0",
		},
		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.`,
	}

	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 {
	return []Tool{
		{
			Name:        "search_options",
			Description: "Search for NixOS configuration options by name or description",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"query": {
						Type:        "string",
						Description: "Search query (matches option names and descriptions)",
					},
					"revision": {
						Type:        "string",
						Description: "Git hash or channel name (e.g., 'nixos-unstable'). Uses default if not specified.",
					},
					"type": {
						Type:        "string",
						Description: "Filter by option type (e.g., 'boolean', 'string', 'list')",
					},
					"namespace": {
						Type:        "string",
						Description: "Filter by namespace prefix (e.g., 'services.nginx')",
					},
					"limit": {
						Type:        "integer",
						Description: "Maximum number of results (default: 50)",
						Default:     50,
					},
				},
				Required: []string{"query"},
			},
		},
		{
			Name:        "get_option",
			Description: "Get full details for a specific NixOS option including its children",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"name": {
						Type:        "string",
						Description: "Full option path (e.g., 'services.nginx.enable')",
					},
					"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: "Fetch the contents of a file from nixpkgs",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"path": {
						Type:        "string",
						Description: "File path relative to nixpkgs root (e.g., 'nixos/modules/services/web-servers/nginx/default.nix')",
					},
					"revision": {
						Type:        "string",
						Description: "Git hash or channel name. Uses default if not specified.",
					},
				},
				Required: []string{"path"},
			},
		},
		{
			Name:        "index_revision",
			Description: "Index a nixpkgs revision to make its options searchable",
			InputSchema: InputSchema{
				Type: "object",
				Properties: map[string]Property{
					"revision": {
						Type:        "string",
						Description: "Git hash (full or short) or channel name (e.g., 'nixos-unstable', 'nixos-24.05')",
					},
				},
				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,
	}
}