diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 1929e7d3fc8..be515392a17 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -61,6 +61,26 @@ class _RoutedActions: return getattr(self._writer_actions, name) +class WriterPinnedClient: + """PrismaClient-shaped view whose `.db` resolves to the writer while it is available. + + Read-after-write paths (e.g. the model reconcile a /model/new triggers to + verify its own just-committed row) must not read through a lagging read + replica: the row is not replayed there yet, so the reconcile concludes the + write is missing and fails the request even though it is durable (#38556). + + While the writer is degraded (`writer_unavailable`), the pin yields to the + routed wrapper so reconcile reads keep working from the replica: a proxy + that starts during a primary outage must still load DB-backed models, and + no read-after-write hazard exists then because writes are failing anyway. + """ + + __slots__ = ("db",) + + def __init__(self, db: "PrismaWrapper | RoutingPrismaWrapper") -> None: + self.db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) and not db.writer_unavailable else db + + class RoutingPrismaWrapper: """ Routes Prisma operations between a writer and a reader Prisma client. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3a70750528f..af0aa9743bc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6764,9 +6764,18 @@ class ProxyConfig: - list: the rows (may be empty if no models exist) - None: signals a DB fetch *failure* — callers must not treat this as "all models deleted" and must not evict existing router deployments. + + Pinned to the writer DB: this read reconciles the router against the rows a + model write just committed, and reading it through a lagging read replica + makes the write-triggered reload report its own durable write as missing + (#38556). It also keeps a stale replica snapshot from evicting a deployment + another pod just added. While the writer is degraded the pin yields to the + replica so reader-only mode keeps loading DB-backed models. """ try: - new_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository(prisma_client).table.find_many() + new_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository( + WriterPinnedClient(prisma_client.db) + ).table.find_many() return new_models except Exception as e: verbose_proxy_logger.exception( @@ -12013,6 +12022,7 @@ async def run_thread( # ) # async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)): from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ( diff --git a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py index dcc0036ff04..966a638f6a4 100644 --- a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py @@ -101,6 +101,49 @@ def test_per_model_reads_route_to_reader_writes_to_writer(): assert actions.delete_many is writer_inner.litellm_usertable.delete_many +def test_writer_pinned_client_bypasses_reader_routing(): + """Regression for #38556: read-after-write reconciles must see the writer's + just-committed rows, so WriterPinnedClient must resolve reads to the writer + even when a read replica is configured.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient + + writer, writer_inner, reader, reader_inner = _make_wrappers() + writer_inner.litellm_proxymodeltable = _model_actions_mock("writer_models") + reader_inner.litellm_proxymodeltable = _model_actions_mock("reader_models") + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + + pinned = WriterPinnedClient(routing) + + assert pinned.db is writer + assert pinned.db.litellm_proxymodeltable.find_many is writer_inner.litellm_proxymodeltable.find_many + + +def test_writer_pinned_client_passes_through_single_db(): + from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient + + writer, _, _, _ = _make_wrappers() + + assert WriterPinnedClient(writer).db is writer + + +def test_writer_pinned_client_yields_to_routed_reads_when_writer_down(): + """The pin must not break reader-only degraded mode: a proxy that starts + during a primary outage still loads DB-backed models from the replica, so + while the writer is degraded the pin resolves to the routed wrapper.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient + + writer, writer_inner, reader, reader_inner = _make_wrappers() + writer_inner.litellm_proxymodeltable = _model_actions_mock("writer_models") + reader_inner.litellm_proxymodeltable = _model_actions_mock("reader_models") + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + routing._writer_unavailable = True + + pinned = WriterPinnedClient(routing) + + assert pinned.db is routing + assert pinned.db.litellm_proxymodeltable.find_many is reader_inner.litellm_proxymodeltable.find_many + + @pytest.mark.asyncio async def test_connect_invokes_both_clients(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 949088ea3ba..71d4d5f4874 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9565,6 +9565,76 @@ class TestDeleteDeploymentSync: assert result is None, f"Expected None on DB failure to signal fetch error, got {result!r}" + @pytest.mark.asyncio + async def test_get_models_from_db_reads_from_writer_not_replica(self): + """ + Regression for #38556: with DATABASE_URL_READ_REPLICA configured, the model + reconcile after /model/new used to read via the replica, so a lagging replica + made the reload miss the just-committed row and fail the request with a 500. + The reconcile read must be pinned to the writer. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + from litellm.proxy.proxy_server import ProxyConfig + + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + committed_row = MagicMock(name="just_committed_model_row") + writer_inner.litellm_proxymodeltable.find_many = AsyncMock(return_value=[committed_row]) + reader_inner.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + mock_prisma = MagicMock() + mock_prisma.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False), + ) + + result = await ProxyConfig()._get_models_from_db(prisma_client=mock_prisma) + + assert result == [committed_row], f"Expected the writer's just-committed row, got {result!r}" + reader_inner.litellm_proxymodeltable.find_many.assert_not_awaited() + + @pytest.mark.asyncio + async def test_get_models_from_db_falls_back_to_replica_when_writer_down(self): + """ + The writer pin must not break reader-only degraded mode: a proxy that + starts during a primary outage (writer connect failed, replica healthy) + must still load DB-backed models through the replica instead of sending + the reconcile read to the unavailable writer. + """ + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + from litellm.proxy.proxy_server import ProxyConfig + + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + replica_row = MagicMock(name="replica_model_row") + writer_inner.litellm_proxymodeltable = SimpleNamespace( + find_many=AsyncMock(side_effect=RuntimeError("writer unreachable")), + create=MagicMock(name="writer_create"), + ) + reader_inner.litellm_proxymodeltable = SimpleNamespace( + find_many=AsyncMock(return_value=[replica_row]), + create=MagicMock(name="reader_create"), + ) + + mock_prisma = MagicMock() + mock_prisma.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False), + ) + mock_prisma.db._writer_unavailable = True + + result = await ProxyConfig()._get_models_from_db(prisma_client=mock_prisma) + + assert result == [replica_row], f"Expected the replica's rows in degraded mode, got {result!r}" + writer_inner.litellm_proxymodeltable.find_many.assert_not_awaited() + def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): """Follow-up to #30223: the flag must be discoverable via /config/list,