fix(proxy): retain unwritten rollups after database disconnects (#51706)

* fix(proxy): retain unwritten rollups after database disconnects

* fix(proxy): initialize auto-router turns per client

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>

litellm/proxy/db/autorouter_session_rollup.py

@@ -21,7 +21,7 @@ import random
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from datetime import datetime, timezone
-from itertools import groupby
+from itertools import chain, groupby
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, NamedTuple
 
@@ -30,6 +30,7 @@ from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
 from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
 from litellm.proxy.db.create_views import SupportsExecuteRaw
 from litellm.proxy.db.db_span import db_span
+from litellm.proxy.db.exception_handler import db_write_never_sent
 
 if TYPE_CHECKING:
     from litellm.proxy._types import SpendLogsPayload
@@ -536,77 +537,110 @@ def _session_partition(transaction: AutoRouterTurnTransaction) -> tuple[str, str
     return (*identity, transaction.session_id, transaction.router_name)
 
 
+class UnwrittenTurns(NamedTuple):
+    """Turns the database provably never received, and the error that proved it."""
+
+    turns: tuple[AutoRouterTurnTransaction, ...]
+    error: Exception
+
+
+class _StoppedPartition(NamedTuple):
+    turns: tuple[AutoRouterTurnTransaction, ...]
+    error: Exception
+
+
 async def _drain_session_partition(
     prisma_client: PrismaClient,
     transactions: tuple[AutoRouterTurnTransaction, ...],
     n_retry_times: int,
     statement: str,
-) -> tuple[AutoRouterTurnTransaction, ...]:
+) -> _StoppedPartition | None:
     for position, transaction in enumerate(transactions):
         try:
             await _upsert_turn_with_retry(prisma_client, transaction, n_retry_times, statement)
-        except Exception as flush_err:  # noqa: BLE001  # stop dependent turns without retrying an ambiguous write
-            verbose_proxy_logger.error(
-                "Spend tracking - auto-router session rollup flush failed for router %s; "
-                "%s of %s turn writes stopped in this partition: %s",
-                transaction.router_name,
-                len(transactions) - position,
-                len(transactions),
-                flush_err,
-            )
-            return transactions[position:]
-    return ()
+        except Exception as flush_err:  # noqa: BLE001  # the caller keeps or drops the stopped turns
+            return _StoppedPartition(transactions[position:], flush_err)
+    return None
+
+
+def _log_dropped_turns(stopped: _StoppedPartition, partition_size: int) -> None:
+    verbose_proxy_logger.error(
+        "Spend tracking - auto-router session rollup flush failed for router %s; "
+        "%s of %s turn writes stopped in this partition: %s",
+        stopped.turns[0].router_name,
+        len(stopped.turns),
+        partition_size,
+        stopped.error,
+    )
 
 
 async def _flush_session_partition(
     prisma_client: PrismaClient,
     transactions: tuple[AutoRouterTurnTransaction, ...],
     n_retry_times: int,
-) -> None:
-    failed_suffix: Final = await _drain_session_partition(
+) -> UnwrittenTurns | None:
+    stopped: Final = await _drain_session_partition(
         prisma_client, transactions, n_retry_times, UPSERT_AUTOROUTER_SESSION_SQL
     )
-    if not failed_suffix or not failed_suffix[0].api_key:
-        return
-    failed_user: Final = failed_suffix[0].user_id
+    if stopped is None:
+        return None
+    if db_write_never_sent(stopped.error):
+        return UnwrittenTurns(stopped.turns, stopped.error)
+    _log_dropped_turns(stopped, len(transactions))
+    if not stopped.turns[0].api_key:
+        return None
+    failed_user: Final = stopped.turns[0].user_id
     other_users: Final = sorted(
         (
             transaction
-            for transaction in failed_suffix[1:]
+            for transaction in stopped.turns[1:]
             if transaction.user_id and transaction.user_id != failed_user
         ),
         key=lambda transaction: transaction.user_id,
     )
-    for _, user_turns in groupby(other_users, key=lambda transaction: transaction.user_id):
-        await _drain_session_partition(
-            prisma_client, tuple(user_turns), n_retry_times, UPSERT_AUTOROUTER_USER_SESSION_SQL
-        )
+    user_partitions: Final = tuple(
+        tuple(user_turns) for _, user_turns in groupby(other_users, key=lambda transaction: transaction.user_id)
+    )
+    for user_partition in user_partitions:
+        if (
+            projection_stopped := await _drain_session_partition(
+                prisma_client, user_partition, n_retry_times, UPSERT_AUTOROUTER_USER_SESSION_SQL
+            )
+        ) is not None:
+            _log_dropped_turns(projection_stopped, len(user_partition))
+    return None
 
 
 async def flush_autorouter_turn_transactions(
     prisma_client: PrismaClient,
     transactions: Sequence[AutoRouterTurnTransaction],
     n_retry_times: int = 3,
-) -> None:
+) -> UnwrittenTurns | None:
     """Drain a queue batch into the rollup, one upsert per turn.
 
     Statements run sequentially in per-session event order: a turn's classification
     depends on the turns before it, and Postgres rejects one multi-row INSERT touching
     the same key twice. Only ConnectError is retried, per statement, because it proves
-    that statement never reached the database. A failed write stops its key and user
-    histories for this batch. Other users sharing that key can still advance their
-    independent user histories, with the key projection disabled and the real key
-    identity preserved. The failed turn is never replayed. Callers must not add their
-    own retry around this function.
+    that statement never reached the database. When a failure proves the database could
+    not take the write at all, the flush stops and returns that turn, the rest of its
+    session and every session not yet attempted, so the caller can queue them again;
+    turns already written are never returned. Any other failed write stops its key and
+    user histories for this batch and is never replayed. Other users sharing that key can
+    still advance their independent user histories, with the key projection disabled and
+    the real key identity preserved. Callers must not add their own retry around this
+    function.
     """
     if not transactions:
