"""
FastMCP server entry point for Prometheus MCP.
This module provides the MCP server that exposes PromQL query tools
for AWS Managed Prometheus with SigV4 authentication.
"""
import asyncio
import logging
import os
import sys
from contextlib import asynccontextmanager
from typing import AsyncIterator
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent
from prometheus_mcp.promql.tools import (
init_client,
get_label_values,
list_labels,
list_metrics,
query_instant,
query_range,
)
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("prometheus-mcp")
# Create MCP server
server = Server("prometheus-mcp")
@asynccontextmanager
async def lifespan() -> AsyncIterator[None]:
"""Manage server lifecycle - initialize and cleanup AMP client."""
workspace_id = os.environ.get("PROMETHEUS_WORKSPACE_ID")
region = os.environ.get("AWS_REGION", "us-east-1")
if not workspace_id:
logger.warning(
"PROMETHEUS_WORKSPACE_ID not set. Set it before making queries."
)
try:
client = init_client(workspace_id=workspace_id, region=region)
logger.info(f"Initialized AMP client for workspace: {workspace_id} in {region}")
yield
finally:
if workspace_id:
client = init_client(workspace_id=workspace_id, region=region)
await client.close()
logger.info("Closed AMP client")
@server.list_tools()
async def list_tools() -> list[Tool]:
"""Return the list of available tools."""
return [
Tool(
name="query_instant",
description=(
"Execute an instant PromQL query against AWS Managed Prometheus. "
"Returns current values of matching time series at a single point in time. "
"Use for current metric values or point-in-time snapshots."
),
inputSchema={
"type": "object",
"properties": {
"query": {
"type": "string",
"description": (
"PromQL query expression (e.g., 'up', 'rate(http_requests_total[5m])')"
),
},
"time": {
"type": "string",
"description": (
"Optional evaluation timestamp in RFC3339 format "
"(e.g., '2024-01-15T10:30:00Z') or Unix timestamp. "
"Defaults to current server time."
),
},
},
"required": ["query"],
},
),
Tool(
name="query_range",
description=(
"Execute a range PromQL query to get time series data over a time period. "
"Returns data points at regular intervals. "
"Use for historical data analysis, graphing, and trend analysis."
),
inputSchema={
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "PromQL query expression",
},
"start": {
"type": "string",
"description": (
"Start timestamp in RFC3339 format (e.g., '2024-01-15T00:00:00Z') "
"or Unix timestamp"
),
},
"end": {
"type": "string",
"description": (
"End timestamp in RFC3339 format (e.g., '2024-01-15T12:00:00Z') "
"or Unix timestamp"
),
},
"step": {
"type": "string",
"description": (
"Query resolution step width (e.g., '15s', '1m', '5m', '1h'). "
"Defaults to '1m'. Smaller steps = more data points but slower queries."
),
"default": "1m",
},
},
"required": ["query", "start", "end"],
},
),
Tool(
name="list_labels",
description=(
"Get all label names from AWS Managed Prometheus. "
"Use for discovering what labels are available for filtering queries."
),
inputSchema={
"type": "object",
"properties": {
"match": {
"type": "array",
"items": {"type": "string"},
"description": (
"Optional list of series selectors to filter which series to consider. "
"For example: ['up', 'http_requests_total']"
),
},
},
"required": [],
},
),
Tool(
name="get_label_values",
description=(
"Get all values for a specific label from AWS Managed Prometheus. "
"Use to discover possible filter values for a specific label."
),
inputSchema={
"type": "object",
"properties": {
"label_name": {
"type": "string",
"description": (
"The label name to get values for (e.g., 'job', 'instance', 'namespace')"
),
},
"match": {
"type": "array",
"items": {"type": "string"},
"description": (
"Optional list of series selectors to filter which series to consider"
),
},
},
"required": ["label_name"],
},
),
Tool(
name="list_metrics",
description=(
"Get all metric names from AWS Managed Prometheus. "
"Optionally includes metadata like metric type, help text, and unit."
),
inputSchema={
"type": "object",
"properties": {
"with_metadata": {
"type": "boolean",
"description": (
"If True, fetches full metadata (type, help, unit) for each metric. "
"This is slower but provides more information. Defaults to False."
),
"default": False,
},
},
"required": [],
},
),
]
@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
"""Handle tool calls."""
try:
if name == "query_instant":
result = await query_instant(
query=arguments["query"],
time=arguments.get("time"),
)
elif name == "query_range":
result = await query_range(
query=arguments["query"],
start=arguments["start"],
end=arguments["end"],
step=arguments.get("step", "1m"),
)
elif name == "list_labels":
result = await list_labels(
match=arguments.get("match"),
)
elif name == "get_label_values":
result = await get_label_values(
label_name=arguments["label_name"],
match=arguments.get("match"),
)
elif name == "list_metrics":
result = await list_metrics(
with_metadata=arguments.get("with_metadata", False),
)
else:
return [TextContent(type="text", text=f"Unknown tool: {name}")]
# Convert Pydantic model to JSON string
return [TextContent(type="text", text=result.model_dump_json(indent=2))]
except Exception as e:
logger.exception(f"Error executing tool {name}")
return [TextContent(type="text", text=f"Error: {str(e)}")]
async def run_server() -> None:
"""Run the MCP server with stdio transport."""
async with lifespan():
async with stdio_server() as (read_stream, write_stream):
await server.run(
read_stream,
write_stream,
server.create_initialization_options(),
)
def main() -> None:
"""Main entry point."""
try:
asyncio.run(run_server())
except KeyboardInterrupt:
logger.info("Server stopped by user")
sys.exit(0)
except Exception as e:
logger.exception("Server error")
sys.exit(1)
if __name__ == "__main__":
main()
