diff --git a/docs/api.md b/docs/api.md index ad67884ff..ef0dc624f 100644 --- a/docs/api.md +++ b/docs/api.md @@ -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 diff --git a/docs/implementation-status.md b/docs/implementation-status.md index 7ed7aad52..468a9fdc7 100644 --- a/docs/implementation-status.md +++ b/docs/implementation-status.md @@ -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::` 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` 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` 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. | diff --git a/docs/superpowers/plans/2026-06-11-codex-proxy-rs.md b/docs/superpowers/plans/2026-06-11-codex-proxy-rs.md index 06d2b0844..313ce0b53 100644 --- a/docs/superpowers/plans/2026-06-11-codex-proxy-rs.md +++ b/docs/superpowers/plans/2026-06-11-codex-proxy-rs.md @@ -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. diff --git a/src/accounts/repository.rs b/src/accounts/repository.rs index 4d028278c..83405117f 100644 --- a/src/accounts/repository.rs +++ b/src/accounts/repository.rs @@ -93,6 +93,28 @@ pub struct AccountUsageRecord { pub last_used_at: Option>, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountUsageListRecord { + pub account_id: String, + pub email: Option, + pub label: Option, + pub plan_type: Option, + pub request_count: i64, + pub input_tokens: i64, + pub output_tokens: i64, + pub cached_tokens: i64, + pub last_used_at: Option>, +} + +#[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 Self { + Self { pool } + } + + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> AccountRepositoryResult> { + 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 { + 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 { + 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::, _>("last_used_at"))?, + }) +} + fn metadata_from_row( row: &sqlx::sqlite::SqliteRow, ) -> AccountRepositoryResult { diff --git a/src/app.rs b/src/app.rs index a2ca09a3e..4dbcc5e48 100644 --- a/src/app.rs +++ b/src/app.rs @@ -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)) diff --git a/src/auth/oauth.rs b/src/auth/oauth.rs index f55bfccc5..e1cf2d00d 100644 --- a/src/auth/oauth.rs +++ b/src/auth/oauth.rs @@ -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, +} + +#[async_trait] +impl TokenRefresher for OpenAiOAuthRefresher { + async fn refresh(&self, refresh_token: &str) -> Result { + 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::(&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 +} diff --git a/src/http/admin.rs b/src/http/admin.rs index b4c880e7a..d9ae55de8 100644 --- a/src/http/admin.rs +++ b/src/http/admin.rs @@ -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, } +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UsageStatsQuery { + pub cursor: Option, + pub limit: Option, +} + #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] pub struct LoginRequest { @@ -290,6 +300,58 @@ impl From for AdminAccountData { } } +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminUsageStatsData { + pub account_id: String, + pub email: Option, + pub label: Option, + pub plan_type: Option, + pub request_count: i64, + pub input_tokens: i64, + pub output_tokens: i64, + pub cached_tokens: i64, + pub last_used_at: Option, +} + +impl From 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 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, + Extension(request_id): Extension, + headers: HeaderMap, + Query(query): Query, +) -> 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, + Extension(request_id): Extension, + 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, Extension(request_id): Extension, diff --git a/src/http/v1.rs b/src/http/v1.rs index 73681a2df..fe56a1afa 100644 --- a/src/http/v1.rs +++ b/src/http/v1.rs @@ -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 { + 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 { + 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 { + 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 { + 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 { + 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, diff --git a/src/main.rs b/src/main.rs index 2b01a48f6..f023af366 100644 --- a/src/main.rs +++ b/src/main.rs @@ -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); diff --git a/src/state.rs b/src/state.rs index 4ad42b3b6..2b0d97632 100644 --- a/src/state.rs +++ b/src/state.rs @@ -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, pub secret_box: Option, pub api_key_hasher: Option, + pub token_refresher: Option>, pub account_pool: Arc>, } @@ -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( + 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 { + Some(AccountUsageRepository::new(self.db()?.clone())) + } + pub fn cookie_repository(&self) -> Option { 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> { + self.services.token_refresher.clone() + } + pub async fn reload_account_pool_from_repository(&self) -> AccountRepositoryResult { let Some(repo) = self.account_repository() else { return Ok(0); diff --git a/tests/admin_usage_stats_route_test.rs b/tests/admin_usage_stats_route_test.rs new file mode 100644 index 000000000..e20600fe4 --- /dev/null +++ b/tests/admin_usage_stats_route_test.rs @@ -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() +} diff --git a/tests/oauth_refresh_test.rs b/tests/oauth_refresh_test.rs new file mode 100644 index 000000000..9090dd7b1 --- /dev/null +++ b/tests/oauth_refresh_test.rs @@ -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")); +} diff --git a/tests/v1_upstream_route_test.rs b/tests/v1_upstream_route_test.rs index 96ff9daa6..8553b9568 100644 --- a/tests/v1_upstream_route_test.rs +++ b/tests/v1_upstream_route_test.rs @@ -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, +} + +#[async_trait] +impl TokenRefresher for StaticTokenRefresher { + async fn refresh(&self, _refresh_token: &str) -> Result { + 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(