codex-lb 0.1.3__tar.gz → 0.1.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (124) hide show
  1. codex_lb-0.1.5/.github/release-please-manifest.json +3 -0
  2. {codex_lb-0.1.3 → codex_lb-0.1.5}/.github/workflows/release.yml +42 -36
  3. {codex_lb-0.1.3 → codex_lb-0.1.5}/CHANGELOG.md +14 -0
  4. {codex_lb-0.1.3 → codex_lb-0.1.5}/PKG-INFO +1 -1
  5. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/balancer/logic.py +5 -12
  6. codex_lb-0.1.5/app/core/utils/retry.py +30 -0
  7. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/db/session.py +27 -9
  8. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/dependencies.py +11 -10
  9. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/proxy/load_balancer.py +3 -3
  10. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/proxy/service.py +36 -15
  11. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/request_logs/repository.py +17 -3
  12. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/static/index.js +3 -11
  13. {codex_lb-0.1.3 → codex_lb-0.1.5}/pyproject.toml +1 -1
  14. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_accounts_api_extended.py +83 -0
  15. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/unit/test_load_balancer.py +7 -3
  16. codex_lb-0.1.3/.github/release-please-manifest.json +0 -3
  17. codex_lb-0.1.3/app/core/utils/retry.py +0 -16
  18. {codex_lb-0.1.3 → codex_lb-0.1.5}/.dockerignore +0 -0
  19. {codex_lb-0.1.3 → codex_lb-0.1.5}/.env.example +0 -0
  20. {codex_lb-0.1.3 → codex_lb-0.1.5}/.github/release-please-config.json +0 -0
  21. {codex_lb-0.1.3 → codex_lb-0.1.5}/.github/workflows/ci.yml +0 -0
  22. {codex_lb-0.1.3 → codex_lb-0.1.5}/.github/workflows/release-please.yml +0 -0
  23. {codex_lb-0.1.3 → codex_lb-0.1.5}/.gitignore +0 -0
  24. {codex_lb-0.1.3 → codex_lb-0.1.5}/.pre-commit-config.yaml +0 -0
  25. {codex_lb-0.1.3 → codex_lb-0.1.5}/AGENTS.md +0 -0
  26. {codex_lb-0.1.3 → codex_lb-0.1.5}/Dockerfile +0 -0
  27. {codex_lb-0.1.3 → codex_lb-0.1.5}/LICENSE +0 -0
  28. {codex_lb-0.1.3 → codex_lb-0.1.5}/README.md +0 -0
  29. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/__init__.py +0 -0
  30. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/cli.py +0 -0
  31. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/__init__.py +0 -0
  32. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/auth/__init__.py +0 -0
  33. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/auth/models.py +0 -0
  34. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/auth/refresh.py +0 -0
  35. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/balancer/__init__.py +0 -0
  36. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/balancer/types.py +0 -0
  37. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/clients/__init__.py +0 -0
  38. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/clients/http.py +0 -0
  39. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/clients/oauth.py +0 -0
  40. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/clients/proxy.py +0 -0
  41. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/clients/usage.py +0 -0
  42. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/config/__init__.py +0 -0
  43. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/config/settings.py +0 -0
  44. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/crypto.py +0 -0
  45. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/errors.py +0 -0
  46. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/openai/__init__.py +0 -0
  47. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/openai/models.py +0 -0
  48. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/openai/parsing.py +0 -0
  49. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/openai/requests.py +0 -0
  50. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/types.py +0 -0
  51. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/usage/__init__.py +0 -0
  52. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/usage/logs.py +0 -0
  53. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/usage/models.py +0 -0
  54. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/usage/pricing.py +0 -0
  55. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/usage/types.py +0 -0
  56. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/utils/__init__.py +0 -0
  57. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/utils/request_id.py +0 -0
  58. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/utils/sse.py +0 -0
  59. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/core/utils/time.py +0 -0
  60. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/db/__init__.py +0 -0
  61. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/db/models.py +0 -0
  62. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/main.py +0 -0
  63. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/__init__.py +0 -0
  64. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/accounts/__init__.py +0 -0
  65. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/accounts/api.py +0 -0
  66. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/accounts/repository.py +0 -0
  67. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/accounts/schemas.py +0 -0
  68. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/accounts/service.py +0 -0
  69. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/health/__init__.py +0 -0
  70. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/health/api.py +0 -0
  71. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/oauth/__init__.py +0 -0
  72. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/oauth/api.py +0 -0
  73. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/oauth/schemas.py +0 -0
  74. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/oauth/service.py +0 -0
  75. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/oauth/templates/oauth_success.html +0 -0
  76. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/proxy/__init__.py +0 -0
  77. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/proxy/api.py +0 -0
  78. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/proxy/auth_manager.py +0 -0
  79. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/proxy/schemas.py +0 -0
  80. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/proxy/types.py +0 -0
  81. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/proxy/usage_updater.py +0 -0
  82. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/request_logs/__init__.py +0 -0
  83. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/request_logs/api.py +0 -0
  84. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/request_logs/schemas.py +0 -0
  85. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/request_logs/service.py +0 -0
  86. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/shared/__init__.py +0 -0
  87. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/shared/schemas.py +0 -0
  88. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/usage/__init__.py +0 -0
  89. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/usage/api.py +0 -0
  90. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/usage/repository.py +0 -0
  91. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/usage/schemas.py +0 -0
  92. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/modules/usage/service.py +0 -0
  93. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/static/7.css +0 -0
  94. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/static/index.css +0 -0
  95. {codex_lb-0.1.3 → codex_lb-0.1.5}/app/static/index.html +0 -0
  96. {codex_lb-0.1.3 → codex_lb-0.1.5}/docker-compose.yml +0 -0
  97. {codex_lb-0.1.3 → codex_lb-0.1.5}/docs/screenshots/accounts.jpeg +0 -0
  98. {codex_lb-0.1.3 → codex_lb-0.1.5}/docs/screenshots/dashboard.jpeg +0 -0
  99. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/__init__.py +0 -0
  100. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/conftest.py +0 -0
  101. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_accounts_api.py +0 -0
  102. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_codex_usage_api.py +0 -0
  103. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_db_models.py +0 -0
  104. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_health_and_errors.py +0 -0
  105. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_load_balancer_integration.py +0 -0
  106. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_oauth_flow.py +0 -0
  107. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_proxy_api_extended.py +0 -0
  108. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_proxy_compact.py +0 -0
  109. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_proxy_responses.py +0 -0
  110. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_repositories.py +0 -0
  111. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_request_logs_api.py +0 -0
  112. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_request_logs_filters.py +0 -0
  113. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_usage_api.py +0 -0
  114. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/integration/test_usage_summary.py +0 -0
  115. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/unit/test_auth.py +0 -0
  116. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/unit/test_auth_refresh.py +0 -0
  117. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/unit/test_oauth_client.py +0 -0
  118. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/unit/test_pricing.py +0 -0
  119. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/unit/test_proxy_utils.py +0 -0
  120. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/unit/test_retry.py +0 -0
  121. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/unit/test_sse.py +0 -0
  122. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/unit/test_usage.py +0 -0
  123. {codex_lb-0.1.3 → codex_lb-0.1.5}/tests/unit/test_usage_client.py +0 -0
  124. {codex_lb-0.1.3 → codex_lb-0.1.5}/uv.lock +0 -0
@@ -0,0 +1,3 @@
1
+ {
2
+ ".": "0.1.5"
3
+ }
@@ -120,6 +120,9 @@ jobs:
120
120
  - name: Set up Docker Buildx
121
121
  uses: docker/setup-buildx-action@v3
122
122
 
123
+ - name: Set up QEMU
124
+ uses: docker/setup-qemu-action@v3
125
+
123
126
  - name: Log in to GHCR
124
127
  uses: docker/login-action@v3
125
128
  with:
@@ -144,6 +147,7 @@ jobs:
144
147
  context: .
145
148
  file: Dockerfile
146
149
  push: true
150
+ platforms: linux/amd64,linux/arm64
147
151
  tags: ${{ steps.meta.outputs.tags }}
148
152
  labels: ${{ steps.meta.outputs.labels }}
149
153
  cache-from: type=gha
@@ -170,62 +174,64 @@ jobs:
170
174
  name: dist
171
175
  path: dist
172
176
 
173
- - name: Generate release notes
174
- id: notes
177
+ - name: Resolve previous tag
178
+ id: prev_tag
175
179
  shell: bash
176
180
  run: |
177
181
  set -euo pipefail
178
182
  TAG_NAME="${RELEASE_TAG}"
179
- TAG_VERSION="${TAG_NAME#v}"
180
- RELEASE_DATE="$(date -u +%Y-%m-%d)"
181
-
182
183
  PREV_TAG=""
