Skip to content

fix(proxy): retry websocket EOF before visible output #1388

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

Open
Komzpa wants to merge 5 commits into
base: main
Choose a base branch
from
Open

fix(proxy): retry websocket EOF before visible output #1388

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a3c95d0
fix(websocket): preserve current backoff lifecycle state
Komzpa Aug 6, 2026
1676031
fix(api-keys): record last-used timestamp
Komzpa Aug 7, 2026
2bdd2b6
fix websocket reconnect cancellation ownership
Komzpa Aug 7, 2026
f5692a8
style: format api key usage update
Komzpa Aug 7, 2026
b73e361
docs: clarify reconnect drain ownership gate
Komzpa Aug 7, 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
5 changes: 5 additions & 0 deletions app/modules/api_keys/repository.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,11 @@ class ApiKeysRepository:
def __init__(self, session: AsyncSession) -> None:
self._session = session

async def update_last_used(self, key_id: str, *, commit: bool = True) -> None:
await self._session.execute(update(ApiKey).where(ApiKey.id == key_id).values(last_used_at=utcnow()))
if commit:
await self._session.commit()

@staticmethod
def _build_account_costs(rows: Sequence[object]) -> list[ApiKeyAccountCost]:
account_costs: list[ApiKeyAccountCost] = []
Expand Down
18 changes: 16 additions & 2 deletions app/modules/proxy/_service/websocket/mixin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1676,8 +1676,21 @@ def take_reader_replay_request_state() -> _WebSocketRequestState | None:
and upstream_control is not None
and upstream_control.reconnect_requested
and upstream_reader is not None
and (
# Keep receiving downstream controls while a live
# reader still owns pending work; only replay or a
# completed/empty drain needs immediate settlement.
upstream_reader.done()
or upstream_control.replay_request_state is not None
or not pending_requests
)
):
await upstream_reader
try:
await upstream_reader
except asyncio.CancelledError:
current_task = asyncio.current_task()
if current_task is not None and current_task.cancelling():
raise
if replay_request_state is None:
replay_request_state = upstream_control.replay_request_state
upstream_reader = None
Expand Down Expand Up @@ -1727,6 +1740,7 @@ def take_reader_replay_request_state() -> _WebSocketRequestState | None:
("previous response", previous_response_owner_account_id),
)
except ProxyResponseError as exc:
response_create_request_state = request_state
error = _parse_openai_error(exc.payload)
error_code = _normalize_error_code(
error.code if error else None,
Expand All @@ -1735,7 +1749,7 @@ def take_reader_replay_request_state() -> _WebSocketRequestState | None:
error_message = error.message if error and error.message else "Upstream error"
error_type = error.type if error and error.type else "server_error"
error_param = error.param if error else None
await proxy._release_websocket_request_state_reservation(request_state)
await proxy._release_websocket_request_state_reservation(response_create_request_state)
await proxy._write_websocket_connect_failure(
account_id=None,
api_key=api_key,
Expand Down
Loading