fix: harden QueuePool hygiene with atomic last-used debounce

Use Redis SET NX for last-used debounce, require the request session in
resolve_current_user, and avoid double-commit when TierService owns the session.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: suguanYang

apps/api/app/services/auth/api_key_authentication_service.py

@@ -174,21 +174,28 @@ async def _update_last_used_best_effort(
     ) -> None:
         """Update last_used_at on the request session without a nested checkout.
 
-        Redis debounce skips redundant writes within the debounce window so job
-        polls do not compete for QueuePool capacity via create_task+get_db_context.
+        Redis SET NX debounce skips redundant writes within the debounce window
+        so job polls do not compete for QueuePool via create_task+get_db_context.
         """
         debounce_key = self._get_last_used_debounce_key(api_key_id)
         try:
-            if await redis_service.exists(debounce_key):
+            acquired = await redis_service.set_nx(
+                debounce_key,
+                "1",
+                ex=_LAST_USED_DEBOUNCE_SECONDS,
+            )
+            if not acquired:
                 return
         except Exception:
             logger.warning(
-                "api_key_authentication: failed to read last-used debounce for api_key_id={}",
+                "api_key_authentication: last-used debounce failed for api_key_id={}; "
+                "updating anyway",
                 api_key_id,
             )
 
         try:
             await self._repository.update_last_used(session, api_key_id)
+            # Request-scoped sessions do not auto-commit.
             await session.commit()
         except Exception as exc:
             logger.warning(
@@ -198,16 +205,3 @@ async def _update_last_used_best_effort(
                 await session.rollback()
             except Exception:
                 pass
-            return
-
-        try:
-            await redis_service.set(
-                debounce_key,
-                "1",
-                ttl=_LAST_USED_DEBOUNCE_SECONDS,
-            )
-        except Exception:
-            logger.warning(
-                "api_key_authentication: failed to write last-used debounce for api_key_id={}",
-                api_key_id,
-            )

apps/api/app/services/rate_limit/job_admission_service.py

@@ -33,7 +33,7 @@ async def resolve_current_user(
         *,
         route_context: RouteAdmissionContext,
         user_id: str,
-        db: AsyncSession | None = None,
+        db: AsyncSession,
     ) -> CurrentUser:
         user_tier = await TierService.get_tier(user_id, session=db)
         self._route_policy_service.enforce_guest_api_key_scope(

apps/api/app/services/rate_limit/tier_service.py

@@ -49,19 +49,29 @@ async def get_tier(
             return cached_tier
 
         if session is not None:
-            user_tier = await TierService._resolve_tier_from_db(session, user_id)
+            user_tier = await TierService._resolve_tier_from_db(
+                session,
+                user_id,
+                commit_on_initialize=True,
+            )
         else:
             async with get_db_context() as owned_session:
                 user_tier = await TierService._resolve_tier_from_db(
                     owned_session,
                     user_id,
+                    commit_on_initialize=False,
                 )
 
         await TierService._set_cached_tier(redis_service, user_id, user_tier)
         return user_tier
 
     @staticmethod
-    async def _resolve_tier_from_db(session: AsyncSession, user_id: str) -> str:
+    async def _resolve_tier_from_db(
+        session: AsyncSession,
+        user_id: str,
+        *,
+        commit_on_initialize: bool,
+    ) -> str:
         """Load tier from DB, initializing missing first-use billing state."""
         try:
             return await TierService._get_tier_from_db(session, user_id)
@@ -70,9 +80,9 @@ async def _resolve_tier_from_db(session: AsyncSession, user_id: str) -> str:
                 session,
                 user_id,
             )
-            # Persist first-use init when reusing the request session
-            # (get_db_context commits on exit when we own the session).
-            await session.commit()
+            # Request-scoped sessions do not auto-commit; get_db_context does.
+            if commit_on_initialize:
+                await session.commit()
             return user_tier
 
     @staticmethod

apps/api/tests/unit/test_job_poll_session_hygiene.py

@@ -98,7 +98,7 @@ async def test_validate_api_key_updates_last_used_on_same_session() -> None:
     session.commit = AsyncMock()
     redis_service = AsyncMock()
     redis_service.get = AsyncMock(return_value=None)
-    redis_service.exists = AsyncMock(return_value=False)
+    redis_service.set_nx = AsyncMock(return_value=True)
     redis_service.set = AsyncMock()
     redis_service.sadd = AsyncMock()
     redis_service.ttl = AsyncMock(return_value=-2)
@@ -125,10 +125,6 @@ async def test_validate_api_key_updates_last_used_on_same_session() -> None:
             "app.services.auth.api_key_authentication_service.hash_api_key",
             return_value="hash-1",
         ),
-        patch(
-            "app.services.auth.api_key_authentication_service.get_db_context",
-            create=True,
-        ) as get_db_context_mock,
         patch("asyncio.create_task") as create_task_mock,
     ):
         user_id = await service.validate_api_key(session, "kw_test_key")
@@ -137,11 +133,10 @@ async def test_validate_api_key_updates_last_used_on_same_session() -> None:
     repository.update_last_used.assert_awaited_once_with(session, "key-1")
     session.commit.assert_awaited_once()
     create_task_mock.assert_not_called()
-    get_db_context_mock.assert_not_called()
-    redis_service.set.assert_any_await(
+    redis_service.set_nx.assert_awaited_once_with(
         "api-key:last-used-debounce:key-1",
         "1",
-        ttl=300,
+        ex=300,
     )
 
 
@@ -151,7 +146,7 @@ async def test_validate_api_key_skips_last_used_when_debounced() -> None:
     session.commit = AsyncMock()
     redis_service = AsyncMock()
     redis_service.get = AsyncMock(return_value=None)
-    redis_service.exists = AsyncMock(return_value=True)
+    redis_service.set_nx = AsyncMock(return_value=False)
     redis_service.set = AsyncMock()
     redis_service.sadd = AsyncMock()
     redis_service.ttl = AsyncMock(return_value=-2)
@@ -203,7 +198,7 @@ async def __aexit__(self, *args: object) -> None:
 
     redis_service = AsyncMock()
     redis_service.get = AsyncMock(side_effect=[None, None])  # api-key miss, tier miss
-    redis_service.exists = AsyncMock(return_value=False)
+    redis_service.set_nx = AsyncMock(return_value=True)
     redis_service.set = AsyncMock()
     redis_service.sadd = AsyncMock()
     redis_service.ttl = AsyncMock(return_value=-2)

packages/shared-python/shared/services/redis/redis_service.py

@@ -111,6 +111,28 @@ async def _operation():
                 original_exception=e,
             )
 
+    async def set_nx(self, key: str, value: str, ex: int) -> bool:
+        """Atomic SET NX EX — set only if the key does not already exist.
+
+        Returns ``True`` if the key was written, ``False`` if it already existed.
+        Does not JSON-encode the value and does not fall back to a default TTL.
+        """
+        try:
+            client = await self._get_client()
+            full_key = self._build_key(key)
+
+            async def _operation():
+                return await client.set(full_key, value, nx=True, ex=ex)
+
+            return bool(await self._execute_with_retry(_operation))
+        except Exception as e:
+            logger.error(f"Redis SET NX operation failed: {e}")
+            raise RedisOperationError(
+                internal_message=f"SET NX operation failed: {str(e)}",
+                operation="SET_NX",
+                original_exception=e,
+            )
+
     async def get(self, key: str, default: Any = None) -> Any:
         """Get a key value."""
         try: