mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 16:06:08 +00:00
feat: add token refresh retry and usage stats
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
Codex
+43
-11
@@ -47,6 +47,8 @@ When the client omits `stream` or sets `"stream": false`, the Rust route collect
|
||||
|
||||
When the client sets `"stream": true`, the Rust route returns `text/event-stream` and passes through upstream SSE frames while collecting usage for the account after the stream finishes. Both modes capture upstream `Set-Cookie`, replay encrypted account-scoped Cookies, record account usage when usage appears in SSE, and write a `v1.response` event log with `requestId`, `accountId`, `route`, `model`, `statusCode`, `latencyMs`, and non-secret metadata.
|
||||
|
||||
If Codex returns `401` and the account has a stored refresh token, the route refreshes the account through the OpenAI OAuth token endpoint once, persists the rotated access token, preserves the old refresh token when the server omits a new one, updates the in-process account pool, and retries the same request once.
|
||||
|
||||
`previous_response_id` and explicit WebSocket-only requests are rejected until the WebSocket transport is implemented and verified.
|
||||
|
||||
## `/admin/*`
|
||||
@@ -157,6 +159,47 @@ Example item:
|
||||
}
|
||||
```
|
||||
|
||||
### `GET /admin/usage-stats`
|
||||
|
||||
Returns cursor-paginated account usage counters recorded after `/v1/*` calls. The endpoint never returns account access tokens, refresh tokens, Cookie values, or client API keys.
|
||||
|
||||
Example item:
|
||||
|
||||
```json
|
||||
{
|
||||
"accountId": "acct_01",
|
||||
"email": "user@example.com",
|
||||
"label": "primary",
|
||||
"planType": "plus",
|
||||
"requestCount": 2,
|
||||
"inputTokens": 14,
|
||||
"outputTokens": 8,
|
||||
"cachedTokens": 3,
|
||||
"lastUsedAt": "2026-06-11T12:00:00Z"
|
||||
}
|
||||
```
|
||||
|
||||
### `GET /admin/usage-stats/summary`
|
||||
|
||||
Returns global usage totals across accounts.
|
||||
|
||||
Success response:
|
||||
|
||||
```json
|
||||
{
|
||||
"code": 200,
|
||||
"message": "OK",
|
||||
"data": {
|
||||
"accountCount": 1,
|
||||
"requestCount": 2,
|
||||
"inputTokens": 14,
|
||||
"outputTokens": 8,
|
||||
"cachedTokens": 3
|
||||
},
|
||||
"requestId": "req_01"
|
||||
}
|
||||
```
|
||||
|
||||
### `GET /admin/api-keys`
|
||||
|
||||
Returns local client API keys with cursor pagination. The plaintext key and hash are never returned by list endpoints.
|
||||
@@ -234,17 +277,6 @@ Success response:
|
||||
}
|
||||
```
|
||||
|
||||
Success body:
|
||||
|
||||
```json
|
||||
{
|
||||
"code": 200,
|
||||
"message": "OK",
|
||||
"data": {},
|
||||
"requestId": "req_01"
|
||||
}
|
||||
```
|
||||
|
||||
List body:
|
||||
|
||||
```json
|
||||
|
||||
@@ -10,17 +10,17 @@ Update this file before each feature commit.
|
||||
| SQLite storage | Completed | `ffec8b5` + `b0c10f2` | `cargo test sqlite_schema_creates_accounts_and_event_tables`; `cargo test connect_sqlite_creates_parent_directory_for_file_database`; full `fmt/test/clippy` before commit | Creates SQLite WAL connector, parent directory bootstrap, sqlx-managed schema for accounts/API keys/sessions/cookies/fingerprints/events, and event/account indexes. |
|
||||
| Admin auth and client API keys | Partial | `24d5799` + `4e0d1e5` + `e9fdea0` | `cargo test client_api_key_has_proxy_prefix_and_verifies_against_hash`; `cargo test admin_password_hash_is_not_a_client_api_key`; `cargo test --test admin_login_route_test`; `cargo test --test admin_api_keys_route_test`; full `fmt/test/clippy` before commit | Admin passwords use Argon2id, client API keys use `cpr_` prefix plus HMAC-SHA256 with server-side pepper; `/admin/login` issues HttpOnly sessions and rejects client API keys as admin credentials. `/admin/api-keys` can create/list local client keys, and `/v1/*` now verifies enabled keys against SQLite instead of accepting any `cpr_` string. Still missing delete/batch-delete/label/status/export/import key management. |
|
||||
| Secret encryption | Completed | `105aa58` + `228febd` | `cargo test secret_box_encrypts_and_decrypts_without_storing_plaintext`; `cargo test secret_box_load_or_create_should_reuse_existing_master_key_file`; full `fmt/test/clippy` before commit | AES-256-GCM `SecretBox` stores upstream secrets as `v1:<nonce>:<ciphertext>` without plaintext and can load/create the runtime master key file used by repositories. |
|
||||
| Logging and pagination | Completed | `eb0341e` + `fc16c48` + `b0c10f2` + `ee0ed90` | `cargo test event_logs_are_cursor_paginated`; `cargo test admin_logs_are_cursor_paginated_and_include_request_id`; `cargo test rotating_writer_splits_files_when_size_limit_is_reached`; `cargo test --test v1_upstream_route_test`; full `fmt/test/clippy` before commit | Adds cursor pagination, `Page<T>` camelCase serialization, admin `/admin/logs` query API, event log repository, runtime tracing initialization, HTTP trace layer, date/size-based file rotation, and v1 response lifecycle events with request/account/route/model/status/latency/non-secret metadata. |
|
||||
| Logging and pagination | Completed | `eb0341e` + `fc16c48` + `b0c10f2` + `ee0ed90` + `pending` | `cargo test event_logs_are_cursor_paginated`; `cargo test admin_logs_are_cursor_paginated_and_include_request_id`; `cargo test rotating_writer_splits_files_when_size_limit_is_reached`; `cargo test --test admin_usage_stats_route_test`; `cargo test --test v1_upstream_route_test`; full `fmt/test/clippy` before commit | Adds cursor pagination, `Page<T>` camelCase serialization, admin `/admin/logs` query API, admin `/admin/usage-stats` and `/admin/usage-stats/summary`, event log repository, runtime tracing initialization, HTTP trace layer, date/size-based file rotation, and v1 response lifecycle events with request/account/route/model/status/latency/non-secret metadata. |
|
||||
| TLS headers and fingerprint | Completed | `1bdd08c` | `cargo test codex_headers_include_desktop_identity_and_turn_state`; `cargo tree \| rg 'reqwest\|rustls'`; full `fmt/test/clippy` before commit | Adds Codex Desktop fingerprint model, exact identity/request headers, and pinned reqwest/rustls client builder. |
|
||||
| Cookie persistence | Completed | `4710926` + `5d3cc6e` | `cargo test cookie_jar_captures_and_replays_account_scoped_cookies`; `cargo test cookie_repository_encrypts_and_replays_account_scoped_cookies`; full `fmt/test/clippy` before commit | Adds account-scoped in-memory Cookie jar and encrypted SQLite Cookie repository; Cookies are replayed only for the same account/domain. |
|
||||
| Account pool and refresh | Partial | `0f97c5c` + `2f25db8` + `4dc0c9a` + `c2cf7a6` + `636fa79` | `cargo test account_pool_skips_expired_disabled_banned_and_quota_exhausted_accounts`; `cargo test --test account_pool_scheduling_test`; `cargo test --test account_refresh_test`; `cargo test --test v1_upstream_route_test`; `cargo test --test account_pool_startup_test`; full `fmt/test/clippy` before commit | Implements base account status filtering, refresh scheduler semantics, max concurrent slots, stale slot cleanup, least-used/round-robin/sticky strategies, tier priority, quota-cache filtering, Cloudflare cooldown skip, request staggering metadata, preferred account selection, exclude IDs, capacity summary, settings-backed pool construction, same-process import-to-pool sync, and startup pool snapshots from SQLite with usage-based last-used metadata. Still missing durable release-time slot state, upstream retry/fallback wiring, and admin account CRUD. |
|
||||
| Account pool and refresh | Partial | `0f97c5c` + `2f25db8` + `4dc0c9a` + `c2cf7a6` + `636fa79` + `pending` | `cargo test account_pool_skips_expired_disabled_banned_and_quota_exhausted_accounts`; `cargo test --test account_pool_scheduling_test`; `cargo test --test account_refresh_test`; `cargo test --test oauth_refresh_test`; `cargo test --test v1_upstream_route_test`; `cargo test --test account_pool_startup_test`; full `fmt/test/clippy` before commit | Implements base account status filtering, refresh scheduler semantics, real OpenAI OAuth refresh-token exchange, 401-triggered refresh and single retry for `/v1/responses`, max concurrent slots, stale slot cleanup, least-used/round-robin/sticky strategies, tier priority, quota-cache filtering, Cloudflare cooldown skip, request staggering metadata, preferred account selection, exclude IDs, capacity summary, settings-backed pool construction, same-process import-to-pool sync, and startup pool snapshots from SQLite with usage-based last-used metadata. Still missing durable release-time slot state, broader upstream retry/fallback wiring, and admin account CRUD. |
|
||||
| Account repository | Completed | `14fe706` | `cargo test --test account_repository_test`; `cargo fmt --check`; `cargo clippy --all-targets --all-features --locked -- -D warnings`; `cargo test` | Adds encrypted SQLite account token storage, decrypting reads, cursor pagination, status/label updates, release-time usage counter persistence, and refresh-token-preserving token updates. |
|
||||
| Translation | Completed | `eff182e` | `cargo test chat_completion_translates_to_codex_response_request`; full `fmt/test/clippy` before commit | Adds OpenAI Chat request types, Codex Responses request type, translation, and OpenAI-compatible error body helper. |
|
||||
| HTTP routes | Partial | `ab0ad8f` + `3830b21` + `6179d44` + `cfec8ba` + `6341336` + `228febd` + `c2cf7a6` + `e9fdea0` + `ee0ed90` | `cargo test v1_requires_client_api_key_not_admin_cookie`; `cargo test models_route_returns_openai_compatible_codex_model_list`; `cargo test --test admin_settings_route_test`; `cargo test --test model_catalog_test`; `cargo test --test admin_accounts_route_test`; `cargo test --test admin_api_keys_route_test`; `cargo test --test v1_upstream_route_test`; full `fmt/test/clippy` before commit | Adds Axum router, health route, `/v1/models`, `/v1/models/catalog`, `/v1/models/{id}`, `/v1/models/{id}/info`, `/debug/models`, `/admin/accounts`, `/admin/accounts/import`, `/admin/api-keys`, `/v1/responses` HTTP SSE collection and `stream:true` passthrough through imported Codex accounts, `/v1` unauthorized OpenAI error body, global request-id response header, binary listener, `gpt-5.5` as the default/test model baseline, and read-only `/admin/settings` with admin-session auth and lower camelCase envelope. `/admin/accounts` reads metadata only and does not decrypt tokens for list views. Still missing OAuth/device/CLI account login flows, admin account mutation/health/quota/Cookie routes, full admin API-key mutation/export/import routes, settings mutation/persistence, diagnostics, usage stats, backend model refresh, model cache persistence, and WebSocket Responses routes. |
|
||||
| HTTP routes | Partial | `ab0ad8f` + `3830b21` + `6179d44` + `cfec8ba` + `6341336` + `228febd` + `c2cf7a6` + `e9fdea0` + `ee0ed90` + `pending` | `cargo test v1_requires_client_api_key_not_admin_cookie`; `cargo test models_route_returns_openai_compatible_codex_model_list`; `cargo test --test admin_settings_route_test`; `cargo test --test model_catalog_test`; `cargo test --test admin_accounts_route_test`; `cargo test --test admin_api_keys_route_test`; `cargo test --test admin_usage_stats_route_test`; `cargo test --test v1_upstream_route_test`; full `fmt/test/clippy` before commit | Adds Axum router, health route, `/v1/models`, `/v1/models/catalog`, `/v1/models/{id}`, `/v1/models/{id}/info`, `/debug/models`, `/admin/accounts`, `/admin/accounts/import`, `/admin/api-keys`, `/admin/usage-stats`, `/admin/usage-stats/summary`, `/v1/responses` HTTP SSE collection and `stream:true` passthrough through imported Codex accounts, `/v1` unauthorized OpenAI error body, global request-id response header, binary listener, `gpt-5.5` as the default/test model baseline, and read-only `/admin/settings` with admin-session auth and lower camelCase envelope. `/admin/accounts` reads metadata only and does not decrypt tokens for list views. Still missing OAuth/device/CLI account login flows, admin account mutation/health/quota/Cookie routes, full admin API-key mutation/export/import routes, settings mutation/persistence, diagnostics, backend model refresh, model cache persistence, and WebSocket Responses routes. |
|
||||
| Admin/frontend response contract | Completed | `1c34ea7` + `b0c10f2` | `cargo test admin_envelope_serializes_lower_camel_case_request_id`; `cargo test admin_logs_reject_missing_admin_session_cookie`; `cargo test v1_response_has_request_id_header_without_admin_body_field`; full `fmt/test/clippy` before commit | Adds reusable admin envelopes with body codes, lower camelCase JSON, paginated list shape, requestId bodies, admin session cookie enforcement, and keeps `/v1/*` OpenAI-compatible. |
|
||||
| HTTP auth boundary | Completed | `cec4607` + `e9fdea0` | `cargo test --test http_auth_test`; `cargo test --test routes_responses_test`; `cargo test --test admin_api_keys_route_test`; `cargo test --test admin_logs_route_test`; `cargo fmt --check`; `cargo clippy --all-targets --all-features --locked -- -D warnings`; `cargo test` | Adds a real `src/http/auth.rs` boundary for client API key and admin session cookie parsing. `/v1/*` accepts only enabled local `Bearer cpr_...` client API keys verified from SQLite, admin session parsing ignores client API keys, and route handlers no longer own ad hoc credential parsing. |
|
||||
| Model catalog | Partial | `6341336` | `cargo test --test model_catalog_test`; `cargo test model_catalog_route_returns_codex_metadata_without_alias_entries`; full `fmt/test/clippy` before commit | Adds a native static Codex model catalog with `gpt-5.5` default, configured aliases, reasoning/service-tier suffix parsing, OpenAI-compatible `/v1/models`, extended catalog/detail/info routes, and debug summary. Still missing backend Codex model fetching per plan, cache persistence, `/admin/refresh-models`, plan allowlist integration from fetched models, and full model metadata parity. |
|
||||
| Upstream lifecycle | Partial | `88cad87` + `225fbff` + `c2cf7a6` + `636fa79` + `ee0ed90` | `cargo test responses_route_rejects_non_codex_provider_models`; `cargo test v1_requires_client_api_key_not_admin_cookie`; `cargo test --test codex_sse_usage_test`; `cargo test --test codex_client_test`; `cargo test --test v1_upstream_route_test`; `cargo test --test account_pool_startup_test`; full `fmt/test/clippy` before commit | `/v1` enforces client API key shape, rejects non-Codex provider models through the model catalog, returns OpenAI-compatible errors, includes reusable SSE/usage helpers, sends Codex Desktop HTTP SSE requests through imported accounts, replays/captures encrypted Cookies, collects `response.completed` for non-streaming clients, passes through client `stream:true` SSE, records token usage, writes v1 lifecycle event logs, and restores the account pool from SQLite on startup. Still missing rate-limit header capture, retry/fallback across accounts, 401 refresh integration, 429 quota cache, 402/403 classification, and previous-response/session affinity. |
|
||||
| Upstream lifecycle | Partial | `88cad87` + `225fbff` + `c2cf7a6` + `636fa79` + `ee0ed90` + `pending` | `cargo test responses_route_rejects_non_codex_provider_models`; `cargo test v1_requires_client_api_key_not_admin_cookie`; `cargo test --test codex_sse_usage_test`; `cargo test --test codex_client_test`; `cargo test --test oauth_refresh_test`; `cargo test --test v1_upstream_route_test`; `cargo test --test account_pool_startup_test`; full `fmt/test/clippy` before commit | `/v1` enforces client API key shape, rejects non-Codex provider models through the model catalog, returns OpenAI-compatible errors, includes reusable SSE/usage helpers, sends Codex Desktop HTTP SSE requests through imported accounts, replays/captures encrypted Cookies, collects `response.completed` for non-streaming clients, passes through client `stream:true` SSE, refreshes and retries once after upstream 401 when a refresh token is stored, records token usage, writes v1 lifecycle event logs, and restores the account pool from SQLite on startup. Still missing rate-limit header capture, retry/fallback across accounts beyond refresh retry, 429 quota cache, 402/403 classification, and previous-response/session affinity. |
|
||||
| Codex WebSocket boundary | Completed | `cc0516d` | `cargo test --test codex_websocket_test`; `cargo fmt --check`; `cargo clippy --all-targets --all-features --locked -- -D warnings`; `cargo test` | Adds `src/codex/websocket.rs` with tested transport classification for HTTP SSE versus WebSocket-only requests. `previous_response_id` and explicit WebSocket requests now fail through a typed boundary instead of being silently downgraded; real WebSocket transport remains an upstream lifecycle follow-up. |
|
||||
| Fingerprint updates | Partial | `dd26d31` + `e32220a` | `cargo test update_manifest_updates_app_version_and_build_number`; `cargo test fingerprint_updater_fetches_manifest_and_persists_history`; full `fmt/test/clippy` before commit | Parses a JSON update manifest and persists app version/build number history to SQLite. Still missing original appcast XML polling, update-state persistence, runtime header fingerprint selection from the latest stored fingerprint, and optional local Codex.app Chromium extraction. |
|
||||
| Runtime docs and packaging | Completed | `33b63f3` | `cargo fmt --check`; `cargo clippy --all-targets --all-features --locked -- -D warnings`; `cargo test`; `cargo run`; `curl -s -i http://127.0.0.1:8080/health` | README scope, `.env.example`, Dockerfile, compose file, and health smoke test are complete. |
|
||||
|
||||
@@ -2967,6 +2967,8 @@ git commit -m "docs: document codex-proxy-rs scope"
|
||||
- [x] Restore the process account pool from SQLite on startup. The snapshot decrypts account tokens through `SecretBox`, carries usage `last_used_at` into scheduling, and lets previously imported accounts survive service restarts.
|
||||
- [x] Split local client API key auth from admin login. `/admin/api-keys` can create/list local `cpr_` keys, plaintext is returned only once, list responses never expose plaintext/hash, and `/v1/*` verifies enabled keys from SQLite instead of accepting any `cpr_` string.
|
||||
- [x] Add tested `/v1/responses` streaming passthrough and lifecycle event logs. Client `stream:true` returns `text/event-stream`, keeps the account slot until stream completion or disconnect, parses usage after SSE completion, writes account usage, and records non-secret `v1.response` logs with request/account/route/model/status/latency metadata. Non-streaming calls now write the same lifecycle log after `response.completed`.
|
||||
- [x] Add tested admin usage-stat read APIs. `/admin/usage-stats` returns cursor-paginated account usage counters, `/admin/usage-stats/summary` returns global totals, and both require the admin session cookie while exposing only non-secret account metadata.
|
||||
- [x] Add real OpenAI OAuth refresh-token exchange and `/v1/responses` 401 retry wiring. The Rust service refreshes access tokens through the OAuth token endpoint, preserves the existing refresh token when the server omits a replacement, persists the rotated access token, updates the in-process account pool, and retries the failed non-streaming or streaming request once.
|
||||
|
||||
## 2026-06-11 Original Project Gap Audit
|
||||
|
||||
@@ -2981,10 +2983,10 @@ The Rust project has no legacy compatibility requirement, but the in-scope behav
|
||||
- [ ] **Client API key management:** implement admin CRUD for local client API keys: list/create/delete/batch-delete/label/status/export/import. Create/list and `/v1/*` SQLite verification are complete. Remove provider/model binding semantics from the TypeScript version; Rust client API keys authorize only local `/v1/*` access.
|
||||
- [ ] **Settings:** implement admin settings for server/runtime fields that remain in scope: default model, reasoning effort, service tier, model aliases, refresh enabled/margin/concurrency, max concurrent per account, request interval, rotation strategy, tier priority, quota refresh/thresholds/skip exhausted, logs state/capacity/body capture, and usage history retention. Read-only `/admin/settings` is complete; persistent mutation and config-store writeback remain. Remove proxy URL, Ollama, third-party provider, OpenAI official key, and Electron/self-update settings.
|
||||
- [ ] **Diagnostics:** expand health/admin diagnostics to include authenticated state, pool summary, capacity summary, transport/fingerprint status, paths, runtime metadata, and test-connection checks. Keep production-local gating for sensitive debug endpoints.
|
||||
- [ ] **Logs and usage stats:** add `/admin/logs/state`, clear/detail, error-log equivalents if retained, and `/admin/usage-stats/summary` plus paginated or bounded history queries. Keep log file rotation and SQLite event pagination.
|
||||
- [ ] **Logs and usage stats:** add `/admin/logs/state`, clear/detail, error-log equivalents if retained, and richer bounded history queries if needed. Keep log file rotation and SQLite event pagination. Initial `/admin/usage-stats` pagination and `/admin/usage-stats/summary` are complete.
|
||||
- [ ] **Model catalog:** replace the hard-coded `/v1/models` list with a model store that supports static defaults, configured aliases, `-low/-medium/-high/-xhigh/-fast/-flex` suffix parsing, backend model fetching per plan, cache persistence, `/v1/models/catalog`, `/v1/models/{id}`, `/v1/models/{id}/info`, `/admin/refresh-models`, and `/debug/models`. Static defaults, aliases, suffix parsing, read-only model routes, and debug summary are complete. Backend fetch/cache, plan snapshots, and admin refresh remain.
|
||||
- [ ] **Account scheduling parity:** extend account acquisition beyond `active + least last_used`. Completed in the in-memory pool: max concurrent slots per account, stale slot cleanup, least-used/round-robin/sticky strategies, quota-cache skip, tier priority, model-plan filtering, exclude IDs, preferred account from session affinity, Cloudflare cooldown skip, request-staggering metadata, capacity summary, settings-backed construction, import-to-pool sync, and startup pool snapshots from SQLite. Remaining: durable release-time slot state and upstream retry/fallback wiring.
|
||||
- [ ] **Upstream Codex lifecycle:** implement account acquire -> Codex request -> SSE/collect-or-stream -> rate-limit capture -> usage release -> retry/fallback. Initial HTTP SSE collection and streaming passthrough are complete: `/v1/responses` acquires imported/restored accounts, sends Codex Desktop headers, replays/captures Cookies, collects `response.completed` for non-streaming clients, passes through `stream:true` SSE, records usage, writes lifecycle logs, and returns the correct OpenAI-compatible JSON or SSE shape. Remaining: rate-limit capture, retry/fallback, 401-triggered refresh, 429 quota cache, 402 quota exhausted, 403 ban/Cloudflare distinction, path-block Cookie clearing, and richer OpenAI-compatible error mapping.
|
||||
- [ ] **Upstream Codex lifecycle:** implement account acquire -> Codex request -> SSE/collect-or-stream -> rate-limit capture -> usage release -> retry/fallback. Initial HTTP SSE collection, streaming passthrough, and 401-triggered refresh retry are complete: `/v1/responses` acquires imported/restored accounts, sends Codex Desktop headers, replays/captures Cookies, collects `response.completed` for non-streaming clients, passes through `stream:true` SSE, refreshes and retries once after upstream 401, records usage, writes lifecycle logs, and returns the correct OpenAI-compatible JSON or SSE shape. Remaining: rate-limit capture, retry/fallback across alternate accounts, 429 quota cache, 402 quota exhausted, 403 ban/Cloudflare distinction, path-block Cookie clearing, and richer OpenAI-compatible error mapping.
|
||||
- [ ] **Session affinity and previous response:** persist an in-memory response-id to account mapping with TTL, prompt-cache/conversation identity, turn-state replay, and preferred account selection. `src/codex/websocket.rs` must either implement the WebSocket path with verified TLS/header behavior or the API must reject WebSocket-only cases explicitly with documented limitations until fingerprint parity is proven.
|
||||
- [ ] **Fingerprint auto-update parity:** poll the real Codex Desktop appcast/update source, parse app version/build number, persist update state, select the latest stored fingerprint for runtime headers, and optionally extract Chromium version from a local Codex.app path when configured.
|
||||
|
||||
|
||||
@@ -93,6 +93,28 @@ pub struct AccountUsageRecord {
|
||||
pub last_used_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountUsageListRecord {
|
||||
pub account_id: String,
|
||||
pub email: Option<String>,
|
||||
pub label: Option<String>,
|
||||
pub plan_type: Option<String>,
|
||||
pub request_count: i64,
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cached_tokens: i64,
|
||||
pub last_used_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct AccountUsageSummary {
|
||||
pub account_count: i64,
|
||||
pub request_count: i64,
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cached_tokens: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AccountRepository {
|
||||
pool: SqlitePool,
|
||||
@@ -438,6 +460,97 @@ fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> AccountRepositoryResult<Acco
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AccountUsageRepository {
|
||||
pool: SqlitePool,
|
||||
}
|
||||
|
||||
impl AccountUsageRepository {
|
||||
pub fn new(pool: SqlitePool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub async fn list(
|
||||
&self,
|
||||
cursor: Option<String>,
|
||||
limit: u32,
|
||||
) -> AccountRepositoryResult<Page<AccountUsageListRecord>> {
|
||||
let fetch_limit = i64::from(limit) + 1;
|
||||
let rows = if let Some(cursor) = cursor {
|
||||
let (last_used_at, account_id) =
|
||||
decode_cursor(&cursor).ok_or(AccountRepositoryError::InvalidCursor)?;
|
||||
sqlx::query(
|
||||
"select account_usage.account_id, accounts.email, accounts.label, accounts.plan_type, account_usage.request_count, account_usage.input_tokens, account_usage.output_tokens, account_usage.cached_tokens, account_usage.last_used_at from account_usage left join accounts on accounts.id = account_usage.account_id where coalesce(account_usage.last_used_at, '') < ? or (coalesce(account_usage.last_used_at, '') = ? and account_usage.account_id < ?) order by coalesce(account_usage.last_used_at, '') desc, account_usage.account_id desc limit ?",
|
||||
)
|
||||
.bind(&last_used_at)
|
||||
.bind(last_used_at)
|
||||
.bind(account_id)
|
||||
.bind(fetch_limit)
|
||||
.fetch_all(&self.pool)
|
||||
.await?
|
||||
} else {
|
||||
sqlx::query(
|
||||
"select account_usage.account_id, accounts.email, accounts.label, accounts.plan_type, account_usage.request_count, account_usage.input_tokens, account_usage.output_tokens, account_usage.cached_tokens, account_usage.last_used_at from account_usage left join accounts on accounts.id = account_usage.account_id order by coalesce(account_usage.last_used_at, '') desc, account_usage.account_id desc limit ?",
|
||||
)
|
||||
.bind(fetch_limit)
|
||||
.fetch_all(&self.pool)
|
||||
.await?
|
||||
};
|
||||
|
||||
let has_next = rows.len() > limit as usize;
|
||||
let take_count = rows.len().min(limit as usize);
|
||||
let mut items = Vec::with_capacity(take_count);
|
||||
for row in rows.into_iter().take(take_count) {
|
||||
items.push(usage_list_from_row(&row)?);
|
||||
}
|
||||
let next_cursor = if has_next {
|
||||
items.last().map(|usage| {
|
||||
encode_cursor(
|
||||
&usage
|
||||
.last_used_at
|
||||
.map(|value| value.to_rfc3339())
|
||||
.unwrap_or_default(),
|
||||
&usage.account_id,
|
||||
)
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
Ok(Page { items, next_cursor })
|
||||
}
|
||||
|
||||
pub async fn summary(&self) -> AccountRepositoryResult<AccountUsageSummary> {
|
||||
let row = sqlx::query(
|
||||
"select count(*) as account_count, coalesce(sum(request_count), 0) as request_count, coalesce(sum(input_tokens), 0) as input_tokens, coalesce(sum(output_tokens), 0) as output_tokens, coalesce(sum(cached_tokens), 0) as cached_tokens from account_usage",
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
Ok(AccountUsageSummary {
|
||||
account_count: row.get("account_count"),
|
||||
request_count: row.get("request_count"),
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_list_from_row(
|
||||
row: &sqlx::sqlite::SqliteRow,
|
||||
) -> AccountRepositoryResult<AccountUsageListRecord> {
|
||||
Ok(AccountUsageListRecord {
|
||||
account_id: row.get("account_id"),
|
||||
email: row.get("email"),
|
||||
label: row.get("label"),
|
||||
plan_type: row.get("plan_type"),
|
||||
request_count: row.get("request_count"),
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
last_used_at: parse_optional_rfc3339(row.get::<Option<String>, _>("last_used_at"))?,
|
||||
})
|
||||
}
|
||||
|
||||
fn metadata_from_row(
|
||||
row: &sqlx::sqlite::SqliteRow,
|
||||
) -> AccountRepositoryResult<StoredAccountMetadata> {
|
||||
|
||||
+6
-1
@@ -7,7 +7,10 @@ use tower_http::trace::TraceLayer;
|
||||
|
||||
use crate::{
|
||||
http::{
|
||||
admin::{accounts, api_keys, create_api_key, import_accounts, login, logs, settings},
|
||||
admin::{
|
||||
accounts, api_keys, create_api_key, import_accounts, login, logs, settings,
|
||||
usage_stats, usage_stats_summary,
|
||||
},
|
||||
health::health,
|
||||
middleware::attach_request_id,
|
||||
v1::{debug_models, model_catalog, model_detail, model_info, models, responses},
|
||||
@@ -28,6 +31,8 @@ pub fn build_router(state: AppState) -> Router {
|
||||
.route("/admin/login", post(login))
|
||||
.route("/admin/logs", get(logs))
|
||||
.route("/admin/settings", get(settings))
|
||||
.route("/admin/usage-stats", get(usage_stats))
|
||||
.route("/admin/usage-stats/summary", get(usage_stats_summary))
|
||||
.route("/admin/accounts", get(accounts))
|
||||
.route("/admin/accounts/import", post(import_accounts))
|
||||
.route("/admin/api-keys", get(api_keys).post(create_api_key))
|
||||
|
||||
@@ -1,6 +1,109 @@
|
||||
use async_trait::async_trait;
|
||||
use reqwest::{Client, StatusCode};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::auth::{
|
||||
refresh::{RefreshFailure, TokenRefresher},
|
||||
token::TokenPair,
|
||||
};
|
||||
|
||||
const CODEX_OAUTH_CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
|
||||
const OPENAI_OAUTH_AUTHORIZE_ENDPOINT: &str = "https://auth.openai.com/oauth/authorize";
|
||||
const OPENAI_OAUTH_TOKEN_ENDPOINT: &str = "https://auth.openai.com/oauth/token";
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OAuthConfig {
|
||||
pub client_id: String,
|
||||
pub auth_endpoint: String,
|
||||
pub token_endpoint: String,
|
||||
}
|
||||
|
||||
impl OAuthConfig {
|
||||
pub fn codex_default() -> Self {
|
||||
Self {
|
||||
client_id: CODEX_OAUTH_CLIENT_ID.to_string(),
|
||||
auth_endpoint: OPENAI_OAUTH_AUTHORIZE_ENDPOINT.to_string(),
|
||||
token_endpoint: OPENAI_OAUTH_TOKEN_ENDPOINT.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct OpenAiOAuthRefresher {
|
||||
client: Client,
|
||||
config: OAuthConfig,
|
||||
}
|
||||
|
||||
impl OpenAiOAuthRefresher {
|
||||
pub fn new(client: Client, config: OAuthConfig) -> Self {
|
||||
Self { client, config }
|
||||
}
|
||||
|
||||
pub fn codex_default(client: Client) -> Self {
|
||||
Self::new(client, OAuthConfig::codex_default())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RefreshTokenResponse {
|
||||
access_token: String,
|
||||
refresh_token: Option<String>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenRefresher for OpenAiOAuthRefresher {
|
||||
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
|
||||
let response = self
|
||||
.client
|
||||
.post(&self.config.token_endpoint)
|
||||
.form(&[
|
||||
("grant_type", "refresh_token"),
|
||||
("client_id", self.config.client_id.as_str()),
|
||||
("refresh_token", refresh_token),
|
||||
])
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| RefreshFailure::Transport)?;
|
||||
let status = response.status();
|
||||
let body = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|_| RefreshFailure::Transport)?;
|
||||
if !status.is_success() {
|
||||
return Err(classify_refresh_failure(status, &body));
|
||||
}
|
||||
let tokens = serde_json::from_str::<RefreshTokenResponse>(&body)
|
||||
.map_err(|_| RefreshFailure::Transport)?;
|
||||
if tokens.access_token.trim().is_empty() {
|
||||
return Err(RefreshFailure::Transport);
|
||||
}
|
||||
Ok(TokenPair {
|
||||
access_token: tokens.access_token,
|
||||
refresh_token: tokens.refresh_token,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn classify_refresh_failure(status: StatusCode, body: &str) -> RefreshFailure {
|
||||
let lower = body.to_ascii_lowercase();
|
||||
if lower.contains("quota") {
|
||||
return RefreshFailure::QuotaExhausted;
|
||||
}
|
||||
if lower.contains("account has been deactivated")
|
||||
|| lower.contains("refresh_token_reused")
|
||||
|| lower.contains("banned")
|
||||
{
|
||||
return RefreshFailure::Banned;
|
||||
}
|
||||
if status == StatusCode::BAD_REQUEST
|
||||
|| status == StatusCode::UNAUTHORIZED
|
||||
|| lower.contains("invalid_grant")
|
||||
|| lower.contains("invalid_token")
|
||||
|| lower.contains("access_denied")
|
||||
|| lower.contains("refresh_token_expired")
|
||||
|| lower.contains("token_revoked")
|
||||
{
|
||||
return RefreshFailure::InvalidGrant;
|
||||
}
|
||||
RefreshFailure::Transport
|
||||
}
|
||||
|
||||
+163
-1
@@ -16,7 +16,10 @@ use uuid::Uuid;
|
||||
use crate::{
|
||||
accounts::{
|
||||
model::{Account, AccountStatus},
|
||||
repository::{NewAccount, StoredAccountMetadata},
|
||||
repository::{
|
||||
AccountRepositoryError, AccountUsageListRecord, AccountUsageSummary, NewAccount,
|
||||
StoredAccountMetadata,
|
||||
},
|
||||
},
|
||||
auth::admin_session::verify_admin_password,
|
||||
auth::api_key_repository::StoredClientApiKey,
|
||||
@@ -138,6 +141,13 @@ pub struct ApiKeysQuery {
|
||||
pub limit: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UsageStatsQuery {
|
||||
pub cursor: Option<String>,
|
||||
pub limit: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct LoginRequest {
|
||||
@@ -290,6 +300,58 @@ impl From<StoredAccountMetadata> for AdminAccountData {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminUsageStatsData {
|
||||
pub account_id: String,
|
||||
pub email: Option<String>,
|
||||
pub label: Option<String>,
|
||||
pub plan_type: Option<String>,
|
||||
pub request_count: i64,
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cached_tokens: i64,
|
||||
pub last_used_at: Option<String>,
|
||||
}
|
||||
|
||||
impl From<AccountUsageListRecord> for AdminUsageStatsData {
|
||||
fn from(usage: AccountUsageListRecord) -> Self {
|
||||
Self {
|
||||
account_id: usage.account_id,
|
||||
email: usage.email,
|
||||
label: usage.label,
|
||||
plan_type: usage.plan_type,
|
||||
request_count: usage.request_count,
|
||||
input_tokens: usage.input_tokens,
|
||||
output_tokens: usage.output_tokens,
|
||||
cached_tokens: usage.cached_tokens,
|
||||
last_used_at: usage.last_used_at.map(|value| value.to_rfc3339()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminUsageStatsSummaryData {
|
||||
pub account_count: i64,
|
||||
pub request_count: i64,
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cached_tokens: i64,
|
||||
}
|
||||
|
||||
impl From<AccountUsageSummary> for AdminUsageStatsSummaryData {
|
||||
fn from(summary: AccountUsageSummary) -> Self {
|
||||
Self {
|
||||
account_count: summary.account_count,
|
||||
request_count: summary.request_count,
|
||||
input_tokens: summary.input_tokens,
|
||||
output_tokens: summary.output_tokens,
|
||||
cached_tokens: summary.cached_tokens,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AdminSettingsData {
|
||||
fn from_config(config: &AppConfig) -> Self {
|
||||
Self {
|
||||
@@ -482,6 +544,106 @@ pub async fn settings(
|
||||
.into_response()
|
||||
}
|
||||
|
||||
pub async fn usage_stats(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<UsageStatsQuery>,
|
||||
) -> Response {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
let Some(pool) = state.db() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "Database is not initialized", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
if let Err(response) = require_admin_session(pool, &headers, &request_id).await {
|
||||
return response;
|
||||
}
|
||||
let Some(repo) = state.account_usage_repository() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(
|
||||
50001,
|
||||
"Account usage repository is not initialized",
|
||||
(),
|
||||
request_id,
|
||||
),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
|
||||
let limit = clamp_limit(query.limit.unwrap_or(50));
|
||||
match repo.list(query.cursor, limit).await {
|
||||
Ok(page) => {
|
||||
let Page { items, next_cursor } = page;
|
||||
let page = Page {
|
||||
items: items.into_iter().map(AdminUsageStatsData::from).collect(),
|
||||
next_cursor,
|
||||
};
|
||||
AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminPageEnvelope::ok(page, limit, request_id),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
Err(AccountRepositoryError::InvalidCursor) => AdminResponse::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
AdminEnvelope::new(40002, "Invalid cursor", (), request_id),
|
||||
)
|
||||
.into_response(),
|
||||
Err(_) => AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "Failed to list usage stats", (), request_id),
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn usage_stats_summary(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
) -> Response {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
let Some(pool) = state.db() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "Database is not initialized", (), request_id),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
if let Err(response) = require_admin_session(pool, &headers, &request_id).await {
|
||||
return response;
|
||||
}
|
||||
let Some(repo) = state.account_usage_repository() else {
|
||||
return AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(
|
||||
50001,
|
||||
"Account usage repository is not initialized",
|
||||
(),
|
||||
request_id,
|
||||
),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
|
||||
match repo.summary().await {
|
||||
Ok(summary) => AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminUsageStatsSummaryData::from(summary), request_id),
|
||||
)
|
||||
.into_response(),
|
||||
Err(_) => AdminResponse::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
AdminEnvelope::new(50001, "Failed to summarize usage stats", (), request_id),
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn accounts(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
|
||||
+123
-6
@@ -13,16 +13,18 @@ use axum::{
|
||||
use chrono::Utc;
|
||||
use futures::{stream, StreamExt};
|
||||
use reqwest::Url;
|
||||
use secrecy::SecretString;
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::{
|
||||
accounts::{
|
||||
model::Account,
|
||||
model::{Account, AccountStatus},
|
||||
pool::{AccountAcquireRequest, AccountPool},
|
||||
repository::UsageDelta,
|
||||
repository::{TokenUpdate, UsageDelta},
|
||||
},
|
||||
auth::refresh::RefreshFailure,
|
||||
codex::{
|
||||
client::{
|
||||
build_reqwest_client, CodexBackendClient, CodexBackendStream, CodexClientError,
|
||||
@@ -112,7 +114,13 @@ pub async fn responses(
|
||||
return responses_stream(state, codex_request, account, log_context).await;
|
||||
}
|
||||
|
||||
let response = send_codex_request(&state, &codex_request, &account, request_id.as_str()).await;
|
||||
let response = send_codex_request_with_refresh_retry(
|
||||
&state,
|
||||
&codex_request,
|
||||
&account,
|
||||
request_id.as_str(),
|
||||
)
|
||||
.await;
|
||||
state.account_pool().lock().await.release(&account.id);
|
||||
|
||||
let response = match response {
|
||||
@@ -251,6 +259,23 @@ fn default_body(default_model: String) -> ResponsesBody {
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_codex_request_with_refresh_retry(
|
||||
state: &AppState,
|
||||
request: &CodexResponsesRequest,
|
||||
account: &Account,
|
||||
request_id: &str,
|
||||
) -> Result<crate::codex::client::CodexBackendResponse, CodexClientError> {
|
||||
match send_codex_request(state, request, account, request_id).await {
|
||||
Err(CodexClientError::Upstream { status, body }) if status == StatusCode::UNAUTHORIZED => {
|
||||
let Some(refreshed) = refresh_account_after_unauthorized(state, account).await else {
|
||||
return Err(CodexClientError::Upstream { status, body });
|
||||
};
|
||||
send_codex_request(state, request, &refreshed, request_id).await
|
||||
}
|
||||
result => result,
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_codex_request(
|
||||
state: &AppState,
|
||||
request: &CodexResponsesRequest,
|
||||
@@ -281,6 +306,23 @@ async fn send_codex_request(
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_codex_stream_request_with_refresh_retry(
|
||||
state: &AppState,
|
||||
request: &CodexResponsesRequest,
|
||||
account: &Account,
|
||||
request_id: &str,
|
||||
) -> Result<CodexBackendStream, CodexClientError> {
|
||||
match send_codex_stream_request(state, request, account, request_id).await {
|
||||
Err(CodexClientError::Upstream { status, body }) if status == StatusCode::UNAUTHORIZED => {
|
||||
let Some(refreshed) = refresh_account_after_unauthorized(state, account).await else {
|
||||
return Err(CodexClientError::Upstream { status, body });
|
||||
};
|
||||
send_codex_stream_request(state, request, &refreshed, request_id).await
|
||||
}
|
||||
result => result,
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_codex_stream_request(
|
||||
state: &AppState,
|
||||
request: &CodexResponsesRequest,
|
||||
@@ -317,9 +359,13 @@ async fn responses_stream(
|
||||
account: Account,
|
||||
log_context: V1LogContext,
|
||||
) -> Response {
|
||||
let stream_response =
|
||||
send_codex_stream_request(&state, &request, &account, log_context.request_id.as_str())
|
||||
.await;
|
||||
let stream_response = send_codex_stream_request_with_refresh_retry(
|
||||
&state,
|
||||
&request,
|
||||
&account,
|
||||
log_context.request_id.as_str(),
|
||||
)
|
||||
.await;
|
||||
let stream_response = match stream_response {
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
@@ -399,6 +445,77 @@ async fn responses_stream(
|
||||
})
|
||||
}
|
||||
|
||||
async fn refresh_account_after_unauthorized(
|
||||
state: &AppState,
|
||||
account: &Account,
|
||||
) -> Option<Account> {
|
||||
if !state.config().auth.refresh_enabled {
|
||||
return None;
|
||||
}
|
||||
let refresh_token = account.refresh_token.as_deref()?;
|
||||
let refresher = state.token_refresher()?;
|
||||
match refresher.refresh(refresh_token).await {
|
||||
Ok(tokens) => persist_refreshed_account(state, account, tokens).await,
|
||||
Err(failure) => {
|
||||
mark_refresh_failure(state, account, failure).await;
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn persist_refreshed_account(
|
||||
state: &AppState,
|
||||
account: &Account,
|
||||
tokens: crate::auth::token::TokenPair,
|
||||
) -> Option<Account> {
|
||||
let repo = state.account_repository()?;
|
||||
let access_token = tokens.access_token;
|
||||
let refresh_token = tokens.refresh_token;
|
||||
repo.update_tokens(
|
||||
&account.id,
|
||||
TokenUpdate {
|
||||
access_token: SecretString::new(access_token.clone().into()),
|
||||
refresh_token: refresh_token
|
||||
.clone()
|
||||
.map(|token| SecretString::new(token.into())),
|
||||
access_token_expires_at: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.ok()?;
|
||||
|
||||
let mut refreshed = account.clone();
|
||||
refreshed.access_token = access_token;
|
||||
if let Some(refresh_token) = refresh_token {
|
||||
refreshed.refresh_token = Some(refresh_token);
|
||||
}
|
||||
refreshed.status = AccountStatus::Active;
|
||||
state.account_pool().lock().await.insert(refreshed.clone());
|
||||
Some(refreshed)
|
||||
}
|
||||
|
||||
async fn mark_refresh_failure(state: &AppState, account: &Account, failure: RefreshFailure) {
|
||||
let Some(status) = status_for_refresh_failure(failure) else {
|
||||
return;
|
||||
};
|
||||
if let Some(repo) = state.account_repository() {
|
||||
let _ = repo.set_status(&account.id, status).await;
|
||||
}
|
||||
let mut updated = account.clone();
|
||||
updated.status = status;
|
||||
state.account_pool().lock().await.insert(updated);
|
||||
}
|
||||
|
||||
fn status_for_refresh_failure(failure: RefreshFailure) -> Option<AccountStatus> {
|
||||
match failure {
|
||||
RefreshFailure::InvalidGrant => Some(AccountStatus::Expired),
|
||||
RefreshFailure::QuotaExhausted => Some(AccountStatus::QuotaExhausted),
|
||||
RefreshFailure::Banned => Some(AccountStatus::Banned),
|
||||
RefreshFailure::Disabled => Some(AccountStatus::Disabled),
|
||||
RefreshFailure::Transport => None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn persist_upstream_cookies(
|
||||
state: &AppState,
|
||||
account_id: &str,
|
||||
|
||||
+14
-3
@@ -1,6 +1,6 @@
|
||||
use codex_proxy_rs::{
|
||||
app::build_router,
|
||||
auth::api_key::ApiKeyHasher,
|
||||
auth::{api_key::ApiKeyHasher, oauth::OpenAiOAuthRefresher},
|
||||
config::AppConfig,
|
||||
crypto::SecretBox,
|
||||
logs::rotation::{init_tracing, RotationConfig},
|
||||
@@ -21,9 +21,20 @@ async fn main() -> anyhow::Result<()> {
|
||||
let port = config.server.port;
|
||||
let secret_box = SecretBox::load_or_create(&config.security.master_key_file)?;
|
||||
let api_key_hasher = ApiKeyHasher::load_or_create(&config.security.api_key_pepper_file)?;
|
||||
let token_refresher = OpenAiOAuthRefresher::codex_default(
|
||||
reqwest::Client::builder()
|
||||
.use_rustls_tls()
|
||||
.no_proxy()
|
||||
.build()?,
|
||||
);
|
||||
let pool = connect_sqlite(&config.database.url).await?;
|
||||
let state =
|
||||
AppState::with_pool_secret_and_api_key_hasher(config, pool, secret_box, api_key_hasher);
|
||||
let state = AppState::with_pool_secret_api_key_hasher_and_token_refresher(
|
||||
config,
|
||||
pool,
|
||||
secret_box,
|
||||
api_key_hasher,
|
||||
token_refresher,
|
||||
);
|
||||
let restored_accounts = state.reload_account_pool_from_repository().await?;
|
||||
tracing::info!(restored_accounts, "account pool restored from sqlite");
|
||||
let app = build_router(state);
|
||||
|
||||
+41
-2
@@ -5,9 +5,11 @@ use tokio::sync::Mutex;
|
||||
|
||||
use crate::accounts::{
|
||||
pool::{AccountPool, AccountPoolOptions, RotationStrategy},
|
||||
repository::{AccountRepository, AccountRepositoryResult},
|
||||
repository::{AccountRepository, AccountRepositoryResult, AccountUsageRepository},
|
||||
};
|
||||
use crate::auth::{
|
||||
api_key::ApiKeyHasher, api_key_repository::ClientApiKeyRepository, refresh::TokenRefresher,
|
||||
};
|
||||
use crate::auth::{api_key::ApiKeyHasher, api_key_repository::ClientApiKeyRepository};
|
||||
use crate::config::AppConfig;
|
||||
use crate::cookies::repository::CookieRepository;
|
||||
use crate::crypto::SecretBox;
|
||||
@@ -24,6 +26,7 @@ pub struct AppServices {
|
||||
pub event_logs: Option<EventLogRepository>,
|
||||
pub secret_box: Option<SecretBox>,
|
||||
pub api_key_hasher: Option<ApiKeyHasher>,
|
||||
pub token_refresher: Option<Arc<dyn TokenRefresher>>,
|
||||
pub account_pool: Arc<Mutex<AccountPool>>,
|
||||
}
|
||||
|
||||
@@ -37,6 +40,7 @@ impl AppState {
|
||||
event_logs: None,
|
||||
secret_box: None,
|
||||
api_key_hasher: None,
|
||||
token_refresher: None,
|
||||
account_pool,
|
||||
}),
|
||||
}
|
||||
@@ -51,6 +55,7 @@ impl AppState {
|
||||
event_logs: Some(EventLogRepository::new(pool)),
|
||||
secret_box: None,
|
||||
api_key_hasher: None,
|
||||
token_refresher: None,
|
||||
account_pool,
|
||||
}),
|
||||
}
|
||||
@@ -69,6 +74,7 @@ impl AppState {
|
||||
event_logs: Some(EventLogRepository::new(pool)),
|
||||
secret_box: Some(secret_box),
|
||||
api_key_hasher: None,
|
||||
token_refresher: None,
|
||||
account_pool,
|
||||
}),
|
||||
}
|
||||
@@ -88,6 +94,31 @@ impl AppState {
|
||||
event_logs: Some(EventLogRepository::new(pool)),
|
||||
secret_box: Some(secret_box),
|
||||
api_key_hasher: Some(api_key_hasher),
|
||||
token_refresher: None,
|
||||
account_pool,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_pool_secret_api_key_hasher_and_token_refresher<C>(
|
||||
config: AppConfig,
|
||||
pool: SqlitePool,
|
||||
secret_box: SecretBox,
|
||||
api_key_hasher: ApiKeyHasher,
|
||||
token_refresher: C,
|
||||
) -> Self
|
||||
where
|
||||
C: TokenRefresher,
|
||||
{
|
||||
let account_pool = account_pool_from_config(&config);
|
||||
Self {
|
||||
services: Arc::new(AppServices {
|
||||
config,
|
||||
db: Some(pool.clone()),
|
||||
event_logs: Some(EventLogRepository::new(pool)),
|
||||
secret_box: Some(secret_box),
|
||||
api_key_hasher: Some(api_key_hasher),
|
||||
token_refresher: Some(Arc::new(token_refresher)),
|
||||
account_pool,
|
||||
}),
|
||||
}
|
||||
@@ -112,6 +143,10 @@ impl AppState {
|
||||
))
|
||||
}
|
||||
|
||||
pub fn account_usage_repository(&self) -> Option<AccountUsageRepository> {
|
||||
Some(AccountUsageRepository::new(self.db()?.clone()))
|
||||
}
|
||||
|
||||
pub fn cookie_repository(&self) -> Option<CookieRepository> {
|
||||
Some(CookieRepository::new(
|
||||
self.db()?.clone(),
|
||||
@@ -131,6 +166,10 @@ impl AppState {
|
||||
self.services.api_key_hasher.as_ref()
|
||||
}
|
||||
|
||||
pub fn token_refresher(&self) -> Option<Arc<dyn TokenRefresher>> {
|
||||
self.services.token_refresher.clone()
|
||||
}
|
||||
|
||||
pub async fn reload_account_pool_from_repository(&self) -> AccountRepositoryResult<usize> {
|
||||
let Some(repo) = self.account_repository() else {
|
||||
return Ok(0);
|
||||
|
||||
@@ -0,0 +1,264 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use axum::{
|
||||
body::{to_bytes, Body},
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use chrono::Utc;
|
||||
use serde_json::Value;
|
||||
use tower::ServiceExt;
|
||||
|
||||
use codex_proxy_rs::{
|
||||
app::build_router,
|
||||
config::{
|
||||
AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig,
|
||||
QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig,
|
||||
UsageStatsConfig,
|
||||
},
|
||||
state::AppState,
|
||||
storage::db::connect_sqlite,
|
||||
};
|
||||
|
||||
fn test_config(database_url: String) -> AppConfig {
|
||||
AppConfig {
|
||||
server: ServerConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 0,
|
||||
},
|
||||
api: ApiConfig {
|
||||
base_url: "https://chatgpt.com/backend-api".to_string(),
|
||||
},
|
||||
model: ModelConfig {
|
||||
default_model: "gpt-5.5".to_string(),
|
||||
default_reasoning_effort: None,
|
||||
service_tier: None,
|
||||
aliases: BTreeMap::new(),
|
||||
},
|
||||
auth: AuthConfig {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_enabled: true,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 3,
|
||||
request_interval_ms: 50,
|
||||
rotation_strategy: "least_used".to_string(),
|
||||
tier_priority: Vec::new(),
|
||||
},
|
||||
quota: QuotaConfig {
|
||||
refresh_interval_minutes: 5,
|
||||
warning_thresholds: QuotaWarningThresholds {
|
||||
primary: vec![80, 90],
|
||||
secondary: vec![80, 90],
|
||||
},
|
||||
skip_exhausted: true,
|
||||
},
|
||||
usage_stats: UsageStatsConfig {
|
||||
history_retention_days: None,
|
||||
},
|
||||
database: DatabaseConfig { url: database_url },
|
||||
security: SecurityConfig {
|
||||
master_key_file: "data/master.key".to_string(),
|
||||
api_key_pepper_file: "data/api-key-pepper.key".to_string(),
|
||||
},
|
||||
tls: TlsConfig {
|
||||
force_http11: false,
|
||||
},
|
||||
admin: AdminConfig {
|
||||
session_ttl_minutes: 1440,
|
||||
},
|
||||
logging: LoggingConfig {
|
||||
directory: "logs".to_string(),
|
||||
max_file_bytes: 10_485_760,
|
||||
retention_days: 14,
|
||||
enabled: false,
|
||||
capacity: 2_000,
|
||||
capture_body: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_usage_stats_should_cursor_page_account_usage() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-usage-stats.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
seed_admin_session(&pool, "session_1").await;
|
||||
seed_usage(&pool).await;
|
||||
let app = build_router(AppState::with_pool(test_config(url), pool));
|
||||
|
||||
let first_response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/admin/usage-stats?limit=1")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_usage")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(first_response.status(), StatusCode::OK);
|
||||
let first_body = response_json(first_response).await;
|
||||
assert_eq!(first_body["code"], 200);
|
||||
assert_eq!(first_body["requestId"], "req_usage");
|
||||
assert_eq!(first_body["data"].as_array().unwrap().len(), 1);
|
||||
assert_eq!(first_body["data"][0]["accountId"], "acct_b");
|
||||
assert_eq!(first_body["data"][0]["requestCount"], 2);
|
||||
assert_eq!(first_body["data"][0]["inputTokens"], 7);
|
||||
assert_eq!(first_body["page"]["limit"], 1);
|
||||
let cursor = first_body["page"]["nextCursor"].as_str().unwrap();
|
||||
|
||||
let second_response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri(format!("/admin/usage-stats?limit=1&cursor={cursor}"))
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(second_response.status(), StatusCode::OK);
|
||||
let second_body = response_json(second_response).await;
|
||||
assert_eq!(second_body["data"][0]["accountId"], "acct_a");
|
||||
assert!(second_body["page"]["nextCursor"].is_null());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_usage_stats_summary_should_return_usage_totals() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-usage-stats.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
seed_admin_session(&pool, "session_1").await;
|
||||
seed_usage(&pool).await;
|
||||
let app = build_router(AppState::with_pool(test_config(url), pool));
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/admin/usage-stats/summary")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_usage_summary")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["code"], 200);
|
||||
assert_eq!(body["data"]["accountCount"], 2);
|
||||
assert_eq!(body["data"]["requestCount"], 5);
|
||||
assert_eq!(body["data"]["inputTokens"], 19);
|
||||
assert_eq!(body["data"]["outputTokens"], 8);
|
||||
assert_eq!(body["data"]["cachedTokens"], 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_usage_stats_should_reject_missing_admin_session_cookie() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-usage-stats.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
let app = build_router(AppState::with_pool(test_config(url), pool));
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/admin/usage-stats")
|
||||
.header("x-request-id", "req_usage")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["code"], 40101);
|
||||
assert_eq!(body["requestId"], "req_usage");
|
||||
}
|
||||
|
||||
async fn seed_admin_session(pool: &sqlx::SqlitePool, session_id: &str) {
|
||||
let now = Utc::now().to_rfc3339();
|
||||
sqlx::query(
|
||||
"insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)",
|
||||
)
|
||||
.bind("admin_1")
|
||||
.bind("hash")
|
||||
.bind(&now)
|
||||
.bind(&now)
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind("admin_1")
|
||||
.bind("2999-01-01T00:00:00Z")
|
||||
.bind(now)
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn seed_usage(pool: &sqlx::SqlitePool) {
|
||||
let now = Utc::now().to_rfc3339();
|
||||
for (id, email, label, plan_type) in [
|
||||
("acct_a", "a@example.com", "primary", "plus"),
|
||||
("acct_b", "b@example.com", "backup", "free"),
|
||||
] {
|
||||
sqlx::query(
|
||||
"insert into accounts (id, email, label, plan_type, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?, 'active', ?, ?)",
|
||||
)
|
||||
.bind(id)
|
||||
.bind(email)
|
||||
.bind(label)
|
||||
.bind(plan_type)
|
||||
.bind("encrypted")
|
||||
.bind(&now)
|
||||
.bind(&now)
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
sqlx::query(
|
||||
"insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind("acct_a")
|
||||
.bind(3_i64)
|
||||
.bind(12_i64)
|
||||
.bind(5_i64)
|
||||
.bind(1_i64)
|
||||
.bind("2026-06-11T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind("acct_b")
|
||||
.bind(2_i64)
|
||||
.bind(7_i64)
|
||||
.bind(3_i64)
|
||||
.bind(2_i64)
|
||||
.bind("2026-06-11T00:01:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn response_json(response: axum::response::Response) -> Value {
|
||||
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
||||
serde_json::from_slice(&bytes).unwrap()
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
use wiremock::{
|
||||
matchers::{method, path},
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
};
|
||||
|
||||
use codex_proxy_rs::auth::{
|
||||
oauth::{OAuthConfig, OpenAiOAuthRefresher},
|
||||
refresh::TokenRefresher,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_oauth_refresher_should_exchange_refresh_token_with_form_body() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"access_token": "new-access",
|
||||
"refresh_token": "new-refresh"
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let refresher = OpenAiOAuthRefresher::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
OAuthConfig {
|
||||
client_id: "codex-client".to_string(),
|
||||
auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(),
|
||||
token_endpoint: format!("{}/oauth/token", server.uri()),
|
||||
},
|
||||
);
|
||||
|
||||
let tokens = refresher.refresh("refresh-secret").await.unwrap();
|
||||
|
||||
assert_eq!(tokens.access_token, "new-access");
|
||||
assert_eq!(tokens.refresh_token.as_deref(), Some("new-refresh"));
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
let body = String::from_utf8(requests[0].body.clone()).unwrap();
|
||||
assert!(body.contains("grant_type=refresh_token"));
|
||||
assert!(body.contains("client_id=codex-client"));
|
||||
assert!(body.contains("refresh_token=refresh-secret"));
|
||||
}
|
||||
@@ -1,10 +1,12 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use axum::{
|
||||
body::{to_bytes, Body},
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use chrono::Utc;
|
||||
use secrecy::ExposeSecret;
|
||||
use serde_json::{json, Value};
|
||||
use tower::ServiceExt;
|
||||
use wiremock::{
|
||||
@@ -13,8 +15,14 @@ use wiremock::{
|
||||
};
|
||||
|
||||
use codex_proxy_rs::{
|
||||
accounts::repository::AccountRepository,
|
||||
app::build_router,
|
||||
auth::{api_key::ApiKeyHasher, api_key_repository::ClientApiKeyRepository},
|
||||
auth::{
|
||||
api_key::ApiKeyHasher,
|
||||
api_key_repository::ClientApiKeyRepository,
|
||||
refresh::{RefreshFailure, TokenRefresher},
|
||||
token::TokenPair,
|
||||
},
|
||||
config::{
|
||||
AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig,
|
||||
QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig,
|
||||
@@ -34,6 +42,22 @@ struct ImportedApp {
|
||||
_tempdir: tempfile::TempDir,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct StaticTokenRefresher {
|
||||
access_token: String,
|
||||
refresh_token: Option<String>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenRefresher for StaticTokenRefresher {
|
||||
async fn refresh(&self, _refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
|
||||
Ok(TokenPair {
|
||||
access_token: self.access_token.clone(),
|
||||
refresh_token: self.refresh_token.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn test_config(database_url: String, base_url: String) -> AppConfig {
|
||||
AppConfig {
|
||||
server: ServerConfig {
|
||||
@@ -144,6 +168,67 @@ async fn build_imported_app(base_url: String) -> ImportedApp {
|
||||
}
|
||||
}
|
||||
|
||||
async fn build_imported_app_with_refresher(base_url: String, access_token: &str) -> ImportedApp {
|
||||
let tempdir = tempfile::tempdir().unwrap();
|
||||
let db = tempdir.path().join("v1-upstream-refresh.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
seed_admin_session(&pool, "session_1").await;
|
||||
let secret_box = SecretBox::new([31u8; 32]);
|
||||
let hasher = ApiKeyHasher::new([32u8; 32]);
|
||||
let generated = hasher.generate_client_api_key("test");
|
||||
ClientApiKeyRepository::new(pool.clone())
|
||||
.insert_generated("test", &generated)
|
||||
.await
|
||||
.unwrap();
|
||||
let app = build_router(
|
||||
AppState::with_pool_secret_api_key_hasher_and_token_refresher(
|
||||
test_config(url, base_url),
|
||||
pool.clone(),
|
||||
secret_box.clone(),
|
||||
hasher,
|
||||
StaticTokenRefresher {
|
||||
access_token: access_token.to_string(),
|
||||
refresh_token: None,
|
||||
},
|
||||
),
|
||||
);
|
||||
|
||||
let import_response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/admin/accounts/import")
|
||||
.header("content-type", "application/json")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"accounts": [{
|
||||
"id": "acct_imported",
|
||||
"accountId": "chatgpt-account",
|
||||
"token": "access-secret",
|
||||
"refreshToken": "refresh-secret",
|
||||
"status": "active"
|
||||
}]
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(import_response.status(), StatusCode::OK);
|
||||
|
||||
ImportedApp {
|
||||
app,
|
||||
pool,
|
||||
secret_box,
|
||||
client_api_key: generated.plaintext,
|
||||
_tempdir: tempdir,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn v1_responses_should_use_imported_account_and_record_usage() {
|
||||
let server = MockServer::start().await;
|
||||
@@ -298,6 +383,137 @@ async fn v1_responses_should_passthrough_stream_and_record_usage_and_log() {
|
||||
assert_eq!(event.4["usage"]["outputTokens"], 5);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn v1_responses_should_refresh_after_401_and_retry_non_stream() {
|
||||
let server = MockServer::start().await;
|
||||
let success_sse = concat!(
|
||||
"event: response.completed\n",
|
||||
"data: {\"response\":{\"id\":\"resp_after_refresh\",\"object\":\"response\",\"usage\":{\"input_tokens\":2,\"output_tokens\":1}}}\n",
|
||||
"\n",
|
||||
);
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.and(header("authorization", "Bearer access-secret"))
|
||||
.respond_with(ResponseTemplate::new(401).set_body_json(json!({
|
||||
"error": {"code": "token_revoked"}
|
||||
})))
|
||||
.up_to_n_times(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.and(header("authorization", "Bearer refreshed-access"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(success_sse),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let imported = build_imported_app_with_refresher(server.uri(), "refreshed-access").await;
|
||||
|
||||
let response = imported
|
||||
.app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.header(
|
||||
"authorization",
|
||||
format!("Bearer {}", imported.client_api_key),
|
||||
)
|
||||
.header("content-type", "application/json")
|
||||
.header("x-request-id", "req_refresh_non_stream")
|
||||
.body(Body::from(r#"{"model":"gpt-5.5","input":[]}"#))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["id"], "resp_after_refresh");
|
||||
let repo = AccountRepository::new(imported.pool.clone(), imported.secret_box);
|
||||
let account = repo.get("acct_imported").await.unwrap().unwrap();
|
||||
assert_eq!(account.access_token.expose_secret(), "refreshed-access");
|
||||
assert_eq!(
|
||||
account.refresh_token.unwrap().expose_secret(),
|
||||
"refresh-secret"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn v1_responses_should_refresh_after_401_and_retry_stream() {
|
||||
let server = MockServer::start().await;
|
||||
let success_sse = concat!(
|
||||
"event: response.output_text.delta\n",
|
||||
"data: {\"delta\":\"pong\"}\n",
|
||||
"\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"response\":{\"id\":\"resp_stream_after_refresh\",\"object\":\"response\",\"usage\":{\"input_tokens\":4,\"output_tokens\":2}}}\n",
|
||||
"\n",
|
||||
);
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.and(header("authorization", "Bearer access-secret"))
|
||||
.respond_with(ResponseTemplate::new(401).set_body_json(json!({
|
||||
"error": {"code": "token_revoked"}
|
||||
})))
|
||||
.up_to_n_times(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.and(header("authorization", "Bearer refreshed-stream-access"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(success_sse),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let imported = build_imported_app_with_refresher(server.uri(), "refreshed-stream-access").await;
|
||||
|
||||
let response = imported
|
||||
.app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.header(
|
||||
"authorization",
|
||||
format!("Bearer {}", imported.client_api_key),
|
||||
)
|
||||
.header("content-type", "application/json")
|
||||
.header("x-request-id", "req_refresh_stream")
|
||||
.body(Body::from(
|
||||
r#"{"model":"gpt-5.5","input":[],"stream":true}"#,
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get("content-type")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
assert!(content_type.starts_with("text/event-stream"));
|
||||
let body = response_text(response).await;
|
||||
assert!(body.contains("event: response.completed"));
|
||||
let usage: (i64, i64, i64, i64) = sqlx::query_as(
|
||||
"select request_count, input_tokens, output_tokens, cached_tokens from account_usage where account_id = ?",
|
||||
)
|
||||
.bind("acct_imported")
|
||||
.fetch_one(&imported.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(usage, (1, 4, 2, 0));
|
||||
}
|
||||
|
||||
async fn seed_admin_session(pool: &sqlx::SqlitePool, session_id: &str) {
|
||||
let now = Utc::now().to_rfc3339();
|
||||
sqlx::query(
|
||||
|
||||
Reference in New Issue
Block a user