mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-10-03 00:55:14 -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:
|
def create_app() -> FastAPI:
|
||||||
"""Create and configure the FastAPI application."""
|
"""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
|
from .mcp_server.context import ClientIdMiddleware
|
||||||
|
|
||||||
# Build the MCP app up-front so we can wire its lifespan into FastAPI's —
|
# 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")
|
mcp_app = mcp.http_app(path="/", transport="http")
|
||||||
|
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def lifespan(app: FastAPI):
|
async def voicebox_lifespan(app: FastAPI):
|
||||||
await _run_startup(app)
|
await _run_startup(app)
|
||||||
async with mcp_app.router.lifespan_context(app):
|
try:
|
||||||
try:
|
yield
|
||||||
yield
|
finally:
|
||||||
finally:
|
# Paired with _run_startup via try/finally: runs whether or
|
||||||
await _run_shutdown()
|
# 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(
|
application = FastAPI(
|
||||||
title="voicebox API",
|
title="voicebox API",
|
||||||
|
|||||||
Reference in New Issue
Block a user