App/mcp/server.py

"""
MCP server exposing DSE equity values from the Uwekezaji backend.

Remote customers connect over HTTP (Streamable HTTP) to the deployed URL.
Local developers can still run `python -m App.mcp.server` over stdio.
"""
from __future__ import annotations

import asyncio
import json
import os
from contextlib import asynccontextmanager
from typing import Any, Optional

from fastmcp import FastMCP
from starlette.applications import Starlette
from starlette.responses import JSONResponse
from starlette.routing import Route

from App.mcp.api_client import UwekezajiApiClient, get_api_base_url
from App.mcp.auth import McpApiKeyMiddleware, get_mcp_api_key
from App.mcp.equities_service import (
    data_source_label,
    get_equity,
    get_equity_price_history,
    list_equities,
)
from db import close_db, init_db

_db_lock = asyncio.Lock()
_db_ready = False
_api_client: Optional[UwekezajiApiClient] = None
_EMBEDDED = os.getenv("UWEKEZAJI_MCP_EMBEDDED", "").strip().lower() in {"1", "true", "yes"}


async def _ensure_db() -> None:
    global _db_ready
    async with _db_lock:
        if _db_ready:
            return
        if _EMBEDDED:
            # Mounted inside main FastAPI; Tortoise is already initialized.
            _db_ready = True
            return
        await init_db()
        _db_ready = True


async def _get_api_client() -> Optional[UwekezajiApiClient]:
    global _api_client
    api_base = get_api_base_url()
    if not api_base:
        return None
    if _api_client is None:
        _api_client = UwekezajiApiClient(api_base)
    return _api_client


@asynccontextmanager
async def _data_session():
    client = await _get_api_client()
    if client is None:
        await _ensure_db()
    try:
        yield client
    finally:
        pass


mcp = FastMCP(
    name="Uwekezaji Equities",
    instructions=(
        "Provides latest market values for Dar es Salaam Stock Exchange listed "
        "equities and ETFs. Government and corporate bonds are excluded."
    ),
)


@mcp.tool
async def list_equity_values() -> dict[str, Any]:
    """List all DSE equities and ETFs with latest prices and market values."""
    async with _data_session() as client:
        return await list_equities(client)


@mcp.tool
async def get_equity_value(symbol: str) -> dict[str, Any]:
    """Get the latest market value snapshot for one equity symbol (e.g. CRDB)."""
    async with _data_session() as client:
        payload = await get_equity(symbol, client)
        if payload is None:
            return {
                "found": False,
                "symbol": symbol.upper(),
                "source": data_source_label(),
                "message": "Equity not found or excluded as a bond.",
            }
        return {"found": True, "source": data_source_label(), "equity": payload}


@mcp.tool
async def get_equity_price_history(symbol: str, days: int = 90) -> dict[str, Any]:
    """Get historical OHLCV price rows for an equity symbol over the last N days."""
    async with _data_session() as client:
        payload = await get_equity_price_history(symbol, days=days, client=client)
        if payload is None:
            return {
                "found": False,
                "symbol": symbol.upper(),
                "source": data_source_label(),
                "message": "Equity not found or excluded as a bond.",
            }
        return {"found": True, **payload}


@mcp.resource("equities://latest")
async def equities_latest_resource() -> str:
    """JSON snapshot of all latest equity values (bonds excluded)."""
    async with _data_session() as client:
        return json.dumps(await list_equities(client), indent=2)


async def _health(_request):
    return JSONResponse(
        {
            "status": "ok",
            "service": "uwekezaji-equities-mcp",
            "transport": "streamable-http",
            "auth_required": bool(get_mcp_api_key()),
            "data_source": data_source_label(),
        }
    )


def create_mcp_http_app():
    """ASGI app for remote MCP clients. Mount at /mcp on the main FastAPI app."""
    mcp_app = mcp.http_app(path="/", transport="streamable-http")
    routes = [Route("/health", _health), Route("/healthz", _health)]
    app = Starlette(routes=routes)
    app.mount("/", mcp_app)

    api_key = get_mcp_api_key()
    if api_key:
        app.add_middleware(McpApiKeyMiddleware, api_key=api_key)
    return app


def main() -> None:
    transport = os.getenv("UWEKEZAJI_MCP_TRANSPORT", "stdio").strip().lower()
    if transport in {"http", "streamable-http", "streamable_http"}:
        host = os.getenv("UWEKEZAJI_MCP_HOST", "0.0.0.0")
        port = int(os.getenv("UWEKEZAJI_MCP_PORT", "8787"))
        mcp.run(transport="streamable-http", host=host, port=port)
        return

    try:
        mcp.run()
    finally:
        if _db_ready and not _EMBEDDED:
            asyncio.run(close_db())
        if _api_client is not None:
            asyncio.run(_api_client.close())


if __name__ == "__main__":
    main()