Sourcegraph MCP Server

import asyncio import logging import pathlib import signal from typing import Any, List import requests from dotenv import load_dotenv from fastmcp import FastMCP from starlette.requests import Request from starlette.responses import JSONResponse, Response from .backends import SourcegraphClient, SourcegraphContentFetcher from .backends.models import FormattedResult from .config import ServerConfig from .core import PromptManager from .exceptions import ContentFetchError, SearchError, ServerShutdownError logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) load_dotenv() class SourcegraphMCPServer: def __init__(self, config: ServerConfig) -> None: self.config = config self.server = FastMCP(sse_path="/sourcegraph/sse", message_path="/sourcegraph/messages/") self._shutdown_requested = False self._setup_clients() self._load_prompts() def _setup_clients(self) -> None: self.search_client = SourcegraphClient( endpoint=self.config.sourcegraph_endpoint, token=self.config.sourcegraph_token ) self.content_fetcher = SourcegraphContentFetcher( endpoint=self.config.sourcegraph_endpoint, token=self.config.sourcegraph_token ) logger.info("Using Sourcegraph backend") def _load_prompts(self) -> None: prompt_manager = PromptManager(file_path=pathlib.Path(__file__).parent / "prompts" / "prompts.yaml") self.codesearch_guide = prompt_manager._load_prompt("guides.codesearch_guide") self.search_tool_description = prompt_manager._load_prompt("tools.search") self.search_prompt_guide_description = prompt_manager._load_prompt("tools.search_prompt_guide") self.fetch_content_description = prompt_manager._load_prompt("tools.fetch_content") try: self.org_guide = prompt_manager._load_prompt("guides.org_guide") except Exception: self.org_guide = "" def signal_handler(self, sig: int, frame: Any = None) -> None: """Handle termination signals for graceful shutdown.""" logger.info(f"Received signal {sig}, initiating graceful shutdown...") self._shutdown_requested = True async def fetch_content(self, repo: str, path: str) -> str: if self._shutdown_requested: logger.info("Shutdown in progress, declining new requests") raise ServerShutdownError("Server is shutting down") try: result = await asyncio.to_thread(self.content_fetcher.get_content, repo, path) return result except ValueError as e: logger.warning(f"Error fetching content from {repo}: {str(e)}") raise ContentFetchError("Invalid arguments: path or repository does not exist") from e except Exception as e: logger.error(f"Unexpected error fetching content: {e}") raise ContentFetchError("Error fetching content") from e async def search(self, query: str, limit: int = 30) -> List[FormattedResult]: if self._shutdown_requested: logger.info("Shutdown in progress, declining new requests") raise ServerShutdownError("Server is shutting down") num_results = min(max(1, limit), 100) logger.info(f"Search query: {query}, limit: {num_results}") try: results = await asyncio.to_thread(self.search_client.search, query, num_results) formatted_results = await asyncio.to_thread(self.search_client.format_results, results, num_results) return formatted_results except requests.exceptions.HTTPError as exc: logger.error(f"Search HTTP error: {exc}") raise SearchError(f"HTTP error during search: {exc}") from exc except Exception as exc: logger.error(f"Unexpected error during search: {exc}") raise SearchError(f"Unexpected error during search: {exc}") from exc async def search_prompt_guide(self, objective: str) -> str: if self._shutdown_requested: logger.info("Shutdown in progress, declining new prompt guide requests") raise ServerShutdownError("Server is shutting down") prompt_parts = [] if self.org_guide: prompt_parts.append(self.org_guide) prompt_parts.append("\n\n") prompt_parts.append(self.codesearch_guide) prompt_parts.append( f"\nGiven this guide create a Sourcegraph query for {objective} and call the search tool accordingly." ) return "".join(prompt_parts) async def _safe_fetch_content(self, repo: str, path: str) -> str: """Safe wrapper for fetch_content that handles exceptions.""" try: return await self.fetch_content(repo, path) except ServerShutdownError: return "" except ContentFetchError as e: return str(e) except Exception as e: logger.error(f"Unexpected error in fetch_content: {e}") return "error fetching content" async def _safe_search(self, query: str, limit: int = 30) -> List[FormattedResult]: """Safe wrapper for search that handles exceptions.""" try: return await self.search(query, limit) except ServerShutdownError: return [] except SearchError as e: logger.error(f"Search error: {e}") return [] except Exception as e: logger.error(f"Unexpected error in search: {e}") return [] async def _safe_search_prompt_guide(self, objective: str) -> str: """Safe wrapper for search_prompt_guide that handles exceptions.""" try: return await self.search_prompt_guide(objective) except ServerShutdownError: return "Server is shutting down" except Exception as e: logger.error(f"Unexpected error in search_prompt_guide: {e}") return "Error generating search guide" def _register_tools(self) -> None: """Register MCP tools with the server.""" tools = [ (self._safe_search, "search", self.search_tool_description), (self._safe_search_prompt_guide, "search_prompt_guide", self.search_prompt_guide_description), (self._safe_fetch_content, "fetch_content", self.fetch_content_description), ] for tool_func, tool_name, description in tools: self.server.tool(tool_func, name=tool_name, description=description) logger.info(f"Registered tool: {tool_name}") def _register_health_endpoints(self) -> None: """Register health check endpoints.""" @self.server.custom_route("/health", methods=["GET"]) async def health_check(request: Request) -> Response: """Simple health check endpoint for liveness probe.""" return JSONResponse({"status": "ok", "service": "sourcegraph-mcp"}) @self.server.custom_route("/ready", methods=["GET"]) async def readiness_check(request: Request) -> Response: """Readiness check endpoint that verifies the service is ready.""" try: # Check if search client is available if not hasattr(self, "search_client") or self.search_client is None: return JSONResponse({"status": "not_ready", "reason": "search_client_unavailable"}, status_code=503) # Check if content fetcher is available if not hasattr(self, "content_fetcher") or self.content_fetcher is None: return JSONResponse( {"status": "not_ready", "reason": "content_fetcher_unavailable"}, status_code=503 ) return JSONResponse({"status": "ready", "service": "sourcegraph-mcp", "backend": "sourcegraph"}) except Exception as e: logger.error(f"Readiness check failed: {e}") return JSONResponse({"status": "error", "reason": str(e)}, status_code=503) async def _run_server(self) -> None: """Run the FastMCP server with both HTTP and SSE transports.""" tasks = [ self.server.run_http_async( transport="streamable-http", host="0.0.0.0", path="/sourcegraph/mcp", port=self.config.streamable_http_port, ), self.server.run_http_async(transport="sse", host="0.0.0.0", port=self.config.sse_port), ] await asyncio.gather(*tasks) async def run(self) -> None: """Start the search server.""" signal.signal(signal.SIGINT, lambda sig, frame: self.signal_handler(sig, frame)) signal.signal(signal.SIGTERM, lambda sig, frame: self.signal_handler(sig, frame)) self._register_tools() self._register_health_endpoints() try: logger.info("Starting Sourcegraph MCP server...") await self._run_server() except KeyboardInterrupt: logger.info("Received keyboard interrupt (CTRL+C)") except Exception as exc: logger.error(f"Server error: {exc}") raise finally: logger.info("Server has shut down.") def main() -> None: config = ServerConfig() server = SourcegraphMCPServer(config) asyncio.run(server.run()) if __name__ == "__main__": main()