fix(fips): retry the startup marker and state reads once after a dropped connection

A dropped connection during the non-FIPS marker read used to count as an unmarked
database, so a blip at boot let a non-FIPS proxy start on a FIPS database. Both reads
now go through call_with_db_reconnect_retry, the reconnect-and-retry the other startup
reads use. A read that still fails keeps its old outcome.

litellm/repositories/fips_repository.py

@@ -4,6 +4,7 @@ from dataclasses import dataclass
 from typing import Final, Protocol, cast  # noqa: TID251  # Prisma delegates query_raw and execute_raw dynamically
 
 from litellm.proxy.db.create_views import SupportsRawQueries
+from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
 
 FIPS_MARKER: Final = "fips_mode"
 _IS_MARKED: Final = 'SELECT 1 FROM "LiteLLM_Config" WHERE param_name = $1'
@@ -25,6 +26,10 @@ class _WithWriter(Protocol):
     @property
     def writer_db(self) -> object: ...
 
+    async def attempt_db_reconnect(
+        self, *, reason: str, timeout_seconds: float | None = None, lock_timeout_seconds: float | None = None
+    ) -> bool: ...
+
 
 @dataclass(frozen=True, slots=True)
 class FipsDatabaseState:
@@ -33,7 +38,8 @@ class FipsDatabaseState:
 
 
 class FipsRepository:
-    """Reads and writes go to the writer, since a replica may not have the marker yet."""
+    """Reads and writes go to the writer, since a replica may not have the marker yet. A read that loses its
+    connection reconnects and runs once more."""
 
     def __init__(self, prisma_client: _WithWriter) -> None:
         self._prisma_client: Final = prisma_client
@@ -44,10 +50,20 @@ class FipsRepository:
         return cast(SupportsRawQueries, writer)  # cast-ok: Prisma delegates database methods dynamically
 
     async def is_marked_fips(self) -> bool:
-        return bool(await self._writer.query_raw(_IS_MARKED, FIPS_MARKER))
+        return await call_with_db_reconnect_retry(
+            self._prisma_client, self._read_marker, reason="fips_startup_read_marker_failure"
+        )
 
     async def state(self) -> FipsDatabaseState:
         """The marker and whether any other LiteLLM row exists, read in one statement so they agree."""
+        return await call_with_db_reconnect_retry(
+            self._prisma_client, self._read_state, reason="fips_startup_read_state_failure"
+        )
+
+    async def _read_marker(self) -> bool:
+        return bool(await self._writer.query_raw(_IS_MARKED, FIPS_MARKER))
+
+    async def _read_state(self) -> FipsDatabaseState:
         writer: Final = self._writer
         tables: Final = tuple(str(row["table_name"]) for row in await writer.query_raw(_OTHER_TABLES))
         any_row_in_other_tables: Final = "".join(f" OR EXISTS (SELECT 1 FROM {_quoted(table)})" for table in tables)

tests/unit/proxy/common_utils/test_fips.py

@@ -4,6 +4,7 @@ from contextlib import nullcontext
 from typing import Final
 from unittest.mock import Mock
 
+import httpx
 import pytest
 from cryptography.exceptions import InternalError
 from cryptography.hazmat.bindings._rust import openssl as rust_openssl
@@ -135,12 +136,16 @@ def test_enforce_fips_mode(
 class _Database:
     """Prisma's raw-query edge over a database whose only LiteLLM table besides LiteLLM_Config is the key table."""
 
-    def __init__(self, *, marked_fips: bool, holds_rows: bool, readable: bool = True) -> None:
+    def __init__(self, *, marked_fips: bool, holds_rows: bool, readable: bool = True, dropped_reads: int = 0) -> None:
         self.marked_fips = marked_fips
         self.holds_rows = holds_rows
         self.readable = readable
+        self.dropped_reads = dropped_reads
 
     async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]:
+        if self.dropped_reads:
+            self.dropped_reads -= 1
+            raise httpx.ReadError("connection dropped")
         if not self.readable:
             raise RuntimeError('relation "LiteLLM_Config" does not exist')
         if "information_schema.tables" in query:
@@ -158,6 +163,11 @@ class _Prisma:
     def __init__(self, database: _Database) -> None:
         self.writer_db: Final = database
 
+    async def attempt_db_reconnect(
+        self, *, reason: str, timeout_seconds: float | None = None, lock_timeout_seconds: float | None = None
+    ) -> bool:
+        return True
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
@@ -170,6 +180,19 @@ class _Prisma:
         (False, _Database(marked_fips=True, holds_rows=True), "set up with general_settings.fips_mode: true", True),
         (False, _Database(marked_fips=False, holds_rows=True), None, False),
         (False, _Database(marked_fips=True, holds_rows=True, readable=False), None, True),
+        (
+            True,
+            _Database(marked_fips=False, holds_rows=True, dropped_reads=1),
+            "already holds data written without it",
+            False,
+        ),
+        (
+            False,
+            _Database(marked_fips=True, holds_rows=True, dropped_reads=1),
+            "set up with general_settings.fips_mode: true",
+            True,
+        ),
+        (False, _Database(marked_fips=True, holds_rows=True, dropped_reads=2), None, True),
     ],
     ids=[
         "fips-on-new-database",
@@ -179,6 +202,9 @@ class _Prisma:
         "fips-off-fips-database",
         "fips-off-used-database",
         "fips-off-unreadable-database",
+        "fips-on-used-database-after-a-dropped-connection",
+        "fips-off-fips-database-after-a-dropped-connection",
+        "fips-off-fips-database-dropped-twice",
     ],
 )
 async def test_enforce_fips_database(