183
184
  if git describe --tags --abbrev=0 "${TAG_NAME}^" >/dev/null 2>&1; then
184
185
  PREV_TAG="$(git describe --tags --abbrev=0 "${TAG_NAME}^")"
185
186
  fi
186
-
187
187
  if [ -n "${PREV_TAG}" ]; then
188
188
  RANGE="${PREV_TAG}..${TAG_NAME}"
189
189
  else
190
190
  RANGE="${TAG_NAME}"
191
191
  fi
192
+ echo "tag_name=${TAG_NAME}" >> "$GITHUB_OUTPUT"
193
+ echo "prev_tag=${PREV_TAG}" >> "$GITHUB_OUTPUT"
194
+ echo "range=${RANGE}" >> "$GITHUB_OUTPUT"
192
195
 
193
- COMMIT_COUNT="$(git rev-list --count "${RANGE}")"
194
-
196
+ - name: Generate release notes (GitHub)
197
+ id: notes
198
+ uses: actions/github-script@v7
199
+ env:
200
+ TAG_NAME: ${{ steps.prev_tag.outputs.tag_name }}
201
+ PREV_TAG: ${{ steps.prev_tag.outputs.prev_tag }}
202
+ with:
203
+ script: |
204
+ const fs = require("fs");
205
+ const tag = process.env.TAG_NAME;
206
+ const prev = process.env.PREV_TAG || undefined;
207
+ if (!tag) {
208
+ core.setFailed("Missing TAG_NAME");
209
+ return;
210
+ }
211
+ const { data } = await github.rest.repos.generateReleaseNotes({
212
+ owner: context.repo.owner,
213
+ repo: context.repo.repo,
214
+ tag_name: tag,
215
+ previous_tag_name: prev,
216
+ });
217
+ fs.writeFileSync("release_notes.md", `${data.body.trimEnd()}\\n`);
218
+ core.setOutput("tag_name", tag);
219
+
220
+ - name: Append install and contributors
221
+ shell: bash
222
+ env:
223
+ TAG_NAME: ${{ steps.prev_tag.outputs.tag_name }}
224
+ run: |
225
+ set -euo pipefail
226
+ TAG_VERSION="${TAG_NAME#v}"
195
227
  {
196
- echo "## What's New in ${TAG_NAME}"
197
- echo
198
- echo "- Release date: ${RELEASE_DATE}"
199
- echo "- Commits: ${COMMIT_COUNT}"
200
- echo
201
- if [ -n "${PREV_TAG}" ]; then
202
- echo "### Changes"
203
- git log --pretty=format:"- %s" "${RANGE}"
204
- else
205
- echo "### Changes"
206
- git log --pretty=format:"- %s" "${RANGE}"
207
- fi
208
228
  echo
209
- echo
210
- echo "### Install / Run"
229
+ echo "## Install / Run"
211
230
  echo '```bash'
212
231
  echo 'uvx codex-lb'
213
- echo 'pip install codex-lb'
214
232
  echo "docker pull ghcr.io/${GITHUB_REPOSITORY}:${TAG_VERSION}"
215
233
  echo '```'
216
- echo
217
- echo "### Artifacts"
218
- if ls dist >/dev/null 2>&1; then
219
- ls -1 dist | sed 's|^|- dist/|'
220
- else
221
- echo "- dist/ (not available)"
222
- fi
223
- echo
224
- echo "### Contributors"
225
- git shortlog -s "${RANGE}" | awk '{$1=$1}1' | sed 's/^/- /'
226
- } > release_notes.md
227
-
228
- echo "tag_name=${TAG_NAME}" >> "$GITHUB_OUTPUT"
234
+ } >> release_notes.md
229
235
 
230
236
  - name: Create GitHub Release
231
237
  uses: softprops/action-gh-release@v2
@@ -1,5 +1,19 @@
1
1
  # Changelog
2
2
 
