at-mintel/packages/umami-mcp/src/index.ts

import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
import express from 'express';
import crypto from 'crypto';
import {
    CallToolRequestSchema,
    ListToolsRequestSchema,
    Tool,
} from "@modelcontextprotocol/sdk/types.js";
import axios from "axios";
import https from "https";

const UMAMI_BASE_URL = process.env.UMAMI_BASE_URL || "https://umami.infra.mintel.me";
const UMAMI_USERNAME = process.env.UMAMI_USERNAME;
const UMAMI_PASSWORD = process.env.UMAMI_PASSWORD;
const UMAMI_API_KEY = process.env.UMAMI_API_KEY; // optional if using API key auth

const httpsAgent = new https.Agent({
    rejectUnauthorized: false,
});

if (!UMAMI_USERNAME && !UMAMI_API_KEY) {
    console.error("Warning: Neither UMAMI_USERNAME/PASSWORD nor UMAMI_API_KEY is set.");
}

// Token cache to avoid logging in on every request
let cachedToken: string | null = null;

async function getAuthHeaders(): Promise<Record<string, string>> {
    try {
        if (UMAMI_API_KEY) {
            return { 'x-umami-api-key': UMAMI_API_KEY };
        }
        if (!cachedToken) {
            console.error(`Attempting Umami login for ${UMAMI_USERNAME} at ${UMAMI_BASE_URL}`);
            const res = await axios.post(`${UMAMI_BASE_URL}/api/auth/login`, {
                username: UMAMI_USERNAME,
                password: UMAMI_PASSWORD,
            }, { 
                httpsAgent,
                timeout: 5000 // Add timeout to prevent hangs
            });
            cachedToken = res.data.token;
            console.error("Umami login successful");
        }
        return { Authorization: `Bearer ${cachedToken}` };
    } catch (e: any) {
        const msg = e.response?.data?.message || e.message;
        console.error(`Umami Auth Error: ${msg}`);
        throw new Error(`Authentication failed: ${msg}`);
    }
}


// --- Tool Definitions ---
const LIST_WEBSITES_TOOL: Tool = {
    name: "umami_list_websites",
    description: "List all websites tracked in Umami",
    inputSchema: { type: "object", properties: {} },
};

const GET_WEBSITE_STATS_TOOL: Tool = {
    name: "umami_get_website_stats",
    description: "Get summary statistics for a website for a time range",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
            start_at: { type: "number", description: "Start timestamp in ms (e.g., Date.now() - 7 days)" },
            end_at: { type: "number", description: "End timestamp in ms (default: now)" },
        },
        required: ["website_id", "start_at"],
    },
};

const GET_PAGE_VIEWS_TOOL: Tool = {
    name: "umami_get_pageviews",
    description: "Get pageview/session time series for a website",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
            start_at: { type: "number", description: "Start timestamp in ms" },
            end_at: { type: "number", description: "End timestamp in ms (default: now)" },
            unit: { type: "string", description: "Time unit: 'hour', 'day', 'month' (default: day)" },
            timezone: { type: "string", description: "Timezone (default: Europe/Berlin)" },
        },
        required: ["website_id", "start_at"],
    },
};

const GET_TOP_PAGES_TOOL: Tool = {
    name: "umami_get_top_pages",
    description: "Get the most visited pages/URLs for a website",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
            start_at: { type: "number", description: "Start timestamp in ms" },
            end_at: { type: "number", description: "End timestamp in ms" },
            limit: { type: "number", description: "Number of results (default: 20)" },
        },
        required: ["website_id", "start_at"],
    },
};

const GET_TOP_REFERRERS_TOOL: Tool = {
    name: "umami_get_top_referrers",
    description: "Get the top traffic referrers for a website",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
            start_at: { type: "number", description: "Start timestamp in ms" },
            end_at: { type: "number", description: "End timestamp in ms" },
            limit: { type: "number", description: "Number of results (default: 10)" },
        },
        required: ["website_id", "start_at"],
    },
};

const GET_COUNTRY_STATS_TOOL: Tool = {
    name: "umami_get_country_stats",
    description: "Get visitor breakdown by country",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
            start_at: { type: "number", description: "Start timestamp in ms" },
            end_at: { type: "number", description: "End timestamp in ms" },
        },
        required: ["website_id", "start_at"],
    },
};

const GET_ACTIVE_VISITORS_TOOL: Tool = {
    name: "umami_get_active_visitors",
    description: "Get the number of visitors currently active on a website (last 5 minutes)",
    inputSchema: {
        type: "object",
        properties: {
            website_id: { type: "string", description: "Umami website UUID" },
        },
        required: ["website_id"],
    },
};

// --- Server Setup ---
const server = new Server(
    { name: "umami-mcp", version: "1.0.0" },
    { capabilities: { tools: {} } }
);

server.setRequestHandler(ListToolsRequestSchema, async () => ({
    tools: [
        LIST_WEBSITES_TOOL,
        GET_WEBSITE_STATS_TOOL,
        GET_PAGE_VIEWS_TOOL,
        GET_TOP_PAGES_TOOL,
        GET_TOP_REFERRERS_TOOL,
        GET_COUNTRY_STATS_TOOL,
        GET_ACTIVE_VISITORS_TOOL,
    ],
}));

