From dedc219f8ede94239d5dff19ecbbce447e9b46b7 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrish+github@berri.ai>
Date: Mon, 20 Apr 2026 07:38:51 -0700
Subject: [PATCH] fix: minor improvements

---
 .../adaptive_router_update_queue.py           | 35 ++++++-------------
 litellm/proxy/proxy_server.py                 |  4 ++-
 .../adaptive_router/adaptive_router.py        | 11 ++++--
 3 files changed, 23 insertions(+), 27 deletions(-)

diff --git a/litellm/proxy/db/db_transaction_queue/adaptive_router_update_queue.py b/litellm/proxy/db/db_transaction_queue/adaptive_router_update_queue.py
index 3a76370e7d..d1e275a076 100644
--- a/litellm/proxy/db/db_transaction_queue/adaptive_router_update_queue.py
+++ b/litellm/proxy/db/db_transaction_queue/adaptive_router_update_queue.py
@@ -107,24 +107,11 @@ class AdaptiveRouterUpdateQueue:
             router, rt, model = key
             payload = batch[key]
             try:
-                existing = (
-                    await prisma_client.db.litellm_adaptiverouterstate.find_unique(
-                        where={
-                            "router_name_request_type_model_name": {
-                                "router_name": router,
-                                "request_type": rt,
-                                "model_name": model,
-                            }
-                        }
-                    )
-                )
-                new_alpha = (existing.alpha if existing else 0.0) + payload[
-                    "delta_alpha"
-                ]
-                new_beta = (existing.beta if existing else 0.0) + payload["delta_beta"]
-                new_samples = (existing.total_samples if existing else 0) + int(
-                    payload["samples_added"]
-                )
+                # Atomic increment: push the delta directly into the DB so
+                # concurrent flushers from multiple pods don't overwrite each
+                # other. The upsert creates the row with the delta as the
+                # initial value on first write, then increments on subsequent
+                # writes — no read-modify-write race.
                 await prisma_client.db.litellm_adaptiverouterstate.upsert(
                     where={
                         "router_name_request_type_model_name": {
@@ -138,14 +125,14 @@ class AdaptiveRouterUpdateQueue:
                             "router_name": router,
                             "request_type": rt,
                             "model_name": model,
-                            "alpha": new_alpha,
-                            "beta": new_beta,
-                            "total_samples": new_samples,
+                            "alpha": payload["delta_alpha"],
+                            "beta": payload["delta_beta"],
+                            "total_samples": int(payload["samples_added"]),
                         },
                         "update": {
-                            "alpha": new_alpha,
-                            "beta": new_beta,
-                            "total_samples": new_samples,
+                            "alpha": {"increment": payload["delta_alpha"]},
+                            "beta": {"increment": payload["delta_beta"]},
+                            "total_samples": {"increment": int(payload["samples_added"])},
                         },
                     },
                 )
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 67a3414d0b..8d5e9bc0fe 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -952,8 +952,10 @@ async def proxy_startup_event(app: FastAPI):  # noqa: PLR0915
             _run_background_health_check()
         )  # start the background health check coroutine.
 
-    # Start adaptive-router queue flusher if any AdaptiveRouter is configured.
+    # Start adaptive-router queue flusher and load persisted state if any AdaptiveRouter is configured.
     if llm_router is not None and getattr(llm_router, "adaptive_routers", None):
+        for _ar in llm_router.adaptive_routers.values():
+            await _ar.load_state_from_db(prisma_client)
         asyncio.create_task(_adaptive_router_flusher_loop())
 
     ## [Optional] Initialize dd tracer
diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py
index 2f3adccad7..b7f7722e0a 100644
--- a/litellm/router_strategy/adaptive_router/adaptive_router.py
+++ b/litellm/router_strategy/adaptive_router/adaptive_router.py
@@ -297,7 +297,9 @@ class AdaptiveRouter:
         """Apply one turn, push session snapshot + bandit deltas to the queue."""
         state = self.get_or_create_session_state(session_id, model_name, request_type)
         delta = apply_turn(state, turn)
-        print("CALLS DELTA", delta)
+        verbose_router_logger.debug(
+            "AdaptiveRouter[%s]: record_turn delta=%s", self.router_name, delta
+        )
 
         snapshot = asdict(state)
         await self.queue.add_session_state(
@@ -305,7 +307,12 @@ class AdaptiveRouter:
         )
 
         d_alpha, d_beta = self._compute_bandit_delta(delta)
-        print("CALLS D_ALPHA", d_alpha)
+        verbose_router_logger.debug(
+            "AdaptiveRouter[%s]: bandit delta alpha=%.2f beta=%.2f",
+            self.router_name,
+            d_alpha,
+            d_beta,
+        )
         if d_alpha != 0 or d_beta != 0:
             # For non-GENERAL turns, attribute to the current-turn classification
             # so genuine mid-session topic shifts (e.g. code → math) update the