Skip to content

Commit 1e2d3f7

Browse files
shancwclaude
andcommitted
feat: Add lifespan handler for automatic heartbeat startup/shutdown
Integrate FastMCP lifespan context manager to automatically start heartbeat on server startup and gracefully stop on shutdown, avoiding event loop conflicts and following async service best practices. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent b5b637f commit 1e2d3f7

File tree

1 file changed

+32
-5
lines changed

1 file changed

+32
-5
lines changed

perplexity/server/app.py

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,24 +3,54 @@
33
"""
44

55
import os
6+
from contextlib import asynccontextmanager
67
from typing import Any, Dict, Iterable, List, Optional, Union
78

89
from fastmcp import FastMCP
910
from fastmcp.server.middleware import Middleware, MiddlewareContext
1011
from fastmcp.server.dependencies import get_http_headers
12+
from starlette.applications import Starlette
1113

1214
from .client_pool import ClientPool
1315
from ..client import Client
1416
from ..config import SEARCH_LANGUAGES
1517
from ..exceptions import ValidationError
18+
from ..logger import get_logger
1619

1720
from .utils import (
1821
sanitize_query, validate_file_data, validate_query_limits, validate_search_params,
1922
)
2023

24+
logger = get_logger("server.app")
25+
2126
# API 密钥配置(从环境变量读取,默认为 sk-123456)
2227
MCP_TOKEN = os.getenv("MCP_TOKEN", "sk-123456")
2328

29+
# 全局 ClientPool 实例
30+
_pool: Optional[ClientPool] = None
31+
32+
33+
def get_pool() -> ClientPool:
34+
"""Get or create the singleton ClientPool instance."""
35+
global _pool
36+
if _pool is None:
37+
_pool = ClientPool()
38+
return _pool
39+
40+
41+
@asynccontextmanager
42+
async def app_lifespan(server: FastMCP):
43+
"""Application lifespan handler for startup/shutdown events."""
44+
# Startup: Initialize pool and start heartbeat
45+
pool = get_pool()
46+
if pool.is_heartbeat_enabled():
47+
pool.start_heartbeat()
48+
logger.info("Heartbeat started via lifespan")
49+
yield
50+
# Shutdown: Stop heartbeat gracefully
51+
pool.stop_heartbeat()
52+
logger.info("Heartbeat stopped via lifespan")
53+
2454

2555
class AuthMiddleware(Middleware):
2656
"""Bearer Token 认证中间件"""
@@ -38,15 +68,12 @@ async def on_request(self, context: MiddlewareContext, call_next):
3868
return await call_next(context)
3969

4070

41-
# Create FastMCP instance
42-
mcp = FastMCP("perplexity-mcp")
71+
# Create FastMCP instance with lifespan
72+
mcp = FastMCP("perplexity-mcp", lifespan=app_lifespan)
4373

4474
# 添加认证中间件
4575
mcp.add_middleware(AuthMiddleware(MCP_TOKEN))
4676

47-
# 全局 ClientPool 实例
48-
_pool: Optional[ClientPool] = None
49-
5077

5178
def get_pool() -> ClientPool:
5279
"""Get or create the singleton ClientPool instance."""

0 commit comments

Comments
 (0)