-        return
+        return None
     ordered: Final = sorted(
         transactions,
         key=lambda transaction: (*_session_partition(transaction), transaction.turn_at),
     )
-    for _, session_group in groupby(
-        ordered,
-        key=_session_partition,
-    ):
-        await _flush_session_partition(prisma_client, tuple(session_group), n_retry_times)
+    partitions: Final = tuple(tuple(group) for _, group in groupby(ordered, key=_session_partition))
+    for position, partition in enumerate(partitions):
+        if (unwritten := await _flush_session_partition(prisma_client, partition, n_retry_times)) is not None:
+            return UnwrittenTurns(
+                (*unwritten.turns, *chain.from_iterable(partitions[position + 1 :])),
+                unwritten.error,
+            )
+    return None

litellm/proxy/db/exception_handler.py

@@ -776,8 +776,17 @@ def db_write_failure_is_requeue_safe(e: Exception) -> bool:
 
 
 def db_write_never_sent(e: Exception) -> bool:
-    """Whether the database could not take the write at all, because it was unreachable (P1001),
-    refusing connections (57P03) or full (53300, P2024), so even a non-idempotent write can be sent again."""
+    """Whether the database could not take the write at all, because the client had no query engine to send it
+    to (ClientNotConnectedError, NotConnectedError, ConnectError), it was unreachable (P1001), refusing connections
+    (57P03) or full (53300, P2024), so even a non-idempotent write can be sent again."""
+    import prisma
+    import prisma.engine.errors
+
+    unsent_types: Final = _exception_types(
+        prisma.errors.ClientNotConnectedError, prisma.engine.errors.NotConnectedError
+    )
+    if isinstance(e, (*DB_RETRY_SAFE_ERROR_TYPES, *unsent_types)):
+        return True
     if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e):
         return True
     if PrismaDBExceptionHandler.is_database_capacity_error(e):

litellm/proxy/utils.py

@@ -337,7 +337,7 @@ if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
     from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
     from litellm.models.team import LiteLLM_TeamTableCachedObj
-    from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
+    from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction, UnwrittenTurns
     from litellm.proxy.db.baseline_accounting import BaselineAccountingRecord
     from litellm.proxy.db.model_usage_rollup import ModelUsageTransaction
     from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
@@ -5705,7 +5705,6 @@ class PrismaClient:
     _tool_usage_transactions_lock = asyncio.Lock()
     model_usage_transactions: ClassVar[list["ModelUsageTransaction"]] = []
     _model_usage_transactions_lock = asyncio.Lock()
-    autorouter_turn_transactions: ClassVar[list["AutoRouterTurnTransaction"]] = []
     _autorouter_turn_transactions_lock = asyncio.Lock()
 
     # How long a health probe failure waits for an in-flight planned engine
@@ -5724,6 +5723,7 @@ class PrismaClient:
     ):
         ## init logging object
         self.baseline_accounting_transactions: list[BaselineAccountingRecord] = []
+        self.autorouter_turn_transactions: list[AutoRouterTurnTransaction] = []
         self.baseline_accounting_lock: Final = asyncio.Lock()
         self.proxy_logging_obj = proxy_logging_obj
         self.token_auth: DatabaseTokenAuth | None = resolve_database_token_auth()
