diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index dd91bb6eef3..feebc3b5ee1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -124,6 +124,7 @@ from .types_utils.utils import get_instance_fn, validate_custom_validate_return_ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from litellm.repositories.database import Database from litellm.tracing import TraceReceiver Span = _Span | Any @@ -133,6 +134,7 @@ else: class ProxyLifespanState(TypedDict): tracing_receiver: ReadOnly["TraceReceiver | None"] + database: ReadOnly["Database | None"] class ReconcileOutcome(NamedTuple): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bc79becf7bb..48e3cd35546 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -168,6 +168,7 @@ from litellm.proxy.common_utils.token_count_settings import ( ) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields +from litellm.repositories import database as database_runtime from litellm.router_utils.access_windows import access_windows_config_error from litellm.router_utils.add_retry_fallback_headers import ( get_fallback_errors_from_headers, @@ -1288,13 +1289,7 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N if worker_heartbeat is not None and prisma_client: await worker_heartbeat.deregister() if prisma_client: - # Drain the SGR fold first: it lives in memory, so an un-drained interval - # is lost, and a write attempted after disconnect raises - # ClientNotConnectedError rather than persisting anything. Ordering this - # inside the same guard is what keeps the two from drifting apart. await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client)) - verbose_proxy_logger.debug("Disconnecting from Prisma") - await prisma_client.disconnect() else: await _drain_audit_tasks_on_shutdown() @@ -1425,8 +1420,19 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries: return client.writer_db +def _forget_prisma_client() -> None: + global prisma_client # noqa: PLW0603 # routes still read this global until they move to DatabaseDependency + prisma_client = None # rebind-ok: a closed client must not survive into the next lifespan entry + + @asynccontextmanager async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: + async with AsyncExitStack() as resources, _proxy_lifespan(app, resources) as state: + yield state + + +@asynccontextmanager +async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGenerator[ProxyLifespanState, None]: global \ prisma_client, \ master_key, \ @@ -1550,6 +1556,11 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState ) ) + resources.callback(_forget_prisma_client) + database: Final = await resources.enter_async_context( + database_runtime.lifespan(None if prisma_client is None else database_runtime.Database(prisma=prisma_client)) + ) + await migrate_if_requested( environ=os.environ, master_key=master_key, @@ -1784,7 +1795,7 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState enabled=tracing_enabled, settings=tracing_settings, ) as receiver: - state: Final[ProxyLifespanState] = {"tracing_receiver": receiver} + state: Final[ProxyLifespanState] = {"tracing_receiver": receiver, "database": database} from litellm.proxy.admin_mcp import admin_mcp_lifespan telemetry_settings: Final = load_telemetry_settings() diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py new file mode 100644 index 00000000000..f4518dbceb3 --- /dev/null +++ b/litellm/repositories/database.py @@ -0,0 +1,27 @@ +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from litellm._logging import verbose_proxy_logger + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +@dataclass(frozen=True, slots=True) +class Database: + prisma: "PrismaClient" + + async def aclose(self) -> None: + verbose_proxy_logger.debug("Disconnecting from Prisma") + await self.prisma.disconnect() + + +@asynccontextmanager +async def lifespan(database: Database | None) -> AsyncGenerator[Database | None, None]: + try: + yield database + finally: + if database is not None: + await database.aclose() diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 11a44a0ec16..5cdc8f9782e 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -33,9 +33,9 @@ from pydantic import BaseModel from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps +from litellm.proxy.management_helpers.audit_logs import track_audit_task from litellm.proxy.proxy_server import ( ProxyStartupEvent, - initialize_shared_aiohttp_session, _resolve_pydantic_type, _resolve_typed_dict_type, cleanup_router_config_variables, @@ -43,14 +43,14 @@ from litellm.proxy.proxy_server import ( get_litellm_model_info, initialize, initialize_from_worker_config, + initialize_shared_aiohttp_session, load_from_azure_key_vault, proxy_shutdown_event, proxy_startup_event, run_ollama_serve, save_worker_config, ) - -from litellm.proxy.management_helpers.audit_logs import track_audit_task +from litellm.repositories import database as database_runtime from .conftest import normalize @@ -101,7 +101,7 @@ def test_cleanup_router_config_variables_fails_on_unknown_attr_raises(): @pytest.mark.asyncio -async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch): +async def test_database_lifespan_disconnects_prisma_after_proxy_shutdown(monkeypatch): fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) @@ -117,7 +117,9 @@ async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch): monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - await proxy_shutdown_event() + async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)): + await proxy_shutdown_event() + fake_prisma.disconnect.assert_not_awaited() observed = { "disconnect_called": fake_prisma.disconnect.await_count == 1, @@ -222,7 +224,8 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - await proxy_shutdown_event() + async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)): + await proxy_shutdown_event() assert calls == ["flush", "disconnect"] @@ -262,7 +265,8 @@ async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_befor monkeypatch.setattr(litellm, "success_callback", [], raising=False) track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts())) - await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) + async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)): + await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) assert calls == ["flush", "audit", "disconnect"] @@ -318,7 +322,7 @@ async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_t @pytest.mark.asyncio -async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): +async def test_database_lifespan_propagates_disconnect_error_after_shutdown(monkeypatch): fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone")) monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) @@ -333,7 +337,8 @@ async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): monkeypatch.setattr(litellm, "success_callback", [], raising=False) with pytest.raises(RuntimeError, match="db gone"): - await proxy_shutdown_event() + async with database_runtime.lifespan(database_runtime.Database(prisma=fake_prisma)): + await proxy_shutdown_event() # --------------------------------------------------------------------------- @@ -1060,7 +1065,7 @@ def test_otel_global_provider_published_after_callback_init(): span was orphaned. The publish (``publish_global_otel_v2_provider``) must therefore appear after ``_initialize_startup_logging`` in the lifespan source. """ - wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event) + wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan) source = inspect.getsource(wrapped) init_pos = source.find("ProxyStartupEvent.initialize_startup_logging(") publish_pos = source.find("publish_global_otel_v2_provider(") @@ -1146,7 +1151,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): The call must follow Prisma setup so DB-backed deployments do not false-positive. Direct ``_warn_budget_without_db`` tests cover the warning behavior itself. """ - wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event) + wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan) source = inspect.getsource(wrapped) budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:") warn_pos = source.find("warn_budget_without_db(") diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index a26a64c6fb0..cb7e12720b5 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -113,7 +113,7 @@ async def test_admin_mcp_failure_still_closes_proxy_resources( message: Final = "LITELLM_LICENSE" if failure_phase == "license" else f"{failure_phase} failed" async def run_lifespan() -> None: async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state: - assert state == {"tracing_receiver": None} + assert state == {"tracing_receiver": None, "database": None} if failure_phase == "serving": raise RuntimeError("serving failed") if failure_phase == "cancelled": @@ -162,7 +162,7 @@ async def test_proxy_shutdown_drains_active_admin_tool_before_closing_connector( transport=httpx2.ASGITransport(app=admin_mcp_proxy), base_url="http://localhost:4000" ) as client: async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state: - assert state == {"tracing_receiver": None} + assert state == {"tracing_receiver": None, "database": None} request: Final = asyncio.create_task(client.post( "/admin/mcp", headers={"Authorization": "Bearer admin", "Accept": "application/json, text/event-stream"}, diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 8c62b225992..1669ebea63f 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -1838,6 +1838,12 @@ async def test_proxy_startup_says_a_lingering_migrate_from_variable_can_be_delet class _PrismaClientWhoseDatabaseRejectsQueries: + def __init__(self) -> None: + self.disconnect_calls = 0 + + async def disconnect(self) -> None: + self.disconnect_calls += 1 + class _Database: async def query_raw(self, query, *args): raise RuntimeError("permission denied for table LiteLLM_CredentialsTable") @@ -1852,7 +1858,8 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc from litellm.proxy.proxy_server import proxy_startup_event _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries()) + prisma: Final = _PrismaClientWhoseDatabaseRejectsQueries() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY) with ( @@ -1865,6 +1872,25 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc notices = [record.getMessage() for record in caplog.records if "LITELLM_MIGRATE_FROM_MASTER_KEY" in record.message] assert len(notices) == 1 assert "Could not migrate stored values" in notices[0] + assert prisma.disconnect_calls == 1 + + +@pytest.mark.asyncio +async def test_proxy_startup_after_a_failed_migration_does_not_reuse_the_closed_client(monkeypatch, tmp_path): + from fastapi import FastAPI + + from litellm.proxy.proxy_server import proxy_startup_event + + _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries()) + monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY) + with pytest.raises(RuntimeError, match="permission denied"): + async with proxy_startup_event(FastAPI()): + pass + + monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", "") + async with proxy_startup_event(FastAPI()) as state: + assert state["database"] is None @pytest.mark.asyncio diff --git a/tests/unit/repositories/test_database.py b/tests/unit/repositories/test_database.py new file mode 100644 index 00000000000..5a84688575f --- /dev/null +++ b/tests/unit/repositories/test_database.py @@ -0,0 +1,71 @@ +import asyncio +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.repositories.database import Database, lifespan + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [RuntimeError, asyncio.CancelledError]) +async def test_lifespan_closes_database_when_owner_fails(failure): + prisma: Final = MagicMock() + prisma.disconnect = AsyncMock() + database: Final = Database(prisma=prisma) + + async def fail(): + async with lifespan(database) as current: + assert current is database + prisma.disconnect.assert_not_awaited() + raise failure("owner failed") + + with pytest.raises(failure): + await fail() + prisma.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_lifespan_without_database_preserves_owner_error(): + async def fail(): + async with lifespan(None) as database: + assert database is None + raise RuntimeError("owner failed") + + with pytest.raises(RuntimeError, match="owner failed"): + await fail() + + +@pytest.mark.asyncio +async def test_cancelling_owner_waits_for_database_close() -> None: + serving: Final = asyncio.Event() + closing: Final = asyncio.Event() + release_close: Final = asyncio.Event() + prisma: Final = MagicMock() + + async def disconnect() -> None: + closing.set() + await release_close.wait() + + prisma.disconnect = AsyncMock(side_effect=disconnect) + + async def serve() -> None: + async with lifespan(Database(prisma=prisma)): + serving.set() + await asyncio.Event().wait() + + owner: Final = asyncio.create_task(serve()) + try: + await asyncio.wait_for(serving.wait(), timeout=2) + owner.cancel() + await asyncio.wait_for(closing.wait(), timeout=2) + assert not owner.done() + release_close.set() + with pytest.raises(asyncio.CancelledError): + await owner + finally: + release_close.set() + if not owner.done(): + owner.cancel() + await asyncio.gather(owner, return_exceptions=True) + prisma.disconnect.assert_awaited_once()