server.setRequestHandler(CallToolRequestSchema, async (request) => {
    try {
        const headers = await getAuthHeaders();
        const api = axios.create({ baseURL: `${UMAMI_BASE_URL}/api`, headers, httpsAgent, timeout: 10000 });
        const now = Date.now();

    if (request.params.name === "umami_list_websites") {
        const res = await api.get('/websites');
        const sites = (res.data.data || res.data || []).map((s: any) => ({
            id: s.id, name: s.name, domain: s.domain
        }));
        return { content: [{ type: "text", text: JSON.stringify(sites, null, 2) }] };
    }

    if (request.params.name === "umami_get_website_stats") {
        const { website_id, start_at, end_at = now } = request.params.arguments as any;
        const res = await api.get(`/websites/${website_id}/stats`, { params: { startAt: start_at, endAt: end_at } });
        return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
    }

    if (request.params.name === "umami_get_pageviews") {
        const { website_id, start_at, end_at = now, unit = 'day', timezone = 'Europe/Berlin' } = request.params.arguments as any;
        const res = await api.get(`/websites/${website_id}/pageviews`, {
            params: { startAt: start_at, endAt: end_at, unit, timezone }
        });
        return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
    }

    if (request.params.name === "umami_get_top_pages") {
        const { website_id, start_at, end_at = now, limit = 20 } = request.params.arguments as any;
        const res = await api.get(`/websites/${website_id}/metrics`, {
            params: { startAt: start_at, endAt: end_at, type: 'url', limit }
        });
        return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
    }

    if (request.params.name === "umami_get_top_referrers") {
        const { website_id, start_at, end_at = now, limit = 10 } = request.params.arguments as any;
        const res = await api.get(`/websites/${website_id}/metrics`, {
            params: { startAt: start_at, endAt: end_at, type: 'referrer', limit }
        });
        return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
    }

    if (request.params.name === "umami_get_country_stats") {
        const { website_id, start_at, end_at = now } = request.params.arguments as any;
        const res = await api.get(`/websites/${website_id}/metrics`, {
            params: { startAt: start_at, endAt: end_at, type: 'country' }
        });
        return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
    }

    if (request.params.name === "umami_get_active_visitors") {
        const { website_id } = request.params.arguments as any;
        const res = await api.get(`/websites/${website_id}/active`);
        return { content: [{ type: "text", text: JSON.stringify(res.data, null, 2) }] };
    }

    throw new Error(`Unknown tool: ${request.params.name}`);


    } catch (e: any) {
        return { 
            isError: true, 
            content: [{ type: "text", text: `MCP Service Error: ${e.message}` }] 
        };
    }
});

// --- Express / SSE Server ---
async function run() {
    const isStdio = process.argv.includes('--stdio');

    if (isStdio) {
        const { StdioServerTransport } = await import('@modelcontextprotocol/sdk/server/stdio.js');
        const transport = new StdioServerTransport();
        await server.connect(transport);
        console.error('Umami MCP server is running on stdio');
    } else {
        const app = express();
        const transports = new Map<string, SSEServerTransport>();

        // Manual CORS middleware
        app.use((req, res, next) => {
            res.setHeader('Access-Control-Allow-Origin', '*');
            res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
            res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Authorization, x-mcp-protocol-version');
            if (req.method === 'OPTIONS') {
                res.status(204).end();
                return;
            }
            next();
        });

        app.use((req, _res, next) => {
            console.error(`${req.method} ${req.url}`);
            next();
        });

        app.get('/sse', async (req, res) => {
            const transport = new SSEServerTransport('/message', res as any);
            const sessionId = transport.sessionId;
            console.error(`New SSE connection: ${sessionId}`);
            transports.set(sessionId, transport);

            const heartbeatInterval = setInterval(() => {
                res.write(": heartbeat\n\n");
            }, 15000);

            req.on('close', () => {
                console.error(`SSE connection closed: ${sessionId}`);
                clearInterval(heartbeatInterval);
                transports.delete(sessionId);
            });

            await server.connect(transport);
            
            await new Promise((resolve) => {
                req.on('close', resolve);
            });
        });

        app.post('/message', express.json(), async (req, res) => {
            const sessionId = req.query.sessionId as string;
            const transport = transports.get(sessionId);

            console.error(`POST /message/${sessionId} - Active transports: ${transports.size}`);

            if (!transport) {
                console.error(`No transport found for session: ${sessionId}`);
                res.status(400).send('No active SSE connection for this session');
                return;
            }
            await transport.handlePostMessage(req, res);
        });

        const PORT = Number(process.env.UMAMI_MCP_PORT) || 3003;
        app.listen(PORT, '0.0.0.0', () => {
            console.error(`Umami MCP server running on http://0.0.0.0:${PORT}/sse`);
        });
    }
}

run().catch((err) => {
    console.error("Fatal error:", err);
    process.exit(1);
});