@@ -9220,25 +9220,27 @@ async def _run_spend_logs_job(
 
     await flush_baseline_accounting(prisma_client)
 
-    async with prisma_client._autorouter_turn_transactions_lock:
+    autorouter_queue_lock: Final = prisma_client._autorouter_turn_transactions_lock
+    async with autorouter_queue_lock:
         autorouter_turns_to_process: Final = prisma_client.autorouter_turn_transactions[:MAX_LOGS_PER_INTERVAL]
         remaining_autorouter_turns: Final = prisma_client.autorouter_turn_transactions[
             len(autorouter_turns_to_process) :
         ]
         prisma_client.autorouter_turn_transactions = remaining_autorouter_turns  # rebind-ok: drain under lock
-    try:
-        from litellm.proxy.db.autorouter_session_rollup import flush_autorouter_turn_transactions
-
-        await flush_autorouter_turn_transactions(
-            prisma_client=prisma_client,
-            transactions=autorouter_turns_to_process,
-        )
-    except Exception as autorouter_tracking_err:  # noqa: BLE001  # a drain bug must not abort the spend job
-        verbose_proxy_logger.error(
-            "Spend tracking - auto-router session rollup drain failed; %s turn transactions dropped: %s",
-            len(autorouter_turns_to_process),
-            autorouter_tracking_err,
+    unwritten_autorouter_turns: Final = await _flush_autorouter_turns(prisma_client, autorouter_turns_to_process)
+    if unwritten_autorouter_turns is not None:
+        async with autorouter_queue_lock:
+            prisma_client.autorouter_turn_transactions = [  # rebind-ok: requeue under lock
+                *unwritten_autorouter_turns.turns,
+                *prisma_client.autorouter_turn_transactions,
+            ]
+        verbose_proxy_logger.warning(
+            "Spend tracking - database unavailable during auto-router session rollup flush, "
+            "requeued %s unwritten turn transactions for the next flush: %s",
+            len(unwritten_autorouter_turns.turns),
+            unwritten_autorouter_turns.error,
         )
+        raise unwritten_autorouter_turns.error
 
     try:
         from litellm.proxy.db.shadow_eval_funnel import flush_shadow_eval_funnel
@@ -9248,6 +9250,22 @@ async def _run_spend_logs_job(
         verbose_proxy_logger.error("Spend tracking - shadow eval funnel drain failed: %s", funnel_err)
 
 
+async def _flush_autorouter_turns(
+    prisma_client: PrismaClient, turns: "Sequence[AutoRouterTurnTransaction]"
+) -> "UnwrittenTurns | None":
+    from litellm.proxy.db.autorouter_session_rollup import flush_autorouter_turn_transactions
+
+    try:
+        return await flush_autorouter_turn_transactions(prisma_client=prisma_client, transactions=turns)
+    except Exception as autorouter_tracking_err:  # noqa: BLE001  # a drain bug must not abort the spend job
+        verbose_proxy_logger.error(
+            "Spend tracking - auto-router session rollup drain failed; %s turn transactions dropped: %s",
+            len(turns),
+            autorouter_tracking_err,
+        )
+        return None
+
+
 MAX_SPEND_LOG_DRAIN_ITERATIONS: Final = 20
 
 

tests/unit/proxy/db/test_autorouter_session_rollup.py

@@ -26,6 +26,13 @@ from litellm.proxy.db.autorouter_session_rollup import (
 )
 from tests.unit.proxy.db.fake_prisma_engine import engine_call
 
+
+def _client_not_connected() -> Exception:
+    from prisma.errors import ClientNotConnectedError
+
+    return ClientNotConnectedError()
+
+
 ROUTING_DECISION = {"router_model_name": "live-auto", "router_type": "complexity", "routed_model": "haiku"}
 
 
@@ -272,7 +279,7 @@ class TestBuildTransaction:
 class _FakeDB:
     def __init__(
         self,
-        failures: "list[Exception] | None" = None,
+        failures: "list[Exception | None] | None" = None,
         poison_session: str | None = None,
         poison_user: str | None = None,
         commit_then_error_users: frozenset[str] = frozenset(),
@@ -290,8 +297,8 @@ class _FakeDB:
             raise RuntimeError("index row size exceeds btree maximum")
         if self._poison_user is not None and params[19] == self._poison_user:
             raise RuntimeError("index row size exceeds btree maximum")
-        if self._failures:
-            raise self._failures.pop(0)
+        if self._failures and (failure := self._failures.pop(0)) is not None:
+            raise failure
         self.calls.append((sql, params))
         if params[19] in self._commit_then_error_users:
             raise RuntimeError("commit succeeded but acknowledgement was lost")
@@ -301,7 +308,7 @@ class _FakeDB:
 class _FakeClient:
     def __init__(
         self,
-        failures: "list[Exception] | None" = None,
+        failures: "list[Exception | None] | None" = None,
         poison_session: str | None = None,
         poison_user: str | None = None,
         commit_then_error_users: frozenset[str] = frozenset(),
@@ -411,7 +418,7 @@ class TestFlush:
                 ("z-other", 10, "k2"),
             )
         )
-        asyncio.run(flush_autorouter_turn_transactions(client, tuple(reversed(turns))))
+        assert asyncio.run(flush_autorouter_turn_transactions(client, tuple(reversed(turns)))) is None
 
         assert client.db.attempts == client.db.calls
         assert [
@@ -444,9 +451,38 @@ class TestFlush:
             _transaction(session_id="s1", at=datetime(2026, 8, 1, 12, 0, 10)),
             _transaction(session_id="s2", at=datetime(2026, 8, 1, 12, 0, 5)),
         ]
-        asyncio.run(flush_autorouter_turn_transactions(client, transactions))
+        assert asyncio.run(flush_autorouter_turn_transactions(client, transactions)) is None
         assert [params[1] for _, params in client.db.calls] == ["s2"]
 
+    @pytest.mark.parametrize(
+        "error",
+        [_client_not_connected(), httpx.ConnectError("connection refused")],
+        ids=["client-not-connected", "connect-error-retries-exhausted"],
+    )
+    def test_an_unsent_turn_is_returned_with_every_turn_after_it_and_never_double_counts(
+        self, error: Exception
+    ) -> None:
+        client: Final = _FakeClient(failures=[None, error])
+        written: Final = _transaction(user_id="a-user", at=datetime(2026, 8, 1, 12, 0, 0))
+        unsent: Final = _transaction(user_id="b-user", at=datetime(2026, 8, 1, 12, 0, 1))
+        same_key_other_user: Final = _transaction(user_id="a-user", at=datetime(2026, 8, 1, 12, 0, 2))
+        other_key: Final = _transaction(api_key="k2", user_id="c-user", at=datetime(2026, 8, 1, 12, 0, 3))
+
+        unwritten: Final = asyncio.run(
+            flush_autorouter_turn_transactions(
+                client, [other_key, same_key_other_user, unsent, written], n_retry_times=0
+            )
+        )
+
+        assert unwritten is not None
+        assert (unwritten.turns, unwritten.error) == ((unsent, same_key_other_user, other_key), error)
+        assert [params[5] for _, params in client.db.calls] == ["2026-08-01T12:00:00"]
+
+        assert asyncio.run(flush_autorouter_turn_transactions(client, unwritten.turns)) is None
+        assert [(sql, params[0], params[5]) for sql, params in client.db.calls] == [
+            (UPSERT_AUTOROUTER_SESSION_SQL, "k1", f"2026-08-01T12:00:0{second}") for second in range(3)
+        ] + [(UPSERT_AUTOROUTER_SESSION_SQL, "k2", "2026-08-01T12:00:03")]
+
     def test_an_empty_batch_writes_nothing(self):
         client = _FakeClient()
         asyncio.run(flush_autorouter_turn_transactions(client, []))

tests/unit/proxy/db/test_exception_handler.py

@@ -6,7 +6,7 @@ from unittest.mock import MagicMock, patch
 import httpx
 import pytest
 from prisma import errors as prisma_errors
-from prisma.engine.errors import BinaryNotFoundError, EngineConnectionError, EngineRequestError
+from prisma.engine.errors import BinaryNotFoundError, EngineConnectionError, EngineRequestError, NotConnectedError
 from prisma.errors import (
     ClientNotConnectedError,
     DataError,
@@ -1100,7 +1100,11 @@ def test_postgres_lifecycle_errors_mean_the_database_is_unavailable(error: Excep
             ),
             True,
         ),
+        (ClientNotConnectedError(), True),
+        (NotConnectedError("Not connected to the query engine"), True),
+        (httpx.ConnectError("connection refused"), True),
         (httpx.ReadTimeout("read timed out"), False),
+        (httpx.ReadError("connection reset"), False),
         (_localized_error("57014"), False),
     ],
 )

tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py

@@ -21,6 +21,11 @@ from unittest.mock import AsyncMock, MagicMock
 import pytest
 
 from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY
+from litellm.proxy.db.autorouter_session_rollup import (
+    UPSERT_AUTOROUTER_SESSION_SQL,
+    UPSERT_AUTOROUTER_USER_SESSION_SQL,
+    AutoRouterTurnTransaction,
+)
 from litellm.proxy.utils import (
     MAX_SPEND_LOG_DRAIN_ITERATIONS,
     drain_spend_logs_queue,
@@ -1001,6 +1006,12 @@ def _postgres_restarting() -> Exception:
     )
 
 
+def _client_not_connected() -> Exception:
+    from prisma.errors import ClientNotConnectedError
+
+    return ClientNotConnectedError()
+
+
 def _postgres_unreachable() -> Exception:
     from prisma.errors import DataError
 
@@ -1257,7 +1268,11 @@ async def test_model_usage_flush_requeues_and_stops_the_job_when_postgres_is_out
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_unreachable()], ids=["57P03", "P1001"])
+@pytest.mark.parametrize(
+    "error",
+    [_postgres_restarting(), _postgres_unreachable(), _client_not_connected()],
+    ids=["57P03", "P1001", "client-not-connected"],
+)
 async def test_rollups_requeue_and_stop_the_job_while_postgres_cannot_take_writes(
     error: Exception,
     mock_prisma_client: MagicMock,
@@ -1284,6 +1299,72 @@ async def test_rollups_requeue_and_stop_the_job_while_postgres_cannot_take_write
     } == {"queue_after": [first, second], "autorouter_queue_after": 1}
 
 
