starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-07 00:05:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add scoped admin account operations

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-06-12 17:18:33 +08:00
co-authored by Codex
parent e1f35e0f55
commit 93cfee362c
8 changed files with 1404 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/implementation-status.md
+4 -4
View File
File diff suppressed because one or more lines are too long
docs/superpowers/plans/2026-06-12-openai-gpt-codex-parity.md
+10 -7
View File
@@ -8,7 +8,7 @@
**Tech Stack:** Rust, Axum, Reqwest/rustls, serde, sqlx/SQLite, wiremock, tokio, futures, optional WebSocket transport dependency.
**Checkpoint 2026-06-12:** Commit `b9c30ba` completed Tasks 1-3 for the in-scope OpenAI GPT/Codex path: Chat Completions route/translation/output, Responses request field parity, default Responses streaming, Codex context headers, and local-only WebSocket flag serialization cleanup. Commit `350ffbe` completed Task 4: `previous_response_id` requests use WebSocket `response.create` and return SSE-compatible output. Commit `3ecd5b6` partially completed Task 5 for non-streaming `/v1/responses`: 429/402/403 account-state classification and fallback retry across imported accounts. Commit `3960a35` extends the same fallback path to HTTP SSE Responses setup and Chat Completions. Commit `1738f24` adds Task 6 cache groundwork: persisted backend model snapshots and cached `/v1/models*` reads. Commit `0b6cc7e` adds Task 7 import compatibility for Sub2API OpenAI OAuth exports and marks the detected source format while ignoring proxy-only fields. Commit `f75db60` adds `/admin/refresh-models` backed by imported accounts. Commit `4b19846` syncs refreshed model-plan allowlists into the runtime account pool. Commit `db6cb98` adds admin account label/status/delete mutations with runtime pool synchronization. Commit `16eb4ed` adds batch delete/status mutations. Commit `8af350d` adds local client API key status/delete mutations. Commit `203270f` adds encrypted account Cookie get/set/delete. Commit `960bbb2` adds local client API key label and batch-delete mutations. Commit `72438c6` adds native and Sub2API account export. Commit `16938e6` adds manual account creation for Codex OAuth tokens; the current correction narrows `POST /admin/accounts` to TS-like `token`/`refreshToken` import semantics for OpenAI GPT/Codex accounts and adds deterministic Codex CLI `auth.json` import through the same validated JWT-claim path. Remaining major work continues with WebSocket/history fallback policy, PKCE/device login, and scoped admin operations.
**Checkpoint 2026-06-12:** Commit `b9c30ba` completed Tasks 1-3 for the in-scope OpenAI GPT/Codex path: Chat Completions route/translation/output, Responses request field parity, default Responses streaming, Codex context headers, and local-only WebSocket flag serialization cleanup. Commit `350ffbe` completed Task 4: `previous_response_id` requests use WebSocket `response.create` and return SSE-compatible output. Commit `3ecd5b6` partially completed Task 5 for non-streaming `/v1/responses`: 429/402/403 account-state classification and fallback retry across imported accounts. Commit `3960a35` extends the same fallback path to HTTP SSE Responses setup and Chat Completions. Commit `1738f24` adds Task 6 cache groundwork: persisted backend model snapshots and cached `/v1/models*` reads. Commit `0b6cc7e` adds Task 7 import compatibility for Sub2API OpenAI OAuth exports and marks the detected source format while ignoring proxy-only fields. Commit `f75db60` adds `/admin/refresh-models` backed by imported accounts. Commit `4b19846` syncs refreshed model-plan allowlists into the runtime account pool. Commit `db6cb98` adds admin account label/status/delete mutations with runtime pool synchronization. Commit `16eb4ed` adds batch delete/status mutations. Commit `8af350d` adds local client API key status/delete mutations. Commit `203270f` adds encrypted account Cookie get/set/delete. Commit `960bbb2` adds local client API key label and batch-delete mutations. Commit `72438c6` adds native and Sub2API account export. Commit `16938e6` adds manual account creation for Codex OAuth tokens; the current correction narrows `POST /admin/accounts` to TS-like `token`/`refreshToken` import semantics for OpenAI GPT/Codex accounts, adds deterministic Codex CLI `auth.json` import through the same validated JWT-claim path, and adds admin health-check, per-account refresh, reset-usage, and quota routes for imported Codex accounts. Remaining major work continues with WebSocket/history fallback policy plus PKCE/device login and auth-status/logout variants.
---
@@ -19,7 +19,7 @@ This audit is the correction point against the TypeScript reference at `/home/zy
| Area / commit | TypeScript reference behavior | Rust current behavior | Classification | Action |
| --- | --- | --- | --- | --- |
| Sub2API OpenAI OAuth import / `0b6cc7e` | Account import accepts token/refresh-token account exports and Sub2API-shaped OpenAI OAuth data through `AccountImportService`; proxy fields can exist in source data. | `/admin/accounts/import` accepts native and Sub2API OpenAI OAuth payloads, returns `sourceFormat`, and ignores proxy-only fields. | Original parity, scoped. | Keep. Use the real Sub2API export only for private smoke verification; never print or commit secrets. |
| Account label/status/delete and batch mutations / `db6cb98`, `16eb4ed` | `/auth/accounts/:id/label`, `/status`, delete, batch-delete, and batch-status mutate the account pool. | `/admin/accounts/*` mutates SQLite and synchronizes the runtime account pool with the admin envelope contract. | Partial parity, adapted. | Keep. Finish quota, health, reset-usage, and login/import variants separately. |
| Account label/status/delete, health, refresh, reset, quota, and batch mutations / `db6cb98`, `16eb4ed`, current Task 7 pass | `/auth/accounts/:id/label`, `/status`, delete, batch-delete, batch-status, health-check, refresh/probe, reset-usage, and quota mutate or inspect the account pool. | `/admin/accounts/*` mutates SQLite and synchronizes the runtime account pool with the admin envelope contract. Health/quota call the Codex backend through imported account tokens, refresh uses the existing OAuth token refresher, reset-usage clears local counters and pool last-used state, and quota stores normalized quota JSON. | Partial parity, adapted. | Keep. Finish login/import variants separately. |
| Account Cookie get/set/delete / `203270f` | Account routes can store browser Cookie headers for account-scoped replay. | `/admin/accounts/{id}/cookies` stores, reads, and clears encrypted per-account Cookies. | Original parity, adapted. | Keep; preserve account-scoped replay and encryption invariants. |
| Account export / `72438c6` | `/auth/accounts/export` supports `full`, `minimal`, `cockpit_tools`, `sub2api`, and `cpa`. | `/admin/accounts/export` supports native/full Rust export and Sub2API OpenAI OAuth export only, without proxy fields. | Partial parity, scoped. | Keep native and Sub2API. Treat `minimal`, `cockpit_tools`, and `cpa` as omitted unless an OpenAI/Codex operation needs them. |
| Manual account creation / `16938e6` + current correction | `POST /auth/accounts` accepts only `token` and/or `refreshToken`, runs `AccountImportService.importOne`, exchanges refresh-token-only imports, derives account identity/profile from JWT claims, updates an existing `chatgpt_account_id` + `chatgpt_user_id` entry, and never returns tokens. | `POST /admin/accounts` accepts only `token` and/or `refreshToken`, rejects missing/invalid/expired JWTs and tokens without `https://api.openai.com/auth.chatgpt_account_id`, derives email/accountId/userId/planType/expiresAt from JWT claims, ignores caller metadata/status/label, exchanges refresh-token-only imports with `AppState`'s OpenAI refresher, preserves or rotates refresh tokens, encrypts secrets, and synchronizes the runtime account pool. It does not yet cache quota from a verification probe. | Corrected scoped parity. | Keep. Finish quota fetch/health/login variants separately without adding proxy pools or non-OpenAI provider support. |
@@ -27,12 +27,12 @@ This audit is the correction point against the TypeScript reference at `/home/zy
| Local client API key label/batch-delete / `960bbb2` | TS label/batch-delete applies to third-party provider key pool entries. | Rust label/batch-delete applies to local HMAC-hashed `cpr_` client keys. | Rust local extension, possible drift if strict parity is required. | Reclassify in docs as local admin utility. Remove only if strict original parity excludes local client-key management beyond create/list/status/delete. |
| API key export/import variants | TS can export/import third-party provider keys for reimport. | Rust local `cpr_` key plaintext is shown only at creation and then HMAC-hashed; provider key pools are not implemented. | Intentional scope cut. | Do not port TS provider key import/export. Document that local client keys cannot be safely exported after creation. |
| Model refresh and plan allowlist sync / `f75db60`, `4b19846` | TS has provider/key model catalog behavior mixed with account/provider management. | Rust probes Codex backend models via imported accounts and syncs plan allowlists into the account pool. | Partial parity, scoped Rust adaptation. | Keep; avoid third-party provider catalog behavior. Finish full Codex model metadata parity later if needed. |
| OAuth/device/CLI account login/import variants | Original scope includes ChatGPT/Codex OAuth login, device login, CLI token import, manual validated token import, logout/status, quota, and health checks. | Rust now supports deterministic admin CLI import from Codex `auth.json` through `/admin/accounts/import-cli`, storing tokens encrypted and syncing the runtime account pool. Rust still lacks OAuth PKCE login-start/code-relay/callback, device login/poll, validated manual token import, logout/auth status equivalents, quota fetch, and health check routes. | Partial OpenAI/Codex parity. | Continue after correcting manual account creation. |
| OAuth/device/CLI account login/import variants | Original scope includes ChatGPT/Codex OAuth login, device login, CLI token import, manual validated token import, logout/status, quota, and health checks. | Rust now supports validated manual `token`/`refreshToken` import, deterministic admin CLI import from Codex `auth.json`, admin quota fetch, health check, per-account refresh, and reset-usage for imported accounts. Rust still lacks OAuth PKCE login-start/code-relay/callback, device login/poll, and logout/auth status equivalents. | Partial OpenAI/Codex parity. | Continue with login/status variants without adding proxy pools or non-OpenAI providers. |
| Proxy/provider metadata and routes | TS project includes per-account proxy assignment, proxy health, and non-OpenAI provider routes. | Rust imports ignore proxy fields and has no outbound proxy/provider pools. | Intentional scope cut. | Keep omitted. Do not add proxy routing, proxy UI/API, Anthropic/Gemini/Ollama/custom provider support, or Electron/frontend flows. |
Immediate priority after this audit:
1. Finish in-scope account operation parity: OAuth PKCE, device login, quota fetch, health check, logout/status/reset-usage.
1. Finish in-scope account operation parity: OAuth PKCE, device login, logout/auth status.
2. Finish Task 5 WebSocket/history fallback policy without breaking `previous_response_id` account affinity.
3. Keep local client API key utilities documented as Rust-local auth, not TypeScript provider-key parity.
@@ -164,13 +164,16 @@ Immediate priority after this audit:
- [x] Write failing tests for TS-like manual `token`/`refreshToken` import semantics, JWT claim derivation, refresh-token-only exchange, duplicate account update, encrypted storage, runtime pool sync, and sanitized responses.
- [x] Correct manual account creation to TS-like `token`/`refreshToken` import semantics for OpenAI GPT/Codex accounts.
- [x] Run: `cargo test --test admin_accounts_route_test`
- [ ] Reclassify local client API key label/batch operations as Rust-local admin utilities, and avoid treating TS provider API key export/import as in-scope parity.
- [x] Write failing tests for Codex CLI `auth.json` import through admin account import.
- [x] Implement `/admin/accounts/import-cli` using a reusable CLI auth parser and the shared validated JWT-claim account import path, without frontend/proxy/provider behavior.
- [x] Run: `cargo test --test cli_auth_import_test`; `cargo test --test admin_accounts_route_test admin_accounts_import_cli_should_read_codex_auth_file_store_encrypted_and_sync_pool`
- [ ] Write failing tests for quota fetch, health check, account login/import variants, and any accepted local client-key utilities.
- [x] Write failing tests for admin account health-check, per-account refresh, reset-usage, and quota routes.
- [x] Implement `/admin/accounts/health-check`, `/admin/accounts/{id}/refresh`, `/admin/accounts/{id}/reset-usage`, and `/admin/accounts/{id}/quota` for imported Codex accounts only.
- [x] Run: `cargo test --test admin_accounts_route_test`
- [ ] Reclassify local client API key label/batch operations as Rust-local admin utilities, and avoid treating TS provider API key export/import as in-scope parity.
- [ ] Write failing tests for account login/status variants and any accepted local client-key utilities.
- [ ] Run: `cargo test --test admin_accounts_route_test --test admin_api_keys_route_test`
- [ ] Implement scoped admin operations needed to operate Codex-backed OpenAI GPT routes.
- [ ] Implement remaining scoped admin operations needed to operate Codex-backed OpenAI GPT routes.
- [ ] Run: `cargo test --test admin_accounts_route_test --test admin_api_keys_route_test`
### Task 8: Final Verification And Docs
src/accounts/pool.rs
+8
View File
@@ -137,6 +137,14 @@ impl AccountPool {
true
}
pub fn reset_usage(&mut self, account_id: &str) -> bool {
let Some(account) = self.accounts.get_mut(account_id) else {
return false;
};
account.last_used_at = None;
true
}
pub fn mark_quota_limited_until(
&mut self,
account_id: &str,
src/accounts/repository.rs
+28
View File
@@ -486,6 +486,24 @@ impl AccountRepository {
row.map(|row| usage_from_row(&row)).transpose()
}
pub async fn update_quota_json(
&self,
account_id: &str,
quota_json: &str,
) -> AccountRepositoryResult<bool> {
let now = Utc::now().to_rfc3339();
let result = sqlx::query(
"update accounts set quota_json = ?, quota_fetched_at = ?, updated_at = ? where id = ?",
)
.bind(quota_json)
.bind(&now)
.bind(now)
.bind(account_id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
fn account_from_row(
&self,
row: &sqlx::sqlite::SqliteRow,
@@ -630,6 +648,16 @@ impl AccountUsageRepository {
cached_tokens: row.get("cached_tokens"),
})
}
pub async fn reset_account(&self, account_id: &str) -> AccountRepositoryResult<()> {
sqlx::query(
"insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, 0, 0, 0, 0, null) on conflict(account_id) do update set request_count = 0, input_tokens = 0, output_tokens = 0, cached_tokens = 0, last_used_at = null",
)
.bind(account_id)
.execute(&self.pool)
.await?;
Ok(())
}
}
fn usage_list_from_row(
src/app.rs
+14 -3
View File
@@ -8,11 +8,12 @@ use tower_http::trace::TraceLayer;
use crate::{
http::{
admin::{
accounts, api_keys, batch_delete_accounts, batch_delete_api_keys,
account_quota, accounts, api_keys, batch_delete_accounts, batch_delete_api_keys,
batch_update_account_status, create_account, create_api_key, delete_account,
delete_account_cookies, delete_api_key, export_accounts, get_account_cookies,
import_accounts, import_cli_auth, login, logs, refresh_models, set_account_cookies,
settings, update_account_label, update_account_status, update_api_key_label,
health_check_accounts, import_accounts, import_cli_auth, login, logs, refresh_account,
refresh_models, reset_account_usage, set_account_cookies, settings,
update_account_label, update_account_status, update_api_key_label,
update_api_key_status, usage_stats, usage_stats_summary,
},
health::health,
@@ -42,12 +43,22 @@ pub fn build_router(state: AppState) -> Router {
.route("/admin/usage-stats", get(usage_stats))
.route("/admin/usage-stats/summary", get(usage_stats_summary))
.route("/admin/accounts", get(accounts).post(create_account))
.route("/admin/accounts/health-check", post(health_check_accounts))
.route("/admin/accounts/batch-delete", post(batch_delete_accounts))
.route(
"/admin/accounts/batch-status",
post(batch_update_account_status),
)
.route("/admin/accounts/export", get(export_accounts))
.route(
"/admin/accounts/{account_id}/refresh",
post(refresh_account),
)
.route(
"/admin/accounts/{account_id}/reset-usage",
post(reset_account_usage),
)
.route("/admin/accounts/{account_id}/quota", get(account_quota))
.route(
"/admin/accounts/{account_id}/cookies",
get(get_account_cookies)
src/codex/client.rs
+54
View File
@@ -195,6 +195,42 @@ impl CodexBackendClient {
Err(CodexClientError::ModelsUnavailable)
}
pub async fn fetch_usage(&self, context: CodexRequestContext<'_>) -> CodexClientResult<Value> {
let mut last_invalid_body = None;
for endpoint in self.usage_endpoints() {
let mut headers = self.request_headers(context)?;
headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
let response = self.client.get(endpoint).headers(headers).send().await?;
let status = response.status();
let retry_after_seconds = retry_after_seconds(response.headers(), None);
let body = response.text().await?;
if status == StatusCode::NOT_FOUND {
last_invalid_body = Some(body);
continue;
}
if !status.is_success() {
return Err(CodexClientError::Upstream {
status,
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
});
}
match serde_json::from_str::<Value>(&body) {
Ok(parsed) if parsed.get("rate_limit").is_some() => return Ok(parsed),
_ => last_invalid_body = Some(body),
}
}
Err(CodexClientError::Upstream {
status: StatusCode::BAD_GATEWAY,
retry_after_seconds: None,
body: last_invalid_body
.map(|body| format!("invalid usage response: {}", truncate_for_error(&body)))
.unwrap_or_else(|| "usage endpoint is unavailable".to_string()),
})
}
async fn send_response_request(
&self,
request: &CodexResponsesRequest,
@@ -211,6 +247,20 @@ impl CodexBackendClient {
.await?)
}
fn usage_endpoints(&self) -> Vec<String> {
if self.base_url.contains("/backend-api") {
vec![
format!("{}/wham/usage", self.base_url),
format!("{}/codex/usage", self.base_url),
]
} else {
vec![
format!("{}/api/codex/usage", self.base_url),
format!("{}/codex/usage", self.base_url),
]
}
}
fn request_headers(&self, context: CodexRequestContext<'_>) -> CodexClientResult<HeaderMap> {
let mut headers = HeaderMap::new();
for (name, value) in build_codex_headers(
@@ -354,3 +404,7 @@ fn retry_after_seconds_from_body(body: &str) -> Option<u64> {
.as_secs();
(resets_at > now).then_some(resets_at - now)
}
fn truncate_for_error(body: &str) -> String {
body.chars().take(200).collect()
}
src/http/admin.rs
+944 -2
View File
@@ -1,3 +1,5 @@
use std::time::{Duration as StdDuration, Instant};
use axum::{
body::Bytes,
extract::{Path, Query, State},
@@ -10,10 +12,13 @@ use axum::{
};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use chrono::{DateTime, Duration, SecondsFormat, Utc};
use futures::{stream, StreamExt};
use reqwest::Url;
use secrecy::{ExposeSecret, SecretString};
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use sqlx::Row;
use tokio::time::sleep;
use uuid::Uuid;
use crate::{
@@ -21,7 +26,7 @@ use crate::{
model::{Account, AccountStatus},
repository::{
AccountClaimsUpdate, AccountRepositoryError, AccountUsageListRecord,
AccountUsageSummary, NewAccount, StoredAccount, StoredAccountMetadata,
AccountUsageSummary, NewAccount, StoredAccount, StoredAccountMetadata, TokenUpdate,
},
},
auth::api_key_repository::StoredClientApiKey,
@@ -29,8 +34,11 @@ use crate::{
admin_session::verify_admin_password,
cli_import::{default_codex_home, read_cli_auth_from_home},
refresh::RefreshFailure,
token::TokenPair,
},
codex::client::{
build_reqwest_client, CodexBackendClient, CodexClientError, CodexRequestContext,
},
codex::client::{build_reqwest_client, CodexBackendClient, CodexRequestContext},
config::{AppConfig, QuotaWarningThresholds},
fingerprint::model::Fingerprint,
http::{auth::admin_session_id, middleware::RequestId},
@@ -257,6 +265,57 @@ pub struct ImportCliAuthRequest {
pub codex_home: Option<String>,
}
#[derive(Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct HealthCheckRequest {
pub ids: Option<Vec<String>>,
#[serde(alias = "stagger_ms")]
pub stagger_ms: Option<u64>,
pub concurrency: Option<u8>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct HealthCheckData {
pub summary: HealthCheckSummary,
pub results: Vec<AccountProbeData>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct HealthCheckSummary {
pub total: usize,
pub alive: usize,
pub dead: usize,
pub skipped: usize,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountProbeData {
pub id: String,
pub email: Option<String>,
pub previous_status: String,
pub result: String,
pub status: Option<String>,
pub error: Option<String>,
pub duration_ms: Option<u128>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ResetAccountUsageData {
pub id: String,
pub reset: bool,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountQuotaData {
pub quota: Value,
pub raw: Value,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AccountExportFormat {
Native,
@@ -1150,6 +1209,430 @@ pub async fn create_account(
.into_response()
}
pub async fn health_check_accounts(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
headers: HeaderMap,
payload: Option<Json<HealthCheckRequest>>,
) -> Response {
let request_id = request_id.as_str().to_string();
let payload = payload.map(|Json(payload)| payload).unwrap_or_default();
if payload.ids.as_ref().is_some_and(Vec::is_empty) {
return AdminResponse::new(
StatusCode::BAD_REQUEST,
AdminEnvelope::new(40001, "Account ids must not be empty", (), request_id),
)
.into_response();
}
if payload
.stagger_ms
.is_some_and(|value| !(500..=30_000).contains(&value))
{
return AdminResponse::new(
StatusCode::BAD_REQUEST,
AdminEnvelope::new(
40001,
"staggerMs must be between 500 and 30000",
(),
request_id,
),
)
.into_response();
}
if payload
.concurrency
.is_some_and(|value| !(1..=10).contains(&value))
{
return AdminResponse::new(
StatusCode::BAD_REQUEST,
AdminEnvelope::new(
40001,
"concurrency must be between 1 and 10",
(),
request_id,
),
)
.into_response();
}
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_repository() else {
return AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(
50001,
"Account repository is not initialized",
(),
request_id,
),
)
.into_response();
};
let accounts = match repo.list_all().await {
Ok(accounts) => filter_health_check_accounts(accounts, payload.ids.as_deref()),
Err(_) => {
return AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(50001, "Failed to list accounts", (), request_id),
)
.into_response();
}
};
let concurrency = usize::from(payload.concurrency.unwrap_or(2));
let stagger_ms = payload.stagger_ms.unwrap_or(3_000);
let results = stream::iter(accounts.into_iter().enumerate())
.map(|(index, account)| {
let state = state.clone();
let repo = repo.clone();
let request_id = request_id.clone();
async move {
if stagger_ms > 0 && index > 0 {
let multiplier = index.min(concurrency);
sleep(StdDuration::from_millis(
stagger_ms.saturating_mul(multiplier as u64),
))
.await;
}
probe_account_with_codex_backend(&state, &repo, account, &request_id).await
}
})
.buffer_unordered(concurrency)
.collect::<Vec<_>>()
.await;
let summary = HealthCheckSummary {
total: results.len(),
alive: results
.iter()
.filter(|result| result.result == "alive")
.count(),
dead: results
.iter()
.filter(|result| result.result == "dead")
.count(),
skipped: results
.iter()
.filter(|result| result.result == "skipped")
.count(),
};
AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(HealthCheckData { summary, results }, request_id),
)
.into_response()
}
pub async fn refresh_account(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
headers: HeaderMap,
Path(account_id): Path<String>,
) -> 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_repository() else {
return AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(
50001,
"Account repository is not initialized",
(),
request_id,
),
)
.into_response();
};
let account = match repo.get(&account_id).await {
Ok(Some(account)) => account,
Ok(None) => {
return AdminResponse::new(
StatusCode::NOT_FOUND,
AdminEnvelope::new(40401, "Account not found", (), request_id),
)
.into_response();
}
Err(_) => {
return AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(50001, "Failed to load account", (), request_id),
)
.into_response();
}
};
if account.status == AccountStatus::Disabled {
return AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(
skipped_probe_result(&account, "manually disabled"),
request_id,
),
)
.into_response();
}
let Some(refresh_token) = account.refresh_token.as_ref() else {
return AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(
skipped_probe_result(&account, "no refresh token"),
request_id,
),
)
.into_response();
};
let Some(refresher) = state.token_refresher() else {
return AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(50001, "Token refresher is not initialized", (), request_id),
)
.into_response();
};
let started_at = Instant::now();
let previous_status = account_status_value(account.status).to_string();
match refresher.refresh(refresh_token.expose_secret()).await {
Ok(tokens) => {
match persist_admin_refreshed_account(&state, &repo, &account.id, tokens).await {
Ok(updated) => AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(
AccountProbeData {
id: updated.id,
email: updated.email,
previous_status,
result: "alive".to_string(),
status: Some(account_status_value(updated.status).to_string()),
error: None,
duration_ms: Some(started_at.elapsed().as_millis()),
},
request_id,
),
)
.into_response(),
Err(_) => AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(50001, "Failed to store refreshed account", (), request_id),
)
.into_response(),
}
}
Err(failure) => {
let status = apply_refresh_failure_status(&state, &repo, &account, failure).await;
AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(
AccountProbeData {
id: account.id,
email: account.email,
previous_status,
result: "dead".to_string(),
status: status.map(account_status_value).map(ToString::to_string),
error: Some(public_refresh_failure(failure).to_string()),
duration_ms: Some(started_at.elapsed().as_millis()),
},
request_id,
),
)
.into_response()
}
}
}
pub async fn reset_account_usage(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
headers: HeaderMap,
Path(account_id): Path<String>,
) -> 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(account_repo) = state.account_repository() else {
return AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(
50001,
"Account repository is not initialized",
(),
request_id,
),
)
.into_response();
};
match account_repo.exists(&account_id).await {
Ok(true) => {}
Ok(false) => {
return AdminResponse::new(
StatusCode::NOT_FOUND,
AdminEnvelope::new(40401, "Account not found", (), request_id),
)
.into_response();
}
Err(_) => {
return AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(50001, "Failed to inspect account", (), request_id),
)
.into_response();
}
}
let Some(usage_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();
};
if usage_repo.reset_account(&account_id).await.is_err() {
return AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(50001, "Failed to reset account usage", (), request_id),
)
.into_response();
}
state.account_pool().lock().await.reset_usage(&account_id);
AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(
ResetAccountUsageData {
id: account_id,
reset: true,
},
request_id,
),
)
.into_response()
}
pub async fn account_quota(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
headers: HeaderMap,
Path(account_id): Path<String>,
) -> 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_repository() else {
return AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(
50001,
"Account repository is not initialized",
(),
request_id,
),
)
.into_response();
};
let account = match repo.get(&account_id).await {
Ok(Some(account)) => account,
Ok(None) => {
return AdminResponse::new(
StatusCode::NOT_FOUND,
AdminEnvelope::new(40401, "Account not found", (), request_id),
)
.into_response();
}
Err(_) => {
return AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(50001, "Failed to load account", (), request_id),
)
.into_response();
}
};
if account.status != AccountStatus::Active {
return AdminResponse::new(
StatusCode::CONFLICT,
AdminEnvelope::new(
40901,
format!(
"Account is {}, cannot query quota",
account_status_value(account.status)
),
(),
request_id,
),
)
.into_response();
}
match fetch_account_usage(&state, &account, &request_id).await {
Ok(raw) => {
let quota = quota_from_usage(&raw);
if repo
.update_quota_json(&account.id, &quota.to_string())
.await
.is_err()
{
return AdminResponse::new(
StatusCode::INTERNAL_SERVER_ERROR,
AdminEnvelope::new(50001, "Failed to store account quota", (), request_id),
)
.into_response();
}
AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountQuotaData { quota, raw }, request_id),
)
.into_response()
}
Err(error) => {
apply_codex_account_error(&state, &repo, &account, &error).await;
AdminResponse::new(
StatusCode::BAD_GATEWAY,
AdminEnvelope::new(
50201,
"Failed to fetch quota from Codex API",
json!({ "error": public_codex_error(&error) }),
request_id,
),
)
.into_response()
}
}
}
pub async fn update_account_label(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
@@ -2376,6 +2859,465 @@ async fn require_admin_session(
}
}
fn filter_health_check_accounts(
accounts: Vec<StoredAccount>,
ids: Option<&[String]>,
) -> Vec<StoredAccount> {
let Some(ids) = ids else {
return accounts;
};
accounts
.into_iter()
.filter(|account| ids.iter().any(|id| id == &account.id))
.collect()
}
fn skipped_probe_result(account: &StoredAccount, error: &str) -> AccountProbeData {
AccountProbeData {
id: account.id.clone(),
email: account.email.clone(),
previous_status: account_status_value(account.status).to_string(),
result: "skipped".to_string(),
status: Some(account_status_value(account.status).to_string()),
error: Some(error.to_string()),
duration_ms: None,
}
}
async fn probe_account_with_codex_backend(
state: &AppState,
repo: &crate::accounts::repository::AccountRepository,
account: StoredAccount,
request_id: &str,
) -> AccountProbeData {
if account.status == AccountStatus::Disabled {
return skipped_probe_result(&account, "manually disabled");
}
let started_at = Instant::now();
let previous_status = account_status_value(account.status).to_string();
match fetch_account_usage(state, &account, request_id).await {
Ok(raw) => {
let quota = quota_from_usage(&raw);
let _ = repo
.update_quota_json(&account.id, &quota.to_string())
.await;
if account.status != AccountStatus::Active {
let _ = repo.set_status(&account.id, AccountStatus::Active).await;
state
.account_pool()
.lock()
.await
.set_status(&account.id, AccountStatus::Active);
}
AccountProbeData {
id: account.id,
email: account.email,
previous_status,
result: "alive".to_string(),
status: Some(account_status_value(AccountStatus::Active).to_string()),
error: None,
duration_ms: Some(started_at.elapsed().as_millis()),
}
}
Err(error) => {
let status = apply_codex_account_error(state, repo, &account, &error).await;
AccountProbeData {
id: account.id,
email: account.email,
previous_status,
result: "dead".to_string(),
status: status.map(account_status_value).map(ToString::to_string),
error: Some(public_codex_error(&error)),
duration_ms: Some(started_at.elapsed().as_millis()),
}
}
}
}
async fn fetch_account_usage(
state: &AppState,
account: &StoredAccount,
request_id: &str,
) -> Result<Value, CodexClientError> {
let cookie_header = account_cookie_header(state, &account.id).await;
let client = CodexBackendClient::new(
build_reqwest_client(state.config().tls.force_http11)?,
state.config().api.base_url.clone(),
Fingerprint::default_codex_desktop(),
);
client
.fetch_usage(CodexRequestContext {
access_token: account.access_token.expose_secret(),
account_id: account.account_id.as_deref(),
request_id,
turn_state: None,
turn_metadata: None,
beta_features: None,
include_timing_metrics: None,
version: None,
codex_window_id: None,
parent_thread_id: None,
cookie_header: cookie_header.as_deref(),
})
.await
}
async fn account_cookie_header(state: &AppState, account_id: &str) -> Option<String> {
let domain = request_domain(&state.config().api.base_url)?;
state
.cookie_repository()?
.cookie_header(account_id, &domain)
.await
.ok()
.flatten()
}
async fn persist_admin_refreshed_account(
state: &AppState,
repo: &crate::accounts::repository::AccountRepository,
account_id: &str,
tokens: TokenPair,
) -> Result<Account, AccountRepositoryError> {
let access_token = tokens.access_token;
let refresh_token = tokens.refresh_token;
repo.update_tokens(
account_id,
TokenUpdate {
access_token: SecretString::new(access_token.into()),
refresh_token: refresh_token.map(|token| SecretString::new(token.into())),
access_token_expires_at: None,
},
)
.await?;
let account = repo
.get(account_id)
.await?
.ok_or(AccountRepositoryError::Database(sqlx::Error::RowNotFound))?;
let account = pool_account_from_stored(account);
state.account_pool().lock().await.insert(account.clone());
Ok(account)
}
async fn apply_refresh_failure_status(
state: &AppState,
repo: &crate::accounts::repository::AccountRepository,
account: &StoredAccount,
failure: RefreshFailure,
) -> Option<AccountStatus> {
let status = status_for_refresh_failure(failure)?;
let _ = repo.set_status(&account.id, status).await;
state
.account_pool()
.lock()
.await
.set_status(&account.id, status);
Some(status)
}
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,
}
}
fn public_refresh_failure(failure: RefreshFailure) -> &'static str {
match failure {
RefreshFailure::InvalidGrant => "invalidGrant",
RefreshFailure::QuotaExhausted => "quotaExhausted",
RefreshFailure::Banned => "banned",
RefreshFailure::Disabled => "disabled",
RefreshFailure::Transport => "transport",
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CodexAccountErrorAction {
SetStatus(AccountStatus),
RateLimited { retry_after_seconds: u64 },
CloudflareChallenge { cooldown_seconds: u64 },
}
async fn apply_codex_account_error(
state: &AppState,
repo: &crate::accounts::repository::AccountRepository,
account: &StoredAccount,
error: &CodexClientError,
) -> Option<AccountStatus> {
match classify_codex_account_error(error) {
Some(CodexAccountErrorAction::SetStatus(status)) => {
let _ = repo.set_status(&account.id, status).await;
state
.account_pool()
.lock()
.await
.set_status(&account.id, status);
Some(status)
}
Some(CodexAccountErrorAction::RateLimited {
retry_after_seconds,
}) => {
let cooldown_until = Utc::now() + Duration::seconds(retry_after_seconds as i64);
state
.account_pool()
.lock()
.await
.mark_quota_limited_until(&account.id, cooldown_until);
None
}
Some(CodexAccountErrorAction::CloudflareChallenge { cooldown_seconds }) => {
let cooldown_until = Utc::now() + Duration::seconds(cooldown_seconds as i64);
state
.account_pool()
.lock()
.await
.set_cloudflare_cooldown_until(&account.id, cooldown_until);
None
}
None => None,
}
}
fn classify_codex_account_error(error: &CodexClientError) -> Option<CodexAccountErrorAction> {
const DEFAULT_RATE_LIMIT_BACKOFF_SECONDS: u64 = 60;
const MAX_RATE_LIMIT_BACKOFF_SECONDS: u64 = 3_600;
const CLOUDFLARE_CHALLENGE_COOLDOWN_SECONDS: u64 = 120;
let CodexClientError::Upstream {
status,
body,
retry_after_seconds,
} = error
else {
return None;
};
let lower = body.to_ascii_lowercase();
if *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 Some(CodexAccountErrorAction::SetStatus(AccountStatus::Expired));
}
if *status == StatusCode::PAYMENT_REQUIRED || lower.contains("quota") {
return Some(CodexAccountErrorAction::SetStatus(
AccountStatus::QuotaExhausted,
));
}
if *status == StatusCode::TOO_MANY_REQUESTS {
return Some(CodexAccountErrorAction::RateLimited {
retry_after_seconds: retry_after_seconds
.unwrap_or(DEFAULT_RATE_LIMIT_BACKOFF_SECONDS)
.min(MAX_RATE_LIMIT_BACKOFF_SECONDS),
});
}
if *status == StatusCode::FORBIDDEN {
if is_cloudflare_challenge(body) {
return Some(CodexAccountErrorAction::CloudflareChallenge {
cooldown_seconds: CLOUDFLARE_CHALLENGE_COOLDOWN_SECONDS,
});
}
return Some(CodexAccountErrorAction::SetStatus(AccountStatus::Banned));
}
if lower.contains("account has been deactivated")
|| lower.contains("deactivated")
|| lower.contains("banned")
|| lower.contains("suspended")
{
return Some(CodexAccountErrorAction::SetStatus(AccountStatus::Banned));
}
None
}
fn public_codex_error(error: &CodexClientError) -> String {
match error {
CodexClientError::Upstream { status, .. } => {
format!("upstream returned status {}", status.as_u16())
}
CodexClientError::Http(_) => "upstream transport failed".to_string(),
CodexClientError::InvalidHeaderName(_) | CodexClientError::InvalidHeaderValue(_) => {
"invalid upstream request headers".to_string()
}
CodexClientError::UnsupportedTransport(_)
| CodexClientError::WebSocket(_)
| CodexClientError::InvalidSse(_)
| CodexClientError::ModelsUnavailable => "Codex backend request failed".to_string(),
}
}
fn is_cloudflare_challenge(body: &str) -> bool {
let lower = body.to_ascii_lowercase();
lower.contains("cf-mitigated")
|| lower.contains("cf-chl-bypass")
|| lower.contains("_cf_chl")
|| lower.contains("cf_chl")
|| lower.contains("attention required")
|| lower.contains("just a moment")
}
fn quota_from_usage(usage: &Value) -> Value {
let additional = usage
.get("additional_rate_limits")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
let mut rate_limits_by_limit_id = Map::new();
for item in &additional {
let Some(limit_id) = item
.get("metered_feature")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
let quota = quota_from_rate_limit(item.get("rate_limit"));
if quota.is_null() {
continue;
}
rate_limits_by_limit_id.insert(
limit_id.to_string(),
json!({
"limit_id": limit_id,
"limit_name": item.get("limit_name").cloned().unwrap_or(Value::Null),
"allowed": quota.get("allowed").cloned().unwrap_or(Value::Null),
"limit_reached": quota.get("limit_reached").cloned().unwrap_or(Value::Null),
"used_percent": quota.get("used_percent").cloned().unwrap_or(Value::Null),
"remaining_percent": quota.get("remaining_percent").cloned().unwrap_or(Value::Null),
"reset_at": quota.get("reset_at").cloned().unwrap_or(Value::Null),
"limit_window_seconds": quota.get("limit_window_seconds").cloned().unwrap_or(Value::Null),
"secondary_rate_limit": secondary_quota_from_rate_limit(item.get("rate_limit")),
}),
);
}
let additional_review = additional.iter().find(|item| {
is_review_limit_id(item.get("metered_feature").and_then(Value::as_str))
|| is_review_limit_id(item.get("limit_name").and_then(Value::as_str))
});
let code_review_rate_limit = match quota_from_rate_limit(usage.get("code_review_rate_limit")) {
Value::Null => {
quota_from_rate_limit(additional_review.and_then(|item| item.get("rate_limit")))
}
quota => quota,
};
json!({
"plan_type": usage.get("plan_type").cloned().unwrap_or(Value::Null),
"rate_limit": quota_from_rate_limit(usage.get("rate_limit")),
"secondary_rate_limit": secondary_quota_from_rate_limit(usage.get("rate_limit")),
"code_review_rate_limit": code_review_rate_limit,
"rate_limits_by_limit_id": if rate_limits_by_limit_id.is_empty() {
Value::Null
} else {
Value::Object(rate_limits_by_limit_id)
},
"credits": normalize_quota_credits(usage.get("credits")),
})
}
fn quota_from_rate_limit(rate_limit: Option<&Value>) -> Value {
let Some(rate_limit) = rate_limit.filter(|value| !value.is_null()) else {
return Value::Null;
};
let primary = rate_limit.get("primary_window");
let used_percent = primary
.and_then(|window| window.get("used_percent"))
.cloned()
.unwrap_or(Value::Null);
json!({
"allowed": rate_limit.get("allowed").cloned().unwrap_or(Value::Null),
"limit_reached": rate_limit.get("limit_reached").cloned().unwrap_or(Value::Null),
"used_percent": used_percent,
"remaining_percent": remaining_percent(primary.and_then(|window| window.get("used_percent"))),
"reset_at": primary.and_then(|window| window.get("reset_at")).cloned().unwrap_or(Value::Null),
"limit_window_seconds": primary.and_then(|window| window.get("limit_window_seconds")).cloned().unwrap_or(Value::Null),
})
}
fn secondary_quota_from_rate_limit(rate_limit: Option<&Value>) -> Value {
let Some(secondary) = rate_limit
.and_then(|rate_limit| rate_limit.get("secondary_window"))
.filter(|value| !value.is_null())
else {
return Value::Null;
};
let used_percent = secondary
.get("used_percent")
.cloned()
.unwrap_or(Value::Null);
let limit_reached = secondary
.get("used_percent")
.and_then(Value::as_f64)
.map(|used| used >= 100.0)
.map(Value::Bool)
.or_else(|| {
rate_limit
.and_then(|rate_limit| rate_limit.get("limit_reached"))
.cloned()
})
.unwrap_or(Value::Null);
json!({
"limit_reached": limit_reached,
"used_percent": used_percent,
"remaining_percent": remaining_percent(secondary.get("used_percent")),
"reset_at": secondary.get("reset_at").cloned().unwrap_or(Value::Null),
"limit_window_seconds": secondary.get("limit_window_seconds").cloned().unwrap_or(Value::Null),
})
}
fn normalize_quota_credits(raw: Option<&Value>) -> Value {
let Some(raw) = raw.filter(|value| !value.is_null()) else {
return Value::Null;
};
let Some(balance) = raw
.get("balance")
.and_then(Value::as_str)
.and_then(|value| value.parse::<f64>().ok())
.filter(|value| value.is_finite())
else {
return Value::Null;
};
json!({
"has_credits": raw.get("has_credits").and_then(Value::as_bool).unwrap_or(false),
"unlimited": raw.get("unlimited").and_then(Value::as_bool).unwrap_or(false),
"overage_limit_reached": raw.get("overage_limit_reached").and_then(Value::as_bool).unwrap_or(false),
"balance": balance,
})
}
fn remaining_percent(used_percent: Option<&Value>) -> Value {
let Some(used_percent) = used_percent.and_then(Value::as_f64) else {
return Value::Null;
};
json!((100.0 - used_percent.clamp(0.0, 100.0)).round() as i64)
}
fn is_review_limit_id(value: Option<&str>) -> bool {
let normalized = value
.unwrap_or_default()
.trim()
.to_ascii_lowercase()
.replace(['-', ' '], "_");
normalized == "review"
|| normalized == "code_review"
|| normalized == "codex_review"
|| normalized == "codex_code_review"
|| normalized.contains("code_review")
|| normalized.contains("codex_review")
}
fn request_domain(base_url: &str) -> Option<String> {
Url::parse(base_url)
.ok()
.and_then(|url| url.host_str().map(ToString::to_string))
}
async fn sync_runtime_account_status(
state: &AppState,
repo: &crate::accounts::repository::AccountRepository,
tests/admin_accounts_route_test.rs
+342 -17
View File
@@ -13,6 +13,10 @@ use serde_json::{json, Value};
use sqlx::Row;
use tokio::sync::Mutex;
use tower::ServiceExt;
use wiremock::{
matchers::{header, method, path},
Mock, MockServer, ResponseTemplate,
};
use codex_proxy_rs::{
accounts::repository::AccountRepository,
@@ -32,6 +36,32 @@ use codex_proxy_rs::{
storage::db::connect_sqlite,
};
#[derive(Clone)]
struct StaticTokenRefresher {
result: Result<TokenPair, RefreshFailure>,
calls: Arc<Mutex<Vec<String>>>,
}
#[derive(Clone)]
struct FailingTokenRefresher {
failure: RefreshFailure,
}
#[async_trait]
impl TokenRefresher for StaticTokenRefresher {
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
self.calls.lock().await.push(refresh_token.to_string());
self.result.clone()
}
}
#[async_trait]
impl TokenRefresher for FailingTokenRefresher {
async fn refresh(&self, _refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
Err(self.failure)
}
}
fn test_config(database_url: String) -> AppConfig {
AppConfig {
server: ServerConfig {
@@ -89,6 +119,13 @@ fn test_config(database_url: String) -> AppConfig {
}
}
fn test_config_with_base_url(database_url: String, base_url: String) -> AppConfig {
AppConfig {
api: ApiConfig { base_url },
..test_config(database_url)
}
}
#[tokio::test]
async fn admin_accounts_import_should_require_admin_session_cookie() {
let dir = tempfile::tempdir().unwrap();
@@ -1438,9 +1475,311 @@ async fn admin_accounts_import_cli_should_read_codex_auth_file_store_encrypted_a
assert_eq!(acquired.id, stored.0);
}
#[tokio::test]
async fn admin_account_refresh_should_update_tokens_and_runtime_pool_without_returning_secrets() {
let (app, state, pool, _dir) = admin_accounts_test_app_with_refresher(
"admin-account-refresh.sqlite",
28,
StaticTokenRefresher {
result: Ok(TokenPair {
access_token: "new-admin-refresh-access".to_string(),
refresh_token: Some("new-admin-refresh-rt".to_string()),
}),
calls: Arc::new(Mutex::new(Vec::new())),
},
)
.await;
import_test_account(&app, "session_1", "acct_refresh").await;
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/admin/accounts/acct_refresh/refresh")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_refresh_account")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(body["data"]["id"], "acct_refresh");
assert_eq!(body["data"]["result"], "alive");
assert_eq!(body["data"]["previousStatus"], "active");
let serialized = serde_json::to_string(&body).unwrap();
assert!(!serialized.contains("new-admin-refresh-access"));
assert!(!serialized.contains("new-admin-refresh-rt"));
let stored: (String, String) = sqlx::query_as(
"select access_token_cipher, refresh_token_cipher from accounts where id = ?",
)
.bind("acct_refresh")
.fetch_one(&pool)
.await
.unwrap();
assert!(stored.0.starts_with("v1:"));
assert!(!stored.0.contains("new-admin-refresh-access"));
assert!(stored.1.starts_with("v1:"));
assert!(!stored.1.contains("new-admin-refresh-rt"));
let acquired = state
.account_pool()
.lock()
.await
.acquire("gpt-5.5")
.unwrap();
assert_eq!(acquired.access_token, "new-admin-refresh-access");
}
#[tokio::test]
async fn admin_account_refresh_should_mark_invalid_refresh_token_as_expired() {
let (app, state, pool, _dir) = admin_accounts_test_app_with_refresher(
"admin-account-refresh-invalid.sqlite",
29,
FailingTokenRefresher {
failure: RefreshFailure::InvalidGrant,
},
)
.await;
import_test_account(&app, "session_1", "acct_refresh_invalid").await;
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/admin/accounts/acct_refresh_invalid/refresh")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_refresh_invalid")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(body["data"]["result"], "dead");
assert_eq!(body["data"]["status"], "expired");
let status: (String,) = sqlx::query_as("select status from accounts where id = ?")
.bind("acct_refresh_invalid")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(status.0, "expired");
assert!(state
.account_pool()
.lock()
.await
.acquire("gpt-5.5")
.is_none());
}
#[tokio::test]
async fn admin_account_reset_usage_should_clear_local_counters_and_pool_last_used() {
let (app, state, pool, _dir) =
admin_accounts_test_app("admin-account-reset-usage.sqlite", 30).await;
import_test_account(&app, "session_1", "acct_reset_usage").await;
sqlx::query(
"insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, 7, 11, 13, 17, ?)",
)
.bind("acct_reset_usage")
.bind("2026-06-12T12:00:00Z")
.execute(&pool)
.await
.unwrap();
state.reload_account_pool_from_repository().await.unwrap();
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/admin/accounts/acct_reset_usage/reset-usage")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_reset_usage")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(body["data"]["id"], "acct_reset_usage");
assert_eq!(body["data"]["reset"], true);
let usage: (i64, i64, i64, i64, Option<String>) = sqlx::query_as(
"select request_count, input_tokens, output_tokens, cached_tokens, last_used_at from account_usage where account_id = ?",
)
.bind("acct_reset_usage")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(usage, (0, 0, 0, 0, None));
let acquired = state
.account_pool()
.lock()
.await
.acquire("gpt-5.5")
.unwrap();
assert_eq!(acquired.last_used_at, None);
}
#[tokio::test]
async fn admin_account_quota_should_fetch_usage_store_quota_and_not_return_secrets() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.and(header("authorization", "Bearer access-acct_quota"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"plan_type": "plus",
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {
"used_percent": 25,
"limit_window_seconds": 300,
"reset_after_seconds": 120,
"reset_at": 1770000000
},
"secondary_window": null
},
"code_review_rate_limit": null,
"additional_rate_limits": null,
"credits": {
"has_credits": false,
"unlimited": false,
"overage_limit_reached": false,
"balance": "0"
}
})))
.expect(1)
.mount(&server)
.await;
let (app, _state, pool, _dir) =
admin_accounts_test_app_with_base_url("admin-account-quota.sqlite", 31, server.uri()).await;
import_test_account(&app, "session_1", "acct_quota").await;
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/admin/accounts/acct_quota/quota")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_quota")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(body["data"]["quota"]["plan_type"], "plus");
assert_eq!(body["data"]["quota"]["rate_limit"]["remaining_percent"], 75);
assert_eq!(
body["data"]["raw"]["rate_limit"]["primary_window"]["used_percent"],
25
);
let serialized = serde_json::to_string(&body).unwrap();
assert!(!serialized.contains("access-acct_quota"));
assert!(!serialized.contains("refresh-acct_quota"));
let stored: (String,) = sqlx::query_as("select quota_json from accounts where id = ?")
.bind("acct_quota")
.fetch_one(&pool)
.await
.unwrap();
assert!(stored.0.contains("\"remaining_percent\":75"));
assert!(!stored.0.contains("access-acct_quota"));
}
#[tokio::test]
async fn admin_accounts_health_check_should_probe_backend_and_mark_invalid_accounts() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.and(header("authorization", "Bearer access-acct_health_alive"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"plan_type": "plus",
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": null,
"secondary_window": null
},
"code_review_rate_limit": null
})))
.expect(1)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.and(header("authorization", "Bearer access-acct_health_dead"))
.respond_with(ResponseTemplate::new(401).set_body_json(json!({
"error": {"message": "invalid token"}
})))
.expect(1)
.mount(&server)
.await;
let (app, _state, pool, _dir) =
admin_accounts_test_app_with_base_url("admin-account-health.sqlite", 32, server.uri())
.await;
import_test_account(&app, "session_1", "acct_health_alive").await;
import_test_account(&app, "session_1", "acct_health_dead").await;
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/admin/accounts/health-check")
.header("content-type", "application/json")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_health")
.body(Body::from(
json!({
"ids": ["acct_health_alive", "acct_health_dead"],
"concurrency": 2
})
.to_string(),
))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(body["data"]["summary"]["total"], 2);
assert_eq!(body["data"]["summary"]["alive"], 1);
assert_eq!(body["data"]["summary"]["dead"], 1);
assert_eq!(body["data"]["summary"]["skipped"], 0);
let serialized = serde_json::to_string(&body).unwrap();
assert!(!serialized.contains("access-acct_health_alive"));
assert!(!serialized.contains("access-acct_health_dead"));
let status: (String,) = sqlx::query_as("select status from accounts where id = ?")
.bind("acct_health_dead")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(status.0, "expired");
}
async fn admin_accounts_test_app(
db_name: &str,
key_byte: u8,
) -> (Router, AppState, sqlx::SqlitePool, tempfile::TempDir) {
admin_accounts_test_app_with_base_url(
db_name,
key_byte,
"https://chatgpt.com/backend-api".to_string(),
)
.await
}
async fn admin_accounts_test_app_with_base_url(
db_name: &str,
key_byte: u8,
base_url: String,
) -> (Router, AppState, sqlx::SqlitePool, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join(db_name);
@@ -1448,7 +1787,7 @@ async fn admin_accounts_test_app(
let pool = connect_sqlite(&url).await.unwrap();
seed_admin_session(&pool, "session_1").await;
let state = AppState::with_pool_and_secret_box(
test_config(url),
test_config_with_base_url(url, base_url),
pool.clone(),
SecretBox::new([key_byte; 32]),
);
@@ -1459,7 +1798,7 @@ async fn admin_accounts_test_app(
async fn admin_accounts_test_app_with_refresher<C>(
db_name: &str,
key_byte: u8,
refresher: C,
token_refresher: C,
) -> (Router, AppState, sqlx::SqlitePool, tempfile::TempDir)
where
C: TokenRefresher,
@@ -1474,7 +1813,7 @@ where
pool.clone(),
SecretBox::new([key_byte; 32]),
ApiKeyHasher::new([key_byte; 32]),
refresher,
token_refresher,
);
let app = build_router(state.clone());
(app, state, pool, dir)
@@ -1593,17 +1932,3 @@ fn test_jwt(
fn jwt_part(value: &Value) -> String {
URL_SAFE_NO_PAD.encode(serde_json::to_vec(value).unwrap())
}
#[derive(Clone)]
struct StaticTokenRefresher {
result: Result<TokenPair, RefreshFailure>,
calls: Arc<Mutex<Vec<String>>>,
}
#[async_trait]
impl TokenRefresher for StaticTokenRefresher {
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
self.calls.lock().await.push(refresh_token.to_string());
self.result.clone()
}
}