Skip to content

Commit bf0345f

Browse files
KomzpaSoju06claude
authored
fix(proxy): reserve stream capacity for recovery (#1188)
* fix(proxy): reserve stream capacity for recovery * fix(proxy): keep reattach streams alive * fix(proxy): reserve stream capacity for true recovery * fix(proxy): read stream recovery reserve from env-backed settings The reserve was read via getattr on the dashboard settings cache result, which lacks the field, so CODEX_LB_PROXY_ACCOUNT_STREAM_RECOVERY_RESERVE was dead and the reserve was pinned to the default of 1. Read get_settings().proxy_account_stream_recovery_reserve instead, and make the regression test exercise the real env-backed settings with a non-default value. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test: add stream recovery reserve to ttft settings stub Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test: add stream recovery reserve to remaining settings stubs The service now reads proxy_account_stream_recovery_reserve explicitly from get_settings(), so every SimpleNamespace stub patched onto the proxy service module's get_settings needs the field. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Soju06 <qlskssk@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent cea5240 commit bf0345f

16 files changed

Lines changed: 215 additions & 65 deletions

File tree

‎app/core/config/settings.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -323,6 +323,7 @@ def upstream_websocket_proxy_env(self) -> Mapping[str, str | None]:
323323
proxy_admission_wait_timeout_seconds: float = Field(default=10.0, gt=0)
324324
proxy_account_response_create_limit: int = Field(default=4, ge=0)
325325
proxy_account_stream_limit: int = Field(default=8, ge=0)
326+
proxy_account_stream_recovery_reserve: int = Field(default=1, ge=0)
326327
proxy_account_inflight_penalty_pct: float = Field(default=2.5, ge=0)
327328
proxy_account_lease_token_weight: float = Field(default=1.0, ge=0)
328329
proxy_account_lease_ttl_seconds: float = Field(default=900.0, gt=0)

‎app/modules/proxy/_service/http_bridge/streaming.py‎

Lines changed: 12 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -1105,7 +1105,9 @@ def prepare_bridge_request(
11051105
allow_previous_response_recovery_rebind=should_attempt_previous_response_recovery,
11061106
allow_bootstrap_owner_rebind=should_attempt_bootstrap_rebind,
11071107
durable_lookup=durable_lookup,
1108-
request_stage="reattach",
1108+
request_stage=(
1109+
"reattach" if should_attempt_previous_response_recovery else "bootstrap_rebind"
1110+
),
11091111
preferred_account_id=request_state.preferred_account_id,
11101112
request_usage_budget=request_state.request_usage_budget,
11111113
session_header_fallback_key=session_header_fallback_key,
@@ -1196,60 +1198,15 @@ def prepare_bridge_request(
11961198
retry_request_state.request_stage = "reattach"
11971199
retry_request_state.preferred_account_id = request_state.preferred_account_id
11981200

1199-
while True:
1200-
try:
1201-
await self._submit_http_bridge_request(
1202-
session,
1203-
request_state=retry_request_state,
1204-
text_data=retry_text_data,
1205-
queue_limit=queue_limit,
1206-
)
1207-
except ProxyResponseError as capacity_exc:
1208-
wait_plan = _http_bridge_capacity_wait_plan(
1209-
capacity_exc,
1210-
request_deadline=request_deadline,
1211-
)
1212-
if wait_plan is None:
1213-
raise
1214-
bounded_wait_seconds, account_capacity_wait_seconds, message = wait_plan
1215-
logger.info(
1216-
"Waiting for account capacity before retrying HTTP bridge owner-forward recovery "
1217-
"submit request_id=%s model=%s account_id=%s sleep_seconds=%.1f "
1218-
"recovery_hint_seconds=%.1f error=%s",
1219-
retry_request_state.request_id,
1220-
retry_request_state.model,
1221-
session.account.id,
1222-
bounded_wait_seconds,
1223-
account_capacity_wait_seconds,
1224-
message,
1225-
)
1226-
async for line in _iter_account_capacity_wait_sse(
1227-
request_id=retry_request_state.request_id,
1228-
reason=message,
1229-
sleep_seconds=bounded_wait_seconds,
1230-
emit_keepalives=not propagate_http_errors,
1231-
):
1232-
yield line
1233-
if _service_time().monotonic() >= request_deadline:
1234-
raise
1235-
continue
1236-
break
1237-
_signal_propagated_capacity_startup_ready()
1238-
if downstream_turn_state is not None:
1239-
await self._register_http_bridge_turn_state(session, downstream_turn_state)
1240-
event_queue = retry_request_state.event_queue
1241-
assert event_queue is not None
1242-
while True:
1243-
event_block = await event_queue.get()
1244-
if event_block is None:
1245-
break
1246-
if retry_request_state.latency_first_token_ms is None:
1247-
block_payload = parse_sse_data_json(event_block)
1248-
block_event_type = _event_type_from_payload(None, block_payload)
1249-
if block_event_type in _TEXT_DELTA_EVENT_TYPES:
1250-
retry_request_state.latency_first_token_ms = int(
1251-
(_service_time().monotonic() - retry_request_state.started_at) * 1000
1252-
)
1201+
async for event_block in self._stream_http_bridge_session_events(
1202+
session,
1203+
request_state=retry_request_state,
1204+
text_data=retry_text_data,
1205+
queue_limit=queue_limit,
1206+
propagate_http_errors=propagate_http_errors,
1207+
downstream_turn_state=downstream_turn_state,
1208+
request_deadline=request_deadline,
1209+
):
12531210
yield event_block
12541211
except BaseException:
12551212
if retry_reservation_reacquired and retry_api_key_reservation is not None:

‎app/modules/proxy/_service/websocket/mixin.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -681,6 +681,7 @@ async def release_current_account_lease() -> None:
681681
if replay_request_state is not None:
682682
request_state = replay_request_state
683683
replay_request_state = None
684+
request_state.request_stage = "reattach"
684685
request_affinity = request_state.affinity_policy
685686
text_data = request_state.request_text
686687
if text_data is None:
@@ -1783,6 +1784,7 @@ async def _select_websocket_connect_account(
17831784
preferred_account_id=preferred_account_id,
17841785
require_security_work_authorized=require_security_work_authorized,
17851786
lease_kind="stream",
1787+
request_stage=request_state.request_stage,
17861788
estimated_lease_tokens=_facade()._estimated_lease_tokens_from_request_usage_budget(
17871789
request_state.request_usage_budget
17881790
),

‎app/modules/proxy/load_balancer.py‎

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -240,13 +240,20 @@ def _acquire_account_lease_locked(
240240
_record_account_inflight_leases(account_id, runtime)
241241
return lease
242242

243-
def _account_lease_allowed_locked(self, account_id: str, *, kind: AccountLeaseKind) -> bool:
243+
def _account_lease_allowed_locked(
244+
self,
245+
account_id: str,
246+
*,
247+
kind: AccountLeaseKind,
248+
stream_reserve_slots: int = 0,
249+
) -> bool:
244250
runtime = self._runtime.setdefault(account_id, RuntimeState())
245251
if kind == "response_create":
246252
cap = get_settings().proxy_account_response_create_limit
247253
return cap <= 0 or runtime.inflight_response_creates < cap
248254
cap = get_settings().proxy_account_stream_limit
249-
return cap <= 0 or runtime.inflight_streams < cap
255+
effective_cap = max(1, cap - max(0, stream_reserve_slots))
256+
return cap <= 0 or runtime.inflight_streams < effective_cap
250257

251258
def _release_account_lease_locked(self, lease: AccountLease, *, reason: str) -> bool:
252259
runtime = self._runtime.get(lease.account_id)
@@ -310,6 +317,7 @@ async def select_account(
310317
routing_costs_by_account_id: RoutingCostsByAccount | None = None,
311318
lease_kind: AccountLeaseKind | None = None,
312319
estimated_lease_tokens: float = 0.0,
320+
stream_reserve_slots: int = 0,
313321
traffic_class: TrafficClass = TRAFFIC_CLASS_FOREGROUND,
314322
) -> AccountSelection:
315323
excluded_ids = set(exclude_account_ids or ())
@@ -427,7 +435,11 @@ async def load_selection_inputs() -> _SelectionInputs:
427435
now=datetime.now(timezone.utc),
428436
)
429437
)
430-
selection_states = _filter_states_for_account_caps(states, lease_kind=lease_kind)
438+
selection_states = _filter_states_for_account_caps(
439+
states,
440+
lease_kind=lease_kind,
441+
stream_reserve_slots=stream_reserve_slots,
442+
)
431443
if not selection_states and states:
432444
selection_error_code = _account_cap_error_code(lease_kind)
433445
error_message = _account_cap_error_message(lease_kind)
@@ -618,7 +630,13 @@ async def load_selection_inputs() -> _SelectionInputs:
618630
sticky_existing_account_id, str
619631
)
620632
selection_states = (
621-
states if hard_sticky else _filter_states_for_account_caps(states, lease_kind=lease_kind)
633+
states
634+
if hard_sticky
635+
else _filter_states_for_account_caps(
636+
states,
637+
lease_kind=lease_kind,
638+
stream_reserve_slots=stream_reserve_slots,
639+
)
622640
)
623641
if not selection_states and states:
624642
selection_error_code = _account_cap_error_code(lease_kind)
@@ -683,7 +701,11 @@ async def load_selection_inputs() -> _SelectionInputs:
683701
selected_snapshot.deactivation_reason = result.account.deactivation_reason
684702
selected_snapshot.reset_at = selected_reset_at
685703
if lease_kind is not None:
686-
if not self._account_lease_allowed_locked(selected.id, kind=lease_kind):
704+
if not self._account_lease_allowed_locked(
705+
selected.id,
706+
kind=lease_kind,
707+
stream_reserve_slots=stream_reserve_slots,
708+
):
687709
selected_snapshot = None
688710
selection_error_code = _account_cap_error_code(lease_kind)
689711
error_message = _account_cap_error_message(lease_kind)
@@ -1709,6 +1731,7 @@ def _filter_states_for_account_caps(
17091731
states: Iterable[AccountState],
17101732
*,
17111733
lease_kind: AccountLeaseKind | None,
1734+
stream_reserve_slots: int = 0,
17121735
) -> list[AccountState]:
17131736
if lease_kind is None:
17141737
return list(states)
@@ -1721,7 +1744,8 @@ def _filter_states_for_account_caps(
17211744
continue
17221745
else:
17231746
cap = settings.proxy_account_stream_limit
1724-
if cap > 0 and state.inflight_streams >= cap:
1747+
effective_cap = max(1, cap - max(0, stream_reserve_slots))
1748+
if cap > 0 and state.inflight_streams >= effective_cap:
17251749
continue
17261750
filtered.append(state)
17271751
return filtered

‎app/modules/proxy/service.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1731,6 +1731,11 @@ async def _select_account_with_budget(
17311731
try:
17321732
with anyio.fail_after(remaining_budget):
17331733
settings = await get_settings_cache().get()
1734+
stream_reserve_slots = (
1735+
get_settings().proxy_account_stream_recovery_reserve
1736+
if lease_kind == "stream" and request_stage != "reattach"
1737+
else 0
1738+
)
17341739
required_preferred_account = (
17351740
preferred_account_id is not None and not fallback_on_preferred_account_unavailable
17361741
)
@@ -1798,6 +1803,7 @@ async def _select_account_with_budget(
17981803
secondary_budget_threshold_pct=_sticky_reallocation_secondary_budget_threshold_pct(settings),
17991804
lease_kind=lease_kind,
18001805
estimated_lease_tokens=estimated_lease_tokens,
1806+
stream_reserve_slots=stream_reserve_slots,
18011807
traffic_class=effective_traffic_class,
18021808
)
18031809
if preferred_selection.account is not None:
@@ -1841,6 +1847,7 @@ async def _select_account_with_budget(
18411847
secondary_budget_threshold_pct=_sticky_reallocation_secondary_budget_threshold_pct(settings),
18421848
lease_kind=lease_kind,
18431849
estimated_lease_tokens=estimated_lease_tokens,
1850+
stream_reserve_slots=stream_reserve_slots,
18441851
traffic_class=effective_traffic_class,
18451852
)
18461853
if selection.account is not None and selection.account.id in excluded_account_ids_set:

‎openspec/changes/reserve-stream-capacity-for-recovery/.openspec.yaml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
schema: spec-driven
2+
created: 2026-07-10

‎openspec/changes/reserve-stream-capacity-for-recovery/proposal.md‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
## Why
2+
3+
Parallel first turns can consume every account-local stream slot. If an established turn then loses its upstream transport, its continuity-safe reattach competes with new work and can remain stuck even though reserving one existing slot would have let it recover.
4+
5+
## What Changes
6+
7+
- Reserve one account-local stream slot by default from ordinary first-turn and follow-up selection.
8+
- Let reattach selection use the full configured account stream cap.
9+
- Keep the reserve configurable and preserve the hard cap.
10+
11+
## Impact
12+
13+
New fan-out may reach local backpressure one slot earlier per account. Established responses retain capacity to recover after a transient upstream disconnect.

‎openspec/changes/reserve-stream-capacity-for-recovery/specs/proxy-admission-control/spec.md‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
## ADDED Requirements
2+
3+
### Requirement: Account stream capacity reserves recovery headroom
4+
5+
The proxy MUST reserve the configured number of account-local stream slots from ordinary first-turn and follow-up selection, while allowing reattach work to use the full account stream cap. The default recovery reserve MUST be one slot. The reserve MUST NOT increase the configured hard stream cap.
6+
7+
#### Scenario: Fan-out leaves one slot for reattach
8+
9+
- **GIVEN** an account stream cap of eight and a recovery reserve of one
10+
- **AND** seven ordinary streams are active
11+
- **WHEN** another ordinary stream and a reattach stream compete for capacity
12+
- **THEN** the ordinary stream receives local account-cap backpressure
13+
- **AND** the reattach stream may acquire the eighth slot

‎openspec/changes/reserve-stream-capacity-for-recovery/tasks.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
- [x] Add a configurable account stream recovery reserve.
2+
- [x] Apply the reserve to ordinary stream selection while exempting reattach.
3+
- [x] Add concurrency regression coverage for seven ordinary leases plus one recovery lease.
4+
- [x] Run focused tests, lint, type checks, and strict OpenSpec validation.

‎tests/integration/test_proxy_api_extended.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1028,6 +1028,7 @@ async def stream_responses(self, *args, **kwargs):
10281028
settings = SimpleNamespace(
10291029
http_responses_session_bridge_enabled=False,
10301030
sse_keepalive_interval_seconds=0.01,
1031+
proxy_account_stream_recovery_reserve=1,
10311032
)
10321033
monkeypatch.setattr(proxy_api_module, "get_settings", lambda: settings)
10331034
monkeypatch.setattr(proxy_api_module.proxy_service_module, "get_settings", lambda: settings)
@@ -1126,6 +1127,7 @@ async def stream_responses(self, *args, **kwargs):
11261127
settings = SimpleNamespace(
11271128
http_responses_session_bridge_enabled=False,
11281129
sse_keepalive_interval_seconds=0.01,
1130+
proxy_account_stream_recovery_reserve=1,
11291131
)
11301132
monkeypatch.setattr(proxy_api_module, "get_settings", lambda: settings)
11311133
monkeypatch.setattr(proxy_api_module.proxy_service_module, "get_settings", lambda: settings)

0 commit comments

Comments
 (0)