diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py index e9ec040023c..abbeab7f89d 100644 --- a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -55,3 +55,10 @@ class ToolDiscoveryQueue: """Return and clear all pending rows.""" items, self._pending = list(self._pending.values()), {} return items + + def requeue(self, items: Iterable[ToolDiscoveryQueueItem]) -> None: + """Put rows from a failed write back, ahead of and merged with anything queued since the flush.""" + since_flush: Final = self._pending + self._pending = {} + for item in (*items, *since_flush.values()): + self.add_update(item) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 8fbc8708f8c..dbf636e1630 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -31,6 +31,24 @@ _CONNECTION_CAPACITY_PHRASES: Final = ( "too many connections for database", "remaining connection slots are reserved", ) +_NOT_ACCEPTING_CONNECTIONS_SQLSTATE: Final = "57P03" +_NOT_ACCEPTING_CONNECTIONS_PHRASES: Final = ( + "the database system is starting up", + "the database system is shutting down", + "the database system is not yet accepting connections", + "the database system is not accepting connections", + "the database system is in recovery mode", +) +_SERVER_TERMINATED_SQLSTATES: Final = frozenset({"57P01", "57P02", "57P05"}) +_SERVER_TERMINATED_PHRASES: Final = ( + "terminating connection due to administrator command", + "terminating connection due to immediate shutdown command", + "terminating connection because of unexpected sigquit signal", + "terminating connection because of crash of another server process", + "terminating connection due to idle-session timeout", +) +_UNREACHABLE_PRISMA_CODE: Final = "P1001" +_UNREACHABLE_PHRASES: Final = ("can't reach database server", "cannot reach database server") _PRISMA_POOL_TIMEOUT_CODE: Final = "P2024" _LOCK_TIMEOUT_SQLSTATE: Final = "55P03" _LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout" @@ -139,6 +157,8 @@ class PrismaDBExceptionHandler: return True if PrismaDBExceptionHandler.is_database_capacity_error(e): return True + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True return isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection @staticmethod @@ -289,6 +309,39 @@ class PrismaDBExceptionHandler: error_message: Final = str(e).lower() return any(phrase in error_message for phrase in _CONNECTION_CAPACITY_PHRASES) + @staticmethod + def is_database_not_accepting_connections_error(e: Exception) -> bool: + """True iff Postgres refused the connection because it is starting up, shutting + down or in crash recovery (SQLSTATE 57P03). The refusal comes before any statement + runs, so nothing the write carried was applied.""" + import prisma + + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + return False + if PrismaDBExceptionHandler.postgres_sqlstate(e) == _NOT_ACCEPTING_CONNECTIONS_SQLSTATE: + return True + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _NOT_ACCEPTING_CONNECTIONS_PHRASES) + + @staticmethod + def is_database_lifecycle_error(e: Exception) -> bool: + """True iff Postgres failed the statement because of a server lifecycle event: + a refused connection during startup, shutdown or recovery (57P03), or a session + the server terminated (57P01, 57P02, 57P05). The server comes back on its own, + so the rows are not at fault. A terminated session may have committed before it + was cut, so only ``is_database_not_accepting_connections_error`` proves a + non-idempotent write never applied.""" + import prisma + + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + return False + if PrismaDBExceptionHandler.postgres_sqlstate(e) in _SERVER_TERMINATED_SQLSTATES: + return True + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _SERVER_TERMINATED_PHRASES) + @staticmethod def is_lock_timeout_error(e: Exception) -> bool: """True iff Postgres cancelled a statement for exceeding ``lock_timeout`` @@ -406,6 +459,8 @@ class PrismaDBExceptionHandler: return True if PrismaDBExceptionHandler.is_database_capacity_error(e): return True + if PrismaDBExceptionHandler.is_database_lifecycle_error(e): + return True if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e): return True if "cached plan must not change result type" in str(e).lower(): @@ -680,3 +735,18 @@ def db_write_failure_is_requeue_safe(e: Exception) -> bool: return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES + + +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.""" + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True + if PrismaDBExceptionHandler.is_database_capacity_error(e): + return True + if _prisma_error_code(e) == _UNREACHABLE_PRISMA_CODE: + return True + if not PrismaDBExceptionHandler.is_prisma_error(e): + return False + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _UNREACHABLE_PHRASES) diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index f7ec88945a1..9600fb40ad3 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -23,7 +23,11 @@ from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( ToolDiscoveryQueue, aggregate_tool_discoveries, ) -from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry, db_write_failure_is_requeue_safe +from litellm.proxy.db.exception_handler import ( + call_with_db_reconnect_retry, + db_write_failure_is_requeue_safe, + db_write_never_sent, +) from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions @@ -134,8 +138,19 @@ async def flush_tool_discovery_queue(prisma_client: "PrismaClient", tool_discove return try: await batch_upsert_tools(prisma_client=prisma_client, items=items) - except Exception as e: # noqa: BLE001 # a lost tool row only delays its discovery, the scheduler job must not fail - verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e) + except Exception as e: # noqa: BLE001 # the rows go back to the queue or are dropped, the scheduler job must not fail + if db_write_never_sent(e): + tool_discovery_queue.requeue(items) + verbose_proxy_logger.warning( + "Tool discovery - %d tool rows went back to the queue for the next tick. Error: %s", len(items), str(e) + ) + return + verbose_proxy_logger.error( + "Tool discovery - dropped %d tool rows: the failed statement may have applied or the database " + "refused the data, so re-sending it is not safe. Error: %s", + len(items), + str(e), + ) async def commit_tool_discoveries( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 145c8e38229..8b77e39b3e5 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -213,6 +213,7 @@ from litellm.proxy.db.db_url_settings import ( from litellm.proxy.db.exception_handler import ( PrismaDBExceptionHandler, call_with_db_reconnect_retry, + db_write_never_sent, ) from litellm.proxy.db.health_check_latest import ( LatestHealthCheckRow, @@ -8803,9 +8804,20 @@ async def update_daily_tag_spend( def _rollup_batch_never_applied(e: Exception) -> bool: """True when a drained rollup batch provably never reached its rows: Postgres had no - connection for it (53300) or cancelled it under ``lock_timeout`` before it took the row - lock (55P03). Either way re-sending the batch cannot double-count, so it is requeued.""" - return PrismaDBExceptionHandler.is_database_capacity_error(e) or PrismaDBExceptionHandler.is_lock_timeout_error(e) + connection for it (53300), was unreachable or refusing connections (P1001, 57P03), or + cancelled it under ``lock_timeout`` before it took the row lock (55P03). In each case + re-sending the batch cannot double-count, so it is requeued.""" + return ( + PrismaDBExceptionHandler.is_database_capacity_error(e) + or db_write_never_sent(e) + or PrismaDBExceptionHandler.is_lock_timeout_error(e) + ) + + +def _rollup_failure_stops_job(e: Exception) -> bool: + """True when the database cannot take any write, so the job stops instead of draining + the queues after this one into the same failure.""" + return PrismaDBExceptionHandler.is_database_capacity_error(e) or db_write_never_sent(e) async def update_spend_logs_job( @@ -8896,12 +8908,12 @@ async def _run_spend_logs_job( *prisma_client.tool_usage_transactions, ] verbose_proxy_logger.warning( - "Spend tracking - database out of connections or rollup row lock timed out during tool usage " + "Spend tracking - database unavailable or rollup row lock timed out during tool usage " "flush, requeued %s tool usage transactions for the next flush: %s", len(tool_usage_to_process), tool_tracking_err, ) - if PrismaDBExceptionHandler.is_database_capacity_error(tool_tracking_err): + if _rollup_failure_stops_job(tool_tracking_err): raise else: verbose_proxy_logger.error( @@ -8928,12 +8940,12 @@ async def _run_spend_logs_job( *prisma_client.model_usage_transactions, ] verbose_proxy_logger.warning( - "Spend tracking - database out of connections or rollup row lock timed out during model usage " + "Spend tracking - database unavailable or rollup row lock timed out during model usage " "flush, requeued %s model usage transactions for the next flush: %s", len(model_usage_to_process), model_usage_err, ) - if PrismaDBExceptionHandler.is_database_capacity_error(model_usage_err): + if _rollup_failure_stops_job(model_usage_err): raise else: verbose_proxy_logger.error( @@ -9110,6 +9122,7 @@ def _is_transient_spend_log_write_error(e: Exception) -> bool: return ( PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_prisma_engine_internal_error(e) + or PrismaDBExceptionHandler.is_database_lifecycle_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e) ) diff --git a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py index 0a015255f2e..814fdcb4eda 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -2,7 +2,6 @@ Unit tests for ToolDiscoveryQueue. """ - import pytest from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( @@ -63,3 +62,19 @@ def test_aggregate_sums_call_counts_across_workers_flushes(): {"tool_name": "get_weather", "origin": "user_defined", "call_count": 5}, {"tool_name": "search", "call_count": 1}, ) + + +def test_requeued_rows_go_ahead_of_and_merge_with_sightings_since_the_flush(queue: ToolDiscoveryQueue) -> None: + queue.add_update({"tool_name": "get_weather", "origin": "mcp_server"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + failed_write = queue.flush() + queue.add_update({"tool_name": "translate", "origin": "user_defined"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + + queue.requeue(failed_write) + + assert queue.flush() == [ + {"tool_name": "get_weather", "origin": "mcp_server", "call_count": 1}, + {"tool_name": "search", "origin": "user_defined", "call_count": 2}, + {"tool_name": "translate", "origin": "user_defined", "call_count": 1}, + ] diff --git a/tests/unit/proxy/db/test_exception_handler.py b/tests/unit/proxy/db/test_exception_handler.py index 46704bf52ef..abc0850194d 100644 --- a/tests/unit/proxy/db/test_exception_handler.py +++ b/tests/unit/proxy/db/test_exception_handler.py @@ -22,7 +22,7 @@ from prisma.errors import ( import litellm from litellm.proxy._types import ProxyErrorTypes, ProxyException -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_never_sent # Test is_database_connection_error method @@ -996,3 +996,137 @@ def test_connection_closure_code_survives_prisma_data_error_wrapping() -> None: ) assert PrismaDBExceptionHandler.is_database_transport_error(error) is True assert PrismaDBExceptionHandler.is_permanent_database_fault(error) is False + + +def _connector_error(message: str) -> DataError: + return DataError( + data={"user_facing_error": {"message": f"Error in connector: Error querying the database: FATAL: {message}"}} + ) + + +def _batched_postgres_error(sqlstate: str, message: str) -> DataError: + return DataError( + data={ + "user_facing_error": { + "message": "Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: " + f'None, kind: QueryError(PostgresError {{ code: "{sqlstate}", message: "{message}", severity: "FATAL" }}) }})' + } + } + ) + + +def _localized_error(sqlstate: str) -> DataError: + return DataError(data={"user_facing_error": {"message": "FATAL: Datenbank", "meta": {"code": sqlstate}}}) + + +_NOT_ACCEPTING_CONNECTIONS: Final = ( + _connector_error("the database system is starting up"), + _connector_error("the database system is shutting down"), + _connector_error("the database system is not yet accepting connections"), + _connector_error("the database system is not accepting connections"), + _connector_error("the database system is in recovery mode"), + _localized_error("57P03"), + _batched_postgres_error("57P03", "the database system is starting up"), + EngineRequestError( + MagicMock(status=500), + '{"is_panic":false,"message":"Error in connector: Error querying the database: FATAL: the database system is ' + 'starting up","backtrace":null}', + ), + RawQueryError( + data={ + "user_facing_error": { + "message": "Raw query failed. Code: `57P03`. Message: `db error: FATAL: the database system is starting up`", + "meta": {"code": "57P03", "message": "FATAL: the database system is starting up"}, + "error_code": "P2010", + } + } + ), +) +_SERVER_TERMINATED: Final = ( + _connector_error("terminating connection due to administrator command"), + _connector_error("terminating connection due to immediate shutdown command"), + _connector_error("terminating connection because of unexpected SIGQUIT signal"), + _connector_error("terminating connection because of crash of another server process"), + _connector_error("terminating connection due to idle-session timeout"), + _localized_error("57P01"), + _localized_error("57P02"), + _localized_error("57P05"), + _batched_postgres_error("57P01", "terminating connection due to administrator command"), +) + + +@pytest.mark.parametrize( + ("error", "connection_error"), + [ + *((error, True) for error in _NOT_ACCEPTING_CONNECTIONS), + *((error, False) for error in _SERVER_TERMINATED), + ], +) +def test_postgres_lifecycle_errors_mean_the_database_is_unavailable(error: Exception, connection_error: bool) -> None: + """A server that is restarting, shutting down or cutting sessions is a transient outage: + auth answers 503 rather than 401, and the spend-log writer must not take the rows for poison. + Only a server refusing connections opens the ``allow_requests_on_db_unavailable`` fallback, + one terminated session does not prove the database is down.""" + assert { + "lifecycle": PrismaDBExceptionHandler.is_database_lifecycle_error(error), + "service_unavailable": PrismaDBExceptionHandler.is_database_service_unavailable_error(error), + "connection_error": PrismaDBExceptionHandler.is_database_connection_error(error), + "permanent": PrismaDBExceptionHandler.is_permanent_database_fault(error), + } == { + "lifecycle": True, + "service_unavailable": True, + "connection_error": connection_error, + "permanent": False, + } + + +@pytest.mark.parametrize( + ("error", "never_sent"), + [ + *((error, True) for error in _NOT_ACCEPTING_CONNECTIONS), + *((error, False) for error in _SERVER_TERMINATED), + (DataError(data={"user_facing_error": {"message": "Can't reach database server at `db`:`5432`"}}), True), + (DataError(data={"user_facing_error": {"error_code": "P1001", "message": "unreachable"}}), True), + (_capacity_error("sorry, too many clients already"), True), + (_localized_error("53300"), True), + ( + DataError( + data={ + "user_facing_error": { + "error_code": "P2024", + "message": "Timed out fetching a new connection from the connection pool", + } + } + ), + True, + ), + (httpx.ReadTimeout("read timed out"), False), + (_localized_error("57014"), False), + ], +) +def test_db_write_never_sent_only_for_failures_before_any_statement_ran(error: Exception, never_sent: bool) -> None: + """A refused, unreachable or full server provably applied nothing, while a terminated session + or a timeout may have committed before the failure, so only the former lets a non-idempotent + write be replayed.""" + assert db_write_never_sent(error) is never_sent + + +@pytest.mark.parametrize( + "error", + [ + _localized_error("57P04"), + _localized_error("57014"), + _localized_error("23505"), + DataError(data={"user_facing_error": {"message": "invalid byte sequence for encoding UTF8: 0x00"}}), + RuntimeError("the database system is starting up"), + ], +) +def test_other_failures_are_not_postgres_lifecycle_errors(error: Exception) -> None: + assert PrismaDBExceptionHandler.is_database_lifecycle_error(error) is False + assert db_write_never_sent(error) is False + + +def test_lifecycle_error_is_seen_through_a_wrapping_exception() -> None: + wrapped: Final = RuntimeError("spend flush failed") + wrapped.__cause__ = _connector_error("the database system is starting up") + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(wrapped) is True diff --git a/tests/unit/proxy/db/test_tool_registry_writer.py b/tests/unit/proxy/db/test_tool_registry_writer.py index 76898b0eeab..a5ad897913e 100644 --- a/tests/unit/proxy/db/test_tool_registry_writer.py +++ b/tests/unit/proxy/db/test_tool_registry_writer.py @@ -69,9 +69,7 @@ def _make_prisma( prisma.db.litellm_tooltable = MagicMock() prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return) prisma.db.execute_raw = AsyncMock(return_value=1) - prisma.db.litellm_tooltable.find_many = AsyncMock( - return_value=find_many_rows if find_many_rows is not None else [] - ) + prisma.db.litellm_tooltable.find_many = AsyncMock(return_value=find_many_rows if find_many_rows is not None else []) prisma.db.litellm_tooltable.find_unique = AsyncMock(return_value=find_unique_row) prisma.db.litellm_verificationtoken = MagicMock() prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=list(key_rows)) @@ -227,9 +225,7 @@ async def test_get_tool_found(): result = await get_tool(prisma, "my_tool") assert result is not None assert result.tool_name == "my_tool" - prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with( - where={"tool_name": "my_tool"} - ) + prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with(where={"tool_name": "my_tool"}) @pytest.mark.asyncio @@ -266,9 +262,7 @@ async def test_update_tool_policy_calls_upsert_then_get_tool(): updated_by="admin", ) prisma = _make_prisma(find_unique_row=row) - result = await update_tool_policy( - prisma, "my_tool", updated_by="admin", input_policy="blocked" - ) + result = await update_tool_policy(prisma, "my_tool", updated_by="admin", input_policy="blocked") assert result is not None assert result.input_policy == "blocked" prisma.db.litellm_tooltable.upsert.assert_awaited_once() @@ -567,3 +561,42 @@ async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])] expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else [] assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "requeued"), + [ + (_postgres_rejection("57P03"), True), + (_postgres_rejection("53300"), True), + (_postgres_rejection("57P01"), False), + (_postgres_rejection("54000"), False), + (_postgres_rejection("23502"), False), + ], + ids=[ + "restarting_server_keeps_the_rows", + "full_server_keeps_the_rows", + "terminated_session_may_have_counted_so_the_rows_are_dropped", + "oversized_value_is_dropped", + "rejected_row_is_dropped", + ], +) +async def test_failed_write_without_redis_keeps_the_rows_in_the_queue_only_when_they_provably_did_not_apply( + failure: Exception, requeued: bool +): + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(side_effect=failure) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=_InMemoryRedis()), + pod_lock_manager=PodLockManager(redis_cache=_InMemoryRedis()), + use_redis_buffer=False, + ) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + + expected = [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2 if requeued else 1}] + assert queue.flush() == expected diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py index fc7ae531c34..293f2d49f43 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -23,9 +23,9 @@ import pytest from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY from litellm.proxy.utils import ( MAX_SPEND_LOG_DRAIN_ITERATIONS, + drain_spend_logs_queue, monitor_spend_logs_queue, raise_failed_update_spend_exception, - drain_spend_logs_queue, recover_parked_spend_logs, update_daily_tag_spend, update_spend, @@ -985,6 +985,85 @@ async def test_update_spend_logs_job_requeues_whole_batch_when_postgres_is_out_o } == {"create_many_calls": 1, "queue_after": ["r1", "r2", "r3"], "backoff_sleeps": []} +def _postgres_restarting() -> Exception: + """prisma's shape for Postgres SQLSTATE 57P03 while the server runs crash recovery after + a restart: a base ``DataError`` carrying only the connector message.""" + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Error in connector: Error querying the database: FATAL: the database system is starting up", + "backtrace": None, + } + } + ) + + +def _postgres_unreachable() -> Exception: + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Can't reach database server at `db.internal`:`5432`", + "error_code": "P1001", + } + } + ) + + +def _postgres_terminated_session() -> Exception: + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Error in connector: Error querying the database: FATAL: terminating connection due to " + "administrator command", + "meta": {"code": "57P01", "message": "terminating connection due to administrator command"}, + } + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_terminated_session()], ids=["57P03", "57P01"]) +async def test_update_spend_logs_job_keeps_every_row_while_postgres_restarts( + error: Exception, + mock_prisma_client: MagicMock, + make_spend_log_row: Callable[..., dict[str, object]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A restarting server refuses or cuts every statement, so the batch is not bisected + into per-row drops: it is retried whole and then goes back to the queue head.""" + monkeypatch.setattr(asyncio, "sleep", AsyncMock()) + proxy_logging: Final = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [ + make_spend_log_row(request_id="r1"), + make_spend_log_row(request_id="r2"), + make_spend_log_row(request_id="r3"), + make_spend_log_row(request_id="r4"), + ] + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=error) + + with suppress(Exception): + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + ) + + assert { + "create_many_calls": mock_prisma_client.db.litellm_spendlogs.create_many.await_count, + "queue_after": [row["request_id"] for row in mock_prisma_client.spend_log_transactions], + } == {"create_many_calls": 4, "queue_after": ["r1", "r2", "r3", "r4"]} + + def _tool_usage_transaction(request_id: str) -> object: from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -1177,6 +1256,53 @@ async def test_model_usage_flush_requeues_and_stops_the_job_when_postgres_is_out } == {"queue_after": [first, second], "autorouter_queue_after": 1} +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_unreachable()], ids=["57P03", "P1001"]) +async def test_rollups_requeue_and_stop_the_job_while_postgres_cannot_take_writes( + error: Exception, + mock_prisma_client: MagicMock, +) -> None: + """A server that is restarting or unreachable never saw the rollup batch, so the tool + usage increments are kept for the next flush, and the job stops before draining the + auto-router queue into the same outage.""" + mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=error) + mock_prisma_client.spend_log_transactions = [] + first, second = _tool_usage_transaction("r1"), _tool_usage_transaction("r2") + mock_prisma_client.tool_usage_transactions = [first, second] + mock_prisma_client.autorouter_turn_transactions = [MagicMock(name="autorouter_turn")] + + with suppress(Exception): + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + + assert { + "queue_after": mock_prisma_client.tool_usage_transactions, + "autorouter_queue_after": len(mock_prisma_client.autorouter_turn_transactions), + } == {"queue_after": [first, second], "autorouter_queue_after": 1} + + +@pytest.mark.asyncio +async def test_rollups_drop_a_batch_whose_session_the_server_terminated( + mock_prisma_client: MagicMock, +) -> None: + """A terminated session (57P01) may have committed before it was cut, and the rollup + increments are not idempotent, so the batch is dropped rather than counted twice.""" + mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=_postgres_terminated_session()) + mock_prisma_client.spend_log_transactions = [] + mock_prisma_client.tool_usage_transactions = [_tool_usage_transaction("r1")] + + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.tool_usage_transactions == [] + + @pytest.mark.asyncio async def test_model_usage_flush_still_drops_a_batch_that_may_have_applied( mock_prisma_client: MagicMock,