3
+ ## [0.1.5](https://github.com/Soju06/codex-lb/compare/v0.1.4...v0.1.5) (2026-01-14)
4
+
5
+
6
+ ### Bug Fixes
7
+
8
+ * align rate-limit backoff and reset handling ([4d59650](https://github.com/Soju06/codex-lb/commit/4d596508e5ad13e68aa6e64f9cb32324bd38f07b))
9
+
10
+ ## [0.1.4](https://github.com/Soju06/codex-lb/compare/v0.1.3...v0.1.4) (2026-01-13)
11
+
12
+
13
+ ### Bug Fixes
14
+
15
+ * **db:** harden session cleanup on cancellation ([dee3916](https://github.com/Soju06/codex-lb/commit/dee3916efa83dedec1d5ad43e1e14950b8c6e4a7))
16
+
3
17
  ## [0.1.3](https://github.com/Soju06/codex-lb/compare/v0.1.2...v0.1.3) (2026-01-12)
4
18
 
5
19
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: codex-lb
3
- Version: 0.1.3
3
+ Version: 0.1.5
4
4
  Summary: Codex load balancer and proxy for ChatGPT accounts with usage dashboard
5
5
  Author-email: Soju06 <qlskssk@gmail.com>
6
6
  Maintainer-email: Soju06 <qlskssk@gmail.com>
@@ -5,7 +5,7 @@ from dataclasses import dataclass
5
5
  from typing import Iterable
6
6
 
7
7
  from app.core.balancer.types import UpstreamError
8
- from app.core.utils.retry import parse_retry_after
8
+ from app.core.utils.retry import backoff_seconds, parse_retry_after
9
9
  from app.db.models import AccountStatus
10
10
 
11
11
  PERMANENT_FAILURE_CODES = {
@@ -22,7 +22,7 @@ class AccountState:
22
22
  account_id: str
23
23
  status: AccountStatus
24
24
  used_percent: float | None = None
25
- reset_at: int | None = None
25
+ reset_at: float | None = None
26
26
  last_error_at: float | None = None
27
27
  last_selected_at: float | None = None
28
28
  error_count: int = 0
@@ -97,18 +97,11 @@ def handle_rate_limit(state: AccountState, error: UpstreamError) -> None:
97
97
  state.status = AccountStatus.RATE_LIMITED
98
98
  state.error_count += 1
99
99
  state.last_error_at = time.time()
100
-
101
- reset_at = _extract_reset_at(error)
102
- if reset_at is not None:
103
- state.reset_at = reset_at
104
- return
105
-
106
100
  message = error.get("message")
107
101
  delay = parse_retry_after(message) if message else None
108
- if delay:
109
- state.reset_at = int(time.time() + delay)
110
- else:
111
- state.reset_at = int(time.time() + 300)
102
+ if delay is None:
103
+ delay = backoff_seconds(state.error_count)
104
+ state.reset_at = time.time() + delay
112
105
 
113
106
 
114
107
  def handle_quota_exceeded(state: AccountState, error: UpstreamError) -> None:
@@ -0,0 +1,30 @@
1
+ from __future__ import annotations
2
+
3
+ import random
4
+ import re
5
+
6
+ _RETRY_PATTERN = re.compile(r"(?i)try again in\s*(\d+(?:\.\d+)?)\s*(s|ms|seconds?)")
7
+ _BACKOFF_INITIAL_DELAY_MS = 200
8
+ _BACKOFF_FACTOR = 2.0
9
+ _BACKOFF_JITTER_MIN = 0.9
10
+ _BACKOFF_JITTER_MAX = 1.1
11
+
12
+
13
+ def parse_retry_after(message: str) -> float | None:
14
+ match = _RETRY_PATTERN.search(message or "")
15
+ if not match:
16
+ return None
17
+ value = float(match.group(1))
18
+ unit = match.group(2).lower()
19
+ if unit == "ms":
20
+ return value / 1000
21
+ return value
22
+
23
+
24
+ def backoff_seconds(attempt: int) -> float:
25
+ if attempt < 1:
26
+ attempt = 1
27
+ exponent = _BACKOFF_FACTOR ** (attempt - 1)
28
+ base_ms = _BACKOFF_INITIAL_DELAY_MS * exponent
29
+ jitter = random.uniform(_BACKOFF_JITTER_MIN, _BACKOFF_JITTER_MAX)
30
+ return (base_ms * jitter) / 1000.0
@@ -1,5 +1,6 @@
1
1
  from __future__ import annotations
2
2
 
3
+ import asyncio
3
4
  from pathlib import Path
4
5
  from typing import AsyncIterator
5
6
 
@@ -23,16 +24,33 @@ def _ensure_sqlite_dir(url: str) -> None:
23
24
  Path(path).expanduser().parent.mkdir(parents=True, exist_ok=True)
24
25
 
25
26
 
27
+ async def _safe_rollback(session: AsyncSession) -> None:
28
+ if not session.in_transaction():
29
+ return
30
+ try:
31
+ await asyncio.shield(session.rollback())
32
+ except Exception:
33
+ return
34
+
35
+
36
+ async def _safe_close(session: AsyncSession) -> None:
37
+ try:
38
+ await asyncio.shield(session.close())
39
+ except Exception:
40
+ return
41
+
42
+
26
43
  async def get_session() -> AsyncIterator[AsyncSession]:
27
- async with SessionLocal() as session:
28
- try:
29
- yield session
30
- except Exception:
31
- await session.rollback()
32
- raise
33
- finally:
34
- if session.in_transaction():
35
- await session.rollback()
44
+ session = SessionLocal()
45
+ try:
46
+ yield session
47
+ except BaseException:
48
+ await _safe_rollback(session)
49
+ raise
50
+ finally:
51
+ if session.in_transaction():
52
+ await _safe_rollback(session)
53
+ await _safe_close(session)
36
54
 
37
55
 
38
56
  async def init_db() -> None:
@@ -7,7 +7,7 @@ from dataclasses import dataclass
7
7
  from fastapi import Depends
8
8
  from sqlalchemy.ext.asyncio import AsyncSession
9
9
 
10
- from app.db.session import SessionLocal, get_session
10
+ from app.db.session import SessionLocal, _safe_close, _safe_rollback, get_session
11
11
  from app.modules.accounts.repository import AccountsRepository
12
12
  from app.modules.accounts.service import AccountsService
13
13
  from app.modules.oauth.service import OauthService
@@ -87,15 +87,16 @@ def get_usage_context(
87
87
 
88
88
  @asynccontextmanager
89
89
  async def _accounts_repo_context() -> AsyncIterator[AccountsRepository]:
90
- async with SessionLocal() as session:
91
- try:
92
- yield AccountsRepository(session)
93
- except Exception:
94
- await session.rollback()
95
- raise
96
- finally:
97
- if session.in_transaction():
98
- await session.rollback()
90
+ session = SessionLocal()
91
+ try:
92
+ yield AccountsRepository(session)
93
+ except BaseException:
94
+ await _safe_rollback(session)
95
+ raise
96
+ finally:
97
+ if session.in_transaction():
98
+ await _safe_rollback(session)
99
+ await _safe_close(session)
99
100
 
100
101
 
101
102
  def get_oauth_context(
@@ -20,7 +20,7 @@ from app.modules.usage.repository import UsageRepository
20
20
 
21
21
  @dataclass
22
22
  class RuntimeState:
23
- reset_at: int | None = None
23
+ reset_at: float | None = None
24
24
  last_error_at: float | None = None
25
25
  last_selected_at: float | None = None
26
26
  error_count: int = 0
@@ -179,10 +179,10 @@ def _apply_secondary_quota(
179
179
  *,
180
180
  status: AccountStatus,
181
181
  primary_used: float | None,
182
- runtime_reset: int | None,
182
+ runtime_reset: float | None,
183
183
  secondary_used: float | None,
184
184
  secondary_reset: int | None,
185
- ) -> tuple[AccountStatus, float | None, int | None]:
185
+ ) -> tuple[AccountStatus, float | None, float | None]:
186
186
  used_percent = primary_used
187
187
  reset_at = runtime_reset
188
188
 
@@ -1,5 +1,6 @@
1
1
  from __future__ import annotations
2
2
 
3
+ import logging
3
4
  import time
4
5
  from datetime import timedelta
5
6
  from typing import AsyncIterator, Iterable, Mapping
@@ -36,6 +37,8 @@ from app.modules.proxy.usage_updater import UsageUpdater
36
37
  from app.modules.request_logs.repository import RequestLogsRepository
37
38
  from app.modules.usage.repository import UsageRepository
38
39
 
40
+ logger = logging.getLogger(__name__)
41
+
39
42
 
40
43
  class ProxyService:
41
44
  def __init__(
@@ -191,6 +194,7 @@ class ProxyService:
191
194
  yield format_sse_event(event)
192
195
  return
193
196
 
197
+ account_id_value = account.id
194
198
  try:
195
199
  account = await self._ensure_fresh(account)
196
200
  async for line in self._stream_once(
@@ -243,7 +247,15 @@ class ProxyService:
243
247
  await self._load_balancer.mark_permanent_failure(account, exc.code)
244
248
  continue
245
249
  except Exception:
246
- await self._load_balancer.record_error(account)
250
+ try:
251
+ await self._load_balancer.record_error(account)
252
+ except Exception:
253
+ logger.warning(
254
+ "Failed to record proxy error account_id=%s request_id=%s",
255
+ account_id_value,
256
+ request_id,
257
+ exc_info=True,
258
+ )
247
259
  if attempt == max_attempts - 1:
248
260
  event = response_failed_event(
249
261
  "upstream_error",
@@ -267,8 +279,9 @@ class ProxyService:
267
279
  request_id: str,
268
280
  allow_retry: bool,
269
281
  ) -> AsyncIterator[str]:
282
+ account_id_value = account.id
270
283
  access_token = self._encryptor.decrypt(account.access_token_encrypted)
271
- account_id = _header_account_id(account.id)
284
+ account_id = _header_account_id(account_id_value)
272
285
  model = payload.model
273
286
  start = time.monotonic()
274
287
  status = "success"
@@ -341,19 +354,27 @@ class ProxyService:
341
354
  reasoning_tokens = (
342
355
  usage.output_tokens_details.reasoning_tokens if usage and usage.output_tokens_details else None
343
356
  )
344
- await self._logs_repo.add_log(
345
- account_id=account.id,
346
- request_id=request_id,
347
- model=model,
348
- input_tokens=input_tokens,
349
- output_tokens=output_tokens,
350
- cached_input_tokens=cached_input_tokens,
351
- reasoning_tokens=reasoning_tokens,
352
- latency_ms=latency_ms,
353
- status=status,
354
- error_code=error_code,
355
- error_message=error_message,
356
- )
357
+ try:
358
+ await self._logs_repo.add_log(
359
+ account_id=account_id_value,
360
+ request_id=request_id,
361
+ model=model,
362
+ input_tokens=input_tokens,
363
+ output_tokens=output_tokens,
364
+ cached_input_tokens=cached_input_tokens,
365
+ reasoning_tokens=reasoning_tokens,
366
+ latency_ms=latency_ms,
367
+ status=status,
368
+ error_code=error_code,
369
+ error_message=error_message,
370
+ )
371
+ except Exception:
372
+ logger.warning(
373
+ "Failed to persist request log account_id=%s request_id=%s",
374
+ account_id_value,
375
+ request_id,
376
+ exc_info=True,
377
+ )
357
378
 
358
379
  async def _refresh_usage(self, accounts: list[Account]) -> None:
359
380
  latest_usage = await self._usage_repo.latest_by_account(window="primary")
@@ -1,5 +1,6 @@
1
1
  from __future__ import annotations
2
2
 
3
+ import asyncio
3
4
  from datetime import datetime
4
5
 
5
6
  from sqlalchemy import and_, select
@@ -49,9 +50,13 @@ class RequestLogsRepository:
49
50
  requested_at=requested_at or utcnow(),
50
51
  )
51
52
  self._session.add(log)
52
- await self._session.commit()
53
- await self._session.refresh(log)
54
- return log
53
+ try:
54
+ await self._session.commit()
55
+ await self._session.refresh(log)
56
+ return log
57
+ except BaseException:
58
+ await _safe_rollback(self._session)
59
+ raise
55
60
 
56
61
  async def list_recent(
57
62
  self,
@@ -84,3 +89,12 @@ class RequestLogsRepository:
84
89
  stmt = stmt.limit(limit)
85
90
  result = await self._session.execute(stmt)
86
91
  return list(result.scalars().all())
92
+
93
+
94
+ async def _safe_rollback(session: AsyncSession) -> None:
95
+ if not session.in_transaction():
96
+ return
97
+ try:
98
+ await asyncio.shield(session.rollback())
99
+ except Exception:
100
+ return
@@ -569,16 +569,9 @@
569
569
  return acc;
570
570
  }, {});
571
571
 
572
- const mergeUsageIntoAccounts = (
573
- accounts,
574
- primaryUsage,
575
- secondaryUsage,
576
- summary,
577
- ) => {
572
+ const mergeUsageIntoAccounts = (accounts, primaryUsage, secondaryUsage) => {
578
573
  const primaryMap = buildUsageIndex(primaryUsage || []);
579
574
  const secondaryMap = buildUsageIndex(secondaryUsage || []);
580
- const resetAtPrimary = summary?.primaryWindow?.resetAt ?? null;
581
- const resetAtSecondary = summary?.secondaryWindow?.resetAt ?? null;
582
575
  return accounts.map((account) => {
583
576
  const primaryRow = primaryMap[account.id];
584
577
  const secondaryRow = secondaryMap[account.id];
@@ -598,8 +591,8 @@
598
591
  account.usage?.secondaryRemainingPercent ??
599
592
  0,
600
593
  },
601
- resetAtPrimary: resetAtPrimary ?? account.resetAtPrimary ?? null,
602
- resetAtSecondary: resetAtSecondary ?? account.resetAtSecondary ?? null,
594
+ resetAtPrimary: account.resetAtPrimary ?? null,
595
+ resetAtSecondary: account.resetAtSecondary ?? null,
603
596
  };
604
597
  });
605
598
  };
@@ -1191,7 +1184,6 @@
1191
1184
  accountsResult.value,
1192
1185
  primaryUsage,
1193
1186
  secondaryUsage,
1194
- summary,
1195
1187
  );
1196
1188
  this.applyData(
1197
1189
  {
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "codex-lb"
3
- version = "0.1.3"
3
+ version = "0.1.5"
4
4
  description = "Codex load balancer and proxy for ChatGPT accounts with usage dashboard"
5
5
  readme = "README.md"
6
6
  license = { file = "LICENSE" }
@@ -2,10 +2,17 @@ from __future__ import annotations
2
2
 
3
3
  import base64
4
4
  import json
5
+ from datetime import datetime, timezone
5
6
 
6
7
  import pytest
7
8
 
8
9
  from app.core.auth import fallback_account_id
10
+ from app.core.crypto import TokenEncryptor
11
+ from app.core.utils.time import utcnow
12
+ from app.db.models import Account, AccountStatus
13
+ from app.db.session import SessionLocal
14
+ from app.modules.accounts.repository import AccountsRepository
15
+ from app.modules.usage.repository import UsageRepository
9
16
 
10
17
  pytestmark = pytest.mark.integration
11
18
 
@@ -33,6 +40,32 @@ def _make_auth_json(account_id: str | None, email: str, plan_type: str = "plus")
33
40
  return {"tokens": tokens}
34
41
 
35
42
 
43
+ def _make_account(account_id: str, email: str, plan_type: str = "plus") -> Account:
44
+ encryptor = TokenEncryptor()
45
+ return Account(
46
+ id=account_id,
47
+ email=email,
48
+ plan_type=plan_type,
49
+ access_token_encrypted=encryptor.encrypt("access"),
50
+ refresh_token_encrypted=encryptor.encrypt("refresh"),
51
+ id_token_encrypted=encryptor.encrypt("id"),
52
+ last_refresh=utcnow(),
53
+ status=AccountStatus.ACTIVE,
54
+ deactivation_reason=None,
55
+ )
56
+
57
+
58
+ def _iso_utc(epoch_seconds: int) -> str:
59
+ return (
60
+ datetime.fromtimestamp(epoch_seconds, tz=timezone.utc)
61
+ .isoformat()
62
+ .replace(
63
+ "+00:00",
64
+ "Z",
65
+ )
66
+ )
67
+
68
+
36
69
  @pytest.mark.asyncio
37
70
  async def test_import_invalid_json_returns_400(async_client):
38
71
  files = {"auth_json": ("auth.json", "not-json", "application/json")}
@@ -78,3 +111,53 @@ async def test_delete_account_removes_from_list(async_client):
78
111
  assert accounts.status_code == 200
79
112
  data = accounts.json()["accounts"]
80
113
  assert all(account["accountId"] != "acc_delete" for account in data)
114
+
115
+
116
+ @pytest.mark.asyncio
117
+ async def test_accounts_list_includes_per_account_reset_times(async_client, db_setup):
118
+ primary_a = 1735689600
119
+ primary_b = 1735693200
120
+ secondary_a = 1736294400
121
+ secondary_b = 1736380800
122
+
123
+ async with SessionLocal() as session:
124
+ accounts_repo = AccountsRepository(session)
125
+ usage_repo = UsageRepository(session)
126
+
127
+ await accounts_repo.upsert(_make_account("acc_reset_a", "a@example.com"))
128
+ await accounts_repo.upsert(_make_account("acc_reset_b", "b@example.com"))
129
+
130
+ await usage_repo.add_entry(
131
+ "acc_reset_a",
132
+ 10.0,
133
+ window="primary",
134
+ reset_at=primary_a,
135
+ )
136
+ await usage_repo.add_entry(
137
+ "acc_reset_b",
138
+ 20.0,
139
+ window="primary",
140
+ reset_at=primary_b,
141
+ )
142
+ await usage_repo.add_entry(
143
+ "acc_reset_a",
144
+ 30.0,
145
+ window="secondary",
146
+ reset_at=secondary_a,
147
+ )
148
+ await usage_repo.add_entry(
149
+ "acc_reset_b",
150
+ 40.0,
151
+ window="secondary",
152
+ reset_at=secondary_b,
153
+ )
154
+
155
+ response = await async_client.get("/api/accounts")
156
+ assert response.status_code == 200
157
+ payload = response.json()
158
+ accounts = {item["accountId"]: item for item in payload["accounts"]}
159
+
160
+ assert accounts["acc_reset_a"]["resetAtPrimary"] == _iso_utc(primary_a)
161
+ assert accounts["acc_reset_b"]["resetAtPrimary"] == _iso_utc(primary_b)
162
+ assert accounts["acc_reset_a"]["resetAtSecondary"] == _iso_utc(secondary_a)
163
+ assert accounts["acc_reset_b"]["resetAtSecondary"] == _iso_utc(secondary_b)
@@ -38,19 +38,23 @@ def test_select_account_skips_rate_limited_until_reset():
38
38
 
39
39
 
40
40
  def test_handle_rate_limit_sets_reset_at_from_message():
41
+ now = time.time()
41
42
  state = AccountState("a", AccountStatus.ACTIVE, used_percent=5.0)
42
43
  handle_rate_limit(state, {"message": "Try again in 1.5s"})
43
44
  assert state.status == AccountStatus.RATE_LIMITED
44
45
  assert state.reset_at is not None
46
+ delay = state.reset_at - now
47
+ assert 1.2 <= delay <= 2.0
45
48
 
46
49
 
47
- def test_handle_rate_limit_uses_resets_in_seconds():
50
+ def test_handle_rate_limit_uses_backoff_when_no_delay():
48
51
  now = time.time()
49
52
  state = AccountState("a", AccountStatus.ACTIVE, used_percent=5.0)
50
- handle_rate_limit(state, {"resets_in_seconds": 10})
53
+ handle_rate_limit(state, {"message": "Rate limit exceeded."})
51
54
  assert state.status == AccountStatus.RATE_LIMITED
52
55
  assert state.reset_at is not None
53
- assert int(now) <= state.reset_at <= int(now) + 15
56
+ delay = state.reset_at - now
57
+ assert 0.15 <= delay <= 0.3
54
58
 
55
59
 
56
60
  def test_handle_quota_exceeded_sets_used_percent():
@@ -1,3 +0,0 @@
1
- {
2
- ".": "0.1.3"
3
- }
@@ -1,16 +0,0 @@
1
- from __future__ import annotations
2
-
3
- import re
4
-
5
- _RETRY_PATTERN = re.compile(r"(?i)try again in\s*(\d+(?:\.\d+)?)\s*(s|ms|seconds?)")
6
-
7
-
8
- def parse_retry_after(message: str) -> float | None:
9
- match = _RETRY_PATTERN.search(message or "")
10
- if not match:
11
- return None
12
- value = float(match.group(1))
13
- unit = match.group(2).lower()
14
- if unit == "ms":
15
- return value / 1000
16
- return value
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes