lifespan shuts as soon as it opens

[MR-GREEN1337]

from contextlib import asynccontextmanager
from typing import AsyncIterator, Union
from mcp.server.lowlevel.server import Server
from mcp.server.fastmcp import FastMCP, Context
from loguru import logger
from datetime import datetime

from src.core.settings import settings
from src.db.postgresql import db
from src.db.qdrant import init_qdrant, get_qdrant, close_qdrant
from playwright.async_api import async_playwright
from motor.motor_asyncio import AsyncIOMotorClient
import httpx
from qdrant_client import AsyncQdrantClient
import asyncpg
from dataclasses import dataclass

@dataclass
class AppContext:
    postgres: asyncpg.Pool
    mongodb: AsyncIOMotorClient
    qdrant: AsyncQdrantClient
    http_client: httpx.AsyncClient
    playwright: any  # Playwright browser instance


@asynccontextmanager
async def server_lifespan(server: Server) -> AsyncIterator[object]:
    """Initialize and manage database and browser connections."""
    start_time = datetime.now()
    logger.info("Starting server lifespan initialization")

    # Initialize all your connections
    mongodb_client = None
    qdrant_client = None
    http_client = None
    playwright = None

    try:
        # Initialize MongoDB
        logger.info("Initializing MongoDB connection")
        mongodb_client = AsyncIOMotorClient(settings.MONGODB_URL)
        await mongodb_client.admin.command('ping')
        logger.success("MongoDB connection established")

        # Initialize Qdrant
        logger.info("Initializing Qdrant connection")
        if settings.QDRANT_URL:
            await init_qdrant()
            qdrant_client = await get_qdrant()
        logger.success("Qdrant connection established")

        # Initialize HTTP client
        logger.info("Initializing HTTP client")
        http_client = httpx.AsyncClient(
            timeout=30.0,
            follow_redirects=True,
            verify=False
        )
        logger.success("HTTP client initialized")

        # Initialize Playwright
        logger.info("Initializing Playwright browser")
        playwright = await async_playwright().start()
        logger.success("Playwright browser launched")

        # Create context
        context = AppContext(
            postgres=db,
            mongodb=mongodb_client,
            qdrant=qdrant_client,
            http_client=http_client,
            playwright=playwright
        )

        initialization_time = datetime.now() - start_time
        logger.info(f"Server lifespan initialization completed in {initialization_time.total_seconds():.2f} seconds")

        # This is where we yield control back to the server
        yield context
        # The server will run until it's explicitly stopped

    except Exception as e:
        logger.error(f"Error during server lifespan initialization: {str(e)}")
        raise
    finally:
        # This cleanup only happens when the server is shutting down
        logger.info("Starting cleanup of server resources")

        if playwright:
            try:
                await playwright.stop()
                logger.info("Playwright stopped")
            except Exception as e:
                logger.error(f"Error stopping playwright: {str(e)}")

        if http_client:
            try:
                await http_client.aclose()
                logger.info("HTTP client closed")
            except Exception as e:
                logger.error(f"Error closing HTTP client: {str(e)}")

        if settings.QDRANT_URL:
            try:
                await close_qdrant()
                logger.info("Qdrant connection closed")
            except Exception as e:
                logger.error(f"Error closing Qdrant connection: {str(e)}")

        if mongodb_client:
            try:
                mongodb_client.close()
                logger.info("MongoDB connection closed")
            except Exception as e:
                logger.error(f"Error closing MongoDB connection: {str(e)}")

        cleanup_time = datetime.now() - start_time
        logger.info(f"Server cleanup completed in {cleanup_time.total_seconds():.2f} seconds")

from mcp.server.lowlevel.server import Server, NotificationOptions
from mcp.server.models import InitializationOptions
import mcp.server.stdio
import mcp.types as types
from src.extractor.server.lifespan import server_lifespan
import mcp

server = Server("extractor", lifespan=server_lifespan)

async def run():
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="extractor",
                server_version="0.1.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                )
            )
        )

if __name__ == "__main__":
    import asyncio
    asyncio.run(run())


2025-02-21 00:37:42.388 | INFO     | src.extractor.server.lifespan:server_lifespan:28 - Starting server lifespan initialization
2025-02-21 00:37:42.388 | INFO     | src.extractor.server.lifespan:server_lifespan:39 - Initializing MongoDB connection
2025-02-21 00:37:43.113 | SUCCESS  | src.extractor.server.lifespan:server_lifespan:43 - MongoDB connection established
2025-02-21 00:37:43.113 | INFO     | src.extractor.server.lifespan:server_lifespan:46 - Initializing Qdrant connection
2025-02-21 00:37:43.113 | SUCCESS  | src.extractor.server.lifespan:server_lifespan:50 - Qdrant connection established
2025-02-21 00:37:43.113 | INFO     | src.extractor.server.lifespan:server_lifespan:53 - Initializing HTTP client
2025-02-21 00:37:43.115 | SUCCESS  | src.extractor.server.lifespan:server_lifespan:59 - HTTP client initialized
2025-02-21 00:37:43.115 | INFO     | src.extractor.server.lifespan:server_lifespan:62 - Initializing Playwright browser
2025-02-21 00:37:43.309 | SUCCESS  | src.extractor.server.lifespan:server_lifespan:65 - Playwright browser launched
2025-02-21 00:37:43.309 | INFO     | src.extractor.server.lifespan:server_lifespan:78 - Server lifespan initialization completed in 0.92 seconds
2025-02-21 00:37:43.310 | INFO     | src.extractor.server.lifespan:server_lifespan:89 - Starting cleanup of server resources
2025-02-21 00:37:43.333 | INFO     | src.extractor.server.lifespan:server_lifespan:95 - Browser closed
2025-02-21 00:37:43.337 | INFO     | src.extractor.server.lifespan:server_lifespan:102 - Playwright stopped
2025-02-21 00:37:43.337 | INFO     | src.extractor.server.lifespan:server_lifespan:109 - HTTP client closed
2025-02-21 00:37:43.386 | INFO     | src.extractor.server.lifespan:server_lifespan:123 - MongoDB connection closed
2025-02-21 00:37:43.386 | INFO     | src.extractor.server.lifespan:server_lifespan:128 - Server cleanup completed in 1.00 seconds

[dsp-ant]

Is there a shorter repro, preferably in form of a test in tests/ for this?