+def _autorouter_turn(second: int) -> AutoRouterTurnTransaction:
+    return AutoRouterTurnTransaction(
+        api_key="k1",
+        session_id="s1",
+        router_name="live-auto",
+        router_type="complexity",
+        model="bedrock/haiku",
+        turn_at=datetime(2026, 10, 2, 12, 0, second),
+        total_tokens=10,
+        spend=0.01,
+        saved_spend=0.0,
+        classifier_cost=0.0,
+        covered=False,
+        cache_hit=False,
+        cache_ttl_seconds=None,
+        cache_touched=True,
+    )
+
+
+@pytest.mark.asyncio
+async def test_unsent_autorouter_turns_go_back_ahead_of_newer_turns_and_stop_the_job(
+    mock_prisma_client: MagicMock,
+) -> None:
+    """A turn the client could not send is kept with the session turns after it, ahead of a turn
+    queued while the flush ran, and the turn already written is not queued again."""
+    written, unsent, dependent, newer = (_autorouter_turn(second) for second in range(4))
+    error: Final = _client_not_connected()
+    turn_statements: Final = frozenset({UPSERT_AUTOROUTER_SESSION_SQL, UPSERT_AUTOROUTER_USER_SESSION_SQL})
+
+    async def execute_raw(sql: str, *params: object) -> int:
+        if sql not in turn_statements:
+            return 0
+        if params[5] == unsent.turn_at.isoformat():
+            raise error
+        mock_prisma_client.autorouter_turn_transactions = [*mock_prisma_client.autorouter_turn_transactions, newer]
+        return 1
+
+    mock_prisma_client.db.execute_raw = AsyncMock(side_effect=execute_raw)
+    mock_prisma_client._autorouter_turn_transactions_lock = asyncio.Lock()
+    mock_prisma_client.autorouter_turn_transactions = [dependent, unsent, written]
+
+    with pytest.raises(type(error)) as raised:
+        await update_spend_logs_job(
+            prisma_client=mock_prisma_client,
+            db_writer_client=None,
+            proxy_logging_obj=MagicMock(),
+        )
+
+    assert {
+        "raised": raised.value,
+        "turn_writes": [
+            (awaited.args[0], awaited.args[6])
+            for awaited in mock_prisma_client.db.execute_raw.await_args_list
+            if awaited.args[0] in turn_statements
+        ],
+        "queue_after": mock_prisma_client.autorouter_turn_transactions,
+    } == {
+        "raised": error,
+        "turn_writes": [
+            (UPSERT_AUTOROUTER_SESSION_SQL, written.turn_at.isoformat()),
+            (UPSERT_AUTOROUTER_SESSION_SQL, unsent.turn_at.isoformat()),
+        ],
+        "queue_after": [unsent, dependent, newer],
+    }
+
+
 @pytest.mark.asyncio
 async def test_rollups_drop_a_batch_whose_session_the_server_terminated(
     mock_prisma_client: MagicMock,