fix(proxy): resolve k8s-era TC regressions (#290)

* fix(shutdown): reset global lifecycle state between app lifespans

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

* fix(proxy): bound balancer retries and bridge reader teardown

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

* test(helm): build chart dependencies before rendering

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

* fix(tests): stabilize bridge errors and postgres rate limits

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

* test(bridge): relax previous-response continuity assertions

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

* fix(proxy): clear stale sticky selections at retry cap

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

* fix(bridge): abort reconnect when reader shutdown stalls

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Author: Soju06

app/core/rate_limiter/db_rate_limiter.py

@@ -2,7 +2,7 @@
 
 from datetime import UTC, datetime, timedelta
 
-from sqlalchemy import delete, func, select, text
+from sqlalchemy import delete, func, insert, literal, select, text
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from app.core.exceptions import DashboardRateLimitError
@@ -71,22 +71,29 @@ async def check_and_increment(self, key: str, session: AsyncSession) -> None:
 
         now = datetime.now(UTC)
         window_start = now - timedelta(seconds=self.window_seconds)
+        key_column = RateLimitAttempt.__table__.c.key
+        type_column = RateLimitAttempt.__table__.c.type
+        attempted_at_column = RateLimitAttempt.__table__.c.attempted_at
 
         raw_result = await session.execute(
-            text(
-                "INSERT INTO rate_limit_attempts (key, type, attempted_at) "
-                "SELECT :key, :type, :now "
-                "WHERE (SELECT COUNT(*) FROM rate_limit_attempts "
-                "       WHERE key = :key AND type = :type AND attempted_at >= :window_start"
-                "      ) < :max_attempts"
-            ),
-            {
-                "key": key,
-                "type": self.type,
-                "now": now,
-                "window_start": window_start,
-                "max_attempts": self.max_attempts,
-            },
+            insert(RateLimitAttempt).from_select(
+                [key_column.name, type_column.name, attempted_at_column.name],
+                select(
+                    literal(key, type_=key_column.type),
+                    literal(self.type, type_=type_column.type),
+                    literal(now, type_=attempted_at_column.type),
+                ).where(
+                    select(func.count())
+                    .select_from(RateLimitAttempt)
+                    .where(
+                        RateLimitAttempt.key == key,
+                        RateLimitAttempt.type == self.type,
+                        RateLimitAttempt.attempted_at >= window_start,
+                    )
+                    .scalar_subquery()
+                    < self.max_attempts
+                ),
+            )
         )
         inserted = raw_result.rowcount > 0  # type: ignore[union-attr]
         await session.commit()

app/core/shutdown.py

@@ -8,6 +8,13 @@
 _in_flight: int = 0
 
 
+def reset() -> None:
+    global _draining, _bridge_drain_active, _in_flight
+    _draining = False
+    _bridge_drain_active = False
+    _in_flight = 0
+
+
 def set_draining(val: bool = True) -> None:
     global _draining
     _draining = val

app/main.py

@@ -97,8 +97,7 @@ async def lifespan(app: FastAPI):
     instance_id = None
 
     startup_module._startup_complete = False
-    shutdown_state.set_draining(False)
-    shutdown_state.set_bridge_drain_active(False)
+    shutdown_state.reset()
     await get_settings_cache().invalidate()
     await get_rate_limit_headers_cache().invalidate()
     reload_additional_quota_registry()
@@ -267,6 +266,7 @@ async def _heartbeat_only(svc: RingMembershipService, iid: str) -> None:
             except Exception:
                 logger.exception("Metrics server stopped with an error")
             finally:
+                shutdown_state.reset()
                 mark_process_dead()
                 await close_db()
 

app/modules/proxy/load_balancer.py

@@ -39,6 +39,8 @@
 
 logger = logging.getLogger(__name__)
 
+_MAX_SELECTION_ATTEMPTS = 4
+
 _STICKY_GRACE_PERIOD_SECONDS = 10.0
 _RECOVERABLE_STATUSES = frozenset(
     {
@@ -140,12 +142,13 @@ async def load_selection_inputs() -> _SelectionInputs:
             )
 
         selected_snapshot: Account | None = None
-        selected_version_at_pick: int = 0
         error_message: str | None = None
         selected_states: list[AccountState] = []
         selected_account_map: dict[str, Account] = {}
         if sticky_key is None:
+            attempt = 0
             while True:
+                attempt += 1
                 self._prune_runtime(selection_inputs.accounts)
                 states, account_map = _build_states(
                     accounts=selection_inputs.accounts,
@@ -189,11 +192,20 @@ async def load_selection_inputs() -> _SelectionInputs:
                         selected_snapshot.status = result.account.status
                         selected_snapshot.deactivation_reason = result.account.deactivation_reason
                         selected_snapshot.reset_at = selected_reset_at
-                        selected_version_at_pick = self._runtime.get(selected.id, RuntimeState()).version
                 else:
                     error_message = result.error_message
 
-                pre_persist_versions = {aid: runtime.version for aid, runtime in self._runtime.items()}
+                pre_persist_runtime_state = {
+                    aid: (
+                        runtime.reset_at,
+                        runtime.cooldown_until,
+                        runtime.error_count,
+                        runtime.last_error_at,
+                    )
+                    for aid, runtime in self._runtime.items()
+                }
+                pre_persist_cache_generation = self._selection_inputs_cache.generation
+
                 async with self._repo_factory() as repos:
                     stale_account_ids = await self._persist_selection_state(
                         repos.accounts,
@@ -202,6 +214,10 @@ async def load_selection_inputs() -> _SelectionInputs:
                     )
                 stale_account_ids = stale_account_ids or set()
                 if selected_snapshot is not None and selected_snapshot.id in stale_account_ids:
+                    if attempt >= _MAX_SELECTION_ATTEMPTS:
+                        selected_snapshot = None
+                        error_message = None
+                        break
                     selection_inputs = await load_selection_inputs()
                     if selection_inputs.error_code is not None and not selection_inputs.accounts:
                         return AccountSelection(
@@ -215,37 +231,53 @@ async def load_selection_inputs() -> _SelectionInputs:
                     selected_account_map = {}
                     continue
 
-                if selected_snapshot is not None:
-                    _sel_runtime = self._runtime.get(selected_snapshot.id)
-                    _sel_pre_ver = selected_version_at_pick
-                    if _sel_runtime is not None and _sel_runtime.version != _sel_pre_ver:
+                if (
+                    selected_snapshot is not None
+                    and self._selection_inputs_cache.generation != pre_persist_cache_generation
+                    and attempt < _MAX_SELECTION_ATTEMPTS
+                ):
+                    selection_inputs = await load_selection_inputs()
+                    if selection_inputs.error_code is not None and not selection_inputs.accounts:
+                        return AccountSelection(
+                            account=None,
+                            error_message=selection_inputs.error_message,
+                            error_code=selection_inputs.error_code,
+                        )
+                    selected_snapshot = None
+                    error_message = None
+                    selected_states = []
+                    selected_account_map = {}
+                    await asyncio.sleep(0)
+                    continue
+
+                if selected_snapshot is None and error_message == "No available accounts":
+                    runtime_recovered = any(
+                        self._runtime.get(account_id, RuntimeState()).reset_at != before[0]
+                        or self._runtime.get(account_id, RuntimeState()).cooldown_until != before[1]
+                        or self._runtime.get(account_id, RuntimeState()).error_count != before[2]
+                        or self._runtime.get(account_id, RuntimeState()).last_error_at != before[3]
+                        for account_id, before in pre_persist_runtime_state.items()
+                    )
+                    if runtime_recovered and attempt < _MAX_SELECTION_ATTEMPTS:
                         selection_inputs = await load_selection_inputs()
                         if selection_inputs.error_code is not None and not selection_inputs.accounts:
                             return AccountSelection(
                                 account=None,
                                 error_message=selection_inputs.error_message,
                                 error_code=selection_inputs.error_code,
                             )
-                        selected_snapshot = None
                         error_message = None
                         selected_states = []
                         selected_account_map = {}
+                        await asyncio.sleep(0)
                         continue
 
-                if selected_snapshot is None and error_message == "No available accounts":
-                    runtime_recovered = any(
-                        self._runtime.get(aid, RuntimeState()).version != pre_persist_versions.get(aid, 0)
-                        for aid in account_map
-                    )
-                    if runtime_recovered:
-                        error_message = None
-                        selected_states = []
-                        selected_account_map = {}
-                        pre_persist_versions = {aid: runtime.version for aid, runtime in self._runtime.items()}
-                        continue
                 break
+
         else:
+            attempt = 0
             while True:
+                attempt += 1
                 self._prune_runtime(selection_inputs.accounts)
                 states, account_map = _build_states(
                     accounts=selection_inputs.accounts,
@@ -304,17 +336,20 @@ async def load_selection_inputs() -> _SelectionInputs:
                     )
                 stale_account_ids = stale_account_ids or set()
                 if selected_snapshot is not None and selected_snapshot.id in stale_account_ids:
+                    selected_snapshot = None
+                    error_message = None
+                    selected_states = []
+                    selected_account_map = {}
+                    if attempt >= _MAX_SELECTION_ATTEMPTS:
+                        break
                     selection_inputs = await load_selection_inputs()
                     if selection_inputs.error_code is not None and not selection_inputs.accounts:
                         return AccountSelection(
                             account=None,
                             error_message=selection_inputs.error_message,
                             error_code=selection_inputs.error_code,
                         )
-                    selected_snapshot = None
-                    error_message = None
-                    selected_states = []
-                    selected_account_map = {}
+                    await asyncio.sleep(0)
                     continue
                 break
 

app/modules/proxy/service.py

@@ -109,6 +109,26 @@
 
 logger = logging.getLogger(__name__)
 
+_TASK_CANCEL_TIMEOUT_SECONDS = 1.0
+
+
+async def _await_cancelled_task(
+    task: asyncio.Task[object] | asyncio.Task[None],
+    *,
+    timeout_seconds: float = _TASK_CANCEL_TIMEOUT_SECONDS,
+    label: str,
+) -> bool:
+    task.cancel()
+    try:
+        await asyncio.wait_for(task, timeout=timeout_seconds)
+    except asyncio.CancelledError:
+        return True
+    except TimeoutError:
+        logger.warning("Timed out waiting for %s cancellation", label)
+        return False
+    return True
+
+
 _TEXT_DELTA_EVENT_TYPES = frozenset({"response.output_text.delta", "response.refusal.delta"})
 _TEXT_DONE_CONTENT_PART_TYPES = frozenset({"output_text", "refusal"})
 _REQUEST_TRANSPORT_HTTP = "http"
@@ -1007,11 +1027,7 @@ async def proxy_responses_websocket(
                         response_create_gate=response_create_gate,
                     )
                     if upstream_reader is not None:
-                        upstream_reader.cancel()
-                        try:
-                            await upstream_reader
-                        except asyncio.CancelledError:
-                            pass
+                        await _await_cancelled_task(upstream_reader, label="proxy websocket upstream reader")
                         upstream_reader = None
                     upstream_control = None
                     if upstream is not None:
@@ -1024,11 +1040,7 @@ async def proxy_responses_websocket(
                     continue
         finally:
             if upstream_reader is not None:
-                upstream_reader.cancel()
-                try:
-                    await upstream_reader
-                except asyncio.CancelledError:
-                    pass
+                await _await_cancelled_task(upstream_reader, label="proxy websocket upstream reader")
             if upstream is not None:
                 try:
                     await upstream.close()
@@ -1780,11 +1792,7 @@ async def _close_http_bridge_session(
         else:
             await self._unregister_http_bridge_turn_states(session)
         if session.upstream_reader is not None:
-            session.upstream_reader.cancel()
-            try:
-                await session.upstream_reader
-            except asyncio.CancelledError:
-                pass
+            await _await_cancelled_task(session.upstream_reader, label="http bridge upstream reader")
         try:
             await session.upstream.close()
         except Exception:
@@ -2319,12 +2327,17 @@ async def _reconnect_http_bridge_session(
         old_upstream = session.upstream
         old_reader = session.upstream_reader if restart_reader else None
         if old_reader is not None:
-            old_reader.cancel()
             if old_reader is not asyncio.current_task():
-                try:
-                    await old_reader
-                except asyncio.CancelledError:
-                    pass
+                cancelled = await _await_cancelled_task(old_reader, label="http bridge upstream reader")
+                if not cancelled:
+                    session.closed = True
+                    raise ProxyResponseError(
+                        502,
+                        openai_error(
+                            "upstream_unavailable",
+                            "HTTP responses session bridge reader did not shut down cleanly",
+                        ),
+                    )
         try:
             await old_upstream.close()
         except Exception: