fix(proxy): keep streams alive while account capacity recovers

.all-contributorsrc

@@ -717,6 +717,16 @@
         "test"
       ]
     },
+    {
+      "login": "softkleenex",
+      "name": "softkleenex",
+      "avatar_url": "https://avatars.githubusercontent.com/u/92619941?v=4",
+      "profile": "https://github.com/softkleenex",
+      "contributions": [
+        "code",
+        "test"
+      ]
+    },
     {
       "login": "plastictaste",
       "name": "plastictaste",

README.md

@@ -578,6 +578,7 @@ Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/e
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/NeoClaw"><img src="https://avatars.githubusercontent.com/u/261536598?v=4?s=100" width="100px;" alt="NeoClaw"/><br /><sub><b>NeoClaw</b></sub></a><br /><a href="https://github.com/Soju06/codex-lb/commits?author=neoclaw" title="Code">💻</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/abarsegov"><img src="https://avatars.githubusercontent.com/u/181010154?v=4?s=100" width="100px;" alt="abarsegov"/><br /><sub><b>abarsegov</b></sub></a><br /><a href="https://github.com/Soju06/codex-lb/commits?author=abarsegov" title="Code">💻</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/KakatkarAkshay"><img src="https://avatars.githubusercontent.com/u/49910222?v=4?s=100" width="100px;" alt="Akshay Kakatkar"/><br /><sub><b>Akshay Kakatkar</b></sub></a><br /><a href="https://github.com/Soju06/codex-lb/commits?author=KakatkarAkshay" title="Code">💻</a> <a href="https://github.com/Soju06/codex-lb/commits?author=KakatkarAkshay" title="Tests">⚠️</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/softkleenex"><img src="https://avatars.githubusercontent.com/u/92619941?v=4?s=100" width="100px;" alt="softkleenex"/><br /><sub><b>softkleenex</b></sub></a><br /><a href="https://github.com/Soju06/codex-lb/commits?author=softkleenex" title="Code">💻</a> <a href="https://github.com/Soju06/codex-lb/commits?author=softkleenex" title="Tests">⚠️</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/plastictaste"><img src="https://avatars.githubusercontent.com/u/261646970?v=4?s=100" width="100px;" alt="plastictaste"/><br /><sub><b>plastictaste</b></sub></a><br /><a href="https://github.com/Soju06/codex-lb/commits?author=plastictaste" title="Code">💻</a> <a href="https://github.com/Soju06/codex-lb/commits?author=plastictaste" title="Tests">⚠️</a> <a href="https://github.com/Soju06/codex-lb/commits?author=plastictaste" title="Documentation">📖</a></td>
     </tr>
   </tbody>

app/modules/proxy/_service/http_bridge/helpers.py

@@ -1034,6 +1034,16 @@ def _http_bridge_runtime_config(
     )
 
 
+def _http_bridge_request_budget_seconds(settings: object) -> float:
+    return float(
+        getattr(
+            settings,
+            "http_responses_session_bridge_request_budget_seconds",
+            getattr(settings, "proxy_request_budget_seconds", 600.0),
+        )
+    )
+
+
 def _http_bridge_owner_check_required(
     key: _HTTPBridgeSessionKey,
     *,

app/modules/proxy/_service/http_bridge/mixin.py

@@ -149,6 +149,7 @@
     _HTTPBridgeOwnerForward,
     _HTTPBridgeSession,
     _HTTPBridgeSessionKey,
+    _sleep_for_account_selection_recovery,
     _WebSocketRequestState,
     _WebSocketUpstreamControl,
 )
@@ -2184,6 +2185,30 @@ async def release_selected_account_lease() -> None:
                 if reuse_current_account_lease and _remaining_budget_seconds(deadline) > 0:
                     preferred_candidate_id = None
                     continue
+                if await _sleep_for_account_selection_recovery(
+                    selection,
+                    request_id=request_state.request_log_id or request_state.request_id,
+                    kind="http_bridge",
+                    request_stage="reattach",
+                    model=session.request_model,
+                    max_sleep_seconds=_remaining_budget_seconds(deadline),
+                    request_state=request_state,
+                ):
+                    excluded_account_ids = set(request_state.excluded_account_ids)
+                    if skip_same_account:
+                        excluded_account_ids.add(session.account.id)
+                    retry_same_account_once = not skip_same_account and session.account.id not in excluded_account_ids
+                    if skip_same_account:
+                        preferred_candidate_id = None
+                    elif forced_refresh_account_id is not None:
+                        preferred_candidate_id = forced_refresh_account_id
+                    elif request_state.preferred_account_id is not None:
+                        preferred_candidate_id = request_state.preferred_account_id
+                    elif session.account.id not in excluded_account_ids:
+                        preferred_candidate_id = session.account.id
+                    else:
+                        preferred_candidate_id = None
+                    continue
                 _record_same_account_takeover(
                     preferred_account_id=session.account.id,
                     selected_account_id=None,