Coverage for server.py: 72%

1#!/usr/bin/env python

2"""

3 NixMCP Server - A MCP server for NixOS resources.

4

5 This implements a comprehensive FastMCP server that provides MCP resources and tools

6 for querying NixOS packages and options using the Model Context Protocol (MCP).

7 The server communicates via standard input/output streams using a JSON-based

8 message format, allowing seamless integration with MCP-compatible AI models.

9"""

10

11 import os

12 import logging

13 import logging.handlers

14 import json

15 import time

16 from typing import Dict, Any

17 import requests

18 from dotenv import load_dotenv

19 from mcp.server.fastmcp import FastMCP

20 from contextlib import asynccontextmanager

21

22# Load environment variables from .env file

23 load_dotenv()

24

25

26# Configure logging

27 def 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")

31

32 # Create logger

33 logger = logging.getLogger("nixmcp")

34

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

39

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

45

46 # Create console handler

47 console_handler = logging.StreamHandler()

48 console_handler.setLevel(getattr(logging, log_level))

49

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)

56

57 # Add handlers to logger

58 logger.addHandler(file_handler)

59 logger.addHandler(console_handler)

60

61 logger.info("Logging initialized")

62

63 return logger

64

65

66# Initialize logging

67 logger = setup_logging()

68

69

70# Simple in-memory cache implementation

71 class SimpleCache:

72 """A simple in-memory cache with TTL expiration."""

73

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

82

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

88

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

95

96 self.hits += 1

97 return value

98

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]

105

106 self.cache[key] = (time.time(), value)

107

108 def clear(self):

109 """Clear all cache entries."""

110 self.cache = {}

111

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 }

126

127

128# Elasticsearch client for accessing NixOS resources

129 class ElasticsearchClient:

130 """Enhanced client for accessing NixOS Elasticsearch API."""

131

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 )

141

142 # Authentication

143 self.es_user = "aWVSALXpZv"

144 self.es_password = "X8gPHnzL52wFEekuxsfQ9cSh"

145 self.es_auth = (self.es_user, self.es_password)

146

147 # AWS Elasticsearch endpoint (for reference)

148 self.es_aws_endpoint = (

149 "https://nixos-search-5886075189.us-east-1.bonsaisearch.net:443"

150 )

151

152 # Initialize cache

153 self.cache = SimpleCache(max_size=500, ttl=600) # 10 minutes TTL

154

155 # Request timeout settings

156 self.connect_timeout = 3.0 # seconds

157 self.read_timeout = 10.0 # seconds

158

159 # Retry settings

160 self.max_retries = 3

161 self.retry_delay = 1.0 # seconds

162

163 logger.info("Elasticsearch client initialized with caching")

164

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)

171

172 if cached_result:

173 logger.debug(f"Cache hit for query: {cache_key[:100]}...")

174 return cached_result

175

176 logger.debug(f"Cache miss for query: {cache_key[:100]}...")

177

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 )

187

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

205

206 response.raise_for_status()

207 result = response.json()

208

209 # Cache successful result

210 self.cache.set(cache_key, result)

211 return result

212

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)}"}

225

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.

231

232 Args:

233 query: Search term

234 limit: Maximum number of results to return

235 offset: Offset for pagination

236

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}")

244

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}")

249

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 }

370

371 # Execute the query

372 data = self.safe_elasticsearch_query(self.es_packages_url, request_data)

373

374 # Check for errors

375 if "error" in data:

376 return data

377

378 # Process the response

379 hits = data.get("hits", {}).get("hits", [])

380 total = data.get("hits", {}).get("total", {}).get("value", 0)

381

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 )

396

397 return {

398 "count": total,

399 "packages": packages,

400 }

401

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.

407

408 Args:

409 query: Search term

410 limit: Maximum number of results to return

411 offset: Offset for pagination

412

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}")

420

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}")

425

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 }

496

497 # Execute the query

498 data = self.safe_elasticsearch_query(self.es_options_url, request_data)

499

500 # Check for errors

501 if "error" in data:

502 return data

503

504 # Process the response

505 hits = data.get("hits", {}).get("hits", [])

506 total = data.get("hits", {}).get("total", {}).get("value", 0)

507

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 )

520

521 return {

522 "count": total,

523 "options": options,

524 }

525

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.

531

532 Args:

533 program: Program name to search for

534 limit: Maximum number of results to return

535 offset: Offset for pagination

536

537 Returns:

538 Dict containing search results and metadata

539 """

540 logger.info(f"Searching for packages providing program: {program}")

541

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 }

579

580 # Execute the query

581 data = self.safe_elasticsearch_query(self.es_packages_url, request_data)

582

583 # Check for errors

584 if "error" in data:

585 return data

586

587 # Process the response

588 hits = data.get("hits", {}).get("hits", [])

589 total = data.get("hits", {}).get("total", {}).get("value", 0)

590

591 packages = []

592 for hit in hits:

593 source = hit.get("_source", {})

594 programs = source.get("package_programs", [])

595

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 ]

608

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 )

619

620 return {

621 "count": total,

622 "packages": packages,

623 }

624

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.

630

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

636

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 )

643

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 }

688

689 # Execute the query

690 data = self.safe_elasticsearch_query(self.es_packages_url, request_data)

691

692 # Check for errors

693 if "error" in data:

694 return data

695

696 # Process the response

697 hits = data.get("hits", {}).get("hits", [])

698 total = data.get("hits", {}).get("total", {}).get("value", 0)

699

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 )

712

713 return {

714 "count": total,

715 "packages": packages,

716 }

717

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.

723

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

729

730 Returns:

731 Dict containing search results and metadata

732 """

733 logger.info(f"Executing advanced query on {index_type}: {query_string}")

734

735 # Determine the endpoint

736 if index_type.lower() == "options":

737 endpoint = self.es_options_url

738 else:

739 endpoint = self.es_packages_url

740

741 request_data = {

742 "from": offset,

743 "size": limit,

744 "query": {

745 "query_string": {"query": query_string, "default_operator": "AND"}

746 },

747 }

748

749 # Execute the query

750 return self.safe_elasticsearch_query(endpoint, request_data)

751

752 def get_package_stats(self, query: str = "*") -> Dict[str, Any]:

753 """

754 Get statistics about NixOS packages.

755

756 Args:

757 query: Optional query to filter packages

758

759 Returns:

760 Dict containing aggregation statistics

761 """

762 logger.info(f"Getting package statistics for query: {query}")

763

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 }

773

774 # Execute the query

775 return self.safe_elasticsearch_query(self.es_packages_url, request_data)

776

777 def get_package(self, package_name: str) -> Dict[str, Any]:

778 """

779 Get detailed information about a specific package.

780

781 Args:

782 package_name: Name of the package

783

784 Returns:

785 Dict containing package details

786 """

787 logger.info(f"Getting detailed information for package: {package_name}")

788

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 }

796

797 # Execute the query

798 data = self.safe_elasticsearch_query(self.es_packages_url, request_data)

799

800 # Check for errors

801 if "error" in data:

802 return {"name": package_name, "error": data["error"], "found": False}

803

804 # Process the response

805 hits = data.get("hits", {}).get("hits", [])

806

807 if not hits:

808 logger.warning(f"Package {package_name} not found")

809 return {"name": package_name, "error": "Package not found", "found": False}

810

811 # Extract package details from the first hit

812 source = hits[0].get("_source", {})

813

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 }

831

832 def get_option(self, option_name: str) -> Dict[str, Any]:

833 """

834 Get detailed information about a specific NixOS option.

835

836 Args:

837 option_name: Name of the option

838

839 Returns:

840 Dict containing option details

841 """

842 logger.info(f"Getting detailed information for option: {option_name}")

843

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 }

849

850 # Execute the query

851 data = self.safe_elasticsearch_query(self.es_options_url, request_data)

852

853 # Check for errors

854 if "error" in data:

855 return {"name": option_name, "error": data["error"], "found": False}

856

857 # Process the response

858 hits = data.get("hits", {}).get("hits", [])

859

860 if not hits:

861 logger.warning(f"Option {option_name} not found")

862 return {"name": option_name, "error": "Option not found", "found": False}

863

864 # Extract option details from the first hit

865 source = hits[0].get("_source", {})

866

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 }

878

879

880# Model Context with app-specific data

881 class NixOSContext:

882 """Provides NixOS resources to AI models."""

883

884 def __init__(self):

885 """Initialize the ModelContext."""

886 self.es_client = ElasticsearchClient()

887 logger.info("NixOSContext initialized")

888

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 }

899

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)

903

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)

907

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)

911

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)

915

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)

919

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 )

927

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)

933

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)

937

938

939# Define the lifespan context manager for app initialization

940@asynccontextmanager

941 async def app_lifespan(mcp_server: FastMCP):

942 logger.info("Initializing NixMCP server")

943 # Set up resources

944 context = NixOSContext()

945

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}")

961

962

963# Helper functions

964 def create_wildcard_query(query: str) -> str:

965 """Create a wildcard query from a regular query string.

966

967 Args:

968 query: The original query string

969

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}*"

981

982

983# Initialize the model context before creating server

984 model_context = NixOSContext()

985

986# Create the MCP server with the lifespan handler

987 logger.info("Creating FastMCP server instance")

988 mcp = FastMCP(

989 "NixMCP",

990 version="1.0.0",

991 description="NixOS HTTP-based Model Context Protocol Server",

992 lifespan=app_lifespan,

993)

994

995# No need for a get_tool method as we're importing tools directly

996

997

998# Define MCP resources for packages

999@mcp.resource("nixos://status")

1000 def status_resource():

1001 """Get the status of the NixMCP server."""

1002 logger.info("Handling status resource request")

1003 return model_context.get_status()

1004

1005

1006@mcp.resource("nixos://package/{package_name}")

1007 def 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)

1011

1012

1013@mcp.resource("nixos://search/packages/{query}")

1014 def 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)

1018

1019

1020@mcp.resource("nixos://search/options/{query}")

1021 def 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)

1025

1026

1027@mcp.resource("nixos://option/{option_name}")

1028 def 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)

1032

1033

1034@mcp.resource("nixos://search/programs/{program}")

1035 def 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)

1039

1040

1041@mcp.resource("nixos://packages/stats")

1042 def package_stats_resource():

1043 """Get statistics about NixOS packages."""

1044 logger.info("Handling package statistics resource request")

1045 return model_context.get_package_stats()

1046

1047

1048# Add MCP tools for searching and retrieving information

1049@mcp.tool()

1050 def search_nixos(query: str, search_type: str = "packages", limit: int = 10) -> str:

1051 """

1052 Search for NixOS packages or options.

1053

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)

1058

1059 Returns:

1060 Results formatted as text

1061 """

1062 logger.info(f"Searching for {search_type} with query '{query}'")

1063

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)}"

1067

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", [])

1074

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}")

1080

1081 try:

1082 results = model_context.search_packages(wildcard_query, limit)

1083 packages = results.get("packages", [])

1084

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)

1090

1091 if not packages:

1092 return f"No packages found for query: '{query}'\n\nTry using wildcards like *{query}* for broader results."

1093

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

1098

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"

1118

1119 return output

1120

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", [])

1126

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}")

1132

1133 try:

1134 results = model_context.search_options(wildcard_query, limit)

1135 options = results.get("options", [])

1136

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)

1142

1143 if not options:

1144 return f"No options found for query: '{query}'\n\nTry using wildcards like *{query}* for broader results."

1145

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

1150

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"

1169

1170 return output

1171

1172 else: # programs

1173 results = model_context.search_programs(query, limit)

1174 packages = results.get("packages", [])

1175

1176 if not packages:

1177 return f"No packages found providing programs matching: '{query}'\n\nTry using wildcards like *{query}* for broader results."

1178

1179 output = f"Found {len(packages)} packages providing programs matching '{query}':\n\n"

1180

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"

1186

1187 # List matching programs

1188 matching_programs = pkg.get("programs", [])

1189 if matching_programs:

1190 output += f" Programs: {', '.join(matching_programs)}\n"

1191

1192 if pkg.get("description"):

1193 output += f" {pkg.get('description')}\n"

1194 output += "\n"

1195

1196 return output

1197

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)}"

1201

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."

1209

1210 return error_message

1211

1212

1213@mcp.tool()

1214 def get_nixos_package(package_name: str) -> str:

1215 """

1216 Get detailed information about a NixOS package.

1217

1218 Args:

1219 package_name: The name of the package

1220

1221 Returns:

1222 Detailed package information formatted as text

1223 """

1224 logger.info(f"Getting detailed information for package: {package_name}")

1225

1226 try:

1227 package_info = model_context.get_package(package_name)

1228

1229 if not package_info.get("found", False):

1230 return f"Package '{package_name}' not found."

1231

1232 # Format the package information

1233 output = f"# {package_info.get('name', package_name)}\n\n"

1234

1235 if package_info.get("version"):

1236 output += f"**Version:** {package_info.get('version')}\n"

1237

1238 if package_info.get("description"):

1239 output += f"\n**Description:** {package_info.get('description')}\n"

1240

1241 if package_info.get("longDescription"):

1242 output += (

1243 f"\n**Long Description:**\n{package_info.get('longDescription')}\n"

1244 )

1245

1246 if package_info.get("license"):

1247 output += f"\n**License:** {package_info.get('license')}\n"

1248

1249 if package_info.get("homepage"):

1250 output += f"\n**Homepage:** {package_info.get('homepage')}\n"

1251

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"

1268

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"

1275

1276 if package_info.get("channel"):

1277 output += f"\n**Channel:** {package_info.get('channel')}\n"

1278

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"

1284

1285 return output

1286

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)}"

1290

1291

1292@mcp.tool()

1293 def get_nixos_option(option_name: str) -> str:

1294 """

