mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-15 04:40:40 -07:00
fix(mcp): correct lifespan shutdown order — drain MCP before unloading models
The inline lifespan ran _run_shutdown inside the MCP context, so the TTS / Whisper / LLM models were unloaded *before* FastMCP's __aexit__ got a chance to cancel its in-flight session tasks. Any MCP request mid-generate at shutdown time would crash on "model unloaded" instead of receiving a clean session-cancelled error. Rewire via compose_lifespan (which was already defined in mcp_server.server for exactly this purpose but never used): AsyncExitStack enters factories in order and exits in LIFO, so MCP teardown fires first — cancelling sessions — and _run_shutdown runs after nothing is holding the models. Smoke test shows the log order flipped as expected: Ready StreamableHTTP session manager started ... running ... StreamableHTTP session manager shutting down ← was last, now first Voicebox server shutting down... ← was first, now last As a side benefit, _run_shutdown is now paired with _run_startup via try/finally inside voicebox_lifespan, so a partial startup (models half-loaded, MCP __aenter__ fails) still unloads whatever was loaded instead of leaking it to process exit. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
co-authored by
Claude Opus 4.7
parent
c0eba9c628
commit
30c2cf1a2c
+17
-7
@@ -69,7 +69,7 @@ def safe_content_disposition(disposition_type: str, filename: str) -> str:
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
"""Create and configure the FastAPI application."""
|
||||
from .mcp_server.server import build_mcp_server
|
||||
from .mcp_server.server import build_mcp_server, compose_lifespan
|
||||
from .mcp_server.context import ClientIdMiddleware
|
||||
|
||||
# Build the MCP app up-front so we can wire its lifespan into FastAPI's —
|
||||
@@ -79,13 +79,23 @@ def create_app() -> FastAPI:
|
||||
mcp_app = mcp.http_app(path="/", transport="http")
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
async def voicebox_lifespan(app: FastAPI):
|
||||
await _run_startup(app)
|
||||
async with mcp_app.router.lifespan_context(app):
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await _run_shutdown()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Paired with _run_startup via try/finally: runs whether or
|
||||
# not the nested MCP lifespan entered cleanly, so a partial
|
||||
# startup still unloads whatever models were loaded.
|
||||
await _run_shutdown()
|
||||
|
||||
# compose_lifespan enters factories in order (voicebox startup →
|
||||
# MCP startup) and exits in LIFO (MCP teardown first → models
|
||||
# unload last). That ordering matters on shutdown: FastMCP's
|
||||
# __aexit__ cancels in-flight session tasks, and we want that to
|
||||
# happen *before* _run_shutdown yanks the TTS / Whisper / LLM
|
||||
# models out from under any MCP request that was still generating.
|
||||
lifespan = compose_lifespan(voicebox_lifespan, mcp_app.router.lifespan_context)
|
||||
|
||||
application = FastAPI(
|
||||
title="voicebox API",
|
||||
|
||||
Reference in New Issue
Block a user