Coverage for server.py: 72%
573 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-03-25 10:15 -0700
« prev ^ index » next coverage.py v7.7.1, created at 2025-03-25 10:15 -0700
1#!/usr/bin/env python
2"""
3NixMCP Server - A MCP server for NixOS resources.
5This implements a comprehensive FastMCP server that provides MCP resources and tools
6for querying NixOS packages and options using the Model Context Protocol (MCP).
7The server communicates via standard input/output streams using a JSON-based
8message format, allowing seamless integration with MCP-compatible AI models.
9"""
11import os
12import logging
13import logging.handlers
14import json
15import time
16from typing import Dict, Any
17import requests
18from dotenv import load_dotenv
19from mcp.server.fastmcp import FastMCP
20from contextlib import asynccontextmanager
22# Load environment variables from .env file
23load_dotenv()
26# Configure logging
27def setup_logging():
28 """Configure logging for the NixMCP server."""
29 log_file = os.environ.get("LOG_FILE", "nixmcp-server.log")
30 log_level = os.environ.get("LOG_LEVEL", "INFO")
32 # Create logger
33 logger = logging.getLogger("nixmcp")
35 # Only configure handlers if they haven't been added yet
36 # This prevents duplicate logging when code is reloaded
37 if not logger.handlers:
38 logger.setLevel(getattr(logging, log_level))
40 # Create file handler with rotation
41 file_handler = logging.handlers.RotatingFileHandler(
42 log_file, maxBytes=10 * 1024 * 1024, backupCount=5
43 )
44 file_handler.setLevel(getattr(logging, log_level))
46 # Create console handler
47 console_handler = logging.StreamHandler()
48 console_handler.setLevel(getattr(logging, log_level))
50 # Create formatter
51 formatter = logging.Formatter(
52 "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
53 )
54 file_handler.setFormatter(formatter)
55 console_handler.setFormatter(formatter)
57 # Add handlers to logger
58 logger.addHandler(file_handler)
59 logger.addHandler(console_handler)
61 logger.info("Logging initialized")
63 return logger
66# Initialize logging
67logger = setup_logging()
70# Simple in-memory cache implementation
71class SimpleCache:
72 """A simple in-memory cache with TTL expiration."""
74 def __init__(self, max_size=1000, ttl=300): # ttl in seconds
75 """Initialize the cache with maximum size and TTL."""
76 self.cache = {}
77 self.max_size = max_size
78 self.ttl = ttl
79 self.hits = 0
80 self.misses = 0
81 logger.info(f"Initialized cache with max_size={max_size}, ttl={ttl}s")
83 def get(self, key):
84 """Retrieve a value from the cache if it exists and is not expired."""
85 if key not in self.cache:
86 self.misses += 1
87 return None
89 timestamp, value = self.cache[key]
90 if time.time() - timestamp > self.ttl:
91 # Expired
92 del self.cache[key]
93 self.misses += 1
94 return None
96 self.hits += 1
97 return value
99 def set(self, key, value):
100 """Store a value in the cache with the current timestamp."""
101 if len(self.cache) >= self.max_size and key not in self.cache:
102 # Simple eviction: remove oldest entry
103 oldest_key = min(self.cache.keys(), key=lambda k: self.cache[k][0])
104 del self.cache[oldest_key]
106 self.cache[key] = (time.time(), value)
108 def clear(self):
109 """Clear all cache entries."""
110 self.cache = {}
112 def get_stats(self):
113 """Get cache statistics."""
114 return {
115 "size": len(self.cache),
116 "max_size": self.max_size,
117 "ttl": self.ttl,
118 "hits": self.hits,
119 "misses": self.misses,
120 "hit_ratio": (
121 self.hits / (self.hits + self.misses)
122 if (self.hits + self.misses) > 0
123 else 0
124 ),
125 }
128# Elasticsearch client for accessing NixOS resources
129class ElasticsearchClient:
130 """Enhanced client for accessing NixOS Elasticsearch API."""
132 def __init__(self):
133 """Initialize the Elasticsearch client with caching."""
134 # Elasticsearch endpoints
135 self.es_packages_url = (
136 "https://search.nixos.org/backend/latest-42-nixos-unstable/_search"
137 )
138 self.es_options_url = (
139 "https://search.nixos.org/backend/latest-42-nixos-unstable-options/_search"
140 )
142 # Authentication
143 self.es_user = "aWVSALXpZv"
144 self.es_password = "X8gPHnzL52wFEekuxsfQ9cSh"
145 self.es_auth = (self.es_user, self.es_password)
147 # AWS Elasticsearch endpoint (for reference)
148 self.es_aws_endpoint = (
149 "https://nixos-search-5886075189.us-east-1.bonsaisearch.net:443"
150 )
152 # Initialize cache
153 self.cache = SimpleCache(max_size=500, ttl=600) # 10 minutes TTL
155 # Request timeout settings
156 self.connect_timeout = 3.0 # seconds
157 self.read_timeout = 10.0 # seconds
159 # Retry settings
160 self.max_retries = 3
161 self.retry_delay = 1.0 # seconds
163 logger.info("Elasticsearch client initialized with caching")
165 def safe_elasticsearch_query(
166 self, endpoint: str, query_data: Dict[str, Any]
167 ) -> Dict[str, Any]:
168 """Execute an Elasticsearch query with robust error handling and retries."""
169 cache_key = f"{endpoint}:{json.dumps(query_data)}"
170 cached_result = self.cache.get(cache_key)
172 if cached_result:
173 logger.debug(f"Cache hit for query: {cache_key[:100]}...")
174 return cached_result
176 logger.debug(f"Cache miss for query: {cache_key[:100]}...")
178 for attempt in range(self.max_retries):
179 try:
180 response = requests.post(
181 endpoint,
182 json=query_data,
183 auth=self.es_auth,
184 headers={"Content-Type": "application/json"},
185 timeout=(self.connect_timeout, self.read_timeout),
186 )
188 # Handle different status codes
189 if response.status_code == 400:
190 logger.warning(f"Bad query: {query_data}")
191 return {"error": "Invalid query syntax", "details": response.json()}
192 elif response.status_code == 401 or response.status_code == 403:
193 logger.error("Authentication failure")
194 return {"error": "Authentication failed"}
195 elif response.status_code >= 500:
196 logger.error(f"Elasticsearch server error: {response.status_code}")
197 if attempt < self.max_retries - 1:
198 wait_time = self.retry_delay * (
199 2**attempt
200 ) # Exponential backoff
201 logger.info(f"Retrying in {wait_time} seconds...")
202 time.sleep(wait_time)
203 continue
204 return {"error": "Elasticsearch server error"}
206 response.raise_for_status()
207 result = response.json()
209 # Cache successful result
210 self.cache.set(cache_key, result)
211 return result
213 except requests.exceptions.ConnectionError:
214 logger.error("Connection error")
215 if attempt < self.max_retries - 1:
216 time.sleep(self.retry_delay)
217 continue
218 return {"error": "Failed to connect to Elasticsearch"}
219 except requests.exceptions.Timeout:
220 logger.error("Request timeout")
221 return {"error": "Request timed out"}
222 except Exception as e:
223 logger.error(f"Error executing query: {str(e)}")
224 return {"error": f"Query error: {str(e)}"}
226 def search_packages(
227 self, query: str, limit: int = 50, offset: int = 0
228 ) -> Dict[str, Any]:
229 """
230 Search for NixOS packages with enhanced query handling and field boosting.
232 Args:
233 query: Search term
234 limit: Maximum number of results to return
235 offset: Offset for pagination
237 Returns:
238 Dict containing search results and metadata
239 """
240 # Check if query contains wildcards
241 if "*" in query:
242 # Use wildcard query for explicit wildcard searches
243 logger.info(f"Using wildcard query for package search: {query}")
245 # Handle special case for queries like *term*
246 if query.startswith("*") and query.endswith("*") and query.count("*") == 2:
247 term = query.strip("*")
248 logger.info(f"Optimizing *term* query to search for: {term}")
250 request_data = {
251 "from": offset,
252 "size": limit,
253 "query": {
254 "bool": {
255 "should": [
256 # Contains match with high boost
257 {
258 "wildcard": {
259 "package_attr_name": {
260 "value": f"*{term}*",
261 "boost": 9,
262 }
263 }
264 },
265 {
266 "wildcard": {
267 "package_pname": {
268 "value": f"*{term}*",
269 "boost": 7,
270 }
271 }
272 },
273 {
274 "match": {
275 "package_description": {
276 "query": term,
277 "boost": 3,
278 }
279 }
280 },
281 {
282 "match": {
283 "package_programs": {"query": term, "boost": 6}
284 }
285 },
286 ],
287 "minimum_should_match": 1,
288 }
289 },
290 }
291 else:
292 # Standard wildcard query
293 request_data = {
294 "from": offset,
295 "size": limit,
296 "query": {
297 "query_string": {
298 "query": query,
299 "fields": [
300 "package_attr_name^9",
301 "package_pname^7",
302 "package_description^3",
303 "package_programs^6",
304 ],
305 "analyze_wildcard": True,
306 }
307 },
308 }
309 else:
310 # For non-wildcard searches, use a more refined approach with field boosting
311 request_data = {
312 "from": offset,
313 "size": limit,
314 "query": {
315 "bool": {
316 "should": [
317 # Exact match with highest boost
318 {
319 "term": {
320 "package_attr_name": {"value": query, "boost": 10}
321 }
322 },
323 {"term": {"package_pname": {"value": query, "boost": 8}}},
324 # Prefix match (starts with)
325 {
326 "prefix": {
327 "package_attr_name": {"value": query, "boost": 7}
328 }
329 },
330 {"prefix": {"package_pname": {"value": query, "boost": 6}}},
331 # Contains match
332 {
333 "wildcard": {
334 "package_attr_name": {
335 "value": f"*{query}*",
336 "boost": 5,
337 }
338 }
339 },
340 {
341 "wildcard": {
342 "package_pname": {"value": f"*{query}*", "boost": 4}
343 }
344 },
345 # Full-text search in description fields
346 {
347 "match": {
348 "package_description": {"query": query, "boost": 3}
349 }
350 },
351 {
352 "match": {
353 "package_longDescription": {
354 "query": query,
355 "boost": 1,
356 }
357 }
358 },
359 # Program search
360 {
361 "match": {
362 "package_programs": {"query": query, "boost": 6}
363 }
364 },
365 ],
366 "minimum_should_match": 1,
367 }
368 },
369 }
371 # Execute the query
372 data = self.safe_elasticsearch_query(self.es_packages_url, request_data)
374 # Check for errors
375 if "error" in data:
376 return data
378 # Process the response
379 hits = data.get("hits", {}).get("hits", [])
380 total = data.get("hits", {}).get("total", {}).get("value", 0)
382 packages = []
383 for hit in hits:
384 source = hit.get("_source", {})
385 packages.append(
386 {
387 "name": source.get("package_attr_name", ""),
388 "pname": source.get("package_pname", ""),
389 "version": source.get("package_version", ""),
390 "description": source.get("package_description", ""),
391 "channel": source.get("package_channel", ""),
392 "score": hit.get("_score", 0),
393 "programs": source.get("package_programs", []),
394 }
395 )
397 return {
398 "count": total,
399 "packages": packages,
400 }
402 def search_options(
403 self, query: str, limit: int = 50, offset: int = 0
404 ) -> Dict[str, Any]:
405 """
406 Search for NixOS options with enhanced query handling.
408 Args:
409 query: Search term
410 limit: Maximum number of results to return
411 offset: Offset for pagination
413 Returns:
414 Dict containing search results and metadata
415 """
416 # Check if query contains wildcards
417 if "*" in query:
418 # Use wildcard query for explicit wildcard searches
419 logger.info(f"Using wildcard query for option search: {query}")
421 # Handle special case for queries like *term*
422 if query.startswith("*") and query.endswith("*") and query.count("*") == 2:
423 term = query.strip("*")
424 logger.info(f"Optimizing *term* query to search for: {term}")
426 request_data = {
427 "from": offset,
428 "size": limit,
429 "query": {
430 "bool": {
431 "should": [
432 # Contains match with high boost
433 {
434 "wildcard": {
435 "option_name": {
436 "value": f"*{term}*",
437 "boost": 9,
438 }
439 }
440 },
441 {
442 "match": {
443 "option_description": {
444 "query": term,
445 "boost": 3,
446 }
447 }
448 },
449 ],
450 "minimum_should_match": 1,
451 }
452 },
453 }
454 else:
455 # Standard wildcard query
456 request_data = {
457 "from": offset,
458 "size": limit,
459 "query": {
460 "query_string": {
461 "query": query,
462 "fields": ["option_name^9", "option_description^3"],
463 "analyze_wildcard": True,
464 }
465 },
466 }
467 else:
468 # For non-wildcard searches, use a more refined approach with field boosting
469 request_data = {
470 "from": offset,
471 "size": limit,
472 "query": {
473 "bool": {
474 "should": [
475 # Exact match with high boost
476 {"term": {"option_name": {"value": query, "boost": 10}}},
477 # Prefix match for option names
478 {"prefix": {"option_name": {"value": query, "boost": 6}}},
479 # Contains match for option names
480 {
481 "wildcard": {
482 "option_name": {"value": f"*{query}*", "boost": 4}
483 }
484 },
485 # Full-text search in description
486 {
487 "match": {
488 "option_description": {"query": query, "boost": 2}
489 }
490 },
491 ],
492 "minimum_should_match": 1,
493 }
494 },
495 }
497 # Execute the query
498 data = self.safe_elasticsearch_query(self.es_options_url, request_data)
500 # Check for errors
501 if "error" in data:
502 return data
504 # Process the response
505 hits = data.get("hits", {}).get("hits", [])
506 total = data.get("hits", {}).get("total", {}).get("value", 0)
508 options = []
509 for hit in hits:
510 source = hit.get("_source", {})
511 options.append(
512 {
513 "name": source.get("option_name", ""),
514 "description": source.get("option_description", ""),
515 "type": source.get("option_type", ""),
516 "default": source.get("option_default", ""),
517 "score": hit.get("_score", 0),
518 }
519 )
521 return {
522 "count": total,
523 "options": options,
524 }
526 def search_programs(
527 self, program: str, limit: int = 50, offset: int = 0
528 ) -> Dict[str, Any]:
529 """
530 Search for packages that provide specific programs.
532 Args:
533 program: Program name to search for
534 limit: Maximum number of results to return
535 offset: Offset for pagination
537 Returns:
538 Dict containing search results and metadata
539 """
540 logger.info(f"Searching for packages providing program: {program}")
542 # Check if program contains wildcards
543 if "*" in program:
544 request_data = {
545 "from": offset,
546 "size": limit,
547 "query": {"wildcard": {"package_programs": {"value": program}}},
548 }
549 else:
550 request_data = {
551 "from": offset,
552 "size": limit,
553 "query": {
554 "bool": {
555 "should": [
556 {
557 "term": {
558 "package_programs": {"value": program, "boost": 10}
559 }
560 },
561 {
562 "prefix": {
563 "package_programs": {"value": program, "boost": 5}
564 }
565 },
566 {
567 "wildcard": {
568 "package_programs": {
569 "value": f"*{program}*",
570 "boost": 3,
571 }
572 }
573 },
574 ],
575 "minimum_should_match": 1,
576 }
577 },
578 }
580 # Execute the query
581 data = self.safe_elasticsearch_query(self.es_packages_url, request_data)
583 # Check for errors
584 if "error" in data:
585 return data
587 # Process the response
588 hits = data.get("hits", {}).get("hits", [])
589 total = data.get("hits", {}).get("total", {}).get("value", 0)
591 packages = []
592 for hit in hits:
593 source = hit.get("_source", {})
594 programs = source.get("package_programs", [])
596 # Filter to only include matching programs in the result
597 matching_programs = []
598 if isinstance(programs, list):
599 if "*" in program:
600 # For wildcard searches, use simple string matching
601 wild_pattern = program.replace("*", "")
602 matching_programs = [p for p in programs if wild_pattern in p]
603 else:
604 # For exact searches, look for exact/partial matches
605 matching_programs = [
606 p for p in programs if program == p or program in p
607 ]
609 packages.append(
610 {
611 "name": source.get("package_attr_name", ""),
612 "version": source.get("package_version", ""),
613 "description": source.get("package_description", ""),
614 "programs": matching_programs,
615 "all_programs": programs,
616 "score": hit.get("_score", 0),
617 }
618 )
620 return {
621 "count": total,
622 "packages": packages,
623 }
625 def search_packages_with_version(
626 self, query: str, version_pattern: str, limit: int = 50, offset: int = 0
627 ) -> Dict[str, Any]:
628 """
629 Search for packages with a specific version pattern.
631 Args:
632 query: Package search term
633 version_pattern: Version pattern to filter by (e.g., "1.*")
634 limit: Maximum number of results to return
635 offset: Offset for pagination
637 Returns:
638 Dict containing search results and metadata
639 """
640 logger.info(
641 f"Searching for packages matching '{query}' with version '{version_pattern}'"
642 )
644 request_data = {
645 "from": offset,
646 "size": limit,
647 "query": {
648 "bool": {
649 "must": [
650 # Basic package search
651 {
652 "bool": {
653 "should": [
654 {
655 "term": {
656 "package_attr_name": {
657 "value": query,
658 "boost": 10,
659 }
660 }
661 },
662 {
663 "wildcard": {
664 "package_attr_name": {
665 "value": f"*{query}*",
666 "boost": 5,
667 }
668 }
669 },
670 {
671 "match": {
672 "package_description": {
673 "query": query,
674 "boost": 2,
675 }
676 }
677 },
678 ],
679 "minimum_should_match": 1,
680 }
681 },
682 # Version filter
683 {"wildcard": {"package_version": version_pattern}},
684 ]
685 }
686 },
687 }
689 # Execute the query
690 data = self.safe_elasticsearch_query(self.es_packages_url, request_data)
692 # Check for errors
693 if "error" in data:
694 return data
696 # Process the response
697 hits = data.get("hits", {}).get("hits", [])
698 total = data.get("hits", {}).get("total", {}).get("value", 0)
700 packages = []
701 for hit in hits:
702 source = hit.get("_source", {})
703 packages.append(
704 {
705 "name": source.get("package_attr_name", ""),
706 "version": source.get("package_version", ""),
707 "description": source.get("package_description", ""),
708 "channel": source.get("package_channel", ""),
709 "score": hit.get("_score", 0),
710 }
711 )
713 return {
714 "count": total,
715 "packages": packages,
716 }
718 def advanced_query(
719 self, index_type: str, query_string: str, limit: int = 50, offset: int = 0
720 ) -> Dict[str, Any]:
721 """
722 Execute an advanced query using Elasticsearch's query string syntax.
724 Args:
725 index_type: Either "packages" or "options"
726 query_string: Elasticsearch query string syntax
727 limit: Maximum number of results to return
728 offset: Offset for pagination
730 Returns:
731 Dict containing search results and metadata
732 """
733 logger.info(f"Executing advanced query on {index_type}: {query_string}")
735 # Determine the endpoint
736 if index_type.lower() == "options":
737 endpoint = self.es_options_url
738 else:
739 endpoint = self.es_packages_url
741 request_data = {
742 "from": offset,
743 "size": limit,
744 "query": {
745 "query_string": {"query": query_string, "default_operator": "AND"}
746 },
747 }
749 # Execute the query
750 return self.safe_elasticsearch_query(endpoint, request_data)
752 def get_package_stats(self, query: str = "*") -> Dict[str, Any]:
753 """
754 Get statistics about NixOS packages.
756 Args:
757 query: Optional query to filter packages
759 Returns:
760 Dict containing aggregation statistics
761 """
762 logger.info(f"Getting package statistics for query: {query}")
764 request_data = {
765 "size": 0, # We only need aggregations, not actual hits
766 "query": {"query_string": {"query": query}},
767 "aggs": {
768 "channels": {"terms": {"field": "package_channel", "size": 10}},
769 "licenses": {"terms": {"field": "package_license", "size": 10}},
770 "platforms": {"terms": {"field": "package_platforms", "size": 10}},
771 },
772 }
774 # Execute the query
775 return self.safe_elasticsearch_query(self.es_packages_url, request_data)
777 def get_package(self, package_name: str) -> Dict[str, Any]:
778 """
779 Get detailed information about a specific package.
781 Args:
782 package_name: Name of the package
784 Returns:
785 Dict containing package details
786 """
787 logger.info(f"Getting detailed information for package: {package_name}")
789 # Build a query to find the exact package by name
790 request_data = {
791 "size": 1, # We only need one result
792 "query": {
793 "bool": {"must": [{"term": {"package_attr_name": package_name}}]}
794 },
795 }
797 # Execute the query
798 data = self.safe_elasticsearch_query(self.es_packages_url, request_data)
800 # Check for errors
801 if "error" in data:
802 return {"name": package_name, "error": data["error"], "found": False}
804 # Process the response
805 hits = data.get("hits", {}).get("hits", [])
807 if not hits:
808 logger.warning(f"Package {package_name} not found")
809 return {"name": package_name, "error": "Package not found", "found": False}
811 # Extract package details from the first hit
812 source = hits[0].get("_source", {})
814 # Return comprehensive package information
815 return {
816 "name": source.get("package_attr_name", package_name),
817 "pname": source.get("package_pname", ""),
818 "version": source.get("package_version", ""),
819 "description": source.get("package_description", ""),
820 "longDescription": source.get("package_longDescription", ""),
821 "license": source.get("package_license", ""),
822 "homepage": source.get("package_homepage", ""),
823 "maintainers": source.get("package_maintainers", []),
824 "platforms": source.get("package_platforms", []),
825 "channel": source.get("package_channel", "nixos-unstable"),
826 "position": source.get("package_position", ""),
827 "outputs": source.get("package_outputs", []),
828 "programs": source.get("package_programs", []),
829 "found": True,
830 }
832 def get_option(self, option_name: str) -> Dict[str, Any]:
833 """
834 Get detailed information about a specific NixOS option.
836 Args:
837 option_name: Name of the option
839 Returns:
840 Dict containing option details
841 """
842 logger.info(f"Getting detailed information for option: {option_name}")
844 # Build a query to find the exact option by name
845 request_data = {
846 "size": 1, # We only need one result
847 "query": {"bool": {"must": [{"term": {"option_name": option_name}}]}},
848 }
850 # Execute the query
851 data = self.safe_elasticsearch_query(self.es_options_url, request_data)
853 # Check for errors
854 if "error" in data:
855 return {"name": option_name, "error": data["error"], "found": False}
857 # Process the response
858 hits = data.get("hits", {}).get("hits", [])
860 if not hits:
861 logger.warning(f"Option {option_name} not found")
862 return {"name": option_name, "error": "Option not found", "found": False}
864 # Extract option details from the first hit
865 source = hits[0].get("_source", {})
867 # Return comprehensive option information
868 return {
869 "name": source.get("option_name", option_name),
870 "description": source.get("option_description", ""),
871 "type": source.get("option_type", ""),
872 "default": source.get("option_default", ""),
873 "example": source.get("option_example", ""),
874 "declarations": source.get("option_declarations", []),
875 "readOnly": source.get("option_readOnly", False),
876 "found": True,
877 }
880# Model Context with app-specific data
881class NixOSContext:
882 """Provides NixOS resources to AI models."""
884 def __init__(self):
885 """Initialize the ModelContext."""
886 self.es_client = ElasticsearchClient()
887 logger.info("NixOSContext initialized")
889 def get_status(self) -> Dict[str, Any]:
890 """Get the status of the NixMCP server."""
891 return {
892 "status": "ok",
893 "version": "1.0.0",
894 "name": "NixMCP",
895 "description": "NixOS HTTP-based Model Context Protocol Server",
896 "server_type": "http",
897 "cache_stats": self.es_client.cache.get_stats(),
898 }
900 def get_package(self, package_name: str) -> Dict[str, Any]:
901 """Get information about a NixOS package."""
902 return self.es_client.get_package(package_name)
904 def search_packages(self, query: str, limit: int = 10) -> Dict[str, Any]:
905 """Search for NixOS packages."""
906 return self.es_client.search_packages(query, limit)
908 def search_options(self, query: str, limit: int = 10) -> Dict[str, Any]:
909 """Search for NixOS options."""
910 return self.es_client.search_options(query, limit)
912 def get_option(self, option_name: str) -> Dict[str, Any]:
913 """Get information about a NixOS option."""
914 return self.es_client.get_option(option_name)
916 def search_programs(self, program: str, limit: int = 10) -> Dict[str, Any]:
917 """Search for packages that provide specific programs."""
918 return self.es_client.search_programs(program, limit)
920 def search_packages_with_version(
921 self, query: str, version_pattern: str, limit: int = 10
922 ) -> Dict[str, Any]:
923 """Search for packages with a specific version pattern."""
924 return self.es_client.search_packages_with_version(
925 query, version_pattern, limit
926 )
928 def advanced_query(
929 self, index_type: str, query_string: str, limit: int = 10
930 ) -> Dict[str, Any]:
931 """Execute an advanced query using Elasticsearch's query string syntax."""
932 return self.es_client.advanced_query(index_type, query_string, limit)
934 def get_package_stats(self, query: str = "*") -> Dict[str, Any]:
935 """Get statistics about NixOS packages."""
936 return self.es_client.get_package_stats(query)
939# Define the lifespan context manager for app initialization
940@asynccontextmanager
941async def app_lifespan(mcp_server: FastMCP):
942 logger.info("Initializing NixMCP server")
943 # Set up resources
944 context = NixOSContext()
946 try:
947 # We yield our context that will be accessible in all handlers
948 yield {"context": context}
949 except Exception as e:
950 logger.error(f"Error in server lifespan: {e}")
951 raise
952 finally:
953 # Cleanup on shutdown
954 logger.info("Shutting down NixMCP server")
955 # Close any open connections or resources
956 try:
957 # Add any cleanup code here if needed
958 pass
959 except Exception as e:
960 logger.error(f"Error during server shutdown cleanup: {e}")
963# Helper functions
964def create_wildcard_query(query: str) -> str:
965 """Create a wildcard query from a regular query string.
967 Args:
968 query: The original query string
970 Returns:
971 A query string with wildcards added
972 """
973 if " " in query:
974 # For multi-word queries, add wildcards around each word
975 words = query.split()
976 wildcard_terms = [f"*{word}*" for word in words]
977 return " ".join(wildcard_terms)
978 else:
979 # For single word queries, just wrap with wildcards
980 return f"*{query}*"
983# Initialize the model context before creating server
984model_context = NixOSContext()
986# Create the MCP server with the lifespan handler
987logger.info("Creating FastMCP server instance")
988mcp = FastMCP(
989 "NixMCP",
990 version="1.0.0",
991 description="NixOS HTTP-based Model Context Protocol Server",
992 lifespan=app_lifespan,
993)
995# No need for a get_tool method as we're importing tools directly
998# Define MCP resources for packages
999@mcp.resource("nixos://status")
1000def status_resource():
1001 """Get the status of the NixMCP server."""
1002 logger.info("Handling status resource request")
1003 return model_context.get_status()
1006@mcp.resource("nixos://package/{package_name}")
1007def package_resource(package_name: str):
1008 """Get information about a NixOS package."""
1009 logger.info(f"Handling package resource request for {package_name}")
1010 return model_context.get_package(package_name)
1013@mcp.resource("nixos://search/packages/{query}")
1014def search_packages_resource(query: str):
1015 """Search for NixOS packages."""
1016 logger.info(f"Handling package search request for {query}")
1017 return model_context.search_packages(query)
1020@mcp.resource("nixos://search/options/{query}")
1021def search_options_resource(query: str):
1022 """Search for NixOS options."""
1023 logger.info(f"Handling option search request for {query}")
1024 return model_context.search_options(query)
1027@mcp.resource("nixos://option/{option_name}")
1028def option_resource(option_name: str):
1029 """Get information about a NixOS option."""
1030 logger.info(f"Handling option resource request for {option_name}")
1031 return model_context.get_option(option_name)
1034@mcp.resource("nixos://search/programs/{program}")
1035def search_programs_resource(program: str):
1036 """Search for packages that provide specific programs."""
1037 logger.info(f"Handling program search request for {program}")
1038 return model_context.search_programs(program)
1041@mcp.resource("nixos://packages/stats")
1042def package_stats_resource():
1043 """Get statistics about NixOS packages."""
1044 logger.info("Handling package statistics resource request")
1045 return model_context.get_package_stats()
1048# Add MCP tools for searching and retrieving information
1049@mcp.tool()
1050def search_nixos(query: str, search_type: str = "packages", limit: int = 10) -> str:
1051 """
1052 Search for NixOS packages or options.
1054 Args:
1055 query: The search term
1056 search_type: Type of search - either "packages", "options", or "programs"
1057 limit: Maximum number of results to return (default: 10)
1059 Returns:
1060 Results formatted as text
1061 """
1062 logger.info(f"Searching for {search_type} with query '{query}'")
1064 valid_types = ["packages", "options", "programs"]
1065 if search_type.lower() not in valid_types:
1066 return f"Error: Invalid search_type. Must be one of: {', '.join(valid_types)}"
1068 try:
1069 # First try the original query as-is
1070 if search_type.lower() == "packages":
1071 logger.info(f"Trying original query first: {query}")
1072 results = model_context.search_packages(query, limit)
1073 packages = results.get("packages", [])
1075 # If no results with original query and it doesn't already have wildcards,
1076 # try with wildcards using the helper function
1077 if not packages and "*" not in query:
1078 wildcard_query = create_wildcard_query(query)
1079 logger.info(f"No results with original query, trying wildcard search: {wildcard_query}")
1081 try:
1082 results = model_context.search_packages(wildcard_query, limit)
1083 packages = results.get("packages", [])
1085 # If we got results with wildcards, note this in the output
1086 if packages:
1087 logger.info(f"Found {len(packages)} results using wildcard search")
1088 except Exception as e:
1089 logger.error(f"Error in wildcard search: {e}", exc_info=True)
1091 if not packages:
1092 return f"No packages found for query: '{query}'\n\nTry using wildcards like *{query}* for broader results."
1094 # Create a flag to track if wildcards were automatically used
1095 used_wildcards = False
1096 if packages and "*" not in query and "wildcard_query" in locals():
1097 used_wildcards = True
1099 # Indicate if wildcards were used to find results
1100 if "*" in query:
1101 output = (
1102 f"Found {len(packages)} packages for wildcard query '{query}':\n\n"
1103 )
1104 elif used_wildcards:
1105 output = f"Found {len(packages)} packages using automatic wildcard search for '{query}':\n\nNote: No exact matches were found, so wildcards were automatically added.\n\n"
1106 else:
1107 output = f"Found {len(packages)} packages for '{query}':\n\n"
1108 for pkg in packages:
1109 output += f"- {pkg.get('name', 'Unknown')}"
1110 if pkg.get("version"):
1111 output += f" ({pkg.get('version')})"
1112 output += "\n"
1113 if pkg.get("description"):
1114 output += f" {pkg.get('description')}\n"
1115 if pkg.get("channel"):
1116 output += f" Channel: {pkg.get('channel')}\n"
1117 output += "\n"
1119 return output
1121 elif search_type.lower() == "options":
1122 # First try the original query as-is
1123 logger.info(f"Trying original query first: {query}")
1124 results = model_context.search_options(query, limit)
1125 options = results.get("options", [])
1127 # If no results with original query and it doesn't already have wildcards,
1128 # try with wildcards using the helper function
1129 if not options and "*" not in query:
1130 wildcard_query = create_wildcard_query(query)
1131 logger.info(f"No results with original query, trying wildcard search: {wildcard_query}")
1133 try:
1134 results = model_context.search_options(wildcard_query, limit)
1135 options = results.get("options", [])
1137 # If we got results with wildcards, note this in the output
1138 if options:
1139 logger.info(f"Found {len(options)} results using wildcard search")
1140 except Exception as e:
1141 logger.error(f"Error in wildcard search: {e}", exc_info=True)
1143 if not options:
1144 return f"No options found for query: '{query}'\n\nTry using wildcards like *{query}* for broader results."
1146 # Create a flag to track if wildcards were automatically used
1147 used_wildcards = False
1148 if options and "*" not in query and "wildcard_query" in locals():
1149 used_wildcards = True
1151 # Indicate if wildcards were used to find results
1152 if "*" in query:
1153 output = (
1154 f"Found {len(options)} options for wildcard query '{query}':\n\n"
1155 )
1156 elif used_wildcards:
1157 output = f"Found {len(options)} options using automatic wildcard search for '{query}':\n\nNote: No exact matches were found, so wildcards were automatically added.\n\n"
1158 else:
1159 output = f"Found {len(options)} options for '{query}':\n\n"
1160 for opt in options:
1161 output += f"- {opt.get('name', 'Unknown')}\n"
1162 if opt.get("description"):
1163 output += f" {opt.get('description')}\n"
1164 if opt.get("type"):
1165 output += f" Type: {opt.get('type')}\n"
1166 if "default" in opt:
1167 output += f" Default: {opt.get('default')}\n"
1168 output += "\n"
1170 return output
1172 else: # programs
1173 results = model_context.search_programs(query, limit)
1174 packages = results.get("packages", [])
1176 if not packages:
1177 return f"No packages found providing programs matching: '{query}'\n\nTry using wildcards like *{query}* for broader results."
1179 output = f"Found {len(packages)} packages providing programs matching '{query}':\n\n"
1181 for pkg in packages:
1182 output += f"- {pkg.get('name', 'Unknown')}"
1183 if pkg.get("version"):
1184 output += f" ({pkg.get('version')})"
1185 output += "\n"
1187 # List matching programs
1188 matching_programs = pkg.get("programs", [])
1189 if matching_programs:
1190 output += f" Programs: {', '.join(matching_programs)}\n"
1192 if pkg.get("description"):
1193 output += f" {pkg.get('description')}\n"
1194 output += "\n"
1196 return output
1198 except Exception as e:
1199 logger.error(f"Error in search_nixos: {e}", exc_info=True)
1200 error_message = f"Error performing search for '{query}': {str(e)}"
1202 # Add helpful suggestions based on the error
1203 if "ConnectionError" in str(e) or "ConnectionTimeout" in str(e):
1204 error_message += "\n\nThere seems to be a connection issue with the Elasticsearch server. Please try again later."
1205 elif "AuthenticationException" in str(e):
1206 error_message += "\n\nAuthentication failed. Please check your Elasticsearch credentials."
1207 else:
1208 error_message += "\n\nTry simplifying your query or using wildcards like *term* for broader results."
1210 return error_message
1213@mcp.tool()
1214def get_nixos_package(package_name: str) -> str:
1215 """
1216 Get detailed information about a NixOS package.
1218 Args:
1219 package_name: The name of the package
1221 Returns:
1222 Detailed package information formatted as text
1223 """
1224 logger.info(f"Getting detailed information for package: {package_name}")
1226 try:
1227 package_info = model_context.get_package(package_name)
1229 if not package_info.get("found", False):
1230 return f"Package '{package_name}' not found."
1232 # Format the package information
1233 output = f"# {package_info.get('name', package_name)}\n\n"
1235 if package_info.get("version"):
1236 output += f"**Version:** {package_info.get('version')}\n"
1238 if package_info.get("description"):
1239 output += f"\n**Description:** {package_info.get('description')}\n"
1241 if package_info.get("longDescription"):
1242 output += (
1243 f"\n**Long Description:**\n{package_info.get('longDescription')}\n"
1244 )
1246 if package_info.get("license"):
1247 output += f"\n**License:** {package_info.get('license')}\n"
1249 if package_info.get("homepage"):
1250 output += f"\n**Homepage:** {package_info.get('homepage')}\n"
1252 if package_info.get("maintainers"):
1253 maintainers = package_info.get("maintainers")
1254 if isinstance(maintainers, list) and maintainers:
1255 # Convert any dictionary items to strings
1256 maintainer_strings = []
1257 for m in maintainers:
1258 if isinstance(m, dict):
1259 if "name" in m:
1260 maintainer_strings.append(m["name"])
1261 elif "email" in m:
1262 maintainer_strings.append(m["email"])
1263 else:
1264 maintainer_strings.append(str(m))
1265 else:
1266 maintainer_strings.append(str(m))
1267 output += f"\n**Maintainers:** {', '.join(maintainer_strings)}\n"
1269 if package_info.get("platforms"):
1270 platforms = package_info.get("platforms")
1271 if isinstance(platforms, list) and platforms:
1272 # Convert any dictionary or complex items to strings
1273 platform_strings = [str(p) for p in platforms]
1274 output += f"\n**Platforms:** {', '.join(platform_strings)}\n"
1276 if package_info.get("channel"):
1277 output += f"\n**Channel:** {package_info.get('channel')}\n"
1279 # Add programs if available
1280 if package_info.get("programs"):
1281 programs = package_info.get("programs")
1282 if isinstance(programs, list) and programs:
1283 output += f"\n**Provided Programs:** {', '.join(programs)}\n"
1285 return output
1287 except Exception as e:
1288 logger.error(f"Error getting package information: {e}")
1289 return f"Error getting information for package '{package_name}': {str(e)}"
1292@mcp.tool()
1293def get_nixos_option(option_name: str) -> str:
1294 """
1295 Get detailed information about a NixOS option.
1297 Args:
1298 option_name: The name of the option
1300 Returns:
1301 Detailed option information formatted as text
1302 """
1303 logger.info(f"Getting detailed information for option: {option_name}")
1305 try:
1306 option_info = model_context.get_option(option_name)
1308 if not option_info.get("found", False):
1309 return f"Option '{option_name}' not found."
1311 # Format the option information
1312 output = f"# {option_info.get('name', option_name)}\n\n"
1314 if option_info.get("description"):
1315 output += f"**Description:** {option_info.get('description')}\n\n"
1317 if option_info.get("type"):
1318 output += f"**Type:** {option_info.get('type')}\n"
1320 if option_info.get("default") is not None:
1321 output += f"**Default:** {option_info.get('default')}\n"
1323 if option_info.get("example"):
1324 output += f"\n**Example:**\n```nix\n{option_info.get('example')}\n```\n"
1326 if option_info.get("declarations"):
1327 declarations = option_info.get("declarations")
1328 if isinstance(declarations, list) and declarations:
1329 output += f"\n**Declared in:**\n"
1330 for decl in declarations:
1331 output += f"- {decl}\n"
1333 if option_info.get("readOnly"):
1334 output += f"\n**Read Only:** Yes\n"
1336 return output
1338 except Exception as e:
1339 logger.error(f"Error getting option information: {e}")
1340 return f"Error getting information for option '{option_name}': {str(e)}"
1343@mcp.tool()
1344def advanced_search(
1345 query_string: str, index_type: str = "packages", limit: int = 20
1346) -> str:
1347 """
1348 Perform an advanced search using Elasticsearch's query string syntax.
1350 Args:
1351 query_string: Elasticsearch query string (e.g. "package_programs:(python OR ruby)")
1352 index_type: Type of index to search ("packages" or "options")
1353 limit: Maximum number of results to return
1355 Returns:
1356 Search results formatted as text
1357 """
1358 logger.info(f"Performing advanced query string search: {query_string}")
1360 if index_type.lower() not in ["packages", "options"]:
1361 return f"Error: Invalid index_type. Must be 'packages' or 'options'."
1363 try:
1364 results = model_context.advanced_query(index_type, query_string, limit)
1366 # Check for errors
1367 if "error" in results:
1368 return f"Error executing query: {results['error']}"
1370 hits = results.get("hits", {}).get("hits", [])
1371 total = results.get("hits", {}).get("total", {}).get("value", 0)
1373 if not hits:
1374 return f"No results found for query: '{query_string}'"
1376 output = f"Found {total} results for query '{query_string}' (showing top {len(hits)}):\n\n"
1378 for hit in hits:
1379 source = hit.get("_source", {})
1380 score = hit.get("_score", 0)
1382 if index_type.lower() == "packages":
1383 # Format package result
1384 name = source.get("package_attr_name", "Unknown")
1385 version = source.get("package_version", "")
1386 description = source.get("package_description", "")
1388 output += f"- {name}"
1389 if version:
1390 output += f" ({version})"
1391 output += f" [score: {score:.2f}]\n"
1392 if description:
1393 output += f" {description}\n"
1394 else:
1395 # Format option result
1396 name = source.get("option_name", "Unknown")
1397 description = source.get("option_description", "")
1399 output += f"- {name} [score: {score:.2f}]\n"
1400 if description:
1401 output += f" {description}\n"
1403 output += "\n"
1405 return output
1407 except Exception as e:
1408 logger.error(f"Error in advanced_search: {e}", exc_info=True)
1409 return f"Error performing advanced search: {str(e)}"
1412@mcp.tool()
1413def package_statistics(query: str = "*") -> str:
1414 """
1415 Get statistics about NixOS packages matching the query.
1417 Args:
1418 query: Search query (default: all packages)
1420 Returns:
1421 Statistics about matching packages
1422 """
1423 logger.info(f"Getting package statistics for query: {query}")
1425 try:
1426 results = model_context.get_package_stats(query)
1428 # Check for errors
1429 if "error" in results:
1430 return f"Error getting statistics: {results['error']}"
1432 # Extract aggregations
1433 aggregations = results.get("aggregations", {})
1435 if not aggregations:
1436 return "No statistics available"
1438 output = f"# NixOS Package Statistics\n\n"
1440 # Channel distribution
1441 channels = aggregations.get("channels", {}).get("buckets", [])
1442 if channels:
1443 output += "## Distribution by Channel\n\n"
1444 for channel in channels:
1445 output += f"- {channel.get('key', 'Unknown')}: {channel.get('doc_count', 0)} packages\n"
1446 output += "\n"
1448 # License distribution
1449 licenses = aggregations.get("licenses", {}).get("buckets", [])
1450 if licenses:
1451 output += "## Distribution by License\n\n"
1452 for license in licenses:
1453 output += f"- {license.get('key', 'Unknown')}: {license.get('doc_count', 0)} packages\n"
1454 output += "\n"
1456 # Platform distribution
1457 platforms = aggregations.get("platforms", {}).get("buckets", [])
1458 if platforms:
1459 output += "## Distribution by Platform\n\n"
1460 for platform in platforms:
1461 output += f"- {platform.get('key', 'Unknown')}: {platform.get('doc_count', 0)} packages\n"
1462 output += "\n"
1464 # Add cache statistics
1465 cache_stats = model_context.es_client.cache.get_stats()
1466 output += "## Cache Statistics\n\n"
1467 output += (
1468 f"- Cache size: {cache_stats['size']}/{cache_stats['max_size']} entries\n"
1469 )
1470 output += f"- Hit ratio: {cache_stats['hit_ratio']*100:.1f}% ({cache_stats['hits']} hits, {cache_stats['misses']} misses)\n"
1472 return output
1474 except Exception as e:
1475 logger.error(f"Error getting package statistics: {e}", exc_info=True)
1476 return f"Error getting package statistics: {str(e)}"
1479@mcp.tool()
1480def version_search(package_query: str, version_pattern: str, limit: int = 10) -> str:
1481 """
1482 Search for packages matching a specific version pattern.
1484 Args:
1485 package_query: Package search term
1486 version_pattern: Version pattern to filter by (e.g., "1.*")
1487 limit: Maximum number of results to return
1489 Returns:
1490 Search results formatted as text
1491 """
1492 logger.info(
1493 f"Searching for packages matching '{package_query}' with version '{version_pattern}'"
1494 )
1496 try:
1497 results = model_context.search_packages_with_version(
1498 package_query, version_pattern, limit
1499 )
1501 # Check for errors
1502 if "error" in results:
1503 return f"Error searching packages: {results['error']}"
1505 packages = results.get("packages", [])
1506 total = results.get("count", 0)
1508 if not packages:
1509 return f"No packages found matching '{package_query}' with version pattern '{version_pattern}'"
1511 output = f"Found {total} packages matching '{package_query}' with version pattern '{version_pattern}' (showing top {len(packages)}):\n\n"
1513 for pkg in packages:
1514 output += (
1515 f"- {pkg.get('name', 'Unknown')} ({pkg.get('version', 'Unknown')})\n"
1516 )
1517 if pkg.get("description"):
1518 output += f" {pkg.get('description')}\n"
1519 if pkg.get("channel"):
1520 output += f" Channel: {pkg.get('channel')}\n"
1521 output += "\n"
1523 return output
1525 except Exception as e:
1526 logger.error(f"Error in version_search: {e}", exc_info=True)
1527 return f"Error searching packages with version pattern: {str(e)}"
1530if __name__ == "__main__":
1531 # This will start the server and keep it running
1532 try:
1533 logger.info("Starting NixMCP server...")
1534 mcp.run()
1535 except KeyboardInterrupt:
1536 logger.info("Server stopped by user")
1537 except Exception as e:
1538 logger.error(f"Error running server: {e}", exc_info=True)