1295 Get detailed information about a NixOS option.

1296

1297 Args:

1298 option_name: The name of the option

1299

1300 Returns:

1301 Detailed option information formatted as text

1302 """

1303 logger.info(f"Getting detailed information for option: {option_name}")

1304

1305 try:

1306 option_info = model_context.get_option(option_name)

1307

1308 if not option_info.get("found", False):

1309 return f"Option '{option_name}' not found."

1310

1311 # Format the option information

1312 output = f"# {option_info.get('name', option_name)}\n\n"

1313

1314 if option_info.get("description"):

1315 output += f"**Description:** {option_info.get('description')}\n\n"

1316

1317 if option_info.get("type"):

1318 output += f"**Type:** {option_info.get('type')}\n"

1319

1320 if option_info.get("default") is not None:

1321 output += f"**Default:** {option_info.get('default')}\n"

1322

1323 if option_info.get("example"):

1324 output += f"\n**Example:**\n```nix\n{option_info.get('example')}\n```\n"

1325

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"

1332

1333 if option_info.get("readOnly"):

1334 output += f"\n**Read Only:** Yes\n"

1335

1336 return output

1337

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)}"

1341

1342

1343@mcp.tool()

1344 def 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.

1349

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

1354

1355 Returns:

1356 Search results formatted as text

1357 """

1358 logger.info(f"Performing advanced query string search: {query_string}")

1359

1360 if index_type.lower() not in ["packages", "options"]:

1361 return f"Error: Invalid index_type. Must be 'packages' or 'options'."

1362

1363 try:

1364 results = model_context.advanced_query(index_type, query_string, limit)

1365

1366 # Check for errors

1367 if "error" in results:

1368 return f"Error executing query: {results['error']}"

1369

1370 hits = results.get("hits", {}).get("hits", [])

1371 total = results.get("hits", {}).get("total", {}).get("value", 0)

1372

1373 if not hits:

1374 return f"No results found for query: '{query_string}'"

1375

1376 output = f"Found {total} results for query '{query_string}' (showing top {len(hits)}):\n\n"

1377

1378 for hit in hits:

1379 source = hit.get("_source", {})

1380 score = hit.get("_score", 0)

1381

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", "")

1387

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", "")

1398

1399 output += f"- {name} [score: {score:.2f}]\n"

1400 if description:

1401 output += f" {description}\n"

1402

1403 output += "\n"

1404

1405 return output

1406

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)}"

1410

1411

1412@mcp.tool()

1413 def package_statistics(query: str = "*") -> str:

1414 """

1415 Get statistics about NixOS packages matching the query.

1416

1417 Args:

1418 query: Search query (default: all packages)

1419

1420 Returns:

1421 Statistics about matching packages

1422 """

1423 logger.info(f"Getting package statistics for query: {query}")

1424

1425 try:

1426 results = model_context.get_package_stats(query)

1427

1428 # Check for errors

1429 if "error" in results:

1430 return f"Error getting statistics: {results['error']}"

1431

1432 # Extract aggregations

1433 aggregations = results.get("aggregations", {})

1434

1435 if not aggregations:

1436 return "No statistics available"

1437

1438 output = f"# NixOS Package Statistics\n\n"

1439

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"

1447

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"

1455

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"

1463

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"

1471

1472 return output

1473

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)}"

1477

1478

1479@mcp.tool()

1480 def version_search(package_query: str, version_pattern: str, limit: int = 10) -> str:

1481 """

1482 Search for packages matching a specific version pattern.

1483

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

1488

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 )

1495

1496 try:

1497 results = model_context.search_packages_with_version(

1498 package_query, version_pattern, limit

1499 )

1500

1501 # Check for errors

1502 if "error" in results:

1503 return f"Error searching packages: {results['error']}"

1504

1505 packages = results.get("packages", [])

1506 total = results.get("count", 0)

1507

1508 if not packages:

1509 return f"No packages found matching '{package_query}' with version pattern '{version_pattern}'"

1510

1511 output = f"Found {total} packages matching '{package_query}' with version pattern '{version_pattern}' (showing top {len(packages)}):\n\n"

1512

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"

1522

1523 return output

1524

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)}"

1528

1529

1530 if __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)