Skip to content

fix(proxy): bind account-bound retries to dispatch owner #1829

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Soju06 merged 2 commits into from
Aug 20, 2026
Merged

fix(proxy): bind account-bound retries to dispatch owner #1829

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f4217b1
fix(proxy): bind account-bound retries to dispatch owner
mastertyko Aug 19, 2026
2229a2f
fix(proxy): defer replay owner until dispatch
mastertyko Aug 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 47 additions & 16 deletions app/modules/proxy/_service/http_bridge/request_submit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,7 @@
_upstream_response_create_max_bytes,
_websocket_auth_failure_permanent_code,
_websocket_auth_failure_requires_reauth,
_websocket_request_text_is_account_neutral_fresh_replay,
)
from app.modules.proxy._service.observability import (
_hash_identifier as _hash_identifier,
Expand Down Expand Up @@ -3094,6 +3095,7 @@ def request_is_retryable(request_state: _WebSocketRequestState) -> bool:
)
if request_state.replay_count >= 1 and not additional_clean_close_retry:
return False
account_bound_replay = False
if request_state.previous_response_id is not None:
require_preferred_reconnect = False
if account_neutral_recovery:
Expand Down Expand Up @@ -3125,11 +3127,26 @@ def request_is_retryable(request_state: _WebSocketRequestState) -> bool:
# Account-scoped uploaded files cannot be replayed on a
# different owner. Keep the preferred account mandatory for
# both silent recovery and clean-close recovery.
require_preferred_reconnect = account_neutral_recovery or request_state.file_required_preferred_account
candidate_text = (
request_state.fresh_upstream_request_text
if request_state.fresh_upstream_request_is_retry_safe and request_state.fresh_upstream_request_text
else request_state.request_text
)
# The send boundary decorates durable operations with
# codex_lb_operation_id after selection. Keep that operation
# identity on its owner unless a dedicated rebind path has
# already replaced the operation ID.
candidate_portable = request_state.operation_id is None and (
_websocket_request_text_is_account_neutral_fresh_replay(candidate_text)
)
request_text = _prepare_websocket_request_state_for_visible_output_replay(request_state)
if request_text is None:
if request_text is None or request_text != candidate_text:
return False
if account_neutral_recovery:
account_bound_replay = not candidate_portable
require_preferred_reconnect = (
account_neutral_recovery or account_bound_replay or request_state.file_required_preferred_account
)
if account_neutral_recovery or account_bound_replay:
request_state.preferred_account_id = session.account.id
elif not request_state.file_required_preferred_account:
if hard_owner_bound and not model_fallback_replay and not fresh_hard_request_account_switch_allowed:
Expand Down Expand Up @@ -3210,7 +3227,7 @@ def request_is_retryable(request_state: _WebSocketRequestState) -> bool:
await self._reconnect_http_bridge_session(
session,
request_state=request_state,
require_same_account=account_neutral_recovery,
require_same_account=account_neutral_recovery or account_bound_replay,
require_preferred_account=True,
**reconnect_reader_kwargs,
)
Expand Down Expand Up @@ -3319,7 +3336,22 @@ async def _retry_http_bridge_precreated_auth_request(
error_message: str | None,
) -> Literal["not_replayable", "retried", "failed"]:
permanent_failure_code = _websocket_auth_failure_permanent_code(error_message)
request_text = _prepare_websocket_request_state_for_auth_replay(request_state)
bound_to_current_account = request_state.replay_required_account_id == session.account.id
if bound_to_current_account and (
_websocket_auth_failure_requires_reauth(error_message)
or request_state.auth_replay_counts_by_account.get(session.account.id, 0) > 0
):
failure_code = permanent_failure_code or _WEBSOCKET_AUTH_INVALIDATED_FAILURE_CODE
await self._load_balancer.mark_permanent_failure(session.account, failure_code)
setattr(request_state, "account_health_error_handled", True)
request_state.force_refresh_account_id = None
request_state.preferred_account_id = None
request_state.excluded_account_ids.add(session.account.id)
return "not_replayable"
request_text = _prepare_websocket_request_state_for_auth_replay(
request_state,
current_account_id=session.account.id,
)
if request_text is None:
await self._load_balancer.mark_permanent_failure(session.account, permanent_failure_code)
setattr(request_state, "account_health_error_handled", True)
Expand Down Expand Up @@ -3368,10 +3400,14 @@ async def _retry_http_bridge_precreated_auth_request(
await self._reconnect_http_bridge_session(
session,
request_state=request_state,
require_same_account=is_http_bridge_account_neutral_replay(
kind=session.key.affinity_kind,
key=session.key.affinity_key,
require_same_account=(
bound_to_current_account
or is_http_bridge_account_neutral_replay(
kind=session.key.affinity_kind,
key=session.key.affinity_key,
)
),
require_preferred_account=bound_to_current_account,
)
request_text = self._http_bridge_text_with_account_installation_id(session, request_state, request_text)
await _send_http_bridge_request_text_with_archive_id(session, request_state, request_text)
Expand Down Expand Up @@ -3413,13 +3449,13 @@ async def _retry_http_bridge_security_work_request(
key=session.key.affinity_key,
):
return False
retry_text = request_state.request_text
if not retry_text:
return False
if request_state.file_required_preferred_account:
return False
if not _websocket_request_can_replay_before_visible_output(request_state):
return False
retry_text = _prepare_websocket_request_state_for_account_switch(request_state)
if retry_text is None:
return False

owner_account_id = session.account.id
previous_replay_count = request_state.replay_count
Expand All @@ -3436,11 +3472,6 @@ async def _retry_http_bridge_security_work_request(
session.turn_state_alias_registration_generations
)
previous_session_headers = session.headers
if request_state.previous_response_id is not None:
retry_text = _prepare_websocket_request_state_for_account_switch(request_state)
if retry_text is None:
return False

request_state.preferred_account_id = None
request_state.excluded_account_ids.add(owner_account_id)
request_state.affinity_policy = replace(
Expand Down
4 changes: 4 additions & 0 deletions app/modules/proxy/_service/http_bridge/service_stubs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -452,6 +452,10 @@ def _prepare_websocket_request_state_for_account_switch(*args: Any, **kwargs: An
return _service_global("_prepare_websocket_request_state_for_account_switch")(*args, **kwargs)


def _websocket_request_text_is_account_neutral_fresh_replay(*args: Any, **kwargs: Any) -> Any:
return _service_global("_websocket_request_text_is_account_neutral_fresh_replay")(*args, **kwargs)


def _matching_websocket_request_states_for_previous_response_error(*args: Any, **kwargs: Any) -> Any:
return _service_global("_matching_websocket_request_states_for_previous_response_error")(*args, **kwargs)

Expand Down
59 changes: 53 additions & 6 deletions app/modules/proxy/_service/streaming/retry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@
is_upstream_model_capacity_error,
)
from app.modules.proxy.load_balancer import AccountLease, AccountSelection
from app.modules.proxy.replay_safety import responses_payload_is_account_neutral_fresh_replay
from app.modules.proxy.selection_errors import USAGE_LIMIT_REACHED, selection_failure_response

_REQUEST_TRANSPORT_HTTP = "http"
Expand Down Expand Up @@ -176,7 +177,10 @@ def _verified_cross_transport_fresh_replay(
stored_fingerprint=continuity_state.last_completed_input_prefix_fingerprint,
):
return None
return payload.model_copy(update={"previous_response_id": None})
fresh_payload = payload.model_copy(update={"previous_response_id": None})
if not responses_payload_is_account_neutral_fresh_replay(fresh_payload.to_replay_safety_payload()):
return None
return fresh_payload


def _effective_http_downstream_transport_policy(
Expand Down Expand Up @@ -396,6 +400,7 @@ async def _stream_with_retry(
deferred_capacity_account: Account | None = None
deferred_capacity_lease: AccountLease | None = None
preferred_account_id: str | None = None
payload_replay_required_account_id: str | None = None
file_preferred_account_id: str | None = rewritten_file_account_id
require_preferred_account = False
last_retryable_stream_error: _RetryableStreamError | None = None
Expand Down Expand Up @@ -577,18 +582,39 @@ async def _settle_process_network_budget_exhaustion(
)
settled = await _settle_stream_usage_before_pending_penalty(settlement)

def _authorize_payload_dispatch(account: Account) -> bool:
required_account_id = payload_replay_required_account_id
if required_account_id is not None and required_account_id != account.id:
raise ProxyResponseError(
502,
openai_error(
"previous_response_owner_unavailable",
"Request payload owner account is unavailable; retry later.",
error_type="server_error",
),
)
return required_account_id is None and not responses_payload_is_account_neutral_fresh_replay(
payload.to_replay_safety_payload()
)

def _move_verified_fresh_replay_from_owner(*, account_id: str, outcome: str) -> bool:
# Only a proxy-injected owner anchor with locally verified full
# input may move; the failed owner stays excluded so sticky
# selection cannot immediately loop back to it.
nonlocal affinity, payload, preferred_account_id, require_preferred_account, verified_fresh_replay_payload
nonlocal affinity, payload, payload_replay_required_account_id
nonlocal preferred_account_id, require_preferred_account, verified_fresh_replay_payload
if not (
require_preferred_account
and preferred_account_id == account_id
and verified_fresh_replay_payload is not None
):
return False
if not responses_payload_is_account_neutral_fresh_replay(
verified_fresh_replay_payload.to_replay_safety_payload()
):
return False
payload = verified_fresh_replay_payload
payload_replay_required_account_id = None
verified_fresh_replay_payload = None
excluded_account_ids.add(account_id)
preferred_account_id = None
Expand Down Expand Up @@ -1037,6 +1063,13 @@ async def _retry_account_model_rejection(
yield format_sse_event(_facade()._proxy_request_timeout_event(request_id))
return
while True:
effective_preferred_account_id = resolve_required_account_id(
("continuation", preferred_account_id),
("dispatched payload", payload_replay_required_account_id),
)
effective_require_preferred_account = (
require_preferred_account or payload_replay_required_account_id is not None
)
try:
selection = await proxy._select_account_with_budget_compatible(
deadline,
Expand All @@ -1050,15 +1083,15 @@ async def _retry_account_model_rejection(
model=payload.model,
service_tier=payload.service_tier,
exclude_account_ids=excluded_account_ids,
preferred_account_id=preferred_account_id,
preferred_account_id=effective_preferred_account_id,
require_security_work_authorized=require_security_work_authorized,
lease_kind="stream",
estimated_lease_tokens=estimated_lease_tokens,
# Keep stored-object and file ownership strict. The
# verified-fresh replay branch below removes its
# anchor before it permits cross-account movement.
fallback_on_preferred_account_unavailable=not (
require_preferred_account or file_required_preferred_account
effective_require_preferred_account or file_required_preferred_account
),
)
except ProxyResponseError as exc:
Expand Down Expand Up @@ -1846,6 +1879,7 @@ async def _retry_account_model_rejection(
)
try:
settlement = _StreamSettlement()
register_payload_owner = _authorize_payload_dispatch(account)
inner_stream = proxy._stream_once(
account,
payload,
Expand Down Expand Up @@ -1887,8 +1921,21 @@ async def _retry_account_model_rejection(
enforce_openai_sdk_contract=enforce_openai_sdk_contract,
)
try:
async for line in inner_stream:
yield line
try:
async for line in inner_stream:
if register_payload_owner:
payload_replay_required_account_id = account.id
register_payload_owner = False
yield line
if register_payload_owner:
payload_replay_required_account_id = account.id
except BaseException as exc:
if register_payload_owner and not (
isinstance(exc, ProxyResponseError)
and is_confirmed_pre_dispatch_transport_error(exc)
):
payload_replay_required_account_id = account.id
raise
finally:
close_task = asyncio.create_task(
inner_stream.aclose(),
Expand Down
3 changes: 3 additions & 0 deletions app/modules/proxy/_service/support.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1066,6 +1066,9 @@ class _WebSocketRequestState:
fresh_upstream_request_responses_lite_model: str | None = None
request_stage: str = "first_turn"
preferred_account_id: str | None = None
# Once an account-bound body has been dispatched, retries remain pinned to
# that owner even when stale-anchor recovery removes previous_response_id.
replay_required_account_id: str | None = None
require_security_work_authorized: bool = False
durable_capability_lineage_required: bool = False
file_required_preferred_account: bool = False
Expand Down
Loading
Loading