labmcp/internal/mcp/server.go

package mcp

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

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

// Server is an MCP server that handles JSON-RPC requests over stdio.
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, reading from r and writing to w.
func (s *Server) Run(ctx context.Context, r io.Reader, w io.Writer) error {
	scanner := bufio.NewScanner(r)
	encoder := json.NewEncoder(w)

	for scanner.Scan() {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		line := scanner.Bytes()
		if len(line) == 0 {
			continue
		}

		var req Request
		if err := json.Unmarshal(line, &req); err != nil {
			s.logger.Printf("Failed to parse request: %v", err)
			resp := Response{
				JSONRPC: "2.0",
				Error: &Error{
					Code:    ParseError,
					Message: "Parse error",
					Data:    err.Error(),
				},
			}
			if err := encoder.Encode(resp); err != nil {
				return fmt.Errorf("failed to write response: %w", err)
			}
			continue
		}

		resp := s.handleRequest(ctx, &req)
		if resp != nil {
			if err := encoder.Encode(resp); err != nil {
				return fmt.Errorf("failed to write response: %w", err)
			}
		}
	}

	if err := scanner.Err(); err != nil {
		return fmt.Errorf("scanner error: %w", err)
	}

	return nil
}

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

	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,
	}
}