diff --git a/.env.example b/.env.example deleted file mode 100644 index e1b2988c6..000000000 --- a/.env.example +++ /dev/null @@ -1,6 +0,0 @@ -CPRS_HOST=127.0.0.1 -CPRS_PORT=8080 -CPRS_DATABASE_URL=sqlite://data/codex-proxy-rs.sqlite -CPRS_MASTER_KEY_FILE=data/master.key -CPRS_API_KEY_PEPPER_FILE=data/api-key-pepper.key -RUST_LOG=codex_proxy_rs=info,tower_http=info diff --git a/.gitignore b/.gitignore index 5b3f219b8..cdee3219c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,11 +1,8 @@ /target/ /data/ /logs/ -/.env -/.env.* -!/.env.example -/config/local.yaml -/config/local.yml +/local.yaml +/local.yml /.worktrees/ /worktrees/ *.sqlite diff --git a/Cargo.lock b/Cargo.lock index 4628fe3c6..315561199 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -376,7 +376,6 @@ dependencies = [ "base64", "chrono", "config", - "dotenvy", "futures", "hmac", "insta", diff --git a/Cargo.toml b/Cargo.toml index c6f72f44b..9a3f69017 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,7 +17,6 @@ axum = "0.8.9" base64 = "0.22.1" chrono = { version = "0.4.45", features = ["serde", "clock"] } config = "0.15.23" -dotenvy = "0.15.7" futures = "0.3.32" hmac = "0.13.0" rand = "0.10.1" diff --git a/Dockerfile b/Dockerfile index dfc585890..91424bb2d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -14,10 +14,9 @@ RUN apt-get update \ WORKDIR /app COPY --from=builder /app/target/release/codex-proxy-rs /usr/local/bin/codex-proxy-rs -COPY config ./config +COPY config.yaml ./config.yaml COPY migrations ./migrations -ENV RUST_LOG=codex_proxy_rs=info,tower_http=info EXPOSE 8080 CMD ["codex-proxy-rs"] diff --git a/README.md b/README.md index 6af27fb18..e09c10c87 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,9 @@ Rust rewrite of Codex Proxy focused only on ChatGPT/Codex accounts and the Codex ## Development +Runtime config is loaded from `config.yaml` in the project root, then optional +`local.yaml` / `local.yml`. + ```bash cargo fmt --check cargo clippy --all-targets --all-features --locked -- -D warnings diff --git a/config/default.yaml b/config.yaml similarity index 98% rename from config/default.yaml rename to config.yaml index 4b3e0a856..6756300ff 100644 --- a/config/default.yaml +++ b/config.yaml @@ -1,5 +1,5 @@ server: - host: 127.0.0.1 + host: 0.0.0.0 port: 8080 api: base_url: https://chatgpt.com/backend-api diff --git a/docker-compose.yml b/docker-compose.yml index 65d737180..e218ff9de 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -3,13 +3,6 @@ services: build: . ports: - "8080:8080" - environment: - CPRS_HOST: 0.0.0.0 - CPRS_PORT: 8080 - CPRS_DATABASE_URL: sqlite://data/codex-proxy-rs.sqlite - CPRS_MASTER_KEY_FILE: data/master.key - CPRS_API_KEY_PEPPER_FILE: data/api-key-pepper.key - RUST_LOG: codex_proxy_rs=info,tower_http=info volumes: - ./data:/app/data - ./logs:/app/logs diff --git a/docs/architecture-reorganization.md b/docs/architecture-reorganization.md new file mode 100644 index 000000000..175a9b8fa --- /dev/null +++ b/docs/architecture-reorganization.md @@ -0,0 +1,1119 @@ +# Architecture Reorganization Roadmap + +## Purpose + +This document defines the next cleanup track for `codex-proxy-rs`. + +The project direction is still valid: a single Rust service exposing OpenAI-compatible +`/v1/*` APIs backed by imported ChatGPT/Codex accounts, plus admin routes for local +operation. The problem is now structural growth. First-stage implementation optimized +for shipping behavior quickly, so HTTP handlers and `AppState` accumulated too many +responsibilities. + +The reorganization goal is not academic Clean Architecture. The goal is to keep feature +work cheap by making boundaries explicit: + +```text +http -> service -> codex / logs / storage / system auth +``` + +`http` parses requests, enforces route-level auth, and shapes responses. Business flow +orchestration moves into `service`. Codex-specific capability groups under the `codex` +domain: imported upstream accounts, Codex transport, OAuth/token refresh, protocol +translation, Codex model catalog, account-scoped cookies, and fingerprinting. Outside +`codex`, modules represent the backend system itself: app startup, HTTP, admin/client +authentication, logs, storage, config, and shared helpers. + +## Initial Signals + +The cleanup started from these pressure points: + +| File | Lines | Signal | +| --- | ---: | --- | +| `src/http/admin.rs` | 5169 | Admin routes, response DTOs, auth helpers, account import/export, account health, quota, cookies, API keys, model refresh, logs, usage, and helper logic are all in one module. | +| `src/http/v1.rs` | 1998 | Responses, Chat Completions, models, fallback/retry, refresh, usage logging, SSE collection, WebSocket bridging, and auth are mixed. | +| `src/state.rs` | 257 | `AppServices` is still a service locator with many optional dependencies. | +| `tests/admin_accounts_route_test.rs` | 2588 | Admin account scenarios are too large to scan or move safely. | +| `tests/v1_upstream_route_test.rs` | 2115 | Responses HTTP SSE, WebSocket, fallback, error mapping, usage, and logging side effects are mixed. | +| `tests/chat_completions_route_test.rs` | 710 | Manageable, but should eventually align with the v1 split. | +| `tests/admin_api_keys_route_test.rs` | 659 | Manageable, but can share common admin test setup. | + +The main coupling pattern is: + +```text +handler -> AppState -> repo / pool / codex client / refresh / logs / translation +``` + +That is acceptable for early delivery. It becomes expensive once new behavior needs to +touch fallback, account state, usage logs, and protocol conversion at the same time. + +## Current Progress + +As of this cleanup pass: + +- `src/http/admin/mod.rs` is the admin module entry, backed by resource-specific files + under `src/http/admin/` and a dedicated admin router. +- `src/http/v1/mod.rs` is the v1 module entry, backed by thin route handlers for auth, + chat, responses, models, errors, and router mounting under `src/http/v1/`. +- `src/service/` now contains backend-facing HTTP use-case services: + `AdminAuthService`, `ChatService`, `ResponsesService`, `ApiKeyService`, + `UsageService`, and `LogService`. +- `src/codex/` is now the upstream Codex integration domain. It contains imported + account state and account service orchestration, transport, protocol translation, + OAuth/token refresh, Codex model catalog and refresh orchestration, desktop + fingerprinting, account-scoped cookies, and v1 upstream dispatch/fallback helpers. +- v1 chat and responses handlers delegate orchestration to services. +- admin API key, usage, and logs handlers delegate repository work to services. +- admin login, logout, and status checks delegate business logic to `AdminAuthService`. +- admin OAuth device login/polling and PKCE start/callback/code-relay flows delegate + session acquisition, OAuth token exchange, account import, and duplicate-completion + handling to `AdminAuthService`. +- admin account list/export/status/label/delete/batch/reset-usage/cookies operations + delegate local repository, cookie repository, and account-pool coordination to + `codex::accounts::service::AccountService`. +- admin manual add, CLI import, OAuth token import, bulk account import, single-account + refresh, quota fetch, and health check now share + `codex::accounts::service::AccountService` storage, token refresh, Codex usage + probing, and account-pool synchronization. +- admin model catalog refresh orchestration lives in `codex::models::service::ModelService`. +- `tests/common/` now contains shared integration-test helpers, including generic + response/session helpers plus admin account and v1 upstream route fixtures. +- `tests/admin_accounts_route_test.rs` has been split into list, mutation, + import/export, OAuth/auth, and cookies/quota/refresh/health scenario files. +- `tests/v1_upstream_route_test.rs` has been split into Responses HTTP/SSE, + Responses WebSocket, upstream fallback/refresh, and upstream error scenario files. +- `src/app.rs` and `src/state.rs` have been moved into `src/app/router.rs` and + `src/app/state.rs`; `src/app/bootstrap.rs` owns production `AppState` construction. +- `src/config.rs` has been split into `src/config/types.rs` and + `src/config/loader.rs`. +- `src/crypto.rs` and `src/pagination.rs` have been moved into `src/utils/`, and the + duplicated JSON string traversal helper now lives in `src/utils/json.rs`. +- `src/error.rs` has been moved to `src/codex/protocol/error.rs`, matching its only + remaining protocol-translation use. +- service modules now use the target names `admin_auth`, `api_key`, `chat`, + `responses`, `usage`, and `log` instead of the temporary `*_service` filenames. + Codex-specific model refresh orchestration moved to `codex/models/service.rs`. +- `codex/upstream/mod.rs` exposes a `CodexUpstreamService` service dependency facade backed + by explicit `CodexUpstreamDependencies`, so request sending, refresh retry, fallback, + cookie persistence, usage recording, stream audit, and v1 logging no longer require + public helper wrappers that accept `AppState`. +- `codex/upstream/dispatch.rs`, `fallback.rs`, `refresh.rs`, `stream.rs`, and + `usage.rs` now own the obvious helper groups for dispatch response helpers, retry + classification, token refresh retry, SSE response collection, and usage recording. +- `codex/accounts/service/import.rs`, `mutation.rs`, `cookies.rs`, `quota.rs`, `health.rs`, + `refresh.rs`, and `runtime_pool.rs` now own account import, mutation, cookie, quota, + health-check, single-account refresh, and runtime-pool workflows instead of leaving + them in `service/account/mod.rs`. +- `ChatService` and `ResponsesService` own a cloned `CodexUpstreamService` instance injected from + `AppState` construction, so v1 orchestration no longer passes raw `AppState` through + upstream helper calls. +- `ChatService` and `ResponsesService` now receive only `ModelConfig`, `ModelService`, + and `CodexUpstreamService` as construction dependencies. v1 handlers no longer pass raw + `AppState` into service `handle` calls. +- admin session validation now runs through `AdminAuthService`; the HTTP auth helper only + extracts the admin session cookie and maps failures to the existing admin envelope + shape. +- old `AppState` raw dependency accessors for repositories, OAuth clients, refreshers, + hashers, event logs, the database, and the runtime account pool have been removed. +- runtime account-pool assertions in integration tests now go through narrow + `AccountService` methods instead of `AppState` exposing the pool. +- `AppServices` no longer stores construction-only raw dependencies such as the + database pool, secret box, API key hasher, token refresher, OAuth client, OAuth + sessions, event-log repository, or account pool. Runtime account-pool restore is now + delegated to `AccountService`. +- `AppState` public constructors now feed a private dependency bundle into one assembly + path, so service wiring is centralized instead of repeated across each test/runtime + constructor. +- v1 Responses, upstream fallback/error, and Chat Completions SSE payloads have been + moved out of route tests into `tests/fixtures/*.sse`. +- v1 Responses WebSocket mocked response payloads have been moved into + `tests/fixtures/*.json`. +- `http/admin/response.rs` now provides `AdminError`, and the admin route modules use + `Result` for standard admin-envelope failures. +- admin auth handlers now propagate session, login, OAuth, and PKCE failures through + `AdminError`; success paths that must set cookies or redirect still build explicit + `Response` values. +- admin account handlers now use `Result` plus focused error mappers for repository, + import, refresh, health-check, quota, and validation failures. The quota route keeps a + small route-specific error type only to preserve the existing upstream fetch-error + detail under `data`. + +## Deferred Follow-Ups + +The core single-crate layering pass is complete. These items are intentionally left as +follow-up choices because doing them now would either change route contracts or split +already manageable files further than the current risk requires: + +- follow-up test cleanup for remaining medium-sized route files, especially Chat + Completions and admin account import/OAuth scenarios if they grow further. +- moving remaining large inline JSON bodies into `tests/fixtures/` only where they are + response/test data rather than request-shape assertions. +- deciding whether to preserve or redesign the admin quota fetch-error `data.error` + body before moving it fully onto the standard null-data `AdminError` contract. + +## Non-Goals + +- Do not split into a Rust workspace yet. +- Do not create `crates/core`, `crates/server`, `crates/store`, or `crates/upstream` in + this phase. +- Do not rewrite behavior while splitting files. +- Do not change route contracts, JSON envelopes, request IDs, auth semantics, or error + bodies during the mechanical split. +- Do not merge test cleanup with service extraction in the same change. + +The project is still moving quickly. A single crate with clearer internal modules is +the lowest-cost structure right now. + +## Rust Implementation Standards + +All code changes in this cleanup track should follow the repository's existing style +and the Rust best-practice rules below. + +### Ownership And Borrowing + +- Prefer borrowed inputs: use `&str` instead of `String`, `&[T]` instead of `Vec`, + and `&T` instead of `T` when the callee only reads data. +- Clone only when a new owned value is actually needed. Acceptable clones include + `Arc` handles, cheap client handles that share internal pools, immutable snapshots, + and owned values required by a downstream API. +- Avoid cloning large request bodies, `Vec`, `HashMap`, account lists, or JSON values in + loops. Prefer iterators, references, or moving ownership at the boundary where the + value is consumed. +- Small `Copy` types can be passed by value. Large structs and enums should be passed + by reference unless ownership transfer is part of the API contract. +- Use `Cow<'_, str>` or `Cow<'_, [T]>` only when an API genuinely accepts either + borrowed or owned data. + +### Error Handling + +- Production code should return `Result` for fallible operations. Do not use + `unwrap()` or `expect()` outside tests unless failure is impossible and the reason is + documented. +- Prefer precise error enums with `thiserror` for service/domain/repository errors. + Reserve `anyhow` for binaries, smoke helpers, or test helpers. +- Use `?` for propagation and `map_err`, `or_else`, or `inspect_err` where conversion, + recovery, or logging is needed. +- Avoid converting meaningful errors into `Option` too early. Preserve enough detail + for service-level classification and HTTP response mapping. +- Service errors should convert into admin envelopes or OpenAI-compatible v1 errors at + the HTTP boundary, not inside repository or domain modules. + +### Options And Control Flow + +- Use `let Some(value) = ... else { ... };` or `let Ok(value) = ... else { ... };` + when early return is the clearest path. +- Use `ok_or_else`, `unwrap_or_else`, and `map_or_else` when constructing fallback values + would allocate or do work. +- Use `match` when every variant matters or when transforming between nested + `Result`/`Option` shapes. + +### Async, Sharing, And Trait Objects + +- Shared runtime state should use `Arc` and explicit synchronization primitives such as + `tokio::sync::Mutex` only where mutation is required. +- Keep lock scopes short. Do not hold account-pool or session locks while performing + network requests, database calls, or expensive serialization. +- Prefer generics/static dispatch for local helper abstractions. Use `Arc` + only at boundaries that need runtime substitution, such as token refreshers or OAuth + clients in tests and production. +- Ensure async errors and trait objects used across tasks are `Send + Sync + 'static` + where the runtime requires it. + +### Performance Discipline + +- Do not optimize by guessing. Keep code simple first, then measure if a route or + service becomes hot. +- Avoid intermediate `collect()` calls when an iterator can be consumed directly. +- Avoid boxing or heap allocation until a type is large, recursive, or must cross a + trait-object boundary. +- Use `cargo clippy --all-targets --all-features --locked -- -D warnings` as the + authoritative lint gate. Pay special attention to `redundant_clone`, + `clone_on_copy`, `needless_collect`, `manual_ok_or`, and `large_enum_variant`. +- If a clippy lint must be suppressed, use `#[expect(clippy::...)]` with a short reason + instead of a broad `#[allow(...)]`. + +### Comments And Documentation + +- Prefer clear names and small functions over explanatory comments. +- Add `//` comments only for non-obvious reasoning: protocol quirks, safety guarantees, + workaround context, or route-contract constraints. +- Public service/domain types should have doc comments when their purpose or invariants + are not obvious from the type name. +- Do not leave untracked `TODO` comments. Any TODO that remains in code should reference + an issue or a concrete follow-up document. + +### Tests + +- Use descriptive test names that read like behavior, for example + `responses_should_retry_next_account_after_retry_after`. +- Prefer one behavior per test. Integration tests may need multiple assertions for a + route contract, but the scenario should still be narrow. +- Put pure logic tests next to the source module when they do not need HTTP setup. +- Use integration tests for route contracts, SQLite side effects, account-pool behavior, + mocked Codex transport, SSE, and WebSocket flows. +- Use fixture files for large JSON/SSE/WebSocket payloads. Keep snapshots small if + snapshot testing is added later. + +### Type-State Pattern + +Type-state is optional, not a default style. Consider it only when it makes invalid +states unrepresentable without excessive generic complexity, such as a future builder +that must prove required OAuth/session fields before a request can be sent. + +## Target Source Layout + +Target shape inside the existing crate: + +```text +src/ + main.rs + lib.rs + + app/ + mod.rs + router.rs + state.rs + bootstrap.rs + + config/ + mod.rs + types.rs + loader.rs + + utils/ + mod.rs + pagination.rs + json.rs + crypto.rs + + http/ + mod.rs + auth.rs + health.rs + middleware.rs + admin/ + mod.rs + router.rs + response.rs + auth.rs + accounts.rs + api_keys.rs + logs.rs + models.rs + settings.rs + usage.rs + v1/ + mod.rs + router.rs + auth.rs + chat.rs + responses.rs + models.rs + errors.rs + service/ + mod.rs + admin_auth.rs + api_key.rs + chat.rs + responses.rs + usage.rs + log.rs + + auth/ + mod.rs + admin_session.rs + api_key.rs + api_key_repository.rs + + codex/ + mod.rs + accounts/ + mod.rs + model.rs + repository.rs + pool.rs + lifecycle.rs + service/ + mod.rs + import.rs + mutation.rs + cookies.rs + quota.rs + health.rs + refresh.rs + runtime_pool.rs + transport/ + mod.rs + client.rs + headers.rs + sse.rs + types.rs + usage.rs + websocket.rs + oauth/ + mod.rs + cli_import.rs + client.rs + refresh.rs + token.rs + protocol/ + mod.rs + openai_to_codex.rs + codex_to_openai.rs + schema.rs + error.rs + models/ + mod.rs + catalog.rs + repository.rs + service.rs + fingerprint/ + mod.rs + model.rs + repository.rs + updater.rs + cookies/ + mod.rs + jar.rs + repository.rs + upstream/ + mod.rs + dispatch.rs + fallback.rs + refresh.rs + stream.rs + usage.rs + + logs/ + storage/ +``` + +This is a target, not a one-commit requirement. The project should stay as one crate +until module responsibilities stabilize. + +### Module Boundary Rules + +- `main.rs` should remain the binary entry point. Startup wiring should move toward + `app/bootstrap.rs`, route assembly toward `app/router.rs`, and runtime state toward + `app/state.rs`. +- `config` owns configuration data, loading, and environment parsing. Runtime business + modules should consume typed config, not environment variables. +- `utils` contains cross-domain primitives only. It can depend on standard crates and + small general-purpose dependencies such as `serde_json`, `base64`, and encryption + crates, but it must not depend on `http`, `service`, `auth`, `codex`, or other domain + modules. +- `utils/pagination.rs` owns cursor encoding, decoding, limit clamping, and the + reusable `Page` shape. +- `utils/json.rs` owns small JSON traversal helpers used by import/parsing code. It + should not contain route-specific DTOs or validation policy. +- `utils/crypto.rs` owns reusable secret encryption primitives such as `SecretBox`, + `CryptoError`, and `CryptoResult`. API-key hashing remains in `auth/api_key.rs`. +- `http` owns route registration, auth extraction, request parsing, and response + mapping. It should not construct repositories, lock account pools, call Codex + clients directly, or perform protocol translation workflows. +- `service` owns backend-facing orchestration for HTTP use cases. It delegates Codex + account, model, upstream retry/fallback, protocol, OAuth, fingerprint, and transport + details into `codex/*`. +- `codex/accounts` owns imported Codex account data shapes, repository contracts, + lifecycle state, and the runtime account pool. It is not a `utils` helper and should + not be named only `accounts`, because this backend can also have its own local admin + users and client API keys. +- `codex/transport` owns direct Codex backend communication: request client, headers, + SSE parsing/encoding, usage extraction, transport types, and WebSocket bridging. +- `codex/oauth` owns Codex/OpenAI account authorization and token refresh mechanics. + System authentication remains in top-level `auth`. +- `codex/protocol` owns OpenAI-compatible request/response translation and + translation-specific errors. Do not recreate a top-level `src/error.rs` unless a + genuinely crate-wide error boundary is introduced. +- `codex/models`, `codex/fingerprint`, and `codex/cookies` own Codex-specific model + catalog, desktop fingerprint, and account-scoped cookie replay. +- `codex/upstream` owns v1 upstream dispatch, fallback, refresh retry, stream + collection, usage recording, and lifecycle logging helpers used by Chat and + Responses services. +- Top-level `logs`, `storage`, `config`, `utils`, `app`, `http`, and system `auth` + remain backend-system modules. + +### Completed Structural Increment + +The current low-risk structural increment moved the remaining top-level helper and +application wiring files into the target directories without changing behavior: + +```text +src/app.rs -> src/app/router.rs +src/state.rs -> src/app/state.rs +src/config.rs -> src/config/{types,loader}.rs +src/crypto.rs -> src/utils/crypto.rs +src/pagination.rs -> src/utils/pagination.rs +src/error.rs -> src/codex/protocol/error.rs +src/service/*_service.rs -> src/service/*.rs +src/service/v1_upstream.rs -> src/codex/upstream/mod.rs +src/service/account_service.rs -> src/codex/accounts/service/mod.rs +src/service/model.rs -> src/codex/models/service.rs +src/service/upstream/* -> src/codex/upstream/* +``` + +### Codex Domain Migration Order + +This structural track makes `codex/` the only top-level home for upstream Codex +concerns. Each step is behavior-preserving and should keep passing the verification +gate: + +1. Move the imported upstream account module into `codex/accounts` and update imports + to `crate::codex::accounts`. +2. Move direct Codex backend files into `codex/transport` and update imports to + `crate::codex::transport`. +3. Move OpenAI/Codex translation into `codex/protocol`. +4. Move Codex OAuth/token refresh/CLI import into `codex/oauth`, leaving system admin + sessions and client API keys in top-level `auth`. +5. Move Codex model catalog, fingerprint, and account-scoped cookies under + `codex/models`, `codex/fingerprint`, and `codex/cookies`. +6. Move remaining Codex-specific orchestration out of top-level `service/account`, + `service/model`, and `service/upstream` into `codex/accounts/service`, + `codex/models/service`, and `codex/upstream`. + +Callers should import cross-domain helpers through: + +```rust +use crate::utils::{ + crypto::{CryptoError, SecretBox}, + pagination::{clamp_limit, decode_cursor, encode_cursor, Page}, +}; +``` + +Integration tests should import public helpers through `codex_proxy_rs::utils::*`. +Do not keep top-level compatibility wrappers for `crypto` or `pagination`; leaving both +paths active would make the target boundary less clear. + +## Phase 1: Split HTTP Files Without Behavior Changes + +### Goal + +Reduce the former `src/http/admin.rs` and `src/http/v1.rs` large files into +resource-focused module directories while keeping handler bodies and helper logic +intact. + +This phase should be mostly mechanical: + +- move related structs/functions into new files; +- keep function names stable where tests or `app.rs` reference them; +- preserve route paths and methods; +- preserve response body shapes; +- preserve existing helper behavior, even when it is not ideally placed yet. + +### Admin Split + +Create: + +```text +src/http/admin/ + mod.rs + router.rs + response.rs + auth.rs + accounts.rs + api_keys.rs + logs.rs + models.rs + settings.rs + usage.rs +``` + +Suggested ownership: + +| Module | Initial contents | +| --- | --- | +| `response.rs` | `AdminEnvelope`, `AdminPageEnvelope`, `AdminResponse`, shared admin DTO helpers that are not resource-specific. | +| `auth.rs` | `login`, `auth_status`, `auth_logout`, OAuth PKCE/device routes, admin session cookie helpers. | +| `accounts.rs` | account list/create/import/export/delete/status/label, health check, refresh, reset usage, quota, cookies, account import parsing helpers. | +| `api_keys.rs` | local client key list/create/import/export/delete/status/label/batch helpers. | +| `logs.rs` | `/admin/logs` handler and log DTOs. | +| `models.rs` | `/admin/refresh-models` handler and model refresh DTOs. | +| `settings.rs` | `/admin/settings` handler and settings DTOs. | +| `usage.rs` | `/admin/usage-stats` and `/admin/usage-stats/summary`. | +| `router.rs` | admin route mounting only. | + +Keep cross-resource helpers private when possible. If a helper is needed by multiple +admin modules during the split, expose it as `pub(super)` and revisit it in Phase 2. + +The admin router should own admin paths: + +```rust +pub fn router() -> Router { + Router::new() + .route("/admin/login", post(auth::login)) + .route("/admin/auth/status", get(auth::auth_status)) + .route("/admin/accounts", get(accounts::accounts).post(accounts::create_account)) + .route("/admin/api-keys", get(api_keys::api_keys).post(api_keys::create_api_key)) +} +``` + +`src/app/router.rs` should then merge admin routes instead of importing every admin handler: + +```rust +Router::new() + .route("/health", get(health)) + .merge(v1::router()) + .merge(admin::router()) +``` + +### V1 Split + +Create: + +```text +src/http/v1/ + mod.rs + router.rs + auth.rs + chat.rs + responses.rs + models.rs + errors.rs + logging.rs + upstream.rs +``` + +Suggested ownership: + +| Module | Initial contents | +| --- | --- | +| `chat.rs` | `/v1/chat/completions` handler and Chat-specific response conversion calls. | +| `responses.rs` | `/v1/responses` handler, SSE collection, non-streaming collection, WebSocket response path. | +| `models.rs` | `/v1/models`, catalog/detail/info/debug model routes. | +| `auth.rs` | client API key route authorization and missing-key response. | +| `errors.rs` | OpenAI-compatible v1 error body helpers such as model-not-found and upstream client errors. | +| `logging.rs` | `CodexRequestLogContext`, Codex upstream response logging, stream metadata helpers. | +| `upstream.rs` | account acquire/release guard, refresh retry, upstream account retry classification, cookie persistence, usage recording. | +| `router.rs` | v1 route mounting only. | + +In Phase 1, `upstream.rs` can still depend on `AppState` directly. The dependency +direction improves in Phase 2 when services become the orchestration boundary. + +### Phase 1 Acceptance + +Run after each small move: + +```bash +cargo fmt --check +cargo test --test admin_accounts_list_test +cargo test --test admin_accounts_mutation_test +cargo test --test admin_accounts_import_export_test +cargo test --test admin_accounts_oauth_test +cargo test --test admin_accounts_cookies_quota_test +cargo test --test admin_api_keys_route_test +cargo test --test v1_responses_http_sse_test +cargo test --test v1_responses_websocket_test +cargo test --test v1_upstream_fallback_test +cargo test --test v1_upstream_errors_test +cargo test --test chat_completions_route_test +``` + +Run before finishing the phase: + +```bash +cargo fmt --check +cargo test +cargo clippy --all-targets --all-features --locked -- -D warnings +``` + +## Phase 2: Introduce Service Layer + +### Goal + +Move business workflow orchestration out of handlers. Handlers should become thin: + +```rust +pub async fn chat_completions( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Json(req): Json, +) -> Response { + match state.services.chat.handle(request_id, headers, req).await { + Ok(resp) => resp.into_response(), + Err(err) => err.into_response(), + } +} +``` + +Do this after Phase 1 so the move is readable and test failures point to business +logic, not file-boundary churn. + +### Initial Services + +Start with v1 because it is the core proxy path: + +| Service | Responsibility | +| --- | --- | +| `ChatService` | Chat Completions flow: auth result input, model resolution, Chat-to-Codex translation, account acquisition, Codex transport, retry/fallback, Chat output conversion, usage/log side effects. | +| `ResponsesService` | Responses flow: model resolution, HTTP SSE/WebSocket transport selection, previous-response affinity, fallback policy, collection/streaming behavior, usage/log side effects. | +| `ModelService` | model catalog reads, cached backend model snapshots, v1 model response shaping. | + +Then move admin orchestration: + +| Service | Responsibility | +| --- | --- | +| `AdminAuthService` | admin login, session lifecycle, OAuth PKCE/device flows, logout/status. | +| `AccountService` | manual import, CLI import, native/Sub2API import/export, status/label/delete, refresh, quota, cookies, runtime pool synchronization. | +| `ApiKeyService` | local client key create/list/update/delete/import/export. | +| `UsageService` | usage list and summary queries. | +| `LogService` | admin log queries and v1 lifecycle log writes if useful. | + +The service layer should depend on repositories and domain modules. HTTP should depend +on services. Domain modules should not depend on HTTP. + +### Phase 2 Acceptance + +Each migrated handler should preserve its current route test. For every service moved, +add or retain focused tests at the cheapest level: + +- pure conversion and classification tests stay near source modules; +- route contract tests stay in `tests/`; +- service tests are useful when the orchestration has many branches and can be tested + without HTTP request boilerplate. + +Do not migrate all handlers at once. Move one vertical flow at a time: + +1. `ModelService` +2. `ResponsesService` +3. `ChatService` +4. `AdminAuthService` +5. `AccountService` +6. `ApiKeyService` + +## Phase 3: Reshape AppState + +### Current Shape + +`AppState` already wraps `Arc`, but `AppServices` still exposes raw, +optional dependencies: + +```rust +pub struct AppServices { + pub config: AppConfig, + pub db: Option, + pub event_logs: Option, + pub secret_box: Option, + pub api_key_hasher: Option, + pub token_refresher: Option>, + pub oauth_client: Option>, + pub oauth_sessions: Arc>, + pub account_pool: Arc>, +} +``` + +This creates two problems: + +- handlers and helpers can reach almost everything; +- every business path has to repeatedly decide what a missing dependency means. + +### Target Shape + +Move toward: + +```rust +pub struct AppState { + pub services: Arc, +} + +pub struct AppServices { + pub config: AppConfig, + pub auth: AdminAuthService, + pub accounts: AccountService, + pub api_keys: ApiKeyService, + pub chat: ChatService, + pub responses: ResponsesService, + pub models: ModelService, + pub usage: UsageService, + pub logs: LogService, +} +``` + +Repositories, Codex clients, account pool handles, OAuth clients, token refreshers, +and hashers should be injected into the service that actually uses them. + +Keep constructors practical. During migration, it is fine to support both old helper +accessors and new service fields. Remove the old accessors only after handlers stop +using them. + +### Dependency Rule + +After this phase: + +- HTTP modules should not construct repositories. +- HTTP modules should not lock the account pool directly. +- HTTP modules should not call Codex clients directly. +- HTTP modules should not update usage or event logs directly. +- Service modules can orchestrate those operations. +- Domain/repository modules should remain HTTP-agnostic. + +## Phase 4: Unify Errors And Response Envelopes + +### Goal + +Make route error handling less repetitive without changing external contracts. + +Targets: + +```text +src/codex/protocol/error.rs +src/http/response.rs +src/http/admin/response.rs +src/http/v1/errors.rs +``` + +Admin and v1 should keep different public contracts: + +- admin routes return the admin envelope and lower camelCase body fields; +- v1 routes return OpenAI-compatible error objects and headers. + +The unification is internal: shared error types, shared conversion points, and fewer +ad hoc `(StatusCode, Json)` helpers in handlers. + +Preferred handler shape: + +```rust +pub async fn handler(...) -> Result { + let data = state.services.accounts.list(...).await?; + Ok(AdminEnvelope::ok(data, request_id)) +} +``` + +Do not force one public response format across admin and v1. + +## Test Reorganization + +The test suite should be split for the same reason as the source: current files prove +behavior, but they are too large to maintain. + +Rust integration tests compile each `tests/*.rs` file as a separate crate, so avoid +extreme fragmentation. Prefer 20-40 focused integration test files rather than one +file per endpoint. + +### Target Test Layout + +Recommended shape: + +```text +tests/ + common/ + mod.rs + app.rs + fixtures.rs + http.rs + mock.rs + assertions.rs + fixtures/ + openai_chat_basic.json + codex_response_basic.sse + codex_response_error.sse + websocket_response_completed.json + admin_auth_test.rs + admin_accounts_list_test.rs + admin_accounts_mutation_test.rs + admin_accounts_import_export_test.rs + admin_accounts_oauth_test.rs + admin_accounts_cookies_quota_test.rs + admin_api_keys_test.rs + admin_logs_models_settings_test.rs + admin_usage_test.rs + v1_chat_test.rs + v1_responses_http_sse_test.rs + v1_responses_websocket_test.rs + v1_upstream_fallback_test.rs + v1_upstream_errors_test.rs + v1_models_test.rs + account_pool_test.rs + account_repository_test.rs + codex_client_test.rs + codex_websocket_test.rs + translation_test.rs +``` + +### `tests/common` + +Move shared setup first, before splitting large tests: + +```rust +pub struct TestApp { + pub router: Router, + pub state: AppState, + pub db: SqlitePool, + pub client_key: Option, + pub admin_cookie: Option, +} +``` + +Useful helpers: + +```rust +pub async fn test_app() -> TestApp; +pub async fn test_app_with_config(config: AppConfig) -> TestApp; +pub async fn create_admin_session(app: &TestApp) -> String; +pub async fn create_client_key(app: &TestApp) -> String; +pub async fn import_test_account(app: &TestApp) -> Account; +pub async fn json_get(app: &TestApp, path: &str) -> Response; +pub async fn json_post(app: &TestApp, path: &str, body: &T) -> Response; +pub fn assert_status(response: &Response, expected: StatusCode); +``` + +If upstream mocking is repeated, wrap it: + +```rust +pub struct MockCodex { + pub base_url: Url, + pub server: MockServer, +} +``` + +### First Test Splits + +Split `tests/admin_accounts_route_test.rs` first: + +| New file | Scenario group | +| --- | --- | +| `admin_accounts_list_test.rs` | list, status read, labels in list, reset usage visibility. | +| `admin_accounts_mutation_test.rs` | create, delete, batch delete, batch status, label/status updates. | +| `admin_accounts_import_export_test.rs` | manual import, native import/export, Sub2API import/export, CLI import. | +| `admin_accounts_oauth_test.rs` | device login, device poll, PKCE login start, callback/code relay. | +| `admin_accounts_cookies_quota_test.rs` | cookies, health check, refresh, quota, reset usage side effects. | + +Then split `tests/v1_upstream_route_test.rs`: + +| New file | Scenario group | +| --- | --- | +| `v1_responses_http_sse_test.rs` | Responses HTTP SSE setup, non-stream collection, passthrough streaming, terminal SSE failures. | +| `v1_responses_websocket_test.rs` | previous response WebSocket, explicit WebSocket, account affinity, handshake failures. | +| `v1_upstream_fallback_test.rs` | 429/402/403 classification, retry-after, fallback account acquisition, exhausted account behavior. | +| `v1_upstream_errors_test.rs` | upstream error bodies, transport error classification, refresh failure mapping, lifecycle log status. | +| `v1_models_test.rs` | model list/catalog/detail/info/debug routes if not already covered elsewhere. | + +Keep `tests/chat_completions_route_test.rs` until v1 helpers are extracted, then rename or +merge into `v1_chat_test.rs`. + +### Unit Tests Near Source + +Move pure logic tests closer to source modules where practical: + +- translation request mapping; +- Codex-to-OpenAI response conversion; +- model catalog suffix parsing; +- upstream error classification; +- retry-after parsing; +- usage extraction; +- quota normalization; +- header construction. + +Integration tests should prove route contracts and side effects. Unit tests should cover +branchy pure functions without HTTP setup. + +### Fixture Files + +Move large JSON/SSE/WebSocket payloads into `tests/fixtures/`: + +```rust +const BASIC_CHAT: &str = include_str!("fixtures/openai_chat_basic.json"); +const CODEX_ERROR_SSE: &str = include_str!("fixtures/codex_response_error.sse"); +``` + +Do this after test setup extraction and initial file splits. Fixture movement is cleanup, +not a prerequisite for architectural splitting. + +## Recommended Execution Order + +The branch has already completed the large HTTP file split, most initial service +extraction, `tests/common` extraction, the two largest integration test splits, +`AdminAuthService`-backed admin session validation, raw `AppState` dependency accessor +removal, v1 upstream dependency injection, the `app/`, `config/`, `utils/`, and +`translation/error.rs` migrations, and the first account/upstream service submodule +splits. From the current state, continue in this order: + +1. Continue moving remaining account/upstream methods from `mod.rs` into the workflow + submodules when a helper group can move without changing route behavior. +2. Finish moving large inline JSON route payloads into `tests/fixtures/` + where they are test data rather than request-shape assertions. +3. Continue unifying internal error/response conversion, starting with the larger admin + `auth` and `accounts` modules after the smaller Result-style handlers have stayed + green. +4. Revisit medium-sized route tests only when they grow or start duplicating setup. + +The original sequence below is kept as a reference for starting from an unorganized +tree or reviewing why earlier moves were ordered that way. + +### Step 1: Establish Test Common Ground + +Create `tests/common` and move duplicated app/database/session/key setup there. Do not +split test files yet. + +Verification: + +```bash +cargo test --test admin_accounts_list_test +cargo test --test admin_accounts_mutation_test +cargo test --test admin_accounts_import_export_test +cargo test --test admin_accounts_oauth_test +cargo test --test admin_accounts_cookies_quota_test +cargo test --test v1_responses_http_sse_test +cargo test --test v1_responses_websocket_test +cargo test --test v1_upstream_fallback_test +cargo test --test v1_upstream_errors_test +cargo test +``` + +### Step 2: Split Admin Tests + +Mechanically move scenarios out of `admin_accounts_route_test.rs` by scenario group. +Keep assertions unchanged unless a shared helper removes duplication. + +Verification after each new file: + +```bash +cargo test --test admin_accounts_list_test +cargo test --test admin_accounts_mutation_test +cargo test --test admin_accounts_import_export_test +cargo test --test admin_accounts_oauth_test +cargo test --test admin_accounts_cookies_quota_test +``` + +### Step 3: Split V1 Tests + +Mechanically move scenarios out of `v1_upstream_route_test.rs`. + +Verification after each new file: + +```bash +cargo test --test v1_responses_http_sse_test +cargo test --test v1_responses_websocket_test +cargo test --test v1_upstream_fallback_test +cargo test --test v1_upstream_errors_test +``` + +### Step 4: Split Former `src/http/admin.rs` + +Move admin code into `src/http/admin/*.rs`, with `src/http/admin/mod.rs` as the module +entry. Preserve behavior. + +Verification: + +```bash +cargo test --test admin_auth_test +cargo test --test admin_accounts_list_test +cargo test --test admin_accounts_mutation_test +cargo test --test admin_accounts_import_export_test +cargo test --test admin_accounts_oauth_test +cargo test --test admin_accounts_cookies_quota_test +cargo test --test admin_api_keys_route_test +cargo test --test admin_logs_route_test +cargo test --test admin_models_route_test +cargo test --test admin_settings_route_test +cargo test --test admin_usage_stats_route_test +``` + +### Step 5: Split Former `src/http/v1.rs` + +Move v1 code into `src/http/v1/*.rs`, with `src/http/v1/mod.rs` as the module entry. +Preserve behavior. + +Verification: + +```bash +cargo test --test chat_completions_route_test +cargo test --test v1_responses_http_sse_test +cargo test --test v1_responses_websocket_test +cargo test --test v1_upstream_fallback_test +cargo test --test v1_upstream_errors_test +cargo test --test routes_responses_test +``` + +### Step 6: Extract Services One Flow At A Time + +Start with low-risk service extraction: + +1. `ModelService` +2. `ResponsesService` +3. `ChatService` +4. `AdminAuthService` +5. `AccountService` +6. `ApiKeyService` + +For each service: + +- move orchestration out of the handler; +- keep HTTP request parsing and response wrapping in HTTP; +- keep repository SQL in repository modules; +- keep Codex transport in `codex`; +- keep translation in `translation`; +- keep route tests green. + +### Step 7: Tighten `AppState` + +After services own dependencies, remove direct handler access to raw repositories, +pool locks, token refreshers, and Codex clients. Keep this as a separate phase so +behavior regressions are easier to isolate. + +### Step 8: Response/Error Cleanup + +Unify internal error conversion after service boundaries are stable. Do not change +public admin or v1 response contracts. + +## Guardrails For Each PR + +Each PR or commit should state which kind of change it is: + +```text +mechanical move only +test split only +service extraction +state dependency cleanup +error/response cleanup +``` + +Avoid mixing categories. + +Before merging a mechanical move, use diff review to confirm: + +- no route path changed; +- no HTTP method changed; +- no JSON field changed; +- no status code changed; +- no auth requirement changed; +- no request-id behavior changed; +- no upstream fallback policy changed; +- no database schema changed. + +Minimum verification for mechanical HTTP moves: + +```bash +cargo fmt --check +cargo test --test admin_accounts_list_test +cargo test --test admin_accounts_mutation_test +cargo test --test admin_accounts_import_export_test +cargo test --test admin_accounts_oauth_test +cargo test --test admin_accounts_cookies_quota_test +cargo test --test admin_api_keys_route_test +cargo test --test v1_responses_http_sse_test +cargo test --test v1_responses_websocket_test +cargo test --test v1_upstream_fallback_test +cargo test --test v1_upstream_errors_test +cargo test --test chat_completions_route_test +``` + +Full verification before declaring a phase complete: + +```bash +cargo fmt --check +cargo test +cargo clippy --all-targets --all-features --locked -- -D warnings +``` + +## Decision Record + +- Keep a single crate for now. +- Split source files before changing architecture. +- Split tests before or alongside source file movement so failures are local. +- Extract v1 services before admin services because v1 is the core proxy path. +- Keep admin and v1 public response contracts separate. +- Treat local `cpr_` API key management as a Rust-local admin utility, not TypeScript + provider-key parity. +- Do not add proxy pools, non-OpenAI providers, frontend/Electron flows, or workspace + crate boundaries as part of this cleanup. diff --git a/docs/implementation-status.md b/docs/implementation-status.md index bd654c019..1e7e13b09 100644 --- a/docs/implementation-status.md +++ b/docs/implementation-status.md @@ -6,7 +6,7 @@ Update this file before each feature commit. | --- | --- | --- | --- | --- | | Scaffold | Completed | `fc47967` | `cargo fmt --check`; `cargo test`; `cargo clippy --all-targets --all-features --locked -- -D warnings` | Rust crate, lint policy, pinned TLS-sensitive dependencies, `.gitignore`, and minimal entry points are in place. | | API contract docs | Completed | `fc47967` | `cargo fmt --check`; `cargo test`; `cargo clippy --all-targets --all-features --locked -- -D warnings` | `/v1/*` OpenAI-compatible body, `/admin/*` frontend envelope, body codes, camelCase, and request ID policy are documented. | -| Configuration | Completed | `9ed8579` + `e0680bf` + `cfec8ba` | `cargo test default_config_keeps_only_codex_backend`; `cargo test config_loader_merges_default_local_and_cprs_env_overrides`; full `fmt/test/clippy` before commit | Defines Codex-only config, default/local/env loading, default YAML, `Arc` state shell, `gpt-5.5` default model, and in-scope model/auth/quota/usage/logging settings fields without proxy/Ollama/provider/OpenAI-key knobs. | +| Configuration | Completed | `9ed8579` + `e0680bf` + `cfec8ba` | `cargo test default_config_keeps_only_codex_backend`; `cargo test config_loader_merges_config_and_local_yaml`; full `fmt/test/clippy` before commit | Defines Codex-only config, root `config.yaml` plus optional local YAML loading, `Arc` state shell, `gpt-5.5` default model, and in-scope model/auth/quota/usage/logging settings fields without proxy/Ollama/provider/OpenAI-key knobs. | | 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` + `8af350d` + `960bbb2` | `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`; `cargo fmt --check`; `cargo test`; `cargo clippy --all-targets --all-features --locked -- -D warnings`; full `fmt/test/clippy` before `8af350d`; full `fmt/test/clippy` before `960bbb2`; current local key import/export targeted `cargo test --test admin_api_keys_route_test` | 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, status-enable/status-disable local client keys, label local client keys, delete local client keys, batch-delete local client keys, export local client key metadata, import metadata by rotating to new local `cpr_` keys, and `/v1/*` verifies only enabled keys against SQLite instead of accepting any `cpr_` string. These routes are Rust-local client auth utilities, not TypeScript third-party provider API-key parity. Provider API key pools and provider key import/export remain out of scope; local `cpr_` key plaintext and key hashes are not exportable after HMAC storage, and import returns newly generated plaintext values only once. | | 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. | @@ -23,4 +23,4 @@ Update this file before each feature commit. | Upstream lifecycle | Partial | `88cad87` + `225fbff` + `c2cf7a6` + `636fa79` + `ee0ed90` + `fdb9374` + `5bf012c` + `9ca5c7d` + `b9c30ba` + `350ffbe` + `3ecd5b6` + `3960a35` + `bfe561d` + `6d18d2c` | `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`; `cargo test --test codex_websocket_test`; `cargo test --test admin_accounts_route_test`; Task 5 targeted `cargo test --test v1_upstream_route_test`, `cargo test --test chat_completions_route_test`, `cargo test --test account_pool_scheduling_test`, and WebSocket route tests; HTTP SSE error mapping targeted `cargo test --test v1_upstream_route_test v1_responses_non_stream_should_return_upstream_error` and `cargo test --test v1_upstream_route_test v1_responses_stream_should_passthrough`; full `cargo fmt --check && cargo test && cargo clippy --all-targets --all-features --locked -- -D warnings` before `b9c30ba` and `350ffbe` | `/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, sends WebSocket `response.create` requests when `previous_response_id` requires upstream history, keeps `previous_response_id` requests account-affine on classified WebSocket handshake failures for streaming and non-streaming clients, replays/captures encrypted Cookies, captures successful WebSocket rate-limit headers, forwards Codex request context headers, fetches Codex usage/quota through imported account tokens for admin health/quota operations, normalizes `service_tier: fast` to upstream `priority`, collects `response.completed`, reconstructs missing assistant output from `response.output_item.done` or text deltas, synchronizes `output_text` for non-streaming clients, maps non-streaming HTTP SSE `error` and `response.failed` terminal frames to OpenAI-compatible `upstream_error` bodies, passes through Responses streaming SSE by default while marking `error` and `response.failed` terminal frames as failed lifecycle logs, adapts Chat Completions responses/chunks, refreshes and retries once after upstream 401 when a refresh token is stored, parses Retry-After plus Codex reset metadata from upstream errors, records 429 request attempts, classifies upstream 429/402/403 into account cooldown/quota/banned/Cloudflare state changes, retries non-streaming `/v1/responses`, HTTP SSE `/v1/responses` setup, non-history WebSocket setup, and `/v1/chat/completions` with another imported account when available, records token usage, writes v1 lifecycle event logs, records refresh failure classification for upstream 401, and restores the account pool from SQLite on startup. Still missing durable 429 quota cooldown persistence because the current account schema has no dedicated cooldown column, and live Chat streaming transform. | | Codex WebSocket boundary | Completed | `cc0516d` + `b9c30ba` + `350ffbe` + `6d18d2c` | `cargo test --test codex_websocket_test`; `cargo test --test v1_upstream_route_test v1_responses_should_use_websocket_for_previous_response_id_streaming`; WebSocket route fallback tests; full `cargo fmt --check && cargo test && cargo clippy --all-targets --all-features --locked -- -D warnings` before `b9c30ba` and `350ffbe` | Adds `src/codex/websocket.rs` with tested transport classification for HTTP SSE versus WebSocket-only requests. `previous_response_id` and explicit WebSocket requests use typed transport selection, local-only `use_websocket` is not serialized into upstream JSON, WebSocket `response.create` carries Codex headers and parity fields, incoming WebSocket messages are emitted as SSE-compatible events, handshake HTTP failures surface as upstream errors before request body send, successful handshake headers capture turn state, cookies, and rate-limit headers, non-history WebSocket setup can fallback to another imported account, and `previous_response_id` WebSocket failures remain on the original account. | | 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. | +| 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, root runtime config, 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 efca34be3..5c0481ca7 100644 --- a/docs/superpowers/plans/2026-06-11-codex-proxy-rs.md +++ b/docs/superpowers/plans/2026-06-11-codex-proxy-rs.md @@ -6,7 +6,7 @@ **Architecture:** Create a new Rust service at `/home/zyy/桌面/Codes/codex-proxy-rs`. The only upstream is `https://chatgpt.com/backend-api`; remove OpenAI official API key passthrough, Anthropic, Gemini, custom providers, Ollama, Electron, and per-account proxy assignment. Treat Codex Desktop impersonation as a first-class subsystem: locked `reqwest` + `rustls`, exact headers, account-scoped Cookie replay, and fingerprint auto-update. Keep `src/main.rs` as the binary/bootstrap layer and keep reusable logic in `src/lib.rs` modules with typed `thiserror` errors. -**Tech Stack:** Rust 2021, `tokio`, `axum`, `tower-http`, `sqlx` + SQLite WAL, `config`, `dotenvy`, `serde`, `serde_json`, `serde_yaml`, `tracing`, `tracing-subscriber`, rotating file logging, `argon2`, `rand`, `uuid`, `chrono`, `thiserror`, `anyhow` only in `main.rs` and tests, `base64`, `aes-gcm`, `hmac`, `sha2`, `secrecy`, `zeroize`, `wiremock`, `insta`, `tempfile`. Use the latest stable crate versions at implementation time except dependencies intentionally pinned for Codex Desktop TLS fingerprint parity: `reqwest = 0.12.28` and `rustls = 0.23.36` until a fingerprint review proves a newer pair is equivalent. +**Tech Stack:** Rust 2021, `tokio`, `axum`, `tower-http`, `sqlx` + SQLite WAL, `config`, `serde`, `serde_json`, `serde_yaml`, `tracing`, `tracing-subscriber`, rotating file logging, `argon2`, `rand`, `uuid`, `chrono`, `thiserror`, `anyhow` only in `main.rs` and tests, `base64`, `aes-gcm`, `hmac`, `sha2`, `secrecy`, `zeroize`, `wiremock`, `insta`, `tempfile`. Use the latest stable crate versions at implementation time except dependencies intentionally pinned for Codex Desktop TLS fingerprint parity: `reqwest = 0.12.28` and `rustls = 0.23.36` until a fingerprint review proves a newer pair is equivalent. **Execution Mode:** Subagent-Driven. Execute one task at a time with implementation, spec review, and code-quality review before moving to the next task. After each completed task, update this plan's checkbox state and `docs/implementation-status.md` before committing. @@ -50,7 +50,7 @@ These adjustments come from the `rust-best-practices` review and should be treat - Hash admin passwords with Argon2id. Hash high-entropy client API keys with HMAC-SHA256 using a server-side pepper, not with the API key display name. - Prefer concrete service structs over early `dyn Trait` abstraction. Use generics for test seams where useful; introduce trait objects only at boundaries that actually need runtime polymorphism. - Keep shared runtime state in an `Arc` container that holds cheap-clone handles (`SqlitePool`, `reqwest::Client`, repositories, secret box, config). Avoid mutable globals and avoid `std::sync::Mutex` across async `.await`; use SQLite, channels, or `tokio::sync` primitives. -- Use the `config` crate with `config/default.yaml`, optional `config/local.yaml`, and `CPRS_*` environment overrides. Keep `.env` support via `dotenvy` for local development only. +- Use the `config` crate with root `config.yaml` plus optional root `local.yaml` or `local.yml`. Do not use environment-variable overrides for runtime config. - Prefer borrowed parameters in public APIs (`&str`, `&[T]`, references to request structs) unless ownership transfer is part of the domain model. - Keep tests descriptive and targeted. Unit tests live beside modules for private behavior; integration tests under `tests/` cover route and storage behavior. Use `insta` only for structured payload snapshots that are worth reviewing. - Add sparse Chinese comments only where the reason is not obvious from names or types. Good targets are protocol/security/business invariants: TLS fingerprint pinning, Codex Desktop header parity, account-scoped Cookie replay, admin-session versus client-API-key boundaries, and token refresh invariants. Do not comment every line, do not restate obvious code, and do not leave untracked TODO comments. @@ -193,7 +193,7 @@ codex-proxy-rs/ Cargo.toml rust-toolchain.toml README.md - config/default.yaml + config.yaml docs/api.md docs/dependency-policy.md docs/implementation-status.md @@ -305,7 +305,7 @@ Expected: `git status --short` prints no tracked files yet. Run these commands so non-fingerprint dependencies resolve to the latest stable versions available at implementation time: ```bash -cargo add anyhow aes-gcm argon2 async-stream base64 bytes dotenvy futures-util hmac http http-body-util rand serde_json serde_yaml sha2 thiserror tracing tracing-appender uuid zeroize +cargo add anyhow aes-gcm argon2 async-stream base64 bytes futures-util hmac http http-body-util rand serde_json serde_yaml sha2 thiserror tracing tracing-appender uuid zeroize cargo add axum --features macros,ws cargo add chrono --features serde cargo add config @@ -340,7 +340,6 @@ base64 = "0.22" bytes = "1" chrono = { version = "0.4", features = ["serde"] } config = "0.15" -dotenvy = "0.15" futures-util = "0.3" hmac = "0.12" http = "1" @@ -715,7 +714,7 @@ git commit -m "docs: define api contracts and status codes" ### Task 2: Configuration and Application State **Files:** -- Create: `/home/zyy/桌面/Codes/codex-proxy-rs/config/default.yaml` +- Create: `/home/zyy/桌面/Codes/codex-proxy-rs/config.yaml` - Create: `/home/zyy/桌面/Codes/codex-proxy-rs/src/config.rs` - Create: `/home/zyy/桌面/Codes/codex-proxy-rs/src/state.rs` - Modify: `/home/zyy/桌面/Codes/codex-proxy-rs/src/lib.rs` @@ -875,7 +874,7 @@ impl AppState { } ``` -Create `/home/zyy/桌面/Codes/codex-proxy-rs/config/default.yaml`: +Create `/home/zyy/桌面/Codes/codex-proxy-rs/config.yaml`: ```yaml server: @@ -1739,8 +1738,8 @@ git commit -m "feat: add paginated event logs" ### Task 6: Codex Fingerprint, Header Generation, and TLS Client **Files:** -- Create: `/home/zyy/桌面/Codes/codex-proxy-rs/src/fingerprint/mod.rs` -- Create: `/home/zyy/桌面/Codes/codex-proxy-rs/src/fingerprint/model.rs` +- Create: `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/fingerprint/mod.rs` +- Create: `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/fingerprint/model.rs` - Create: `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/mod.rs` - Create: `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/headers.rs` - Create: `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/client.rs` @@ -1753,8 +1752,8 @@ Create `/home/zyy/桌面/Codes/codex-proxy-rs/tests/codex_headers_test.rs`: ```rust use codex_proxy_rs::{ - codex::headers::build_codex_headers, - fingerprint::model::Fingerprint, + codex::transport::headers::build_codex_headers, + codex::fingerprint::model::Fingerprint, }; #[test] @@ -1791,21 +1790,21 @@ pub mod codex; pub mod config; pub mod crypto; pub mod error; -pub mod fingerprint; + pub mod logs; pub mod pagination; pub mod state; pub mod storage; ``` -Create `/home/zyy/桌面/Codes/codex-proxy-rs/src/fingerprint/mod.rs`: +Create `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/fingerprint/mod.rs`: ```rust pub mod model; pub mod updater; ``` -Create `/home/zyy/桌面/Codes/codex-proxy-rs/src/fingerprint/model.rs`: +Create `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/fingerprint/model.rs`: ```rust use serde::{Deserialize, Serialize}; @@ -1860,7 +1859,7 @@ Create `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/headers.rs`: ```rust use std::collections::BTreeMap; -use crate::fingerprint::model::Fingerprint; +use crate::codex::fingerprint::model::Fingerprint; pub fn build_codex_headers( fp: &Fingerprint, @@ -1907,7 +1906,7 @@ pub fn build_reqwest_client(force_http11: bool) -> Result AppResult Value { } ``` -Create `/home/zyy/桌面/Codes/codex-proxy-rs/src/translation/schema.rs`: +Create `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/protocol/schema.rs`: ```rust #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -2478,7 +2477,7 @@ Expected: pass. - [x] **Step 5: Commit translation layer** ```bash -git add src/translation src/codex/types.rs src/lib.rs tests/routes_chat_test.rs +git add src/codex/protocol src/codex/types.rs src/lib.rs tests/routes_chat_test.rs git commit -m "feat: add openai to codex translation" ``` @@ -2550,7 +2549,7 @@ Expected: compile failure because app/http modules do not exist. Modify `/home/zyy/桌面/Codes/codex-proxy-rs/src/lib.rs`: ```rust -pub mod accounts; + pub mod app; pub mod auth; pub mod codex; @@ -2558,13 +2557,13 @@ pub mod config; pub mod cookies; pub mod crypto; pub mod error; -pub mod fingerprint; + pub mod http; pub mod logs; pub mod pagination; pub mod state; pub mod storage; -pub mod translation; + ``` Create `/home/zyy/桌面/Codes/codex-proxy-rs/src/app.rs`: @@ -2609,7 +2608,7 @@ Create `/home/zyy/桌面/Codes/codex-proxy-rs/src/http/v1.rs`: ```rust use axum::{http::StatusCode, response::IntoResponse, Json}; -use crate::translation::codex_to_openai::openai_error; +use crate::codex::protocol::codex_to_openai::openai_error; pub async fn responses() -> impl IntoResponse { ( @@ -2723,7 +2722,7 @@ Modify `/home/zyy/桌面/Codes/codex-proxy-rs/src/http/v1.rs`: use axum::{http::{HeaderMap, StatusCode}, response::IntoResponse, Json}; use serde::Deserialize; -use crate::translation::codex_to_openai::openai_error; +use crate::codex::protocol::codex_to_openai::openai_error; #[derive(Deserialize)] struct ResponsesBody { @@ -2774,8 +2773,8 @@ git commit -m "feat: enforce codex-only upstream routing" ### Task 12: Fingerprint Auto-Update **Files:** -- Modify: `/home/zyy/桌面/Codes/codex-proxy-rs/src/fingerprint/updater.rs` -- Modify: `/home/zyy/桌面/Codes/codex-proxy-rs/src/fingerprint/model.rs` +- Modify: `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/fingerprint/updater.rs` +- Modify: `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/fingerprint/model.rs` - Test: `/home/zyy/桌面/Codes/codex-proxy-rs/tests/codex_headers_test.rs` - [x] **Step 1: Add failing updater parse test** @@ -2783,7 +2782,7 @@ git commit -m "feat: enforce codex-only upstream routing" Append to `/home/zyy/桌面/Codes/codex-proxy-rs/tests/codex_headers_test.rs`: ```rust -use codex_proxy_rs::fingerprint::updater::parse_update_manifest; +use codex_proxy_rs::codex::fingerprint::updater::parse_update_manifest; #[test] fn update_manifest_updates_app_version_and_build_number() { @@ -2806,7 +2805,7 @@ Expected: compile failure because `parse_update_manifest` does not exist. - [x] **Step 3: Implement updater parser** -Create `/home/zyy/桌面/Codes/codex-proxy-rs/src/fingerprint/updater.rs`: +Create `/home/zyy/桌面/Codes/codex-proxy-rs/src/codex/fingerprint/updater.rs`: ```rust use serde::Deserialize; @@ -2852,7 +2851,7 @@ Expected: pass. - [x] **Step 5: Commit fingerprint updater parser** ```bash -git add src/fingerprint/updater.rs tests/codex_headers_test.rs +git add src/codex/fingerprint/updater.rs tests/codex_headers_test.rs git commit -m "feat: parse codex desktop fingerprint updates" ``` @@ -2862,7 +2861,6 @@ git commit -m "feat: parse codex desktop fingerprint updates" **Files:** - Modify: `/home/zyy/桌面/Codes/codex-proxy-rs/README.md` -- Create: `/home/zyy/桌面/Codes/codex-proxy-rs/.env.example` - Create: `/home/zyy/桌面/Codes/codex-proxy-rs/Dockerfile` - Create: `/home/zyy/桌面/Codes/codex-proxy-rs/docker-compose.yml` @@ -2901,18 +2899,12 @@ Rust rewrite of Codex Proxy focused only on ChatGPT/Codex accounts and the Codex - Codex account tokens are internal upstream credentials. ``` -- [x] **Step 2: Add environment example** +- [x] **Step 2: Document local YAML overrides** -Create `/home/zyy/桌面/Codes/codex-proxy-rs/.env.example`: +Document that runtime config loads from root `config.yaml`, then optional root +`local.yaml` or `local.yml`. Local overrides stay untracked and replace the old +environment-variable override path. -```dotenv -CPRS_HOST=127.0.0.1 -CPRS_PORT=8080 -CPRS_DATABASE_URL=sqlite://data/codex-proxy-rs.sqlite -CPRS_MASTER_KEY_FILE=data/master.key -CPRS_API_KEY_PEPPER_FILE=data/api-key-pepper.key -RUST_LOG=codex_proxy_rs=info,tower_http=info -``` - [x] **Step 3: Verify full test suite and lint** @@ -2944,7 +2936,7 @@ Expected response: - [x] **Step 5: Commit docs and runtime files** ```bash -git add README.md .env.example Dockerfile docker-compose.yml +git add README.md config.yaml Dockerfile docker-compose.yml git commit -m "docs: document codex-proxy-rs scope" ``` @@ -2979,7 +2971,7 @@ The Rust project has no legacy compatibility requirement, but the in-scope behav ### Keep and Rebuild Natively - [x] **HTTP auth boundary:** restore a real `src/http/auth.rs` module instead of scattered header checks. It validates `Bearer cpr_...` client API keys for `/v1/*`, validates admin session cookies for `/admin/*`, rejects cross-use in both directions, and returns the documented OpenAI/admin response families. -- [x] **Account SQLite repository:** implement `src/accounts/repository.rs` for encrypted access/refresh tokens, account metadata, quota JSON, usage counters, labels, status transitions, pagination, batch-safe mutations, metadata-only auth status reads, and account-wide delete for logout. Do not add TypeScript JSON migration or legacy import shims. +- [x] **Account SQLite repository:** implement `src/codex/accounts/repository.rs` for encrypted access/refresh tokens, account metadata, quota JSON, usage counters, labels, status transitions, pagination, batch-safe mutations, metadata-only auth status reads, and account-wide delete for logout. Do not add TypeScript JSON migration or legacy import shims. - [ ] **Admin account routes:** rebuild `/auth/status`, `/auth/accounts`, import/export, batch delete/status, health-check, per-account refresh, label, delete, reset usage, quota, Cookie CRUD, and quota warnings as `/admin/*` endpoints with admin envelopes and cursor pagination where a list is returned. Encrypted import, manual token add, CLI import, paginated sanitized list, native/Sub2API export, Sub2API import marking, label/status/delete, batch delete/status, Cookie CRUD, health/probe, refresh, quota, reset usage, admin auth status, and admin logout are complete. Quota warnings remain. - [x] **Login routes:** implement OAuth PKCE login-start/code-relay/callback, device-login/device-poll, CLI auth import, manual access-token import, and logout as account-login flows. CLI auth import, validated manual token import, device login/poll, OAuth PKCE login-start/code-relay/callback, admin-session-gated auth status, and logout are complete. Keep this separate from admin password login and client API key creation. - [x] **Client API key management:** implement admin CRUD for local client API keys: list/create/delete/batch-delete/label/status/export/import. Create/list/status/delete/batch-delete/label/export/import and `/v1/*` SQLite verification are complete. Remove provider/model binding semantics from the TypeScript version; Rust client API keys authorize only local `/v1/*` access. @@ -3038,7 +3030,7 @@ This pass re-read the TypeScript routes and core modules under `/home/zyy/Codes/ - Non-fingerprint dependencies are resolved from current stable crate releases at implementation time and `Cargo.lock` is committed. - `reqwest` and `rustls` remain explicitly pinned until a separate TLS fingerprint review approves newer versions. - Library modules use `thiserror` error enums. `anyhow` is limited to `src/main.rs`, tests, and setup utilities. -- Runtime config loads from `config/default.yaml`, optional local config, and `CPRS_*` environment variables. +- Runtime config loads from root `config.yaml`, then optional root `local.yaml` or `local.yml`. - In-memory upstream secrets use `secrecy` types, and persisted access tokens, refresh tokens, Cookies, master key material, and API key pepper material are encrypted or stored outside plaintext database fields. - `/v1/*` and `/admin/*` response envelopes and status codes match `docs/api.md` and `docs/status-codes.md`. - `/v1/*` never uses the admin/frontend `code/message/data/requestId` envelope; it keeps OpenAI-compatible response and error bodies. diff --git a/docs/superpowers/plans/2026-06-12-openai-gpt-codex-parity.md b/docs/superpowers/plans/2026-06-12-openai-gpt-codex-parity.md index 8661a009b..d103c5f7e 100644 --- a/docs/superpowers/plans/2026-06-12-openai-gpt-codex-parity.md +++ b/docs/superpowers/plans/2026-06-12-openai-gpt-codex-parity.md @@ -43,7 +43,7 @@ Immediate priority after this audit: ### Task 1: Chat Completions Translation **Files:** -- Modify: `src/translation/openai_to_codex.rs` +- Modify: `src/codex/protocol/openai_to_codex.rs` - Modify: `src/codex/types.rs` - Test: `tests/routes_chat_test.rs` @@ -57,7 +57,7 @@ Immediate priority after this audit: **Files:** - Modify: `src/app.rs` - Modify: `src/http/v1.rs` -- Modify: `src/translation/codex_to_openai.rs` +- Modify: `src/codex/protocol/codex_to_openai.rs` - Test: `tests/chat_completions_route_test.rs` - [x] Write failing route tests proving `/v1/chat/completions` sends translated Codex payloads and returns OpenAI chat JSON/SSE. @@ -98,8 +98,8 @@ Immediate priority after this audit: **Files:** - Modify: `src/http/v1.rs` -- Modify: `src/accounts/pool.rs` -- Modify: `src/accounts/repository.rs` +- Modify: `src/codex/accounts/pool.rs` +- Modify: `src/codex/accounts/repository.rs` - Test: `tests/v1_upstream_route_test.rs` - Test: `tests/account_pool_scheduling_test.rs` @@ -123,8 +123,8 @@ Immediate priority after this audit: ### Task 6: Model Catalog Refresh **Files:** -- Modify: `src/models/catalog.rs` -- Create: `src/models/repository.rs` +- Modify: `src/codex/models/catalog.rs` +- Create: `src/codex/models/repository.rs` - Modify: `src/http/v1.rs` - Test: `tests/model_catalog_test.rs` @@ -140,8 +140,8 @@ Immediate priority after this audit: **Files:** - Modify: `src/http/admin.rs` -- Modify: `src/accounts/repository.rs` -- Modify: `src/cookies/repository.rs` +- Modify: `src/codex/accounts/repository.rs` +- Modify: `src/codex/cookies/repository.rs` - Test: `tests/admin_accounts_route_test.rs` - Test: `tests/admin_api_keys_route_test.rs` diff --git a/docs/superpowers/plans/2026-06-13-utils-module-migration.md b/docs/superpowers/plans/2026-06-13-utils-module-migration.md new file mode 100644 index 000000000..80e57c43a --- /dev/null +++ b/docs/superpowers/plans/2026-06-13-utils-module-migration.md @@ -0,0 +1,141 @@ +# Utils Module Migration Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Move cross-domain helpers out of crate root into `src/utils/` without changing runtime behavior. + +**Architecture:** Keep the project as one Rust crate and tighten internal module boundaries. `utils` will hold reusable primitives that do not depend on HTTP, services, or domain modules; callers update to the new explicit paths. + +**Tech Stack:** Rust, Cargo, Axum integration tests, `thiserror`, `serde_json`, `base64`, `aes-gcm`, `secrecy`. + +--- + +### Task 1: Document The Utils Boundary + +**Files:** +- Modify: `docs/architecture-reorganization.md` + +- [x] **Step 1: Replace the target source layout with the accepted single-crate layout** + +The architecture document should show `app/`, `config/`, `utils/`, `http/`, `service/`, and the existing domain modules. + +- [x] **Step 2: Add boundary rules** + +The document should state that `utils` may contain `pagination`, `json`, and `crypto`, and that it must not depend on `http`, `service`, `accounts`, `auth`, or `codex`. + +### Task 2: Create `src/utils` + +**Files:** +- Create: `src/utils/mod.rs` +- Create: `src/utils/pagination.rs` +- Create: `src/utils/crypto.rs` +- Create: `src/utils/json.rs` +- Modify: `src/lib.rs` +- Delete: `src/pagination.rs` +- Delete: `src/crypto.rs` + +- [ ] **Step 1: Add the utils module root** + +```rust +pub mod crypto; +pub mod json; +pub mod pagination; +``` + +- [ ] **Step 2: Move pagination helpers** + +Move `Page`, `encode_cursor`, `decode_cursor`, and `clamp_limit` into `src/utils/pagination.rs` unchanged. + +- [ ] **Step 3: Move crypto helpers** + +Move `CryptoError`, `CryptoResult`, and `SecretBox` into `src/utils/crypto.rs` unchanged. + +- [ ] **Step 4: Add JSON string helpers** + +```rust +use serde_json::Value; + +pub fn first_string(value: &Value, paths: &[&[&str]]) -> Option { + paths.iter().find_map(|path| string_at(value, path)) +} + +pub fn string_at(value: &Value, path: &[&str]) -> Option { + let mut current = value; + for segment in path { + current = current.get(*segment)?; + } + current + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) +} +``` + +- [ ] **Step 5: Expose `utils` from the library** + +```rust +pub mod utils; +``` + +Remove the old top-level `pub mod crypto;` and `pub mod pagination;` declarations after all imports are updated. + +### Task 3: Update Callers + +**Files:** +- Modify: `src/codex/accounts/repository.rs` +- Modify: `src/auth/api_key_repository.rs` +- Modify: `src/codex/cookies/repository.rs` +- Modify: `src/http/admin/*.rs` +- Modify: `src/service/*.rs` +- Modify: `src/main.rs` +- Modify: `src/state.rs` +- Modify: `tests/*.rs` +- Modify: `tests/common/*.rs` + +- [ ] **Step 1: Update source imports** + +Use `crate::utils::crypto::{CryptoError, SecretBox}` and +`crate::utils::pagination::{clamp_limit, decode_cursor, encode_cursor, Page}`. + +- [ ] **Step 2: Update integration test imports** + +Use `codex_proxy_rs::utils::crypto::SecretBox` and +`codex_proxy_rs::utils::pagination::Page`. + +- [ ] **Step 3: Remove duplicate JSON helper functions** + +Import `crate::utils::json::first_string` in `src/service/api_key_service.rs` and +`src/http/admin/accounts.rs`, then delete the local `first_string` and `string_at` +definitions. + +### Task 4: Verify + +**Files:** +- No source edits after this task unless verification fails. + +- [ ] **Step 1: Format** + +Run: + +```bash +cargo fmt +``` + +- [ ] **Step 2: Run focused tests** + +Run: + +```bash +cargo test --test crypto_test --test account_repository_test --test admin_accounts_import_export_test --test admin_api_keys_route_test +``` + +- [ ] **Step 3: Run full gate** + +Run: + +```bash +cargo fmt --check && cargo test && cargo clippy --all-targets --all-features --locked -- -D warnings +``` + +Expected: all commands exit successfully. diff --git a/docs/superpowers/specs/2026-06-12-openai-gpt-codex-parity-design.md b/docs/superpowers/specs/2026-06-12-openai-gpt-codex-parity-design.md index 2d3daf625..de9abf24e 100644 --- a/docs/superpowers/specs/2026-06-12-openai-gpt-codex-parity-design.md +++ b/docs/superpowers/specs/2026-06-12-openai-gpt-codex-parity-design.md @@ -22,14 +22,14 @@ Current progress: Tasks 1-4 are implemented for the OpenAI GPT/Codex scope. Comm ## Components -- `src/translation/openai_to_codex.rs`: OpenAI Chat request schema and Chat-to-Codex translation, including messages, tools/functions, multimodal input, response formats, reasoning effort, and service tier. -- `src/translation/codex_to_openai.rs`: Codex SSE to OpenAI Chat JSON/SSE conversion, including text, reasoning, tool calls, usage, and errors. +- `src/codex/protocol/openai_to_codex.rs`: OpenAI Chat request schema and Chat-to-Codex translation, including messages, tools/functions, multimodal input, response formats, reasoning effort, and service tier. +- `src/codex/protocol/codex_to_openai.rs`: Codex SSE to OpenAI Chat JSON/SSE conversion, including text, reasoning, tool calls, usage, and errors. - `src/codex/types.rs`: Expanded Codex request/response types for Responses and WebSocket. - `src/codex/websocket.rs`: WebSocket transport selection and `response.create` conversion into SSE-compatible events for upstream-history requests such as `previous_response_id`. - `src/codex/client.rs`: HTTP SSE and WebSocket Codex transport, context headers, cookies, and rate-limit header capture. - `src/http/v1.rs`: Separate `/v1/responses` and `/v1/chat/completions` handlers sharing upstream lifecycle helpers. -- `src/accounts/pool.rs` and repositories: fallback selection, quota/rate-limit updates, durable usage release, and account state mutation hooks. -- `src/models/catalog.rs`: static catalog plus persisted backend snapshots, plan allowlist derivation, and admin refresh trigger. +- `src/codex/accounts/pool.rs` and repositories: fallback selection, quota/rate-limit updates, durable usage release, and account state mutation hooks. +- `src/codex/models/catalog.rs`: static catalog plus persisted backend snapshots, plan allowlist derivation, and admin refresh trigger. - `src/http/admin.rs`: scoped account/API-key/account-health/quota/cookie/admin model operations needed to operate Codex-backed GPT compatibility, including Sub2API account import normalization without proxy-pool support. ## Data Flow diff --git a/src/app/bootstrap.rs b/src/app/bootstrap.rs new file mode 100644 index 000000000..247fd94f9 --- /dev/null +++ b/src/app/bootstrap.rs @@ -0,0 +1,49 @@ +use thiserror::Error; + +use crate::{ + auth::{api_key::ApiKeyHasher, error::AuthError}, + codex::accounts::repository::AccountRepositoryError, + codex::oauth::OpenAiOAuthRefresher, + config::AppConfig, + storage::db::connect_sqlite, + utils::crypto::{CryptoError, SecretBox}, +}; + +use super::state::AppState; + +#[derive(Debug, Error)] +pub enum BootstrapError { + #[error("secret storage error: {0}")] + Crypto(#[from] CryptoError), + #[error("api key hasher error: {0}")] + Auth(#[from] AuthError), + #[error("http client error: {0}")] + HttpClient(#[from] reqwest::Error), + #[error("database error: {0}")] + Database(#[from] sqlx::Error), + #[error("account repository error: {0}")] + AccountRepository(#[from] AccountRepositoryError), +} + +pub type BootstrapResult = Result; + +pub async fn build_state(config: AppConfig) -> BootstrapResult<(AppState, usize)> { + 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 oauth_client = 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_api_key_hasher_and_oauth_client( + config, + pool, + secret_box, + api_key_hasher, + oauth_client, + ); + let restored_accounts = state.reload_account_pool_from_repository().await?; + Ok((state, restored_accounts)) +} diff --git a/src/app/mod.rs b/src/app/mod.rs new file mode 100644 index 000000000..554fe9924 --- /dev/null +++ b/src/app/mod.rs @@ -0,0 +1,5 @@ +pub mod bootstrap; +pub mod router; +pub mod state; + +pub use router::build_router; diff --git a/src/app/router.rs b/src/app/router.rs new file mode 100644 index 000000000..3140fd938 --- /dev/null +++ b/src/app/router.rs @@ -0,0 +1,16 @@ +use axum::{middleware::from_fn, routing::get, Router}; +use tower_http::trace::TraceLayer; + +use crate::http::{admin, health::health, middleware::attach_request_id, v1}; + +use super::state::AppState; + +pub fn build_router(state: AppState) -> Router { + Router::new() + .route("/health", get(health)) + .merge(v1::router()) + .merge(admin::router()) + .with_state(state) + .layer(TraceLayer::new_for_http()) + .layer(from_fn(attach_request_id)) +} diff --git a/src/app/state.rs b/src/app/state.rs new file mode 100644 index 000000000..f640aedce --- /dev/null +++ b/src/app/state.rs @@ -0,0 +1,356 @@ +use std::sync::Arc; + +use sqlx::SqlitePool; +use tokio::sync::Mutex; + +use crate::auth::{api_key::ApiKeyHasher, api_key_repository::ClientApiKeyRepository}; +use crate::codex::accounts::{ + pool::{AccountPool, AccountPoolOptions, RotationStrategy}, + repository::{AccountRepository, AccountRepositoryResult, AccountUsageRepository}, + service::AccountService, +}; +use crate::codex::cookies::repository::CookieRepository; +use crate::codex::models::repository::ModelSnapshotRepository; +use crate::codex::models::service::ModelService; +use crate::codex::oauth::{OAuthClient, PkceSessionStore, TokenRefresher}; +use crate::codex::upstream::CodexUpstreamService; +use crate::config::AppConfig; +use crate::logs::repository::EventLogRepository; +use crate::service::{ + admin_auth::AdminAuthService, api_key::ApiKeyService, chat::ChatService, log::LogService, + responses::ResponsesService, usage::UsageService, +}; +use crate::utils::crypto::SecretBox; + +#[derive(Clone)] +pub struct AppState { + pub services: Arc, +} + +pub struct AppServices { + pub config: AppConfig, + pub admin_auth: AdminAuthService, + pub accounts: AccountService, + pub api_keys: ApiKeyService, + pub logs: LogService, + pub usage: UsageService, + pub chat: ChatService, + pub responses: ResponsesService, + pub models: ModelService, +} + +#[derive(Default)] +struct AppStateDependencies { + pool: Option, + secret_box: Option, + api_key_hasher: Option, + token_refresher: Option>, + oauth_client: Option>, +} + +impl AppState { + pub fn new(config: AppConfig) -> Self { + Self::from_dependencies(config, AppStateDependencies::default()) + } + + pub fn with_pool(config: AppConfig, pool: SqlitePool) -> Self { + Self::from_dependencies( + config, + AppStateDependencies { + pool: Some(pool), + ..AppStateDependencies::default() + }, + ) + } + + pub fn with_pool_and_secret_box( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + ) -> Self { + Self::from_dependencies( + config, + AppStateDependencies { + pool: Some(pool), + secret_box: Some(secret_box), + ..AppStateDependencies::default() + }, + ) + } + + pub fn with_pool_secret_and_api_key_hasher( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + api_key_hasher: ApiKeyHasher, + ) -> Self { + Self::from_dependencies( + config, + AppStateDependencies { + pool: Some(pool), + secret_box: Some(secret_box), + api_key_hasher: Some(api_key_hasher), + ..AppStateDependencies::default() + }, + ) + } + + 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 token_refresher: Arc = Arc::new(token_refresher); + Self::from_dependencies( + config, + AppStateDependencies { + pool: Some(pool), + secret_box: Some(secret_box), + api_key_hasher: Some(api_key_hasher), + token_refresher: Some(token_refresher), + ..AppStateDependencies::default() + }, + ) + } + + pub fn with_pool_secret_api_key_hasher_and_oauth_client( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + api_key_hasher: ApiKeyHasher, + oauth_client: C, + ) -> Self + where + C: OAuthClient + TokenRefresher, + { + let oauth_client = Arc::new(oauth_client); + let token_refresher: Arc = oauth_client.clone(); + let oauth_client: Arc = oauth_client; + Self::from_dependencies( + config, + AppStateDependencies { + pool: Some(pool), + secret_box: Some(secret_box), + api_key_hasher: Some(api_key_hasher), + token_refresher: Some(token_refresher), + oauth_client: Some(oauth_client), + }, + ) + } + + pub fn config(&self) -> &AppConfig { + &self.services.config + } + + pub async fn reload_account_pool_from_repository(&self) -> AccountRepositoryResult { + self.services + .accounts + .reload_runtime_accounts_from_repository() + .await + } + + fn from_dependencies(config: AppConfig, dependencies: AppStateDependencies) -> Self { + let AppStateDependencies { + pool, + secret_box, + api_key_hasher, + token_refresher, + oauth_client, + } = dependencies; + let pool_ref = pool.as_ref(); + let secret_box_ref = secret_box.as_ref(); + let account_pool = account_pool_from_config(&config); + let oauth_sessions = Arc::new(Mutex::new(PkceSessionStore::default())); + let api_keys = api_key_service(pool_ref, api_key_hasher.as_ref()); + let logs = log_service(pool_ref); + let usage = usage_service(pool_ref); + let accounts = account_service( + &config, + account_repository(pool_ref, secret_box_ref), + pool_ref, + secret_box_ref, + token_refresher.clone(), + account_pool.clone(), + ); + let admin_auth = admin_auth_service( + &config, + pool_ref, + secret_box_ref, + account_pool.clone(), + oauth_client, + oauth_sessions, + accounts.clone(), + ); + let V1Services { + chat, + responses, + models, + } = v1_services( + &config, + pool_ref, + secret_box_ref, + token_refresher, + account_pool, + ); + Self { + services: Arc::new(AppServices { + config, + admin_auth, + accounts, + api_keys, + logs, + usage, + chat, + responses, + models, + }), + } + } +} + +fn account_pool_from_config(config: &AppConfig) -> Arc> { + Arc::new(Mutex::new(AccountPool::with_options(AccountPoolOptions { + max_concurrent_per_account: config.auth.max_concurrent_per_account, + stale_slot_ttl: chrono::Duration::minutes(5), + rotation_strategy: match config.auth.rotation_strategy.as_str() { + "round_robin" => RotationStrategy::RoundRobin, + "sticky" => RotationStrategy::Sticky, + _ => RotationStrategy::LeastUsed, + }, + skip_quota_limited: config.quota.skip_exhausted, + tier_priority: config.auth.tier_priority.clone(), + model_plan_allowlist: std::collections::BTreeMap::new(), + }))) +} + +struct V1Services { + chat: ChatService, + responses: ResponsesService, + models: ModelService, +} + +fn v1_services( + config: &AppConfig, + pool: Option<&SqlitePool>, + secret_box: Option<&SecretBox>, + token_refresher: Option>, + account_pool: Arc>, +) -> V1Services { + let upstream = codex_upstream_service( + config, + pool, + secret_box, + token_refresher, + account_pool.clone(), + ); + let models = model_service(config, pool, secret_box, account_pool); + let model_config = config.model.clone(); + V1Services { + chat: ChatService::new(model_config.clone(), models.clone(), upstream.clone()), + responses: ResponsesService::new(model_config, models.clone(), upstream), + models, + } +} + +fn model_service( + config: &AppConfig, + pool: Option<&SqlitePool>, + secret_box: Option<&SecretBox>, + account_pool: Arc>, +) -> ModelService { + ModelService::new( + config.clone(), + pool.cloned().map(ModelSnapshotRepository::new), + account_repository(pool, secret_box), + account_pool, + ) +} + +fn admin_auth_service( + config: &AppConfig, + pool: Option<&SqlitePool>, + secret_box: Option<&SecretBox>, + account_pool: Arc>, + oauth_client: Option>, + oauth_sessions: Arc>, + accounts: AccountService, +) -> AdminAuthService { + AdminAuthService::new( + config.clone(), + pool.cloned(), + account_repository(pool, secret_box), + account_pool, + oauth_client, + oauth_sessions, + accounts, + ) +} + +fn account_service( + config: &AppConfig, + repository: Option, + pool: Option<&SqlitePool>, + secret_box: Option<&SecretBox>, + token_refresher: Option>, + account_pool: Arc>, +) -> AccountService { + AccountService::new( + config.clone(), + repository, + pool.cloned().map(AccountUsageRepository::new), + cookie_repository(pool, secret_box), + token_refresher, + account_pool, + ) +} + +fn codex_upstream_service( + config: &AppConfig, + pool: Option<&SqlitePool>, + secret_box: Option<&SecretBox>, + token_refresher: Option>, + account_pool: Arc>, +) -> CodexUpstreamService { + CodexUpstreamService::new( + config.clone(), + account_pool, + account_repository(pool, secret_box), + cookie_repository(pool, secret_box), + pool.cloned().map(EventLogRepository::new), + token_refresher, + ) +} + +fn account_repository( + pool: Option<&SqlitePool>, + secret_box: Option<&SecretBox>, +) -> Option { + Some(AccountRepository::new(pool?.clone(), secret_box?.clone())) +} + +fn cookie_repository( + pool: Option<&SqlitePool>, + secret_box: Option<&SecretBox>, +) -> Option { + Some(CookieRepository::new(pool?.clone(), secret_box?.clone())) +} + +fn api_key_service(pool: Option<&SqlitePool>, hasher: Option<&ApiKeyHasher>) -> ApiKeyService { + ApiKeyService::new( + pool.cloned().map(ClientApiKeyRepository::new), + hasher.cloned(), + ) +} + +fn log_service(pool: Option<&SqlitePool>) -> LogService { + LogService::new(pool.cloned().map(EventLogRepository::new)) +} + +fn usage_service(pool: Option<&SqlitePool>) -> UsageService { + UsageService::new(pool.cloned().map(AccountUsageRepository::new)) +} diff --git a/src/auth/api_key_repository.rs b/src/auth/api_key_repository.rs index 91b10be84..f43b07665 100644 --- a/src/auth/api_key_repository.rs +++ b/src/auth/api_key_repository.rs @@ -5,7 +5,7 @@ use uuid::Uuid; use crate::{ auth::api_key::{ApiKeyHasher, GeneratedClientApiKey}, - pagination::{decode_cursor, encode_cursor, Page}, + utils::pagination::{decode_cursor, encode_cursor, Page}, }; #[derive(Debug, Error)] diff --git a/src/auth/mod.rs b/src/auth/mod.rs index 3ea33c700..be4c5cd44 100644 --- a/src/auth/mod.rs +++ b/src/auth/mod.rs @@ -1,8 +1,4 @@ pub mod admin_session; pub mod api_key; pub mod api_key_repository; -pub mod cli_import; pub mod error; -pub mod oauth; -pub mod refresh; -pub mod token; diff --git a/src/accounts/lifecycle.rs b/src/codex/accounts/lifecycle.rs similarity index 100% rename from src/accounts/lifecycle.rs rename to src/codex/accounts/lifecycle.rs diff --git a/src/accounts/mod.rs b/src/codex/accounts/mod.rs similarity index 80% rename from src/accounts/mod.rs rename to src/codex/accounts/mod.rs index 5e3860392..b80507705 100644 --- a/src/accounts/mod.rs +++ b/src/codex/accounts/mod.rs @@ -2,3 +2,4 @@ pub mod lifecycle; pub mod model; pub mod pool; pub mod repository; +pub mod service; diff --git a/src/accounts/model.rs b/src/codex/accounts/model.rs similarity index 100% rename from src/accounts/model.rs rename to src/codex/accounts/model.rs diff --git a/src/accounts/pool.rs b/src/codex/accounts/pool.rs similarity index 99% rename from src/accounts/pool.rs rename to src/codex/accounts/pool.rs index 53fbf7421..6c5f3702c 100644 --- a/src/accounts/pool.rs +++ b/src/codex/accounts/pool.rs @@ -2,7 +2,7 @@ use std::collections::{BTreeMap, VecDeque}; use chrono::{DateTime, Duration, Utc}; -use crate::accounts::model::{Account, AccountStatus}; +use crate::codex::accounts::model::{Account, AccountStatus}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RotationStrategy { diff --git a/src/accounts/repository.rs b/src/codex/accounts/repository.rs similarity index 99% rename from src/accounts/repository.rs rename to src/codex/accounts/repository.rs index 72b49847b..ed01a09db 100644 --- a/src/accounts/repository.rs +++ b/src/codex/accounts/repository.rs @@ -4,9 +4,9 @@ use sqlx::{Row, SqlitePool}; use thiserror::Error; use crate::{ - accounts::model::{Account, AccountStatus}, - crypto::{CryptoError, SecretBox}, - pagination::{decode_cursor, encode_cursor, Page}, + codex::accounts::model::{Account, AccountStatus}, + utils::crypto::{CryptoError, SecretBox}, + utils::pagination::{decode_cursor, encode_cursor, Page}, }; #[derive(Debug, Error)] diff --git a/src/codex/accounts/service/cookies.rs b/src/codex/accounts/service/cookies.rs new file mode 100644 index 000000000..e9432ccee --- /dev/null +++ b/src/codex/accounts/service/cookies.rs @@ -0,0 +1,51 @@ +use reqwest::Url; + +use super::{AccountService, AccountServiceError}; + +impl AccountService { + pub async fn get_cookies( + &self, + account_id: &str, + ) -> Result, AccountServiceError> { + self.ensure_account_exists(account_id).await?; + self.cookie_repository()? + .cookie_header(account_id, "chatgpt.com") + .await + .map_err(|_| AccountServiceError::LoadCookies) + } + + pub async fn set_cookies( + &self, + account_id: &str, + cookie_header: &str, + ) -> Result, AccountServiceError> { + self.ensure_account_exists(account_id).await?; + let cookie_repo = self.cookie_repository()?; + match cookie_repo + .set_cookie_header(account_id, cookie_header) + .await + { + Ok(0) => Err(AccountServiceError::NoValidCookies), + Ok(_) => cookie_repo + .cookie_header(account_id, "chatgpt.com") + .await + .map_err(|_| AccountServiceError::LoadCookies), + Err(_) => Err(AccountServiceError::StoreCookies), + } + } + + pub async fn delete_cookies(&self, account_id: &str) -> Result<(), AccountServiceError> { + self.ensure_account_exists(account_id).await?; + self.cookie_repository()? + .delete_account_cookies(account_id) + .await + .map(|_| ()) + .map_err(|_| AccountServiceError::DeleteCookies) + } +} + +pub(super) fn request_domain(base_url: &str) -> Option { + Url::parse(base_url) + .ok() + .and_then(|url| url.host_str().map(ToString::to_string)) +} diff --git a/src/codex/accounts/service/health.rs b/src/codex/accounts/service/health.rs new file mode 100644 index 000000000..05080a3a8 --- /dev/null +++ b/src/codex/accounts/service/health.rs @@ -0,0 +1,308 @@ +use std::time::Duration as StdDuration; + +use chrono::Utc; +use futures::{stream, StreamExt}; +use reqwest::StatusCode; +use secrecy::ExposeSecret; +use serde_json::Value; +use tokio::time::sleep; + +use crate::{ + codex::accounts::{ + model::AccountStatus, + repository::{AccountRepository, StoredAccount}, + }, + codex::fingerprint::model::Fingerprint, + codex::transport::client::{ + build_reqwest_client, CodexBackendClient, CodexClientError, CodexRequestContext, + }, +}; + +use super::{ + cookies::request_domain, quota::quota_from_usage, AccountProbeOutcome, AccountProbeResult, + AccountService, HealthCheckError, +}; + +impl AccountService { + pub async fn health_check_accounts( + &self, + ids: Option>, + concurrency: usize, + stagger_ms: u64, + request_id: &str, + ) -> Result, HealthCheckError> { + let repo = self + .repository + .as_ref() + .ok_or(HealthCheckError::RepositoryUnavailable)?; + let accounts = repo + .list_all() + .await + .map(|accounts| filter_health_check_accounts(accounts, ids.as_deref())) + .map_err(|_| HealthCheckError::List)?; + let results = stream::iter(accounts.into_iter().enumerate()) + .map(|(index, account)| { + let service = self.clone(); + let repo = repo.clone(); + let request_id = request_id.to_string(); + 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; + } + service + .probe_account_with_codex_backend(&repo, account, &request_id) + .await + } + }) + .buffer_unordered(concurrency) + .collect::>() + .await; + Ok(results) + } + + async fn probe_account_with_codex_backend( + &self, + repo: &AccountRepository, + account: StoredAccount, + request_id: &str, + ) -> AccountProbeResult { + if account.status == AccountStatus::Disabled { + return skipped_probe_result(&account, "manually disabled"); + } + + let started_at = std::time::Instant::now(); + let previous_status = account.status; + match fetch_account_usage(self, &account, request_id).await { + Ok(raw) => { + let quota = quota_from_usage(&raw); + let _ = repo + .update_quota_json(&account.id, "a.to_string()) + .await; + if account.status != AccountStatus::Active { + let _ = repo.set_status(&account.id, AccountStatus::Active).await; + self.account_pool + .lock() + .await + .set_status(&account.id, AccountStatus::Active); + } + AccountProbeResult { + id: account.id, + email: account.email, + previous_status, + outcome: AccountProbeOutcome::Alive, + status: Some(AccountStatus::Active), + error: None, + duration_ms: Some(started_at.elapsed().as_millis()), + } + } + Err(error) => { + let status = apply_codex_account_error(self, repo, &account, &error).await; + AccountProbeResult { + id: account.id, + email: account.email, + previous_status, + outcome: AccountProbeOutcome::Dead, + status, + error: Some(public_codex_error(&error)), + duration_ms: Some(started_at.elapsed().as_millis()), + } + } + } + } +} + +pub(super) async fn fetch_account_usage( + service: &AccountService, + account: &StoredAccount, + request_id: &str, +) -> Result { + let cookie_header = account_cookie_header(service, &account.id).await; + let client = CodexBackendClient::new( + build_reqwest_client(service.config.tls.force_http11)?, + service.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(service: &AccountService, account_id: &str) -> Option { + let domain = request_domain(&service.config.api.base_url)?; + service + .cookie_repository + .as_ref()? + .cookie_header(account_id, &domain) + .await + .ok() + .flatten() +} + +pub(super) async fn apply_codex_account_error( + service: &AccountService, + repo: &AccountRepository, + account: &StoredAccount, + error: &CodexClientError, +) -> Option { + match classify_codex_account_error(error) { + Some(CodexAccountErrorAction::SetStatus(status)) => { + let _ = repo.set_status(&account.id, status).await; + service + .account_pool + .lock() + .await + .set_status(&account.id, status); + Some(status) + } + Some(CodexAccountErrorAction::RateLimited { + retry_after_seconds, + }) => { + let cooldown_until = Utc::now() + chrono::Duration::seconds(retry_after_seconds as i64); + service + .account_pool + .lock() + .await + .mark_quota_limited_until(&account.id, cooldown_until); + None + } + Some(CodexAccountErrorAction::CloudflareChallenge { cooldown_seconds }) => { + let cooldown_until = Utc::now() + chrono::Duration::seconds(cooldown_seconds as i64); + service + .account_pool + .lock() + .await + .set_cloudflare_cooldown_until(&account.id, cooldown_until); + None + } + None => None, + } +} + +pub(super) fn skipped_probe_result(account: &StoredAccount, error: &str) -> AccountProbeResult { + AccountProbeResult { + id: account.id.clone(), + email: account.email.clone(), + previous_status: account.status, + outcome: AccountProbeOutcome::Skipped, + status: Some(account.status), + error: Some(error.to_string()), + duration_ms: None, + } +} + +pub(super) fn filter_health_check_accounts( + accounts: Vec, + ids: Option<&[String]>, +) -> Vec { + let Some(ids) = ids else { + return accounts; + }; + accounts + .into_iter() + .filter(|account| ids.iter().any(|id| id == &account.id)) + .collect() +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum CodexAccountErrorAction { + SetStatus(AccountStatus), + RateLimited { retry_after_seconds: u64 }, + CloudflareChallenge { cooldown_seconds: u64 }, +} + +pub(super) fn classify_codex_account_error( + error: &CodexClientError, +) -> Option { + 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 +} + +pub(super) 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") +} diff --git a/src/codex/accounts/service/import.rs b/src/codex/accounts/service/import.rs new file mode 100644 index 000000000..98cc82aeb --- /dev/null +++ b/src/codex/accounts/service/import.rs @@ -0,0 +1,319 @@ +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use chrono::{DateTime, Utc}; +use secrecy::SecretString; +use serde_json::{Map, Value}; +use uuid::Uuid; + +use crate::{ + codex::accounts::{ + model::{Account, AccountStatus}, + repository::{AccountClaimsUpdate, AccountRepository, NewAccount, StoredAccount}, + }, + codex::oauth::RefreshFailure, +}; + +use super::{ + runtime_pool::pool_account_from_stored, AccountImportCounts, AccountImportEntry, + AccountService, StoreImportAccountError, ValidatedAccountImportError, +}; + +enum StoredImportAccount { + Imported, + Skipped, +} + +impl AccountService { + pub async fn import_validated( + &self, + token: Option, + refresh_token: Option, + ) -> Result { + let repo = self + .repository + .as_ref() + .ok_or(ValidatedAccountImportError::RepositoryUnavailable)?; + let (access_token, refresh_token_update, new_account_refresh_token) = match ( + empty_to_none(token.map(normalize_bearer_token)), + empty_to_none(refresh_token), + ) { + (Some(token), refresh_token) => (token, refresh_token.clone(), refresh_token), + (None, Some(refresh_token)) => { + let Some(refresher) = self.token_refresher.as_ref() else { + return Err(ValidatedAccountImportError::TokenRefresherUnavailable); + }; + let tokens = match refresher.refresh(&refresh_token).await { + Ok(tokens) => tokens, + Err(RefreshFailure::Transport) => { + return Err(ValidatedAccountImportError::RefreshTransport); + } + Err(_) => return Err(ValidatedAccountImportError::RefreshRejected), + }; + let access_token = normalize_bearer_token(tokens.access_token); + let rotated_refresh_token = empty_to_none(tokens.refresh_token); + let new_account_refresh_token = + rotated_refresh_token.clone().or(Some(refresh_token)); + ( + access_token, + rotated_refresh_token, + new_account_refresh_token, + ) + } + (None, None) => return Err(ValidatedAccountImportError::TokenRequired), + }; + + let claims = manual_account_claims(&access_token, Utc::now()) + .map_err(ValidatedAccountImportError::InvalidToken)?; + let existing = repo + .find_by_chatgpt_identity(&claims.account_id, claims.user_id.as_deref()) + .await + .map_err(|_| ValidatedAccountImportError::Inspect)?; + + let account_id = if let Some(existing) = existing { + let updated = repo + .update_from_claims( + &existing.id, + AccountClaimsUpdate { + email: claims.email.clone(), + account_id: Some(claims.account_id.clone()), + user_id: claims.user_id.clone(), + plan_type: claims.plan_type.clone(), + access_token: SecretString::new(access_token.into()), + refresh_token: refresh_token_update + .map(|token| SecretString::new(token.into())), + access_token_expires_at: Some(claims.expires_at), + status: AccountStatus::Active, + }, + ) + .await + .map_err(|_| ValidatedAccountImportError::Update)?; + if !updated { + return Err(ValidatedAccountImportError::NotFound); + } + existing.id + } else { + let id = normalized_account_id(None); + let account = NewAccount { + id: id.clone(), + email: claims.email.clone(), + account_id: Some(claims.account_id.clone()), + user_id: claims.user_id.clone(), + label: None, + plan_type: claims.plan_type.clone(), + access_token: SecretString::new(access_token.into()), + refresh_token: new_account_refresh_token + .map(|token| SecretString::new(token.into())), + access_token_expires_at: Some(claims.expires_at), + status: AccountStatus::Active, + }; + repo.insert(account) + .await + .map_err(|_| ValidatedAccountImportError::Insert)?; + id + }; + + let stored = repo + .get(&account_id) + .await + .map_err(|_| ValidatedAccountImportError::Load)? + .ok_or(ValidatedAccountImportError::NotFound)?; + self.account_pool + .lock() + .await + .insert(pool_account_from_stored(stored.clone())); + Ok(stored) + } + + pub async fn import_entries( + &self, + entries: Vec, + ) -> Result { + let repo = self + .repository + .as_ref() + .ok_or(StoreImportAccountError::RepositoryUnavailable)?; + let mut imported = 0u32; + let mut skipped = 0u32; + for entry in entries { + match self.store_import_entry(repo, entry).await? { + StoredImportAccount::Imported => imported += 1, + StoredImportAccount::Skipped => skipped += 1, + } + } + Ok(AccountImportCounts { imported, skipped }) + } + + async fn store_import_entry( + &self, + repo: &AccountRepository, + entry: AccountImportEntry, + ) -> Result { + let access_token = entry.token.as_deref().unwrap_or_default().trim(); + if access_token.is_empty() { + return Ok(StoredImportAccount::Skipped); + } + if entry + .label + .as_ref() + .is_some_and(|label| label.chars().count() > 64) + { + return Err(StoreImportAccountError::Invalid( + "Account label must be 64 characters or fewer".to_string(), + )); + } + + let id = normalized_account_id(entry.id); + match repo.exists(&id).await { + Ok(true) => return Ok(StoredImportAccount::Skipped), + Ok(false) => {} + Err(_) => return Err(StoreImportAccountError::Inspect), + } + + let status = parse_import_status(entry.status.as_deref()) + .map_err(StoreImportAccountError::Invalid)?; + let email = empty_to_none(entry.email); + let account_id = empty_to_none(entry.account_id); + let user_id = empty_to_none(entry.user_id); + let label = empty_to_none(entry.label); + let plan_type = empty_to_none(entry.plan_type); + let refresh_token = empty_to_none(entry.refresh_token); + let access_token = access_token.to_string(); + let now = Utc::now().to_rfc3339(); + let pool_account = Account { + id: id.clone(), + email: email.clone(), + account_id: account_id.clone(), + user_id: user_id.clone(), + label: label.clone(), + plan_type: plan_type.clone(), + access_token: access_token.clone(), + refresh_token: refresh_token.clone(), + access_token_expires_at: None, + status, + quota_limit_reached: false, + quota_cooldown_until: None, + cloudflare_cooldown_until: None, + added_at: now, + last_used_at: None, + }; + let account = NewAccount { + id, + email, + account_id, + user_id, + label, + plan_type, + access_token: SecretString::new(access_token.into()), + refresh_token: refresh_token.map(|token| SecretString::new(token.into())), + access_token_expires_at: None, + status, + }; + repo.insert(account) + .await + .map_err(|_| StoreImportAccountError::Insert)?; + self.account_pool.lock().await.insert(pool_account); + + Ok(StoredImportAccount::Imported) + } +} + +#[derive(Debug, Clone)] +pub(super) struct ManualAccountClaims { + pub account_id: String, + pub user_id: Option, + pub email: Option, + pub plan_type: Option, + pub expires_at: DateTime, +} + +pub(super) fn manual_account_claims( + token: &str, + now: DateTime, +) -> Result { + let payload = decode_jwt_payload(token).ok_or("Invalid JWT format")?; + let exp = payload + .get("exp") + .and_then(Value::as_i64) + .ok_or("Token is expired")?; + if now.timestamp() >= exp { + return Err("Token is expired"); + } + let expires_at = DateTime::::from_timestamp(exp, 0).ok_or("Invalid JWT exp claim")?; + let auth = payload + .get("https://api.openai.com/auth") + .and_then(Value::as_object) + .ok_or("Token missing chatgpt_account_id claim")?; + let account_id = + string_claim(auth, "chatgpt_account_id").ok_or("Token missing chatgpt_account_id claim")?; + let profile = payload + .get("https://api.openai.com/profile") + .and_then(Value::as_object); + let user_id = string_claim(auth, "chatgpt_user_id") + .or_else(|| profile.and_then(|profile| string_claim(profile, "chatgpt_user_id"))); + let plan_type = string_claim(auth, "chatgpt_plan_type") + .or_else(|| profile.and_then(|profile| string_claim(profile, "chatgpt_plan_type"))); + let email = profile.and_then(|profile| string_claim(profile, "email")); + + Ok(ManualAccountClaims { + account_id, + user_id, + email, + plan_type, + expires_at, + }) +} + +pub(super) fn normalize_bearer_token(value: String) -> String { + value + .trim() + .strip_prefix("Bearer ") + .or_else(|| value.trim().strip_prefix("bearer ")) + .unwrap_or(value.trim()) + .trim() + .to_string() +} + +pub(super) fn empty_to_none(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +pub(super) fn normalized_account_id(id: Option) -> String { + id.and_then(|id| empty_to_none(Some(id))) + .unwrap_or_else(|| format!("acct_{}", Uuid::new_v4().simple())) +} + +pub(super) fn parse_import_status(status: Option<&str>) -> Result { + let normalized = status.unwrap_or("active").trim().to_ascii_lowercase(); + match normalized.as_str() { + "active" => Ok(AccountStatus::Active), + "expired" => Ok(AccountStatus::Expired), + "quota_exhausted" => Ok(AccountStatus::QuotaExhausted), + "refreshing" => Ok(AccountStatus::Refreshing), + "disabled" => Ok(AccountStatus::Disabled), + "banned" => Ok(AccountStatus::Banned), + other => Err(format!("Unsupported account status: {other}")), + } +} + +fn decode_jwt_payload(token: &str) -> Option> { + let mut parts = token.split('.'); + let _header = parts.next()?; + let payload = parts.next()?; + if payload.is_empty() { + return None; + } + let bytes = URL_SAFE_NO_PAD.decode(payload).ok()?; + serde_json::from_slice::(&bytes) + .ok()? + .as_object() + .cloned() +} + +fn string_claim(map: &Map, key: &str) -> Option { + map.get(key) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) +} diff --git a/src/codex/accounts/service/mod.rs b/src/codex/accounts/service/mod.rs new file mode 100644 index 000000000..65fce85f0 --- /dev/null +++ b/src/codex/accounts/service/mod.rs @@ -0,0 +1,262 @@ +pub mod cookies; +pub mod health; +pub mod import; +pub mod mutation; +pub mod quota; +pub mod refresh; +pub mod runtime_pool; + +use std::sync::Arc; + +use serde_json::Value; +use tokio::sync::Mutex; + +use crate::{ + codex::accounts::{ + model::AccountStatus, + pool::AccountPool, + repository::{ + AccountRepository, AccountUsageRepository, StoredAccount, StoredAccountMetadata, + }, + }, + codex::cookies::repository::CookieRepository, + codex::oauth::TokenRefresher, + config::AppConfig, + utils::pagination::Page, +}; + +#[derive(Clone)] +pub struct AccountService { + config: AppConfig, + repository: Option, + usage_repository: Option, + cookie_repository: Option, + token_refresher: Option>, + account_pool: Arc>, +} + +#[derive(Debug)] +pub enum AccountServiceError { + RepositoryUnavailable, + UsageRepositoryUnavailable, + CookieRepositoryUnavailable, + AccountNotFound, + List, + Export, + Inspect, + ResetUsage, + UpdateLabel, + UpdateStatus, + SyncStatus, + Delete, + LoadCookies, + StoreCookies, + DeleteCookies, + NoValidCookies, + EmptyIds, + InvalidStatus(String), + LabelTooLong, +} + +#[derive(Debug)] +pub struct BatchDeleteAccounts { + pub deleted: u32, + pub not_found: Vec, +} + +#[derive(Debug)] +pub struct BatchUpdateAccountStatus { + pub updated: u32, + pub not_found: Vec, +} + +#[derive(Debug)] +pub struct UpdateAccountStatus { + pub id: String, + pub status: AccountStatus, +} + +#[derive(Debug)] +pub struct AccountImportCounts { + pub imported: u32, + pub skipped: u32, +} + +#[derive(Debug, Clone)] +pub struct AccountImportEntry { + pub id: Option, + pub email: Option, + pub account_id: Option, + pub user_id: Option, + pub label: Option, + pub plan_type: Option, + pub token: Option, + pub refresh_token: Option, + pub status: Option, +} + +#[derive(Debug)] +pub enum AccountProbeOutcome { + Alive, + Dead, + Skipped, +} + +impl AccountProbeOutcome { + pub fn as_str(&self) -> &'static str { + match self { + Self::Alive => "alive", + Self::Dead => "dead", + Self::Skipped => "skipped", + } + } +} + +#[derive(Debug)] +pub struct AccountProbeResult { + pub id: String, + pub email: Option, + pub previous_status: AccountStatus, + pub outcome: AccountProbeOutcome, + pub status: Option, + pub error: Option, + pub duration_ms: Option, +} + +#[derive(Debug)] +pub struct AccountQuotaResult { + pub quota: Value, + pub raw: Value, +} + +#[derive(Debug)] +pub enum AccountQuotaError { + RepositoryUnavailable, + Load, + NotFound, + Inactive(AccountStatus), + StoreQuota, + Fetch(String), +} + +#[derive(Debug)] +pub enum HealthCheckError { + RepositoryUnavailable, + List, +} + +#[derive(Debug)] +pub enum RefreshAccountError { + RepositoryUnavailable, + Load, + NotFound, + TokenRefresherUnavailable, + StoreRefreshed, +} + +#[derive(Debug)] +pub enum StoreImportAccountError { + RepositoryUnavailable, + Inspect, + Invalid(String), + Insert, +} + +#[derive(Debug)] +pub enum ValidatedAccountImportError { + RepositoryUnavailable, + TokenRequired, + TokenRefresherUnavailable, + RefreshTransport, + RefreshRejected, + InvalidToken(&'static str), + Inspect, + NotFound, + Update, + Insert, + Load, +} + +impl AccountService { + pub fn new( + config: AppConfig, + repository: Option, + usage_repository: Option, + cookie_repository: Option, + token_refresher: Option>, + account_pool: Arc>, + ) -> Self { + Self { + config, + repository, + usage_repository, + cookie_repository, + token_refresher, + account_pool, + } + } + + pub fn has_repository(&self) -> bool { + self.repository.is_some() + } + + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> Result, AccountServiceError> { + self.repository()? + .list_metadata(cursor, limit) + .await + .map_err(|_| AccountServiceError::List) + } + + pub async fn export( + &self, + ids: Vec, + ) -> Result, AccountServiceError> { + let repo = self.repository()?; + if ids.is_empty() { + return repo + .list_all() + .await + .map_err(|_| AccountServiceError::Export); + } + + let mut accounts = Vec::with_capacity(ids.len()); + for id in ids { + match repo.get(&id).await { + Ok(Some(account)) => accounts.push(account), + Ok(None) => {} + Err(_) => return Err(AccountServiceError::Export), + } + } + Ok(accounts) + } + + fn repository(&self) -> Result<&AccountRepository, AccountServiceError> { + self.repository + .as_ref() + .ok_or(AccountServiceError::RepositoryUnavailable) + } + + fn usage_repository(&self) -> Result<&AccountUsageRepository, AccountServiceError> { + self.usage_repository + .as_ref() + .ok_or(AccountServiceError::UsageRepositoryUnavailable) + } + + fn cookie_repository(&self) -> Result<&CookieRepository, AccountServiceError> { + self.cookie_repository + .as_ref() + .ok_or(AccountServiceError::CookieRepositoryUnavailable) + } + + async fn ensure_account_exists(&self, account_id: &str) -> Result<(), AccountServiceError> { + match self.repository()?.exists(account_id).await { + Ok(true) => Ok(()), + Ok(false) => Err(AccountServiceError::AccountNotFound), + Err(_) => Err(AccountServiceError::Inspect), + } + } +} diff --git a/src/codex/accounts/service/mutation.rs b/src/codex/accounts/service/mutation.rs new file mode 100644 index 000000000..8f44815b6 --- /dev/null +++ b/src/codex/accounts/service/mutation.rs @@ -0,0 +1,169 @@ +use crate::codex::accounts::{model::AccountStatus, repository::AccountRepository}; + +use super::{ + runtime_pool::pool_account_from_stored, AccountService, AccountServiceError, + BatchDeleteAccounts, BatchUpdateAccountStatus, UpdateAccountStatus, +}; + +impl AccountService { + pub async fn reset_usage(&self, account_id: &str) -> Result { + match self.repository()?.exists(account_id).await { + Ok(true) => {} + Ok(false) => return Ok(false), + Err(_) => return Err(AccountServiceError::Inspect), + } + self.usage_repository()? + .reset_account(account_id) + .await + .map_err(|_| AccountServiceError::ResetUsage)?; + self.account_pool.lock().await.reset_usage(account_id); + Ok(true) + } + + pub async fn update_label( + &self, + account_id: &str, + label: Option, + ) -> Result { + if label + .as_ref() + .is_some_and(|label| label.chars().count() > 64) + { + return Err(AccountServiceError::LabelTooLong); + } + match self + .repository()? + .set_label(account_id, label.clone()) + .await + { + Ok(true) => { + self.account_pool.lock().await.set_label(account_id, label); + Ok(true) + } + Ok(false) => Ok(false), + Err(_) => Err(AccountServiceError::UpdateLabel), + } + } + + pub async fn update_status( + &self, + account_id: &str, + status: &str, + ) -> Result, AccountServiceError> { + let status = parse_admin_account_status(status)?; + let repo = self.repository()?; + match repo.set_status(account_id, status).await { + Ok(true) => { + self.sync_runtime_account_status(repo, account_id, status) + .await?; + Ok(Some(UpdateAccountStatus { + id: account_id.to_string(), + status, + })) + } + Ok(false) => Ok(None), + Err(_) => Err(AccountServiceError::UpdateStatus), + } + } + + pub async fn delete(&self, account_id: &str) -> Result { + match self.repository()?.delete(account_id).await { + Ok(true) => { + self.account_pool.lock().await.remove(account_id); + Ok(true) + } + Ok(false) => Ok(false), + Err(_) => Err(AccountServiceError::Delete), + } + } + + pub async fn batch_delete( + &self, + ids: Vec, + ) -> Result { + if ids.is_empty() { + return Err(AccountServiceError::EmptyIds); + } + let repo = self.repository()?; + let mut deleted = 0u32; + let mut not_found = Vec::new(); + for account_id in ids { + match repo.delete(&account_id).await { + Ok(true) => { + self.account_pool.lock().await.remove(&account_id); + deleted += 1; + } + Ok(false) => not_found.push(account_id), + Err(_) => return Err(AccountServiceError::Delete), + } + } + Ok(BatchDeleteAccounts { deleted, not_found }) + } + + pub async fn batch_update_status( + &self, + ids: Vec, + status: &str, + ) -> Result { + if ids.is_empty() { + return Err(AccountServiceError::EmptyIds); + } + let status = parse_admin_account_status(status)?; + let repo = self.repository()?; + let mut updated = 0u32; + let mut not_found = Vec::new(); + for account_id in ids { + match repo.set_status(&account_id, status).await { + Ok(true) => { + self.sync_runtime_account_status(repo, &account_id, status) + .await?; + updated += 1; + } + Ok(false) => not_found.push(account_id), + Err(_) => return Err(AccountServiceError::UpdateStatus), + } + } + Ok(BatchUpdateAccountStatus { updated, not_found }) + } + + async fn sync_runtime_account_status( + &self, + repo: &AccountRepository, + account_id: &str, + status: AccountStatus, + ) -> Result<(), AccountServiceError> { + if status == AccountStatus::Active { + match repo.get(account_id).await { + Ok(Some(account)) => { + self.account_pool + .lock() + .await + .insert(pool_account_from_stored(account)); + } + Ok(None) => { + self.account_pool.lock().await.remove(account_id); + } + Err(_) => return Err(AccountServiceError::SyncStatus), + } + return Ok(()); + } + self.account_pool.lock().await.remove(account_id); + Ok(()) + } +} + +pub(super) fn parse_admin_account_status( + status: &str, +) -> Result { + match status.trim().to_ascii_lowercase().as_str() { + "active" => Ok(AccountStatus::Active), + "expired" => Ok(AccountStatus::Expired), + "quota_exhausted" | "quota-exhausted" => Ok(AccountStatus::QuotaExhausted), + "refreshing" => Ok(AccountStatus::Refreshing), + "disabled" => Ok(AccountStatus::Disabled), + "banned" => Ok(AccountStatus::Banned), + other => Err(AccountServiceError::InvalidStatus(format!( + "Unsupported account status: {other}" + ))), + } +} diff --git a/src/codex/accounts/service/quota.rs b/src/codex/accounts/service/quota.rs new file mode 100644 index 000000000..da877f418 --- /dev/null +++ b/src/codex/accounts/service/quota.rs @@ -0,0 +1,198 @@ +use crate::codex::accounts::model::AccountStatus; + +use super::{ + health::{apply_codex_account_error, fetch_account_usage, public_codex_error}, + AccountQuotaError, AccountQuotaResult, AccountService, +}; + +use serde_json::{json, Map, Value}; + +impl AccountService { + pub async fn account_quota( + &self, + account_id: &str, + request_id: &str, + ) -> Result { + let repo = self + .repository + .as_ref() + .ok_or(AccountQuotaError::RepositoryUnavailable)?; + let account = match repo.get(account_id).await { + Ok(Some(account)) => account, + Ok(None) => return Err(AccountQuotaError::NotFound), + Err(_) => return Err(AccountQuotaError::Load), + }; + if account.status != AccountStatus::Active { + return Err(AccountQuotaError::Inactive(account.status)); + } + + match fetch_account_usage(self, &account, request_id).await { + Ok(raw) => { + let quota = quota_from_usage(&raw); + if repo + .update_quota_json(&account.id, "a.to_string()) + .await + .is_err() + { + return Err(AccountQuotaError::StoreQuota); + } + Ok(AccountQuotaResult { quota, raw }) + } + Err(error) => { + apply_codex_account_error(self, repo, &account, &error).await; + Err(AccountQuotaError::Fetch(public_codex_error(&error))) + } + } + } +} + +pub(super) 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::().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") +} diff --git a/src/codex/accounts/service/refresh.rs b/src/codex/accounts/service/refresh.rs new file mode 100644 index 000000000..d712263b7 --- /dev/null +++ b/src/codex/accounts/service/refresh.rs @@ -0,0 +1,137 @@ +use secrecy::{ExposeSecret, SecretString}; + +use crate::{ + codex::accounts::{ + model::{Account, AccountStatus}, + repository::{AccountRepository, StoredAccount, TokenUpdate}, + }, + codex::oauth::RefreshFailure, +}; + +use super::{ + health::skipped_probe_result, runtime_pool::pool_account_from_stored, AccountProbeOutcome, + AccountProbeResult, AccountService, RefreshAccountError, +}; + +impl AccountService { + pub async fn refresh_account( + &self, + account_id: &str, + ) -> Result { + let repo = self + .repository + .as_ref() + .ok_or(RefreshAccountError::RepositoryUnavailable)?; + let account = match repo.get(account_id).await { + Ok(Some(account)) => account, + Ok(None) => return Err(RefreshAccountError::NotFound), + Err(_) => return Err(RefreshAccountError::Load), + }; + if account.status == AccountStatus::Disabled { + return Ok(skipped_probe_result(&account, "manually disabled")); + } + let Some(refresh_token) = account.refresh_token.as_ref() else { + return Ok(skipped_probe_result(&account, "no refresh token")); + }; + let Some(refresher) = self.token_refresher.as_ref() else { + return Err(RefreshAccountError::TokenRefresherUnavailable); + }; + + let started_at = std::time::Instant::now(); + let previous_status = account.status; + match refresher.refresh(refresh_token.expose_secret()).await { + Ok(tokens) => { + let updated = self + .persist_refreshed_account( + repo, + &account.id, + tokens.access_token, + tokens.refresh_token, + ) + .await + .map_err(|()| RefreshAccountError::StoreRefreshed)?; + Ok(AccountProbeResult { + id: updated.id, + email: updated.email, + previous_status, + outcome: AccountProbeOutcome::Alive, + status: Some(updated.status), + error: None, + duration_ms: Some(started_at.elapsed().as_millis()), + }) + } + Err(failure) => { + let status = self + .apply_refresh_failure_status(repo, &account, failure) + .await; + Ok(AccountProbeResult { + id: account.id, + email: account.email, + previous_status, + outcome: AccountProbeOutcome::Dead, + status, + error: Some(public_refresh_failure(failure).to_string()), + duration_ms: Some(started_at.elapsed().as_millis()), + }) + } + } + } + + async fn persist_refreshed_account( + &self, + repo: &AccountRepository, + account_id: &str, + access_token: String, + refresh_token: Option, + ) -> Result { + 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 + .map_err(|_| ())?; + let account = repo.get(account_id).await.map_err(|_| ())?.ok_or(())?; + let account = pool_account_from_stored(account); + self.account_pool.lock().await.insert(account.clone()); + Ok(account) + } + + async fn apply_refresh_failure_status( + &self, + repo: &AccountRepository, + account: &StoredAccount, + failure: RefreshFailure, + ) -> Option { + let status = status_for_refresh_failure(failure)?; + let _ = repo.set_status(&account.id, status).await; + self.account_pool + .lock() + .await + .set_status(&account.id, status); + Some(status) + } +} + +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, + } +} + +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", + } +} diff --git a/src/codex/accounts/service/runtime_pool.rs b/src/codex/accounts/service/runtime_pool.rs new file mode 100644 index 000000000..4c97dee68 --- /dev/null +++ b/src/codex/accounts/service/runtime_pool.rs @@ -0,0 +1,54 @@ +use secrecy::ExposeSecret; + +use crate::codex::accounts::{ + model::Account, + repository::{AccountRepositoryResult, StoredAccount}, +}; + +use super::AccountService; + +impl AccountService { + pub async fn reload_runtime_accounts_from_repository(&self) -> AccountRepositoryResult { + let Some(repository) = self.repository.as_ref() else { + return Ok(0); + }; + let accounts = repository.list_pool_accounts().await?; + let restored = accounts.len(); + let mut pool = self.account_pool.lock().await; + for account in accounts { + pool.insert(account); + } + Ok(restored) + } + + pub async fn insert_runtime_account(&self, account: Account) { + self.account_pool.lock().await.insert(account); + } + + pub async fn acquire_runtime_account(&self, model: &str) -> Option { + self.account_pool.lock().await.acquire(model) + } +} + +pub(super) fn pool_account_from_stored(account: StoredAccount) -> Account { + Account { + id: account.id, + email: account.email, + account_id: account.account_id, + user_id: account.user_id, + label: account.label, + plan_type: account.plan_type, + access_token: account.access_token.expose_secret().to_string(), + refresh_token: account + .refresh_token + .as_ref() + .map(|token| token.expose_secret().to_string()), + access_token_expires_at: account.access_token_expires_at, + status: account.status, + quota_limit_reached: false, + quota_cooldown_until: None, + cloudflare_cooldown_until: None, + added_at: account.added_at.to_rfc3339(), + last_used_at: None, + } +} diff --git a/src/cookies/jar.rs b/src/codex/cookies/jar.rs similarity index 100% rename from src/cookies/jar.rs rename to src/codex/cookies/jar.rs diff --git a/src/cookies/mod.rs b/src/codex/cookies/mod.rs similarity index 100% rename from src/cookies/mod.rs rename to src/codex/cookies/mod.rs diff --git a/src/cookies/repository.rs b/src/codex/cookies/repository.rs similarity index 99% rename from src/cookies/repository.rs rename to src/codex/cookies/repository.rs index c1c0dbc80..304fa8f08 100644 --- a/src/cookies/repository.rs +++ b/src/codex/cookies/repository.rs @@ -4,7 +4,7 @@ use sqlx::{Row, SqlitePool}; use thiserror::Error; use uuid::Uuid; -use crate::crypto::{CryptoError, SecretBox}; +use crate::utils::crypto::{CryptoError, SecretBox}; #[derive(Debug, Error)] pub enum CookieError { diff --git a/src/fingerprint/mod.rs b/src/codex/fingerprint/mod.rs similarity index 100% rename from src/fingerprint/mod.rs rename to src/codex/fingerprint/mod.rs diff --git a/src/fingerprint/model.rs b/src/codex/fingerprint/model.rs similarity index 100% rename from src/fingerprint/model.rs rename to src/codex/fingerprint/model.rs diff --git a/src/fingerprint/repository.rs b/src/codex/fingerprint/repository.rs similarity index 98% rename from src/fingerprint/repository.rs rename to src/codex/fingerprint/repository.rs index 955275474..d4231f1df 100644 --- a/src/fingerprint/repository.rs +++ b/src/codex/fingerprint/repository.rs @@ -2,7 +2,7 @@ use chrono::Utc; use sqlx::{Row, SqlitePool}; use uuid::Uuid; -use crate::fingerprint::{ +use crate::codex::fingerprint::{ model::Fingerprint, updater::{FingerprintUpdate, CODEX_DESKTOP_UPDATE_SOURCE}, }; diff --git a/src/fingerprint/updater.rs b/src/codex/fingerprint/updater.rs similarity index 96% rename from src/fingerprint/updater.rs rename to src/codex/fingerprint/updater.rs index 171c16649..bc8254bd8 100644 --- a/src/fingerprint/updater.rs +++ b/src/codex/fingerprint/updater.rs @@ -1,7 +1,7 @@ use serde::Deserialize; use thiserror::Error; -use crate::fingerprint::repository::FingerprintRepository; +use crate::codex::fingerprint::repository::FingerprintRepository; pub const CODEX_DESKTOP_UPDATE_SOURCE: &str = "codex_desktop_update_source"; diff --git a/src/codex/mod.rs b/src/codex/mod.rs index 038fae8a2..abf527181 100644 --- a/src/codex/mod.rs +++ b/src/codex/mod.rs @@ -1,6 +1,8 @@ -pub mod client; -pub mod headers; -pub mod sse; -pub mod types; -pub mod usage; -pub mod websocket; +pub mod accounts; +pub mod cookies; +pub mod fingerprint; +pub mod models; +pub mod oauth; +pub mod protocol; +pub mod transport; +pub mod upstream; diff --git a/src/models/catalog.rs b/src/codex/models/catalog.rs similarity index 100% rename from src/models/catalog.rs rename to src/codex/models/catalog.rs diff --git a/src/models/mod.rs b/src/codex/models/mod.rs similarity index 68% rename from src/models/mod.rs rename to src/codex/models/mod.rs index 27e868291..c01714035 100644 --- a/src/models/mod.rs +++ b/src/codex/models/mod.rs @@ -1,2 +1,3 @@ pub mod catalog; pub mod repository; +pub mod service; diff --git a/src/models/repository.rs b/src/codex/models/repository.rs similarity index 96% rename from src/models/repository.rs rename to src/codex/models/repository.rs index 502b660d1..5ec49256b 100644 --- a/src/models/repository.rs +++ b/src/codex/models/repository.rs @@ -2,7 +2,7 @@ use chrono::Utc; use sqlx::{Row, SqlitePool}; use thiserror::Error; -use crate::models::catalog::{CodexModelInfo, ModelPlanSnapshot}; +use crate::codex::models::catalog::{CodexModelInfo, ModelPlanSnapshot}; #[derive(Debug, Error)] pub enum ModelSnapshotRepositoryError { diff --git a/src/codex/models/service.rs b/src/codex/models/service.rs new file mode 100644 index 000000000..a3e4ef037 --- /dev/null +++ b/src/codex/models/service.rs @@ -0,0 +1,184 @@ +use std::sync::Arc; + +use tokio::sync::Mutex; + +use crate::{ + codex::accounts::{ + model::{Account, AccountStatus}, + pool::AccountPool, + repository::AccountRepository, + }, + codex::fingerprint::model::Fingerprint, + codex::models::{ + catalog::{ModelCatalog, ModelPlanSnapshot}, + repository::ModelSnapshotRepository, + }, + codex::transport::client::{build_reqwest_client, CodexBackendClient, CodexRequestContext}, + config::AppConfig, +}; + +#[derive(Clone)] +pub struct ModelService { + config: AppConfig, + snapshot_repository: Option, + account_repository: Option, + account_pool: Arc>, +} + +#[derive(Debug, Clone, Copy)] +pub struct ModelRefreshResult { + pub refreshed_plans: usize, + pub model_count: usize, + pub failed_plans: usize, +} + +#[derive(Debug)] +pub enum ModelServiceError { + AccountRepositoryUnavailable, + ModelRepositoryUnavailable, + ListAccounts, + NoAccounts, + BuildClient, + StoreSnapshot, + AllPlansFailed(ModelRefreshResult), + LoadSnapshots, +} + +impl ModelService { + pub fn new( + config: AppConfig, + snapshot_repository: Option, + account_repository: Option, + account_pool: Arc>, + ) -> Self { + Self { + config, + snapshot_repository, + account_repository, + account_pool, + } + } + + pub async fn catalog(&self) -> ModelCatalog { + let Some(repo) = self.snapshot_repository.as_ref() else { + return ModelCatalog::from_config(&self.config.model); + }; + match repo.list_plan_snapshots().await { + Ok(snapshots) if !snapshots.is_empty() => { + ModelCatalog::from_config_and_snapshots(&self.config.model, &snapshots) + } + Ok(_) => ModelCatalog::from_config(&self.config.model), + Err(error) => { + // Cached model snapshots extend the catalog; fallback keeps proxy routes available. + tracing::warn!(?error, "failed to load cached model snapshots"); + ModelCatalog::from_config(&self.config.model) + } + } + } + + pub async fn refresh_backend_models( + &self, + request_id: &str, + ) -> Result { + let account_repo = self + .account_repository + .as_ref() + .ok_or(ModelServiceError::AccountRepositoryUnavailable)?; + let model_repo = self + .snapshot_repository + .as_ref() + .ok_or(ModelServiceError::ModelRepositoryUnavailable)?; + + let accounts = account_repo + .list_pool_accounts() + .await + .map_err(|_| ModelServiceError::ListAccounts)?; + let plan_accounts = distinct_active_plan_accounts(accounts); + if plan_accounts.is_empty() { + return Err(ModelServiceError::NoAccounts); + } + + let client = build_reqwest_client(self.config.tls.force_http11) + .map(|client| { + CodexBackendClient::new( + client, + self.config.api.base_url.clone(), + Fingerprint::default_codex_desktop(), + ) + }) + .map_err(|_| ModelServiceError::BuildClient)?; + + let mut result = ModelRefreshResult { + refreshed_plans: 0, + model_count: 0, + failed_plans: 0, + }; + for (plan_type, account) in plan_accounts { + let context = CodexRequestContext { + access_token: &account.access_token, + 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: None, + }; + let entries = match client.fetch_models(context).await { + Ok(entries) if !entries.is_empty() => entries, + Ok(_) => { + result.failed_plans += 1; + continue; + } + Err(error) => { + tracing::warn!(?error, plan_type, "failed to refresh backend models"); + result.failed_plans += 1; + continue; + } + }; + let snapshot = ModelPlanSnapshot::from_backend_entries(plan_type, entries); + result.model_count += snapshot.models.len(); + model_repo + .replace_plan_snapshot(&snapshot) + .await + .map_err(|_| ModelServiceError::StoreSnapshot)?; + result.refreshed_plans += 1; + } + + if result.refreshed_plans == 0 { + return Err(ModelServiceError::AllPlansFailed(result)); + } + + let snapshots = model_repo + .list_plan_snapshots() + .await + .map_err(|_| ModelServiceError::LoadSnapshots)?; + let allowlist = ModelCatalog::from_config_and_snapshots(&self.config.model, &snapshots) + .model_plan_allowlist(); + self.account_pool + .lock() + .await + .set_model_plan_allowlist(allowlist); + + Ok(result) + } +} + +fn distinct_active_plan_accounts(accounts: Vec) -> Vec<(String, Account)> { + let mut by_plan = std::collections::BTreeMap::new(); + for account in accounts { + if account.status != AccountStatus::Active { + continue; + } + let plan_type = account + .plan_type + .as_deref() + .unwrap_or("unknown") + .to_string(); + by_plan.entry(plan_type).or_insert(account); + } + by_plan.into_iter().collect() +} diff --git a/src/auth/cli_import.rs b/src/codex/oauth/cli_import.rs similarity index 100% rename from src/auth/cli_import.rs rename to src/codex/oauth/cli_import.rs diff --git a/src/auth/oauth.rs b/src/codex/oauth/client.rs similarity index 99% rename from src/auth/oauth.rs rename to src/codex/oauth/client.rs index a2964e61d..b3eea6471 100644 --- a/src/auth/oauth.rs +++ b/src/codex/oauth/client.rs @@ -9,7 +9,7 @@ use serde::Deserialize; use sha2::{Digest, Sha256}; use thiserror::Error; -use crate::auth::{ +use super::{ refresh::{RefreshFailure, TokenRefresher}, token::TokenPair, }; diff --git a/src/codex/oauth/mod.rs b/src/codex/oauth/mod.rs new file mode 100644 index 000000000..9f465635a --- /dev/null +++ b/src/codex/oauth/mod.rs @@ -0,0 +1,9 @@ +pub mod cli_import; +pub mod client; +pub mod refresh; +pub mod token; + +pub use cli_import::*; +pub use client::*; +pub use refresh::*; +pub use token::*; diff --git a/src/auth/refresh.rs b/src/codex/oauth/refresh.rs similarity index 97% rename from src/auth/refresh.rs rename to src/codex/oauth/refresh.rs index 3e7e12767..3bd8afa14 100644 --- a/src/auth/refresh.rs +++ b/src/codex/oauth/refresh.rs @@ -5,10 +5,9 @@ use chrono::{DateTime, Duration, Utc}; use thiserror::Error; use tokio::sync::Semaphore; -use crate::{ - accounts::model::{Account, AccountStatus}, - auth::token::TokenPair, -}; +use crate::codex::accounts::model::{Account, AccountStatus}; + +use super::token::TokenPair; #[derive(Debug, Clone, Copy)] pub struct RefreshPolicy { diff --git a/src/auth/token.rs b/src/codex/oauth/token.rs similarity index 100% rename from src/auth/token.rs rename to src/codex/oauth/token.rs diff --git a/src/translation/codex_to_openai.rs b/src/codex/protocol/codex_to_openai.rs similarity index 97% rename from src/translation/codex_to_openai.rs rename to src/codex/protocol/codex_to_openai.rs index 3da738f8b..9daeb34ad 100644 --- a/src/translation/codex_to_openai.rs +++ b/src/codex/protocol/codex_to_openai.rs @@ -4,7 +4,7 @@ use chrono::Utc; use serde_json::{json, Value}; use uuid::Uuid; -use crate::codex::sse::{parse_sse_events, SseError}; +use crate::codex::transport::sse::{parse_sse_events, SseError}; pub fn openai_error(message: &str, code: &str) -> Value { json!({ @@ -39,11 +39,9 @@ pub fn chat_completion_from_codex_sse( content.push_str(delta); } } - Some("response.reasoning_summary_text.delta") => { - if include_reasoning { - if let Some(delta) = value.get("delta").and_then(Value::as_str) { - reasoning.push_str(delta); - } + Some("response.reasoning_summary_text.delta") if include_reasoning => { + if let Some(delta) = value.get("delta").and_then(Value::as_str) { + reasoning.push_str(delta); } } Some("response.output_item.added") => { @@ -281,10 +279,8 @@ pub fn chat_completion_stream_from_codex_sse( call_ids_with_deltas.insert(call_id); } } - Some("response.output_item.done") => { - if is_function_call_item(value.get("item")) { - has_tool_calls = true; - } + Some("response.output_item.done") if is_function_call_item(value.get("item")) => { + has_tool_calls = true; } Some("response.completed") => { completed_response = value.get("response").cloned(); diff --git a/src/error.rs b/src/codex/protocol/error.rs similarity index 100% rename from src/error.rs rename to src/codex/protocol/error.rs diff --git a/src/translation/mod.rs b/src/codex/protocol/mod.rs similarity index 81% rename from src/translation/mod.rs rename to src/codex/protocol/mod.rs index 6992ee40a..1ebdcd482 100644 --- a/src/translation/mod.rs +++ b/src/codex/protocol/mod.rs @@ -1,3 +1,4 @@ pub mod codex_to_openai; +pub mod error; pub mod openai_to_codex; pub mod schema; diff --git a/src/translation/openai_to_codex.rs b/src/codex/protocol/openai_to_codex.rs similarity index 98% rename from src/translation/openai_to_codex.rs rename to src/codex/protocol/openai_to_codex.rs index 813a21cd3..3c6016983 100644 --- a/src/translation/openai_to_codex.rs +++ b/src/codex/protocol/openai_to_codex.rs @@ -2,8 +2,8 @@ use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use crate::{ - codex::types::CodexResponsesRequest, - error::{AppError, AppResult}, + codex::protocol::error::{AppError, AppResult}, + codex::transport::types::CodexResponsesRequest, }; #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/src/translation/schema.rs b/src/codex/protocol/schema.rs similarity index 100% rename from src/translation/schema.rs rename to src/codex/protocol/schema.rs diff --git a/src/codex/client.rs b/src/codex/transport/client.rs similarity index 99% rename from src/codex/client.rs rename to src/codex/transport/client.rs index 6e695de03..806e66568 100644 --- a/src/codex/client.rs +++ b/src/codex/transport/client.rs @@ -8,7 +8,9 @@ use serde_json::Value; use thiserror::Error; use crate::{ - codex::{ + codex::fingerprint::model::Fingerprint, + codex::models::catalog::BackendModelEntry, + codex::transport::{ headers::build_codex_headers, sse::SseError, types::CodexResponsesRequest, @@ -18,8 +20,6 @@ use crate::{ CodexTransport, CodexWebSocketError, WebSocketSupportError, }, }, - fingerprint::model::Fingerprint, - models::catalog::BackendModelEntry, }; #[derive(Debug, Error)] diff --git a/src/codex/headers.rs b/src/codex/transport/headers.rs similarity index 95% rename from src/codex/headers.rs rename to src/codex/transport/headers.rs index 871e55874..73bdaadd7 100644 --- a/src/codex/headers.rs +++ b/src/codex/transport/headers.rs @@ -1,6 +1,6 @@ use std::collections::BTreeMap; -use crate::fingerprint::model::Fingerprint; +use crate::codex::fingerprint::model::Fingerprint; pub fn build_codex_headers( fp: &Fingerprint, diff --git a/src/codex/transport/mod.rs b/src/codex/transport/mod.rs new file mode 100644 index 000000000..038fae8a2 --- /dev/null +++ b/src/codex/transport/mod.rs @@ -0,0 +1,6 @@ +pub mod client; +pub mod headers; +pub mod sse; +pub mod types; +pub mod usage; +pub mod websocket; diff --git a/src/codex/sse.rs b/src/codex/transport/sse.rs similarity index 100% rename from src/codex/sse.rs rename to src/codex/transport/sse.rs diff --git a/src/codex/types.rs b/src/codex/transport/types.rs similarity index 100% rename from src/codex/types.rs rename to src/codex/transport/types.rs diff --git a/src/codex/usage.rs b/src/codex/transport/usage.rs similarity index 98% rename from src/codex/usage.rs rename to src/codex/transport/usage.rs index 69dc18108..9e9fa0e5c 100644 --- a/src/codex/usage.rs +++ b/src/codex/transport/usage.rs @@ -1,7 +1,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; -use crate::codex::sse::{parse_sse_events, SseError}; +use crate::codex::transport::sse::{parse_sse_events, SseError}; #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] diff --git a/src/codex/websocket.rs b/src/codex/transport/websocket.rs similarity index 99% rename from src/codex/websocket.rs rename to src/codex/transport/websocket.rs index 41c5058b3..3663594d8 100644 --- a/src/codex/websocket.rs +++ b/src/codex/transport/websocket.rs @@ -12,7 +12,7 @@ use futures::{SinkExt, StreamExt}; use reqwest::{header::HeaderMap, StatusCode}; use serde_json::{json, Value}; -use crate::codex::{sse::encode_sse_event, types::CodexResponsesRequest}; +use crate::codex::transport::{sse::encode_sse_event, types::CodexResponsesRequest}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum CodexTransport { diff --git a/src/codex/upstream/dispatch.rs b/src/codex/upstream/dispatch.rs new file mode 100644 index 000000000..79bf6b173 --- /dev/null +++ b/src/codex/upstream/dispatch.rs @@ -0,0 +1,29 @@ +use axum::{http::StatusCode, Json}; +use reqwest::Url; +use serde_json::Value; + +use crate::codex::protocol::codex_to_openai::openai_error; + +pub(crate) fn no_available_accounts_response() -> (StatusCode, Json) { + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(openai_error( + "No available Codex accounts", + "no_available_accounts", + )), + ) +} + +pub(crate) fn normalize_service_tier_for_upstream(service_tier: String) -> String { + if service_tier == "fast" { + "priority".to_string() + } else { + service_tier + } +} + +pub(super) fn request_domain(base_url: &str) -> Option { + Url::parse(base_url) + .ok() + .and_then(|url| url.host_str().map(ToString::to_string)) +} diff --git a/src/codex/upstream/fallback.rs b/src/codex/upstream/fallback.rs new file mode 100644 index 000000000..f3132c3b2 --- /dev/null +++ b/src/codex/upstream/fallback.rs @@ -0,0 +1,132 @@ +use axum::http::StatusCode; +use serde_json::{json, Value}; + +use crate::codex::transport::client::CodexClientError; + +const DEFAULT_RATE_LIMIT_BACKOFF_SECONDS: u64 = 60; +const MAX_RATE_LIMIT_BACKOFF_SECONDS: u64 = 86_400 * 7; +const CLOUDFLARE_CHALLENGE_COOLDOWN_SECONDS: u64 = 10; + +#[derive(Debug, Clone, Copy)] +pub(crate) enum UpstreamAccountRetry { + RateLimited { retry_after_seconds: u64 }, + QuotaExhausted, + CloudflareChallenge { cooldown_seconds: u64 }, + Banned, +} + +impl UpstreamAccountRetry { + pub(crate) fn status(self) -> StatusCode { + match self { + Self::RateLimited { .. } => StatusCode::TOO_MANY_REQUESTS, + Self::QuotaExhausted => StatusCode::PAYMENT_REQUIRED, + Self::CloudflareChallenge { .. } => StatusCode::FORBIDDEN, + Self::Banned => StatusCode::FORBIDDEN, + } + } + + pub(crate) fn metadata(self, stream: bool) -> Value { + match self { + Self::RateLimited { + retry_after_seconds, + } => json!({ + "stream": stream, + "retry": true, + "reason": "rateLimited", + "retryAfterSeconds": retry_after_seconds, + }), + Self::QuotaExhausted => json!({ + "stream": stream, + "retry": true, + "reason": "quotaExhausted", + }), + Self::CloudflareChallenge { cooldown_seconds } => json!({ + "stream": stream, + "retry": true, + "reason": "cloudflareChallenge", + "cooldownSeconds": cooldown_seconds, + }), + Self::Banned => json!({ + "stream": stream, + "retry": true, + "reason": "banned", + }), + } + } +} + +pub(crate) fn classify_upstream_account_retry( + error: &CodexClientError, +) -> Option { + match error { + CodexClientError::Upstream { + status, + retry_after_seconds, + .. + } if *status == StatusCode::TOO_MANY_REQUESTS => Some(UpstreamAccountRetry::RateLimited { + retry_after_seconds: retry_after_seconds + .unwrap_or(DEFAULT_RATE_LIMIT_BACKOFF_SECONDS) + .min(MAX_RATE_LIMIT_BACKOFF_SECONDS), + }), + CodexClientError::Upstream { status, .. } if *status == StatusCode::PAYMENT_REQUIRED => { + Some(UpstreamAccountRetry::QuotaExhausted) + } + CodexClientError::Upstream { status, body, .. } if *status == StatusCode::FORBIDDEN => { + if is_cloudflare_challenge(body) { + Some(UpstreamAccountRetry::CloudflareChallenge { + cooldown_seconds: CLOUDFLARE_CHALLENGE_COOLDOWN_SECONDS, + }) + } else { + Some(UpstreamAccountRetry::Banned) + } + } + _ => None, + } +} + +pub(crate) fn websocket_history_retry_metadata(retry: UpstreamAccountRetry, stream: bool) -> Value { + match retry { + UpstreamAccountRetry::RateLimited { + retry_after_seconds, + } => json!({ + "stream": stream, + "transport": "websocket", + "retry": false, + "reason": "rateLimited", + "retryAfterSeconds": retry_after_seconds, + "accountAffinity": "previousResponseId", + }), + UpstreamAccountRetry::QuotaExhausted => json!({ + "stream": stream, + "transport": "websocket", + "retry": false, + "reason": "quotaExhausted", + "accountAffinity": "previousResponseId", + }), + UpstreamAccountRetry::CloudflareChallenge { cooldown_seconds } => json!({ + "stream": stream, + "transport": "websocket", + "retry": false, + "reason": "cloudflareChallenge", + "cooldownSeconds": cooldown_seconds, + "accountAffinity": "previousResponseId", + }), + UpstreamAccountRetry::Banned => json!({ + "stream": stream, + "transport": "websocket", + "retry": false, + "reason": "banned", + "accountAffinity": "previousResponseId", + }), + } +} + +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") +} diff --git a/src/codex/upstream/mod.rs b/src/codex/upstream/mod.rs new file mode 100644 index 000000000..934f67a66 --- /dev/null +++ b/src/codex/upstream/mod.rs @@ -0,0 +1,859 @@ +pub mod dispatch; +pub mod fallback; +pub mod refresh; +pub mod stream; +pub mod usage; + +use std::{sync::Arc, time::Instant}; + +use axum::{ + body::Body, + http::{ + header::{CACHE_CONTROL, CONTENT_TYPE}, + StatusCode, + }, + response::{IntoResponse, Response}, + Json, +}; +use chrono::{Duration, Utc}; +use futures::{stream as futures_stream, StreamExt}; +use serde_json::{json, Value}; +use tokio::sync::Mutex; + +use crate::{ + codex::accounts::{ + model::{Account, AccountStatus}, + pool::{AccountAcquireRequest, AccountPool}, + repository::AccountRepository, + }, + codex::cookies::repository::CookieRepository, + codex::fingerprint::model::Fingerprint, + codex::oauth::TokenRefresher, + codex::protocol::codex_to_openai::openai_error, + codex::transport::{ + client::{ + build_reqwest_client, CodexBackendClient, CodexBackendStream, CodexClientError, + CodexRequestContext, + }, + types::CodexResponsesRequest, + usage::{extract_sse_usage, TokenUsage}, + websocket::{transport_for_request, CodexTransport}, + }, + config::AppConfig, + logs::event::{EventLevel, EventLog}, + logs::repository::EventLogRepository, +}; + +use crate::http::v1::errors::codex_client_error_response; + +pub(crate) use self::{ + dispatch::{no_available_accounts_response, normalize_service_tier_for_upstream}, + fallback::{ + classify_upstream_account_retry, websocket_history_retry_metadata, UpstreamAccountRetry, + }, + stream::{completed_response_json, CollectedResponse}, +}; + +use self::{ + dispatch::request_domain, + refresh::refresh_account_after_unauthorized, + stream::{ensure_stream_metadata, responses_sse_failure}, + usage::{record_request_attempt, record_usage_with_deps}, +}; + +#[derive(Clone)] +struct CodexUpstreamDependencies { + config: AppConfig, + account_pool: Arc>, + account_repository: Option, + cookie_repository: Option, + event_logs: Option, + token_refresher: Option>, +} + +#[derive(Clone)] +pub(crate) struct CodexUpstreamService { + deps: CodexUpstreamDependencies, +} + +impl CodexUpstreamService { + pub(crate) fn new( + config: AppConfig, + account_pool: Arc>, + account_repository: Option, + cookie_repository: Option, + event_logs: Option, + token_refresher: Option>, + ) -> Self { + Self { + deps: CodexUpstreamDependencies { + config, + account_pool, + account_repository, + cookie_repository, + event_logs, + token_refresher, + }, + } + } + + pub(crate) async fn acquire_account(&self, model: &str) -> Option { + self.deps + .account_pool + .lock() + .await + .acquire_with(AccountAcquireRequest::new(model, Utc::now())) + .map(|acquired| acquired.account) + } + + pub(crate) async fn release_account(&self, account_id: &str) { + self.deps.account_pool.lock().await.release(account_id); + } + + pub(crate) async fn send_codex_request_with_refresh_retry( + &self, + request: &CodexResponsesRequest, + account: &Account, + request_id: &str, + ) -> Result { + send_codex_request_with_refresh_retry_deps(&self.deps, request, account, request_id).await + } + + pub(crate) async fn apply_retry_and_acquire_fallback( + &self, + account: &Account, + retry: UpstreamAccountRetry, + model: &str, + excluded_account_ids: &mut Vec, + ) -> Option { + apply_upstream_retry_and_acquire_fallback_with_deps( + &self.deps, + account, + retry, + model, + excluded_account_ids, + ) + .await + } + + pub(crate) async fn apply_account_retry(&self, account: &Account, retry: UpstreamAccountRetry) { + apply_upstream_account_retry_with_deps(&self.deps, account, retry).await; + } + + pub(crate) async fn responses_stream( + &self, + request: CodexResponsesRequest, + mut account: Account, + mut log_context: CodexRequestLogContext, + ) -> Response { + let deps = self.deps.clone(); + if transport_for_request(&request) == CodexTransport::WebSocketRequired { + return responses_websocket_stream(deps, request, account, log_context).await; + } + + let mut excluded_account_ids = Vec::new(); + let stream_response = loop { + let stream_response = send_codex_stream_request_with_refresh_retry( + &deps, + &request, + &account, + log_context.request_id.as_str(), + ) + .await; + match stream_response { + Ok(response) => break response, + Err(error) => { + deps.account_pool.lock().await.release(&account.id); + if let Some(retry) = classify_upstream_account_retry(&error) { + let fallback = apply_upstream_retry_and_acquire_fallback_with_deps( + &deps, + &account, + retry, + &request.model, + &mut excluded_account_ids, + ) + .await; + log_codex_upstream_response_with_deps( + &deps, + &log_context, + retry.status(), + EventLevel::Warn, + "v1 responses stream upstream retrying with fallback account", + retry.metadata(true), + ) + .await; + if let Some(fallback) = fallback { + account = fallback; + log_context = log_context.with_account(&account.id); + continue; + } + } + let error_response = codex_client_error_response(error); + log_codex_upstream_response_with_deps( + &deps, + &log_context, + error_response.0, + EventLevel::Error, + "v1 responses stream upstream request failed", + json!({"stream": true}), + ) + .await; + return error_response.into_response(); + } + } + }; + + if persist_upstream_cookies_with_deps( + &deps, + &account.id, + &stream_response.set_cookie_headers, + ) + .await + .is_err() + { + deps.account_pool.lock().await.release(&account.id); + log_codex_upstream_response_with_deps( + &deps, + &log_context, + StatusCode::INTERNAL_SERVER_ERROR, + EventLevel::Error, + "v1 responses stream cookie persistence failed", + json!({"stream": true, "cookieStoreError": true}), + ) + .await; + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(openai_error( + "Failed to persist upstream cookies", + "cookie_store_error", + )), + ) + .into_response(); + } + + let upstream = Box::pin(stream_response.response.bytes_stream()); + let audit = StreamAudit::new(deps, log_context, account.id); + let body_stream = + futures_stream::unfold(Some((upstream, Vec::new(), audit)), |state| async move { + let (mut upstream, mut collected, mut audit) = state?; + match upstream.next().await { + Some(Ok(chunk)) => { + collected.extend_from_slice(&chunk); + Some((Ok(chunk), Some((upstream, collected, audit)))) + } + Some(Err(error)) => { + audit.log_transport_error(&error).await; + Some((Err(error), None)) + } + None => { + audit.complete(&collected).await; + None + } + } + }); + + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream") + .header(CACHE_CONTROL, "no-cache") + .body(Body::from_stream(body_stream)) + .unwrap_or_else(|_| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(openai_error( + "Failed to build stream response", + "stream_response_error", + )), + ) + .into_response() + }) + } + + pub(crate) async fn persist_cookies( + &self, + account_id: &str, + set_cookie_headers: &[String], + ) -> Result<(), ()> { + persist_upstream_cookies_with_deps(&self.deps, account_id, set_cookie_headers).await + } + + pub(crate) async fn record_usage(&self, account_id: &str, usage: TokenUsage) -> Result<(), ()> { + record_usage_with_deps(&self.deps, account_id, usage).await + } + + pub(crate) async fn log_response( + &self, + context: &CodexRequestLogContext, + status: StatusCode, + level: EventLevel, + message: &str, + metadata: Value, + ) { + log_codex_upstream_response_with_deps( + &self.deps, context, status, level, message, metadata, + ) + .await; + } +} + +async fn send_codex_request_with_refresh_retry_deps( + deps: &CodexUpstreamDependencies, + request: &CodexResponsesRequest, + account: &Account, + request_id: &str, +) -> Result { + match send_codex_request(deps, request, account, request_id).await { + Err(CodexClientError::Upstream { + status, + body, + retry_after_seconds, + }) if status == StatusCode::UNAUTHORIZED => { + let Some(refreshed) = + refresh_account_after_unauthorized(deps, request, account, request_id).await + else { + return Err(CodexClientError::Upstream { + status, + body, + retry_after_seconds, + }); + }; + send_codex_request(deps, request, &refreshed, request_id).await + } + result => result, + } +} + +async fn send_codex_request( + deps: &CodexUpstreamDependencies, + request: &CodexResponsesRequest, + account: &Account, + request_id: &str, +) -> Result { + let request_domain = request_domain(&deps.config.api.base_url); + let cookie_header = match (deps.cookie_repository.as_ref(), request_domain.as_deref()) { + (Some(repo), Some(domain)) => repo.cookie_header(&account.id, domain).await.ok().flatten(), + _ => None, + }; + let client = CodexBackendClient::new( + build_reqwest_client(deps.config.tls.force_http11)?, + deps.config.api.base_url.clone(), + Fingerprint::default_codex_desktop(), + ); + client + .create_response( + request, + CodexRequestContext { + access_token: &account.access_token, + account_id: account.account_id.as_deref(), + request_id, + turn_state: request.turn_state.as_deref(), + turn_metadata: request.turn_metadata.as_deref(), + beta_features: request.beta_features.as_deref(), + include_timing_metrics: request.include_timing_metrics.as_deref(), + version: request.version.as_deref(), + codex_window_id: request.codex_window_id.as_deref(), + parent_thread_id: request.parent_thread_id.as_deref(), + cookie_header: cookie_header.as_deref(), + }, + ) + .await +} + +async fn send_codex_stream_request_with_refresh_retry( + deps: &CodexUpstreamDependencies, + request: &CodexResponsesRequest, + account: &Account, + request_id: &str, +) -> Result { + match send_codex_stream_request(deps, request, account, request_id).await { + Err(CodexClientError::Upstream { + status, + body, + retry_after_seconds, + }) if status == StatusCode::UNAUTHORIZED => { + let Some(refreshed) = + refresh_account_after_unauthorized(deps, request, account, request_id).await + else { + return Err(CodexClientError::Upstream { + status, + body, + retry_after_seconds, + }); + }; + send_codex_stream_request(deps, request, &refreshed, request_id).await + } + result => result, + } +} + +async fn send_codex_stream_request( + deps: &CodexUpstreamDependencies, + request: &CodexResponsesRequest, + account: &Account, + request_id: &str, +) -> Result { + let request_domain = request_domain(&deps.config.api.base_url); + let cookie_header = match (deps.cookie_repository.as_ref(), request_domain.as_deref()) { + (Some(repo), Some(domain)) => repo.cookie_header(&account.id, domain).await.ok().flatten(), + _ => None, + }; + let client = CodexBackendClient::new( + build_reqwest_client(deps.config.tls.force_http11)?, + deps.config.api.base_url.clone(), + Fingerprint::default_codex_desktop(), + ); + client + .stream_response( + request, + CodexRequestContext { + access_token: &account.access_token, + account_id: account.account_id.as_deref(), + request_id, + turn_state: request.turn_state.as_deref(), + turn_metadata: request.turn_metadata.as_deref(), + beta_features: request.beta_features.as_deref(), + include_timing_metrics: request.include_timing_metrics.as_deref(), + version: request.version.as_deref(), + codex_window_id: request.codex_window_id.as_deref(), + parent_thread_id: request.parent_thread_id.as_deref(), + cookie_header: cookie_header.as_deref(), + }, + ) + .await +} + +async fn apply_upstream_retry_and_acquire_fallback_with_deps( + deps: &CodexUpstreamDependencies, + account: &Account, + retry: UpstreamAccountRetry, + model: &str, + excluded_account_ids: &mut Vec, +) -> Option { + apply_upstream_account_retry_with_deps(deps, account, retry).await; + excluded_account_ids.push(account.id.clone()); + deps.account_pool + .lock() + .await + .acquire_with( + AccountAcquireRequest::new(model, Utc::now()) + .with_exclude_account_ids(excluded_account_ids.iter().cloned()), + ) + .map(|fallback| fallback.account) +} + +async fn apply_upstream_account_retry_with_deps( + deps: &CodexUpstreamDependencies, + account: &Account, + retry: UpstreamAccountRetry, +) { + match retry { + UpstreamAccountRetry::RateLimited { + retry_after_seconds, + } => { + let cooldown_until = Utc::now() + Duration::seconds(retry_after_seconds as i64); + deps.account_pool + .lock() + .await + .mark_quota_limited_until(&account.id, cooldown_until); + if record_request_attempt(deps, &account.id).await.is_err() { + tracing::warn!( + account_id = %account.id, + "failed to record rate-limited account attempt" + ); + } + } + UpstreamAccountRetry::QuotaExhausted => { + set_account_status(deps, account, AccountStatus::QuotaExhausted).await; + } + UpstreamAccountRetry::CloudflareChallenge { cooldown_seconds } => { + let cooldown_until = Utc::now() + Duration::seconds(cooldown_seconds as i64); + deps.account_pool + .lock() + .await + .set_cloudflare_cooldown_until(&account.id, cooldown_until); + } + UpstreamAccountRetry::Banned => { + set_account_status(deps, account, AccountStatus::Banned).await; + } + } +} + +async fn set_account_status( + deps: &CodexUpstreamDependencies, + account: &Account, + status: AccountStatus, +) { + if let Some(repo) = deps.account_repository.as_ref() { + if repo.set_status(&account.id, status).await.is_err() { + tracing::warn!( + account_id = %account.id, + "failed to persist upstream account status" + ); + } + } + deps.account_pool + .lock() + .await + .set_status(&account.id, status); +} + +async fn responses_websocket_stream( + deps: CodexUpstreamDependencies, + request: CodexResponsesRequest, + mut account: Account, + mut log_context: CodexRequestLogContext, +) -> Response { + let mut excluded_account_ids = Vec::new(); + let response = loop { + let response = send_codex_request_with_refresh_retry_deps( + &deps, + &request, + &account, + log_context.request_id.as_str(), + ) + .await; + deps.account_pool.lock().await.release(&account.id); + + match response { + Ok(response) => break response, + Err(error) => { + if let Some(retry) = classify_upstream_account_retry(&error) { + if request.previous_response_id.is_some() { + // previous_response_id 的历史由上游账号持有，换账号会静默丢失会话上下文。 + apply_upstream_account_retry_with_deps(&deps, &account, retry).await; + log_codex_upstream_response_with_deps( + &deps, + &log_context, + retry.status(), + EventLevel::Warn, + "v1 responses websocket history request kept on original account", + websocket_history_retry_metadata(retry, true), + ) + .await; + } else { + let fallback = apply_upstream_retry_and_acquire_fallback_with_deps( + &deps, + &account, + retry, + &request.model, + &mut excluded_account_ids, + ) + .await; + log_codex_upstream_response_with_deps( + &deps, + &log_context, + retry.status(), + EventLevel::Warn, + "v1 responses websocket upstream retrying with fallback account", + retry.metadata(true), + ) + .await; + if let Some(fallback) = fallback { + account = fallback; + log_context = log_context.with_account(&account.id); + continue; + } + } + } + let error_response = codex_client_error_response(error); + log_codex_upstream_response_with_deps( + &deps, + &log_context, + error_response.0, + EventLevel::Error, + "v1 responses websocket stream upstream request failed", + json!({"stream": true, "transport": "websocket"}), + ) + .await; + return error_response.into_response(); + } + } + }; + + if persist_upstream_cookies_with_deps(&deps, &account.id, &response.set_cookie_headers) + .await + .is_err() + { + log_codex_upstream_response_with_deps( + &deps, + &log_context, + StatusCode::INTERNAL_SERVER_ERROR, + EventLevel::Error, + "v1 responses websocket stream cookie persistence failed", + json!({"stream": true, "transport": "websocket", "cookieStoreError": true}), + ) + .await; + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(openai_error( + "Failed to persist upstream cookies", + "cookie_store_error", + )), + ) + .into_response(); + } + + let rate_limit_headers = response.rate_limit_headers.clone(); + let mut level = EventLevel::Info; + let mut message = "v1 responses websocket stream completed"; + let mut metadata = json!({ + "stream": true, + "transport": "websocket", + "usage": response.usage, + "rateLimitHeaders": rate_limit_headers.clone(), + }); + if let Some(usage) = response.usage { + if record_usage_with_deps(&deps, &account.id, usage) + .await + .is_err() + { + level = EventLevel::Warn; + message = "v1 responses websocket stream completed with usage store error"; + metadata = json!({ + "stream": true, + "transport": "websocket", + "usage": usage, + "rateLimitHeaders": rate_limit_headers.clone(), + "usageStoreError": true, + }); + } + } + log_codex_upstream_response_with_deps( + &deps, + &log_context, + StatusCode::OK, + level, + message, + metadata, + ) + .await; + + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream") + .header(CACHE_CONTROL, "no-cache") + .body(Body::from(response.body)) + .unwrap_or_else(|_| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(openai_error( + "Failed to build stream response", + "stream_response_error", + )), + ) + .into_response() + }) +} + +async fn persist_upstream_cookies_with_deps( + deps: &CodexUpstreamDependencies, + account_id: &str, + set_cookie_headers: &[String], +) -> Result<(), ()> { + let Some(cookie_repo) = deps.cookie_repository.as_ref() else { + return Ok(()); + }; + for cookie in set_cookie_headers { + cookie_repo + .capture_set_cookie(account_id, cookie) + .await + .map_err(|_| ())?; + } + Ok(()) +} + +#[derive(Clone)] +pub(crate) struct CodexRequestLogContext { + request_id: String, + account_id: String, + model: String, + stream: bool, + started_at: Instant, +} + +impl CodexRequestLogContext { + pub(crate) fn new( + request_id: &str, + account_id: &str, + model: &str, + stream: bool, + started_at: Instant, + ) -> Self { + Self { + request_id: request_id.to_string(), + account_id: account_id.to_string(), + model: model.to_string(), + stream, + started_at, + } + } + + fn latency_ms(&self) -> i64 { + self.started_at.elapsed().as_millis().min(i64::MAX as u128) as i64 + } + + pub(crate) fn with_account(&self, account_id: &str) -> Self { + Self { + request_id: self.request_id.clone(), + account_id: account_id.to_string(), + model: self.model.clone(), + stream: self.stream, + started_at: self.started_at, + } + } +} + +struct StreamAudit { + deps: CodexUpstreamDependencies, + context: CodexRequestLogContext, + account_slot: AccountSlotGuard, +} + +impl StreamAudit { + fn new( + deps: CodexUpstreamDependencies, + context: CodexRequestLogContext, + account_id: String, + ) -> Self { + let account_slot = AccountSlotGuard::new(deps.account_pool.clone(), account_id); + Self { + deps, + context, + account_slot, + } + } + + async fn complete(&mut self, body: &[u8]) { + let body = String::from_utf8_lossy(body); + let mut status = StatusCode::OK; + let mut level = EventLevel::Info; + let mut message = "v1 responses stream completed"; + let mut metadata = match extract_sse_usage(&body) { + Ok(usage) => { + if let Some(usage) = usage { + if record_usage_with_deps(&self.deps, &self.context.account_id, usage) + .await + .is_err() + { + level = EventLevel::Warn; + message = "v1 responses stream completed with usage store error"; + json!({"stream": true, "usage": usage, "usageStoreError": true}) + } else { + json!({"stream": true, "usage": usage}) + } + } else { + json!({"stream": true, "usage": null}) + } + } + Err(error) => { + level = EventLevel::Warn; + message = "v1 responses stream completed with invalid SSE usage"; + json!({"stream": true, "sseParseError": error.to_string()}) + } + }; + match responses_sse_failure(&body) { + Ok(Some(failure)) => { + // SSE 响应头已发出，HTTP 状态不能回滚；用终止事件透传给客户端，并在审计里标记上游失败。 + status = StatusCode::BAD_GATEWAY; + level = EventLevel::Error; + message = "v1 responses stream upstream SSE failed"; + failure.extend_metadata(&mut metadata); + } + Ok(None) => {} + Err(error) => { + level = EventLevel::Warn; + message = "v1 responses stream completed with invalid SSE failure metadata"; + metadata = json!({"stream": true, "sseParseError": error.to_string()}); + } + } + ensure_stream_metadata(&mut metadata, true); + log_codex_upstream_response_with_deps( + &self.deps, + &self.context, + status, + level, + message, + metadata, + ) + .await; + self.account_slot.release().await; + } + + async fn log_transport_error(&mut self, error: &reqwest::Error) { + log_codex_upstream_response_with_deps( + &self.deps, + &self.context, + StatusCode::BAD_GATEWAY, + EventLevel::Error, + "v1 responses stream transport failed", + json!({"stream": true, "transportError": error.to_string()}), + ) + .await; + self.account_slot.release().await; + } +} + +struct AccountSlotGuard { + pool: Arc>, + account_id: String, + released: bool, +} + +impl AccountSlotGuard { + fn new(pool: Arc>, account_id: String) -> Self { + Self { + pool, + account_id, + released: false, + } + } + + async fn release(&mut self) { + if self.released { + return; + } + self.pool.lock().await.release(&self.account_id); + self.released = true; + } +} + +impl Drop for AccountSlotGuard { + fn drop(&mut self) { + if self.released { + return; + } + let pool = self.pool.clone(); + let account_id = self.account_id.clone(); + let Ok(handle) = tokio::runtime::Handle::try_current() else { + return; + }; + handle.spawn(async move { + pool.lock().await.release(&account_id); + }); + } +} + +async fn log_codex_upstream_response_with_deps( + deps: &CodexUpstreamDependencies, + context: &CodexRequestLogContext, + status: StatusCode, + level: EventLevel, + message: &str, + mut metadata: Value, +) { + let Some(repo) = deps.event_logs.as_ref() else { + return; + }; + ensure_stream_metadata(&mut metadata, context.stream); + let mut event = EventLog::new("v1.response", level, message); + event.request_id = Some(context.request_id.clone()); + event.account_id = Some(context.account_id.clone()); + event.route = Some("/v1/responses".to_string()); + event.model = Some(context.model.clone()); + event.status_code = Some(i64::from(status.as_u16())); + event.latency_ms = Some(context.latency_ms()); + event.metadata = metadata; + if let Err(error) = repo.insert(event).await { + tracing::warn!(?error, "failed to insert v1 response event log"); + } +} diff --git a/src/codex/upstream/refresh.rs b/src/codex/upstream/refresh.rs new file mode 100644 index 000000000..9c3ffd62c --- /dev/null +++ b/src/codex/upstream/refresh.rs @@ -0,0 +1,147 @@ +use axum::http::StatusCode; +use secrecy::SecretString; +use serde_json::json; + +use crate::{ + codex::accounts::{ + model::{Account, AccountStatus}, + repository::TokenUpdate, + }, + codex::oauth::{RefreshFailure, TokenPair}, + codex::transport::types::CodexResponsesRequest, + logs::event::{EventLevel, EventLog}, +}; + +use super::CodexUpstreamDependencies; + +pub(super) async fn refresh_account_after_unauthorized( + deps: &CodexUpstreamDependencies, + request: &CodexResponsesRequest, + account: &Account, + request_id: &str, +) -> Option { + if !deps.config.auth.refresh_enabled { + return None; + } + let refresh_token = account.refresh_token.as_deref()?; + let refresher = deps.token_refresher.as_ref()?; + match refresher.refresh(refresh_token).await { + Ok(tokens) => persist_refreshed_account(deps, account, tokens).await, + Err(failure) => { + mark_refresh_failure(deps, account, failure, request_id, &request.model).await; + None + } + } +} + +async fn persist_refreshed_account( + deps: &CodexUpstreamDependencies, + account: &Account, + tokens: TokenPair, +) -> Option { + let repo = deps.account_repository.as_ref()?; + 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; + deps.account_pool.lock().await.insert(refreshed.clone()); + Some(refreshed) +} + +async fn mark_refresh_failure( + deps: &CodexUpstreamDependencies, + account: &Account, + failure: RefreshFailure, + request_id: &str, + model: &str, +) { + let status = status_for_refresh_failure(failure); + if let Some(status) = status { + if let Some(repo) = deps.account_repository.as_ref() { + let _ = repo.set_status(&account.id, status).await; + } + let mut updated = account.clone(); + updated.status = status; + deps.account_pool.lock().await.insert(updated); + } + log_account_refresh_failure(deps, account, failure, status, request_id, model).await; +} + +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 log_account_refresh_failure( + deps: &CodexUpstreamDependencies, + account: &Account, + failure: RefreshFailure, + status: Option, + request_id: &str, + model: &str, +) { + let Some(repo) = deps.event_logs.as_ref() else { + return; + }; + let mut event = EventLog::new( + "account.refresh", + EventLevel::Warn, + "account refresh failed after upstream 401", + ); + event.request_id = Some(request_id.to_string()); + event.account_id = Some(account.id.clone()); + event.route = Some("/v1/responses".to_string()); + event.model = Some(model.to_string()); + event.status_code = Some(i64::from(StatusCode::UNAUTHORIZED.as_u16())); + event.metadata = json!({ + "trigger": "upstream_401", + "failure": refresh_failure_value(failure), + "accountStatus": status.map(account_status_value), + }); + if let Err(error) = repo.insert(event).await { + tracing::warn!(?error, "failed to insert account refresh event log"); + } +} + +fn refresh_failure_value(failure: RefreshFailure) -> &'static str { + match failure { + RefreshFailure::InvalidGrant => "invalidGrant", + RefreshFailure::QuotaExhausted => "quotaExhausted", + RefreshFailure::Banned => "banned", + RefreshFailure::Disabled => "disabled", + RefreshFailure::Transport => "transport", + } +} + +fn account_status_value(status: AccountStatus) -> &'static str { + match status { + AccountStatus::Active => "active", + AccountStatus::Expired => "expired", + AccountStatus::QuotaExhausted => "quotaExhausted", + AccountStatus::Refreshing => "refreshing", + AccountStatus::Disabled => "disabled", + AccountStatus::Banned => "banned", + } +} diff --git a/src/codex/upstream/stream.rs b/src/codex/upstream/stream.rs new file mode 100644 index 000000000..3ba2b4329 --- /dev/null +++ b/src/codex/upstream/stream.rs @@ -0,0 +1,207 @@ +use serde_json::{json, Value}; + +use crate::codex::transport::sse::parse_sse_events; + +pub(crate) enum CollectedResponse { + Completed(Value), + Failed(ResponsesSseFailure), + MissingCompleted, +} + +#[derive(Debug, Clone)] +pub(crate) struct ResponsesSseFailure { + event: String, + message: String, + upstream_code: Option, +} + +impl ResponsesSseFailure { + fn from_event(event: &str, value: &Value) -> Self { + Self { + event: event.to_string(), + message: failure_message(value).unwrap_or_else(|| "Codex upstream SSE failed".into()), + upstream_code: failure_code(value), + } + } + + pub(crate) fn openai_error_message(&self) -> String { + format!("Codex upstream SSE failed: {}", self.message) + } + + pub(crate) fn metadata(&self, stream: bool) -> Value { + let mut metadata = json!({"stream": stream}); + self.extend_metadata(&mut metadata); + metadata + } + + pub(super) fn extend_metadata(&self, metadata: &mut Value) { + let Some(object) = metadata.as_object_mut() else { + *metadata = self.metadata(true); + return; + }; + object.insert("failureEvent".to_string(), json!(self.event)); + object.insert("failureMessage".to_string(), json!(self.message)); + if let Some(code) = &self.upstream_code { + object.insert("upstreamCode".to_string(), json!(code)); + } + } +} + +pub(crate) fn completed_response_json( + body: &str, +) -> Result { + let events = parse_sse_events(body)?; + let mut output_text = String::new(); + let mut output_items = Vec::new(); + let mut completed_response = None; + let mut failed_response = None; + for event in events { + let Ok(value) = serde_json::from_str::(&event.data) else { + continue; + }; + match event.event.as_deref() { + Some("response.output_text.delta") => { + if let Some(delta) = value.get("delta").and_then(Value::as_str) { + output_text.push_str(delta); + } + } + Some("response.output_item.done") => { + if let Some(item) = value.get("item") { + output_items.push(item.clone()); + } + } + Some("response.completed") => { + if let Some(response) = value.get("response") { + completed_response = Some(response.clone()); + } + } + Some(event_name @ ("error" | "response.failed")) if failed_response.is_none() => { + failed_response = Some(ResponsesSseFailure::from_event(event_name, &value)); + } + _ => {} + } + } + if let Some(failure) = failed_response { + return Ok(CollectedResponse::Failed(failure)); + } + let Some(mut response) = completed_response else { + return Ok(CollectedResponse::MissingCompleted); + }; + ensure_completed_response_output(&mut response, &output_items, &output_text); + sync_output_text_from_output(&mut response); + Ok(CollectedResponse::Completed(response)) +} + +pub(super) fn responses_sse_failure( + body: &str, +) -> Result, crate::codex::transport::sse::SseError> { + for event in parse_sse_events(body)? { + let Some(event_name @ ("error" | "response.failed")) = event.event.as_deref() else { + continue; + }; + let Ok(value) = serde_json::from_str::(&event.data) else { + continue; + }; + return Ok(Some(ResponsesSseFailure::from_event(event_name, &value))); + } + Ok(None) +} + +pub(super) fn ensure_stream_metadata(metadata: &mut Value, stream_value: bool) { + let Some(object) = metadata.as_object_mut() else { + *metadata = json!({"stream": stream_value}); + return; + }; + object + .entry("stream".to_string()) + .or_insert_with(|| json!(stream_value)); +} + +fn failure_message(value: &Value) -> Option { + value + .pointer("/response/error/message") + .or_else(|| value.pointer("/error/message")) + .or_else(|| value.get("message")) + .and_then(Value::as_str) + .filter(|message| !message.trim().is_empty()) + .map(ToString::to_string) +} + +fn failure_code(value: &Value) -> Option { + value + .pointer("/response/error/code") + .or_else(|| value.pointer("/error/code")) + .or_else(|| value.get("code")) + .and_then(Value::as_str) + .filter(|code| !code.trim().is_empty()) + .map(ToString::to_string) +} + +fn ensure_completed_response_output( + response: &mut Value, + output_items: &[Value], + output_text: &str, +) { + let output_is_empty = response + .get("output") + .and_then(Value::as_array) + .is_none_or(Vec::is_empty); + if !output_is_empty { + return; + } + + if !output_items.is_empty() { + response["output"] = Value::Array(output_items.to_vec()); + return; + } + if output_text.is_empty() { + return; + } + + response["output"] = json!([{ + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{ + "type": "output_text", + "text": output_text, + "annotations": [] + }] + }]); +} + +fn sync_output_text_from_output(response: &mut Value) { + let output_text_is_empty = response + .get("output_text") + .and_then(Value::as_str) + .is_none_or(str::is_empty); + if !output_text_is_empty { + return; + } + let Some(items) = response.get("output").and_then(Value::as_array) else { + return; + }; + let texts = items + .iter() + .filter_map(output_text_from_item) + .collect::>(); + if texts.is_empty() { + return; + } + response["output_text"] = Value::String(texts.join("\n\n")); +} + +fn output_text_from_item(item: &Value) -> Option { + let content = item.get("content")?.as_array()?; + let text = content + .iter() + .filter_map(|part| { + let part_type = part.get("type")?.as_str()?; + if part_type != "output_text" && part_type != "text" { + return None; + } + part.get("text")?.as_str() + }) + .collect::(); + (!text.is_empty()).then_some(text) +} diff --git a/src/codex/upstream/usage.rs b/src/codex/upstream/usage.rs new file mode 100644 index 000000000..67d6404f1 --- /dev/null +++ b/src/codex/upstream/usage.rs @@ -0,0 +1,46 @@ +use crate::{codex::accounts::repository::UsageDelta, codex::transport::usage::TokenUsage}; + +use super::CodexUpstreamDependencies; + +pub(super) async fn record_usage_with_deps( + deps: &CodexUpstreamDependencies, + account_id: &str, + usage: TokenUsage, +) -> Result<(), ()> { + let Some(repo) = deps.account_repository.as_ref() else { + return Ok(()); + }; + repo.record_usage( + account_id, + UsageDelta { + input_tokens: u64_to_i64_saturating(usage.input_tokens), + output_tokens: u64_to_i64_saturating(usage.output_tokens), + cached_tokens: u64_to_i64_saturating(usage.cached_tokens), + }, + ) + .await + .map_err(|_| ()) +} + +pub(super) async fn record_request_attempt( + deps: &CodexUpstreamDependencies, + account_id: &str, +) -> Result<(), ()> { + let Some(repo) = deps.account_repository.as_ref() else { + return Ok(()); + }; + repo.record_usage( + account_id, + UsageDelta { + input_tokens: 0, + output_tokens: 0, + cached_tokens: 0, + }, + ) + .await + .map_err(|_| ()) +} + +fn u64_to_i64_saturating(value: u64) -> i64 { + value.min(i64::MAX as u64) as i64 +} diff --git a/src/config.rs b/src/config.rs deleted file mode 100644 index 7c90b3977..000000000 --- a/src/config.rs +++ /dev/null @@ -1,216 +0,0 @@ -use std::{collections::BTreeMap, path::Path}; - -use serde::{Deserialize, Serialize}; - -pub type ConfigResult = Result; - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct AppConfig { - pub server: ServerConfig, - pub api: ApiConfig, - pub model: ModelConfig, - pub auth: AuthConfig, - pub quota: QuotaConfig, - pub usage_stats: UsageStatsConfig, - pub database: DatabaseConfig, - pub security: SecurityConfig, - pub tls: TlsConfig, - pub admin: AdminConfig, - pub logging: LoggingConfig, -} - -impl AppConfig { - pub fn load() -> ConfigResult { - let _ = dotenvy::dotenv(); - Self::load_from_dir_with_env("config", std::env::vars()) - } - - pub fn load_from_dir_with_env( - config_dir: impl AsRef, - env: I, - ) -> ConfigResult - where - K: AsRef, - V: AsRef, - I: IntoIterator, - { - let config_dir = config_dir.as_ref(); - let mut cfg: Self = config::Config::builder() - .add_source(config::File::from(config_dir.join("default.yaml")).required(true)) - .add_source(config::File::from(config_dir.join("local.yaml")).required(false)) - .add_source(config::File::from(config_dir.join("local.yml")).required(false)) - .build()? - .try_deserialize()?; - cfg.apply_env(env)?; - Ok(cfg) - } - - fn apply_env(&mut self, env: I) -> ConfigResult<()> - where - K: AsRef, - V: AsRef, - I: IntoIterator, - { - for (key, value) in env { - let key = key.as_ref(); - let value = value.as_ref(); - match key { - "CPRS_HOST" => self.server.host = value.to_string(), - "CPRS_PORT" => self.server.port = parse_env(key, value)?, - "CPRS_BASE_URL" => self.api.base_url = value.to_string(), - "CPRS_DEFAULT_MODEL" => self.model.default_model = value.to_string(), - "CPRS_DEFAULT_REASONING_EFFORT" => { - self.model.default_reasoning_effort = parse_optional_string(value); - } - "CPRS_SERVICE_TIER" => self.model.service_tier = parse_optional_string(value), - "CPRS_REFRESH_MARGIN_SECONDS" => { - self.auth.refresh_margin_seconds = parse_env(key, value)?; - } - "CPRS_REFRESH_ENABLED" => self.auth.refresh_enabled = parse_env(key, value)?, - "CPRS_REFRESH_CONCURRENCY" => { - self.auth.refresh_concurrency = parse_env(key, value)?; - } - "CPRS_MAX_CONCURRENT_PER_ACCOUNT" => { - self.auth.max_concurrent_per_account = parse_env(key, value)?; - } - "CPRS_REQUEST_INTERVAL_MS" => { - self.auth.request_interval_ms = parse_env(key, value)?; - } - "CPRS_ROTATION_STRATEGY" => self.auth.rotation_strategy = value.to_string(), - "CPRS_QUOTA_REFRESH_INTERVAL_MINUTES" => { - self.quota.refresh_interval_minutes = parse_env(key, value)?; - } - "CPRS_QUOTA_SKIP_EXHAUSTED" => { - self.quota.skip_exhausted = parse_env(key, value)?; - } - "CPRS_USAGE_HISTORY_RETENTION_DAYS" => { - self.usage_stats.history_retention_days = parse_optional_env(key, value)?; - } - "CPRS_DATABASE_URL" => self.database.url = value.to_string(), - "CPRS_MASTER_KEY_FILE" => self.security.master_key_file = value.to_string(), - "CPRS_API_KEY_PEPPER_FILE" => { - self.security.api_key_pepper_file = value.to_string(); - } - "CPRS_FORCE_HTTP11" => self.tls.force_http11 = parse_env(key, value)?, - "CPRS_ADMIN_SESSION_TTL_MINUTES" => { - self.admin.session_ttl_minutes = parse_env(key, value)?; - } - "CPRS_LOG_DIR" => self.logging.directory = value.to_string(), - "CPRS_LOG_MAX_FILE_BYTES" => self.logging.max_file_bytes = parse_env(key, value)?, - "CPRS_LOG_RETENTION_DAYS" => self.logging.retention_days = parse_env(key, value)?, - "CPRS_LOGS_ENABLED" => self.logging.enabled = parse_env(key, value)?, - "CPRS_LOGS_CAPACITY" => self.logging.capacity = parse_env(key, value)?, - "CPRS_LOGS_CAPTURE_BODY" => self.logging.capture_body = parse_env(key, value)?, - _ => {} - } - } - Ok(()) - } -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct ServerConfig { - pub host: String, - pub port: u16, -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct ApiConfig { - pub base_url: String, -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct ModelConfig { - #[serde(rename = "default")] - pub default_model: String, - pub default_reasoning_effort: Option, - #[serde(rename = "default_service_tier")] - pub service_tier: Option, - pub aliases: BTreeMap, -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct AuthConfig { - pub refresh_margin_seconds: u64, - pub refresh_enabled: bool, - pub refresh_concurrency: u32, - pub max_concurrent_per_account: usize, - pub request_interval_ms: u64, - pub rotation_strategy: String, - pub tier_priority: Vec, -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct QuotaConfig { - pub refresh_interval_minutes: u64, - pub warning_thresholds: QuotaWarningThresholds, - pub skip_exhausted: bool, -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct QuotaWarningThresholds { - pub primary: Vec, - pub secondary: Vec, -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct UsageStatsConfig { - pub history_retention_days: Option, -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct DatabaseConfig { - pub url: String, -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct SecurityConfig { - pub master_key_file: String, - pub api_key_pepper_file: String, -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct TlsConfig { - pub force_http11: bool, -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct AdminConfig { - pub session_ttl_minutes: u64, -} - -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct LoggingConfig { - pub directory: String, - pub max_file_bytes: u64, - pub retention_days: u64, - pub enabled: bool, - pub capacity: u32, - pub capture_body: bool, -} - -fn parse_env(key: &str, value: &str) -> ConfigResult -where - T: std::str::FromStr, - T::Err: std::fmt::Display, -{ - value - .parse::() - .map_err(|error| config::ConfigError::Message(format!("invalid {key}: {error}"))) -} - -fn parse_optional_env(key: &str, value: &str) -> ConfigResult> -where - T: std::str::FromStr, - T::Err: std::fmt::Display, -{ - match parse_optional_string(value) { - Some(value) => parse_env(key, &value).map(Some), - None => Ok(None), - } -} - -fn parse_optional_string(value: &str) -> Option { - let value = value.trim(); - (!value.is_empty() && value != "null").then(|| value.to_string()) -} diff --git a/src/config/loader.rs b/src/config/loader.rs new file mode 100644 index 000000000..3dd9bb937 --- /dev/null +++ b/src/config/loader.rs @@ -0,0 +1,19 @@ +use std::path::Path; + +use super::{types::AppConfig, ConfigResult}; + +impl AppConfig { + pub fn load() -> ConfigResult { + Self::load_from_dir(".") + } + + pub fn load_from_dir(config_dir: impl AsRef) -> ConfigResult { + let config_dir = config_dir.as_ref(); + ::config::Config::builder() + .add_source(::config::File::from(config_dir.join("config.yaml")).required(true)) + .add_source(::config::File::from(config_dir.join("local.yaml")).required(false)) + .add_source(::config::File::from(config_dir.join("local.yml")).required(false)) + .build()? + .try_deserialize() + } +} diff --git a/src/config/mod.rs b/src/config/mod.rs new file mode 100644 index 000000000..447e78c72 --- /dev/null +++ b/src/config/mod.rs @@ -0,0 +1,6 @@ +mod loader; +mod types; + +pub use types::*; + +pub type ConfigResult = Result; diff --git a/src/config/types.rs b/src/config/types.rs new file mode 100644 index 000000000..264a19eb3 --- /dev/null +++ b/src/config/types.rs @@ -0,0 +1,99 @@ +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct AppConfig { + pub server: ServerConfig, + pub api: ApiConfig, + pub model: ModelConfig, + pub auth: AuthConfig, + pub quota: QuotaConfig, + pub usage_stats: UsageStatsConfig, + pub database: DatabaseConfig, + pub security: SecurityConfig, + pub tls: TlsConfig, + pub admin: AdminConfig, + pub logging: LoggingConfig, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct ServerConfig { + pub host: String, + pub port: u16, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct ApiConfig { + pub base_url: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct ModelConfig { + #[serde(rename = "default")] + pub default_model: String, + pub default_reasoning_effort: Option, + #[serde(rename = "default_service_tier")] + pub service_tier: Option, + pub aliases: BTreeMap, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct AuthConfig { + pub refresh_margin_seconds: u64, + pub refresh_enabled: bool, + pub refresh_concurrency: u32, + pub max_concurrent_per_account: usize, + pub request_interval_ms: u64, + pub rotation_strategy: String, + pub tier_priority: Vec, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct QuotaConfig { + pub refresh_interval_minutes: u64, + pub warning_thresholds: QuotaWarningThresholds, + pub skip_exhausted: bool, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct QuotaWarningThresholds { + pub primary: Vec, + pub secondary: Vec, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct UsageStatsConfig { + pub history_retention_days: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct DatabaseConfig { + pub url: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct SecurityConfig { + pub master_key_file: String, + pub api_key_pepper_file: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct TlsConfig { + pub force_http11: bool, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct AdminConfig { + pub session_ttl_minutes: u64, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct LoggingConfig { + pub directory: String, + pub max_file_bytes: u64, + pub retention_days: u64, + pub enabled: bool, + pub capacity: u32, + pub capture_body: bool, +} diff --git a/src/http/admin.rs b/src/http/admin.rs deleted file mode 100644 index 628854f97..000000000 --- a/src/http/admin.rs +++ /dev/null @@ -1,5169 +0,0 @@ -use std::time::{Duration as StdDuration, Instant}; - -use axum::{ - body::Bytes, - extract::{Path, Query, State}, - http::{ - header::{HeaderValue, HOST, LOCATION, SET_COOKIE}, - HeaderMap, StatusCode, - }, - response::{IntoResponse, Response}, - Extension, Json, -}; -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::{ - accounts::{ - model::{Account, AccountStatus}, - repository::{ - AccountClaimsUpdate, AccountRepositoryError, AccountUsageListRecord, - AccountUsageSummary, NewAccount, StoredAccount, StoredAccountMetadata, TokenUpdate, - }, - }, - auth::api_key_repository::StoredClientApiKey, - auth::{ - admin_session::verify_admin_password, - cli_import::{default_codex_home, read_cli_auth_from_home}, - oauth::{DeviceCode, OAuthClient, OAuthConfig, OAuthError, PkceSession}, - refresh::RefreshFailure, - token::TokenPair, - }, - codex::client::{ - build_reqwest_client, CodexBackendClient, CodexClientError, CodexRequestContext, - }, - config::{AppConfig, QuotaWarningThresholds}, - fingerprint::model::Fingerprint, - http::{auth::admin_session_id, middleware::RequestId}, - models::catalog::{ModelCatalog, ModelPlanSnapshot}, - pagination::{clamp_limit, Page}, - state::AppState, -}; - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminEnvelope { - pub code: u32, - pub message: String, - pub data: T, - pub request_id: String, -} - -impl AdminEnvelope { - pub fn new( - code: u32, - message: impl Into, - data: T, - request_id: impl Into, - ) -> Self { - // body code 给前端做业务分支，HTTP status 仍然是传输层真相。 - Self { - code, - message: message.into(), - data, - request_id: request_id.into(), - } - } - - pub fn ok(data: T, request_id: impl Into) -> Self { - Self::new(200, "OK", data, request_id) - } -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct PageMeta { - pub limit: u32, - pub next_cursor: Option, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminPageEnvelope { - pub code: u32, - pub message: String, - pub data: Vec, - pub page: PageMeta, - pub request_id: String, -} - -impl AdminPageEnvelope { - pub fn new( - code: u32, - message: impl Into, - page: Page, - limit: u32, - request_id: impl Into, - ) -> Self { - let Page { items, next_cursor } = page; - Self { - code, - message: message.into(), - data: items, - page: PageMeta { limit, next_cursor }, - request_id: request_id.into(), - } - } - - pub fn ok(page: Page, limit: u32, request_id: impl Into) -> Self { - Self::new(200, "OK", page, limit, request_id) - } -} - -#[derive(Debug, Clone)] -pub struct AdminResponse { - pub status: StatusCode, - pub body: T, -} - -impl AdminResponse { - pub fn new(status: StatusCode, body: T) -> Self { - Self { status, body } - } -} - -impl IntoResponse for AdminResponse -where - T: Serialize, -{ - fn into_response(self) -> Response { - (self.status, Json(self.body)).into_response() - } -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct LogsQuery { - pub cursor: Option, - pub limit: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountsQuery { - pub cursor: Option, - pub limit: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountExportQuery { - pub ids: Option, - pub format: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ApiKeysQuery { - pub cursor: Option, - pub limit: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ApiKeyExportQuery { - pub ids: 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 { - pub password: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct LoginData { - pub expires_at: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminAuthStatusData { - pub authenticated: bool, - pub user: Option, - pub pool: AdminAuthPoolSummaryData, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminAuthUserData { - pub id: String, - pub email: Option, - pub account_id: Option, - pub user_id: Option, - pub label: Option, - pub plan_type: Option, - pub status: String, - pub access_token_expires_at: Option, -} - -#[derive(Debug, Default, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminAuthPoolSummaryData { - pub total: usize, - pub active: usize, - pub expired: usize, - pub quota_exhausted: usize, - pub refreshing: usize, - pub disabled: usize, - pub banned: usize, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminAuthLogoutData { - pub success: bool, - pub deleted: u64, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminAuthDeviceLoginData { - pub user_code: String, - pub verification_uri: String, - pub verification_uri_complete: String, - pub device_code: String, - pub expires_in: u64, - pub interval: u64, -} - -impl From for AdminAuthDeviceLoginData { - fn from(device: DeviceCode) -> Self { - Self { - user_code: device.user_code, - verification_uri: device.verification_uri, - verification_uri_complete: device.verification_uri_complete, - device_code: device.device_code, - expires_in: device.expires_in, - interval: device.interval, - } - } -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminAuthDevicePollData { - pub success: bool, - pub pending: bool, - pub code: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminAuthLoginStartData { - pub auth_url: String, - pub state: String, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminAuthCodeRelayRequest { - pub callback_url: String, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminAuthCallbackQuery { - pub code: Option, - pub state: Option, - pub error: Option, - pub error_description: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminAuthCodeRelayData { - pub success: bool, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminSettingsData { - pub default_model: String, - pub default_reasoning_effort: Option, - pub service_tier: Option, - pub model_aliases: std::collections::BTreeMap, - pub refresh_enabled: bool, - pub refresh_margin_seconds: u64, - pub refresh_concurrency: u32, - pub max_concurrent_per_account: usize, - pub request_interval_ms: u64, - pub rotation_strategy: String, - pub tier_priority: Vec, - pub quota_refresh_interval_minutes: u64, - pub quota_warning_thresholds: QuotaWarningThresholds, - pub quota_skip_exhausted: bool, - pub logs_enabled: bool, - pub logs_capacity: u32, - pub logs_capture_body: bool, - pub usage_history_retention_days: Option, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum AccountImportFormat { - Native, - Sub2Api, - CodexCli, -} - -impl AccountImportFormat { - fn as_str(self) -> &'static str { - match self { - Self::Native => "native", - Self::Sub2Api => "sub2api", - Self::CodexCli => "codex_cli", - } - } -} - -#[derive(Debug, Clone)] -struct ParsedAccountImportPayload { - accounts: Vec, - source_format: AccountImportFormat, -} - -#[derive(Debug, Clone)] -struct AccountImportEntry { - pub id: Option, - pub email: Option, - pub account_id: Option, - pub user_id: Option, - pub label: Option, - pub plan_type: Option, - pub token: Option, - pub refresh_token: Option, - pub status: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountImportData { - pub imported: u32, - pub skipped: u32, - pub source_format: String, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct CreateAccountRequest { - pub token: Option, - pub refresh_token: Option, -} - -#[derive(Debug, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ImportCliAuthRequest { - pub codex_home: Option, -} - -#[derive(Debug, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct HealthCheckRequest { - pub ids: Option>, - #[serde(alias = "stagger_ms")] - pub stagger_ms: Option, - pub concurrency: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct HealthCheckData { - pub summary: HealthCheckSummary, - pub results: Vec, -} - -#[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, - pub previous_status: String, - pub result: String, - pub status: Option, - pub error: Option, - pub duration_ms: Option, -} - -#[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, - Sub2Api, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdateAccountLabelRequest { - pub label: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdateAccountLabelData { - pub id: String, - pub label: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdateAccountStatusRequest { - pub status: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdateAccountStatusData { - pub id: String, - pub status: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct DeleteAccountData { - pub deleted: bool, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct BatchDeleteAccountsRequest { - pub ids: Vec, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct BatchDeleteAccountsData { - pub deleted: u32, - pub not_found: Vec, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct BatchUpdateAccountStatusRequest { - pub ids: Vec, - pub status: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct BatchUpdateAccountStatusData { - pub updated: u32, - pub not_found: Vec, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SetAccountCookiesRequest { - pub cookies: Value, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountCookiesData { - pub cookies: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct DeleteAccountCookiesData { - pub deleted: bool, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct RefreshModelsData { - pub refreshed_plans: usize, - pub model_count: usize, - pub failed_plans: usize, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct CreateApiKeyRequest { - pub name: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct ClientApiKeyData { - pub id: String, - pub name: String, - pub label: Option, - pub prefix: String, - pub enabled: bool, - pub created_at: String, - pub last_used_at: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct CreatedClientApiKeyData { - pub id: String, - pub name: String, - pub label: Option, - pub prefix: String, - pub enabled: bool, - pub created_at: String, - pub last_used_at: Option, - pub plaintext: String, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdateClientApiKeyStatusRequest { - pub status: String, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdateClientApiKeyLabelRequest { - pub label: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdateClientApiKeyStatusData { - pub id: String, - pub enabled: bool, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct BatchDeleteClientApiKeysRequest { - pub ids: Vec, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct BatchDeleteClientApiKeysData { - pub deleted: u32, - pub not_found: Vec, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct DeleteClientApiKeyData { - pub deleted: bool, -} - -#[derive(Debug, Clone)] -struct ClientApiKeyImportEntry { - pub source_id: Option, - pub source_prefix: Option, - pub name: String, - pub label: Option, - pub enabled: bool, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct ClientApiKeyExportData { - pub source_format: &'static str, - pub rotation_required: bool, - pub api_keys: Vec, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct ClientApiKeyExportEntry { - pub id: String, - pub name: String, - pub label: Option, - pub prefix: String, - pub enabled: bool, - pub created_at: String, - pub last_used_at: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct ClientApiKeyImportData { - pub imported: u32, - pub skipped: u32, - pub rotated: bool, - pub keys: Vec, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct ImportedClientApiKeyData { - pub source_id: Option, - pub source_prefix: Option, - pub id: String, - pub name: String, - pub label: Option, - pub prefix: String, - pub enabled: bool, - pub created_at: String, - pub last_used_at: Option, - pub plaintext: String, -} - -impl From for ClientApiKeyData { - fn from(key: StoredClientApiKey) -> Self { - Self { - id: key.id, - name: key.name, - label: key.label, - prefix: key.prefix, - enabled: key.enabled, - created_at: key.created_at, - last_used_at: key.last_used_at, - } - } -} - -impl From for ClientApiKeyExportEntry { - fn from(key: StoredClientApiKey) -> Self { - Self { - id: key.id, - name: key.name, - label: key.label, - prefix: key.prefix, - enabled: key.enabled, - created_at: key.created_at, - last_used_at: key.last_used_at, - } - } -} - -impl CreatedClientApiKeyData { - fn new(key: StoredClientApiKey, plaintext: String) -> Self { - Self { - id: key.id, - name: key.name, - label: key.label, - prefix: key.prefix, - enabled: key.enabled, - created_at: key.created_at, - last_used_at: key.last_used_at, - plaintext, - } - } -} - -impl ImportedClientApiKeyData { - fn new( - key: StoredClientApiKey, - plaintext: String, - source_id: Option, - source_prefix: Option, - ) -> Self { - Self { - source_id, - source_prefix, - id: key.id, - name: key.name, - label: key.label, - prefix: key.prefix, - enabled: key.enabled, - created_at: key.created_at, - last_used_at: key.last_used_at, - plaintext, - } - } -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminAccountData { - pub id: String, - pub email: Option, - pub account_id: Option, - pub user_id: Option, - pub label: Option, - pub plan_type: Option, - pub status: String, - pub access_token_expires_at: Option, - pub added_at: String, - pub updated_at: String, -} - -impl From for AdminAccountData { - fn from(account: StoredAccountMetadata) -> Self { - Self { - id: account.id, - email: account.email, - account_id: account.account_id, - user_id: account.user_id, - label: account.label, - plan_type: account.plan_type, - status: account_status_value(account.status).to_string(), - access_token_expires_at: account - .access_token_expires_at - .map(|value| value.to_rfc3339()), - added_at: account.added_at.to_rfc3339(), - updated_at: account.updated_at.to_rfc3339(), - } - } -} - -fn account_auth_user(account: &StoredAccountMetadata) -> AdminAuthUserData { - AdminAuthUserData { - id: account.id.clone(), - email: account.email.clone(), - account_id: account.account_id.clone(), - user_id: account.user_id.clone(), - label: account.label.clone(), - plan_type: account.plan_type.clone(), - status: account_status_value(account.status).to_string(), - access_token_expires_at: account - .access_token_expires_at - .map(|value| value.to_rfc3339()), - } -} - -fn account_auth_pool_summary(accounts: &[StoredAccountMetadata]) -> AdminAuthPoolSummaryData { - let mut summary = AdminAuthPoolSummaryData { - total: accounts.len(), - ..AdminAuthPoolSummaryData::default() - }; - for account in accounts { - match account.status { - AccountStatus::Active => summary.active += 1, - AccountStatus::Expired => summary.expired += 1, - AccountStatus::QuotaExhausted => summary.quota_exhausted += 1, - AccountStatus::Refreshing => summary.refreshing += 1, - AccountStatus::Disabled => summary.disabled += 1, - AccountStatus::Banned => summary.banned += 1, - } - } - summary -} - -#[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 { - default_model: config.model.default_model.clone(), - default_reasoning_effort: config.model.default_reasoning_effort.clone(), - service_tier: config.model.service_tier.clone(), - model_aliases: config.model.aliases.clone(), - refresh_enabled: config.auth.refresh_enabled, - refresh_margin_seconds: config.auth.refresh_margin_seconds, - refresh_concurrency: config.auth.refresh_concurrency, - max_concurrent_per_account: config.auth.max_concurrent_per_account, - request_interval_ms: config.auth.request_interval_ms, - rotation_strategy: config.auth.rotation_strategy.clone(), - tier_priority: config.auth.tier_priority.clone(), - quota_refresh_interval_minutes: config.quota.refresh_interval_minutes, - quota_warning_thresholds: config.quota.warning_thresholds.clone(), - quota_skip_exhausted: config.quota.skip_exhausted, - logs_enabled: config.logging.enabled, - logs_capacity: config.logging.capacity, - logs_capture_body: config.logging.capture_body, - usage_history_retention_days: config.usage_stats.history_retention_days, - } - } -} - -pub async fn login( - State(state): State, - Extension(request_id): Extension, - Json(payload): Json, -) -> 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(); - }; - - let admin = match load_first_admin(pool).await { - Ok(Some(admin)) => admin, - Ok(None) => { - return AdminResponse::new( - StatusCode::UNAUTHORIZED, - AdminEnvelope::new(40102, "Admin password invalid", (), request_id), - ) - .into_response(); - } - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to load admin user", (), request_id), - ) - .into_response(); - } - }; - - // admin 登录只校验管理员密码；客户端 cpr_ API Key 不能参与后台登录。 - match verify_admin_password(&payload.password, &admin.password_hash) { - Ok(true) => {} - Ok(false) => { - return AdminResponse::new( - StatusCode::UNAUTHORIZED, - AdminEnvelope::new(40102, "Admin password invalid", (), request_id), - ) - .into_response(); - } - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to verify admin password", (), request_id), - ) - .into_response(); - } - } - - let ttl_minutes = state.config().admin.session_ttl_minutes; - let Ok(ttl_minutes_i64) = i64::try_from(ttl_minutes) else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Admin session ttl is invalid", (), request_id), - ) - .into_response(); - }; - let expires_at = Utc::now() + Duration::minutes(ttl_minutes_i64); - let session_id = format!("sess_{}", Uuid::new_v4().simple()); - if create_admin_session(pool, &session_id, &admin.id, expires_at) - .await - .is_err() - { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to create admin session", (), request_id), - ) - .into_response(); - } - - let Some(cookie) = admin_session_set_cookie(&session_id, ttl_minutes) else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "Failed to create admin session cookie", - (), - request_id, - ), - ) - .into_response(); - }; - let mut response = AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - LoginData { - expires_at: expires_at.to_rfc3339(), - }, - request_id, - ), - ) - .into_response(); - response.headers_mut().insert(SET_COOKIE, cookie); - response -} - -pub async fn auth_status( - 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_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "Account repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - - match repo.list_all_metadata().await { - Ok(accounts) => { - let summary = account_auth_pool_summary(&accounts); - let user = accounts.first().map(account_auth_user); - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - AdminAuthStatusData { - authenticated: summary.total > 0, - user, - pool: summary, - }, - request_id, - ), - ) - .into_response() - } - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "Failed to inspect account auth status", - (), - request_id, - ), - ) - .into_response(), - } -} - -pub async fn auth_logout( - 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_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "Account repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - - match repo.delete_all().await { - Ok(deleted) => { - // 账号 logout 要同时清掉调度池；SQLite 外键会级联 usage/cookies 等账号附属数据。 - state.account_pool().lock().await.clear(); - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - AdminAuthLogoutData { - success: true, - deleted, - }, - request_id, - ), - ) - .into_response() - } - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to clear accounts", (), request_id), - ) - .into_response(), - } -} - -pub async fn auth_device_login( - 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(oauth_client) = state.oauth_client() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "OAuth client is not initialized", (), request_id), - ) - .into_response(); - }; - - match oauth_client.request_device_code().await { - Ok(device) => AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AdminAuthDeviceLoginData::from(device), request_id), - ) - .into_response(), - Err(error) => oauth_error_response(error, request_id, "Device code request failed"), - } -} - -pub async fn auth_device_poll( - State(state): State, - Path(device_code): Path, - 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(oauth_client) = state.oauth_client() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "OAuth client is not initialized", (), request_id), - ) - .into_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 device_code = device_code.trim(); - if device_code.is_empty() { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "Device code is required", (), request_id), - ) - .into_response(); - } - - match oauth_client.poll_device_token(device_code).await { - Ok(tokens) => match store_validated_account_import( - &state, - &repo, - Some(tokens.access_token), - tokens.refresh_token, - ) - .await - { - Ok(_) => AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - AdminAuthDevicePollData { - success: true, - pending: false, - code: None, - }, - request_id, - ), - ) - .into_response(), - Err(error) => validated_account_import_error_response(error, request_id), - }, - Err(error) => { - if let Some(code) = error.pending_code() { - return AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - AdminAuthDevicePollData { - success: false, - pending: true, - code: Some(code.to_string()), - }, - request_id, - ), - ) - .into_response(); - } - oauth_error_response(error, request_id, "Device authorization failed") - } - } -} - -pub async fn auth_login_start( - 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 return_host = request_host(&headers, state.config()); - let login = state - .oauth_sessions() - .lock() - .await - .start_login(&return_host, &OAuthConfig::codex_default()); - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - AdminAuthLoginStartData { - auth_url: login.auth_url, - state: login.state, - }, - request_id, - ), - ) - .into_response() -} - -pub async fn auth_code_relay( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Json(payload): Json, -) -> 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 Ok((code, oauth_state)) = parse_callback_url(&payload.callback_url) else { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "Invalid OAuth callback URL", (), request_id), - ) - .into_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 Some(oauth_client) = state.oauth_client() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "OAuth client is not initialized", (), request_id), - ) - .into_response(); - }; - - let session = match acquire_pkce_session(&state, &oauth_state, &request_id).await { - Ok(Some(session)) => session, - Ok(None) => { - return AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AdminAuthCodeRelayData { success: true }, request_id), - ) - .into_response(); - } - Err(response) => return response, - }; - match exchange_pkce_session( - &state, - &repo, - oauth_client.as_ref(), - &oauth_state, - &code, - session, - &request_id, - ) - .await - { - Ok(_) => AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AdminAuthCodeRelayData { success: true }, request_id), - ) - .into_response(), - Err(response) => response, - } -} - -pub async fn auth_callback( - 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; - } - if query.error.is_some() { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new( - 40001, - query - .error_description - .unwrap_or_else(|| "OAuth error".to_string()), - (), - request_id, - ), - ) - .into_response(); - } - let (Some(code), Some(oauth_state)) = (query.code.as_deref(), query.state.as_deref()) else { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "Missing OAuth code or state", (), request_id), - ) - .into_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 Some(oauth_client) = state.oauth_client() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "OAuth client is not initialized", (), request_id), - ) - .into_response(); - }; - let session = match acquire_pkce_session(&state, oauth_state, &request_id).await { - Ok(Some(session)) => session, - Ok(None) => return redirect_to_host(&request_host(&headers, state.config()), request_id), - Err(response) => return response, - }; - let return_host = session.return_host.clone(); - match exchange_pkce_session( - &state, - &repo, - oauth_client.as_ref(), - oauth_state, - code, - session, - &request_id, - ) - .await - { - Ok(_) => redirect_to_host(&return_host, request_id), - Err(response) => response, - } -} - -pub async fn logs( - 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 limit = clamp_limit(query.limit.unwrap_or(50)); - let Some(repo) = state.event_logs() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "Event log repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - match repo.list(query.cursor, limit).await { - Ok(page) => AdminResponse::new( - StatusCode::OK, - AdminPageEnvelope::ok(page, limit, request_id), - ) - .into_response(), - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to list event logs", (), request_id), - ) - .into_response(), - } -} - -pub async fn settings( - 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; - } - - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AdminSettingsData::from_config(state.config()), request_id), - ) - .into_response() -} - -pub async fn refresh_models( - 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(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(); - }; - let Some(model_repo) = state.model_snapshot_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Model repository is not initialized", (), request_id), - ) - .into_response(); - }; - let accounts = match account_repo.list_pool_accounts().await { - Ok(accounts) => accounts, - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to list accounts", (), request_id), - ) - .into_response(); - } - }; - let plan_accounts = distinct_active_plan_accounts(accounts); - if plan_accounts.is_empty() { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new( - 40001, - "No accounts available for model refresh", - (), - request_id, - ), - ) - .into_response(); - } - let client = match build_reqwest_client(state.config().tls.force_http11) { - Ok(client) => CodexBackendClient::new( - client, - state.config().api.base_url.clone(), - Fingerprint::default_codex_desktop(), - ), - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to build Codex client", (), request_id), - ) - .into_response(); - } - }; - - let mut refreshed_plans = 0usize; - let mut model_count = 0usize; - let mut failed_plans = 0usize; - for (plan_type, account) in plan_accounts { - let context = CodexRequestContext { - access_token: &account.access_token, - account_id: account.account_id.as_deref(), - request_id: &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: None, - }; - let entries = match client.fetch_models(context).await { - Ok(entries) if !entries.is_empty() => entries, - Ok(_) => { - failed_plans += 1; - continue; - } - Err(error) => { - tracing::warn!(?error, plan_type, "failed to refresh backend models"); - failed_plans += 1; - continue; - } - }; - let snapshot = ModelPlanSnapshot::from_backend_entries(plan_type, entries); - model_count += snapshot.models.len(); - if model_repo.replace_plan_snapshot(&snapshot).await.is_err() { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to store model snapshot", (), request_id), - ) - .into_response(); - } - refreshed_plans += 1; - } - - let data = RefreshModelsData { - refreshed_plans, - model_count, - failed_plans, - }; - if refreshed_plans == 0 { - return AdminResponse::new( - StatusCode::BAD_GATEWAY, - AdminEnvelope::new(50201, "Failed to refresh backend models", data, request_id), - ) - .into_response(); - } - match model_repo.list_plan_snapshots().await { - Ok(snapshots) => { - let allowlist = - ModelCatalog::from_config_and_snapshots(&state.config().model, &snapshots) - .model_plan_allowlist(); - // 刷新后的 model -> plans 要立即同步给调度器，避免新模型被分配到不支持的账号 plan。 - state - .account_pool() - .lock() - .await - .set_model_plan_allowlist(allowlist); - } - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to load model snapshots", (), request_id), - ) - .into_response(); - } - } - - AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data, request_id)).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, - 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_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "Account repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - - let limit = clamp_limit(query.limit.unwrap_or(50)); - match repo.list_metadata(query.cursor, limit).await { - Ok(page) => { - let Page { items, next_cursor } = page; - let page = Page { - items: items.into_iter().map(AdminAccountData::from).collect(), - next_cursor, - }; - AdminResponse::new( - StatusCode::OK, - AdminPageEnvelope::ok(page, limit, request_id), - ) - .into_response() - } - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to list accounts", (), request_id), - ) - .into_response(), - } -} - -pub async fn export_accounts( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Query(query): Query, -) -> Response { - let request_id = request_id.as_str().to_string(); - let format = match parse_account_export_format(query.format.as_deref()) { - Ok(format) => format, - Err(message) => { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, message, (), 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 ids = account_export_ids(query.ids.as_deref()); - let accounts = if ids.is_empty() { - match repo.list_all().await { - Ok(accounts) => accounts, - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to export accounts", (), request_id), - ) - .into_response(); - } - } - } else { - let mut accounts = Vec::with_capacity(ids.len()); - for id in ids { - match repo.get(&id).await { - Ok(Some(account)) => accounts.push(account), - Ok(None) => {} - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to export accounts", (), request_id), - ) - .into_response(); - } - } - } - accounts - }; - - // 账号导出会返回可重新导入的 OAuth token；只允许 admin session 访问，不写入日志。 - let data = match format { - AccountExportFormat::Native => native_account_export(accounts), - AccountExportFormat::Sub2Api => sub2api_account_export(accounts), - }; - AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data, request_id)).into_response() -} - -pub async fn create_account( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Json(payload): Json, -) -> 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 stored = - match store_validated_account_import(&state, &repo, payload.token, payload.refresh_token) - .await - { - Ok(stored) => stored, - Err(error) => return validated_account_import_error_response(error, request_id), - }; - - // 手动添加账号的响应只返回可展示元数据，OAuth token 永不回显。 - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(admin_account_data_from_stored(stored), request_id), - ) - .into_response() -} - -pub async fn health_check_accounts( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - payload: Option>, -) -> 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::>() - .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, - Extension(request_id): Extension, - headers: HeaderMap, - Path(account_id): Path, -) -> 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, - Extension(request_id): Extension, - headers: HeaderMap, - Path(account_id): Path, -) -> 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, - Extension(request_id): Extension, - headers: HeaderMap, - Path(account_id): Path, -) -> 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, "a.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, - Extension(request_id): Extension, - headers: HeaderMap, - Path(account_id): Path, - Json(payload): Json, -) -> 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(); - }; - if payload - .label - .as_ref() - .is_some_and(|label| label.chars().count() > 64) - { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new( - 40001, - "Account label must be 64 characters or fewer", - (), - request_id, - ), - ) - .into_response(); - } - - match repo.set_label(&account_id, payload.label.clone()).await { - Ok(true) => { - // 管理后台改名后要同步内存调度池，否则列表和实际调度状态会短暂不一致。 - state - .account_pool() - .lock() - .await - .set_label(&account_id, payload.label.clone()); - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - UpdateAccountLabelData { - id: account_id, - label: payload.label, - }, - request_id, - ), - ) - .into_response() - } - Ok(false) => AdminResponse::new( - StatusCode::NOT_FOUND, - AdminEnvelope::new(40401, "Account not found", (), request_id), - ) - .into_response(), - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to update account label", (), request_id), - ) - .into_response(), - } -} - -pub async fn update_account_status( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Path(account_id): Path, - Json(payload): Json, -) -> 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 status = match parse_admin_account_status(&payload.status) { - Ok(status) => status, - Err(message) => { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, message, (), request_id), - ) - .into_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(); - }; - - match repo.set_status(&account_id, status).await { - Ok(true) => { - if sync_runtime_account_status(&state, &repo, &account_id, status) - .await - .is_err() - { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to sync account status", (), request_id), - ) - .into_response(); - } - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - UpdateAccountStatusData { - id: account_id, - status: account_status_value(status).to_string(), - }, - request_id, - ), - ) - .into_response() - } - Ok(false) => AdminResponse::new( - StatusCode::NOT_FOUND, - AdminEnvelope::new(40401, "Account not found", (), request_id), - ) - .into_response(), - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to update account status", (), request_id), - ) - .into_response(), - } -} - -pub async fn delete_account( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Path(account_id): Path, -) -> 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(); - }; - - match repo.delete(&account_id).await { - Ok(true) => { - // DB 外键会级联清理账号关联数据；内存池仍需立即摘除，避免删除后继续被调度。 - state.account_pool().lock().await.remove(&account_id); - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(DeleteAccountData { deleted: true }, request_id), - ) - .into_response() - } - Ok(false) => AdminResponse::new( - StatusCode::NOT_FOUND, - AdminEnvelope::new(40401, "Account not found", (), request_id), - ) - .into_response(), - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to delete account", (), request_id), - ) - .into_response(), - } -} - -pub async fn batch_delete_accounts( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Json(payload): Json, -) -> Response { - let request_id = request_id.as_str().to_string(); - if payload.ids.is_empty() { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "Account ids are required", (), 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 mut deleted = 0u32; - let mut not_found = Vec::new(); - for account_id in payload.ids { - match repo.delete(&account_id).await { - Ok(true) => { - state.account_pool().lock().await.remove(&account_id); - deleted += 1; - } - Ok(false) => not_found.push(account_id), - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to delete account", (), request_id), - ) - .into_response(); - } - } - } - - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(BatchDeleteAccountsData { deleted, not_found }, request_id), - ) - .into_response() -} - -pub async fn batch_update_account_status( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Json(payload): Json, -) -> Response { - let request_id = request_id.as_str().to_string(); - if payload.ids.is_empty() { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "Account ids are required", (), request_id), - ) - .into_response(); - } - let status = match parse_admin_account_status(&payload.status) { - Ok(status) => status, - Err(message) => { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, message, (), 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 mut updated = 0u32; - let mut not_found = Vec::new(); - for account_id in payload.ids { - match repo.set_status(&account_id, status).await { - Ok(true) => { - if sync_runtime_account_status(&state, &repo, &account_id, status) - .await - .is_err() - { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to sync account status", (), request_id), - ) - .into_response(); - } - updated += 1; - } - Ok(false) => not_found.push(account_id), - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to update account status", (), request_id), - ) - .into_response(); - } - } - } - - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - BatchUpdateAccountStatusData { updated, not_found }, - request_id, - ), - ) - .into_response() -} - -pub async fn get_account_cookies( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Path(account_id): Path, -) -> 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(cookie_repo) = state.cookie_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "Cookie repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - - match cookie_repo.cookie_header(&account_id, "chatgpt.com").await { - Ok(cookies) => AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountCookiesData { cookies }, request_id), - ) - .into_response(), - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to load account cookies", (), request_id), - ) - .into_response(), - } -} - -pub async fn set_account_cookies( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Path(account_id): Path, - Json(payload): Json, -) -> Response { - let request_id = request_id.as_str().to_string(); - let cookie_header = match admin_cookie_header(&payload.cookies) { - Ok(cookie_header) => cookie_header, - Err(message) => { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, message, (), 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(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(cookie_repo) = state.cookie_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "Cookie repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - - // 管理端粘贴的是浏览器 Cookie header；保存时逐项加密，读取时只回放当前账号自己的 Cookie。 - match cookie_repo - .set_cookie_header(&account_id, &cookie_header) - .await - { - Ok(0) => AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "No valid cookies found", (), request_id), - ) - .into_response(), - Ok(_) => match cookie_repo.cookie_header(&account_id, "chatgpt.com").await { - Ok(cookies) => AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountCookiesData { cookies }, request_id), - ) - .into_response(), - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to load account cookies", (), request_id), - ) - .into_response(), - }, - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to store account cookies", (), request_id), - ) - .into_response(), - } -} - -pub async fn delete_account_cookies( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Path(account_id): Path, -) -> 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(cookie_repo) = state.cookie_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "Cookie repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - - match cookie_repo.delete_account_cookies(&account_id).await { - Ok(_) => AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(DeleteAccountCookiesData { deleted: true }, request_id), - ) - .into_response(), - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to delete account cookies", (), request_id), - ) - .into_response(), - } -} - -pub async fn import_accounts( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Json(payload): Json, -) -> 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 parsed = parse_account_import_payload(&payload); - if parsed.accounts.is_empty() { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "No importable accounts found", (), request_id), - ) - .into_response(); - } - - let mut imported = 0u32; - let mut skipped = 0u32; - for entry in parsed.accounts { - match store_import_account_entry(&state, &repo, entry).await { - Ok(StoredImportAccount::Imported) => { - imported += 1; - } - Ok(StoredImportAccount::Skipped) => { - skipped += 1; - } - Err(StoreImportAccountError::Inspect) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to inspect account", (), request_id), - ) - .into_response(); - } - Err(StoreImportAccountError::Invalid(message)) => { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, message, (), request_id), - ) - .into_response(); - } - Err(StoreImportAccountError::Insert) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to import account", (), request_id), - ) - .into_response(); - } - } - } - - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - AccountImportData { - imported, - skipped, - source_format: parsed.source_format.as_str().to_string(), - }, - request_id, - ), - ) - .into_response() -} - -pub async fn import_cli_auth( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - body: Bytes, -) -> 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 payload = if body.is_empty() { - ImportCliAuthRequest::default() - } else { - match serde_json::from_slice::(&body) { - Ok(payload) => payload, - Err(_) => { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "Invalid CLI import request", (), request_id), - ) - .into_response(); - } - } - }; - let codex_home = match empty_to_none(payload.codex_home) { - Some(path) => std::path::PathBuf::from(path), - None => match default_codex_home() { - Ok(path) => path, - Err(error) => { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, error.to_string(), (), request_id), - ) - .into_response(); - } - }, - }; - let cli_auth = match read_cli_auth_from_home(&codex_home) { - Ok(auth) => auth, - Err(error) => { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, error.to_string(), (), request_id), - ) - .into_response(); - } - }; - let _stored = match store_validated_account_import( - &state, - &repo, - Some(cli_auth.access_token().to_string()), - cli_auth.refresh_token().map(str::to_string), - ) - .await - { - Ok(stored) => stored, - Err(error) => return validated_account_import_error_response(error, request_id), - }; - - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - AccountImportData { - imported: 1, - skipped: 0, - source_format: AccountImportFormat::CodexCli.as_str().to_string(), - }, - request_id, - ), - ) - .into_response() -} - -pub async fn api_keys( - 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.client_api_key_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "API key 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(ClientApiKeyData::from).collect(), - next_cursor, - }; - AdminResponse::new( - StatusCode::OK, - AdminPageEnvelope::ok(page, limit, request_id), - ) - .into_response() - } - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to list API keys", (), request_id), - ) - .into_response(), - } -} - -pub async fn export_api_keys( - 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.client_api_key_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "API key repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - - let ids = account_export_ids(query.ids.as_deref()); - let keys = if ids.is_empty() { - match repo.list_all().await { - Ok(keys) => keys, - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to export API keys", (), request_id), - ) - .into_response(); - } - } - } else { - let mut keys = Vec::with_capacity(ids.len()); - for id in ids { - match repo.get(&id).await { - Ok(Some(key)) => keys.push(key), - Ok(None) => {} - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to export API keys", (), request_id), - ) - .into_response(); - } - } - } - keys - }; - - // 安全边界：本地 cpr_ key 只导出可展示元数据，绝不导出 plaintext、key_hash 或 pepper。 - let data = ClientApiKeyExportData { - source_format: "rustLocalClientApiKeys", - rotation_required: true, - api_keys: keys - .into_iter() - .map(ClientApiKeyExportEntry::from) - .collect(), - }; - AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data, request_id)).into_response() -} - -pub async fn import_api_keys( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Json(payload): Json, -) -> 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.client_api_key_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "API key repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - let Some(hasher) = state.api_key_hasher() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "API key hasher is not initialized", (), request_id), - ) - .into_response(); - }; - - let entries = parse_client_api_key_import_payload(&payload); - if entries.is_empty() { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "No importable API keys found", (), request_id), - ) - .into_response(); - } - - let mut imported = 0u32; - let mut skipped = 0u32; - let mut keys = Vec::with_capacity(entries.len()); - for entry in entries { - let name = entry.name.trim(); - if name.is_empty() { - skipped += 1; - continue; - } - if entry - .label - .as_ref() - .is_some_and(|label| label.chars().count() > 64) - { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new( - 40001, - "API key label must be 64 characters or fewer", - (), - request_id, - ), - ) - .into_response(); - } - - let generated = hasher.generate_client_api_key(name); - let plaintext = generated.plaintext.clone(); - let source_id = entry.source_id; - let source_prefix = entry.source_prefix; - match repo - .insert_generated_with_metadata(name, entry.label.as_deref(), entry.enabled, &generated) - .await - { - Ok(key) => { - imported += 1; - keys.push(ImportedClientApiKeyData::new( - key, - plaintext, - source_id, - source_prefix, - )); - } - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to import API key", (), request_id), - ) - .into_response(); - } - } - } - - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - ClientApiKeyImportData { - imported, - skipped, - rotated: true, - keys, - }, - request_id, - ), - ) - .into_response() -} - -pub async fn create_api_key( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Json(payload): Json, -) -> 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 name = payload.name.trim(); - if name.is_empty() { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "API key name is required", (), request_id), - ) - .into_response(); - } - let Some(repo) = state.client_api_key_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "API key repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - let Some(hasher) = state.api_key_hasher() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "API key hasher is not initialized", (), request_id), - ) - .into_response(); - }; - - let generated = hasher.generate_client_api_key(name); - let plaintext = generated.plaintext.clone(); - match repo.insert_generated(name, &generated).await { - Ok(key) => AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(CreatedClientApiKeyData::new(key, plaintext), request_id), - ) - .into_response(), - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to create API key", (), request_id), - ) - .into_response(), - } -} - -pub async fn batch_delete_api_keys( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Json(payload): Json, -) -> Response { - let request_id = request_id.as_str().to_string(); - if payload.ids.is_empty() { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "API key ids are required", (), 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.client_api_key_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "API key repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - - let mut deleted = 0u32; - let mut not_found = Vec::new(); - for key_id in payload.ids { - match repo.delete(&key_id).await { - Ok(true) => deleted += 1, - Ok(false) => not_found.push(key_id), - Err(_) => { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to delete API key", (), request_id), - ) - .into_response(); - } - } - } - - AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - BatchDeleteClientApiKeysData { deleted, not_found }, - request_id, - ), - ) - .into_response() -} - -pub async fn update_api_key_label( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Path(key_id): Path, - Json(payload): Json, -) -> Response { - let request_id = request_id.as_str().to_string(); - if payload - .label - .as_ref() - .is_some_and(|label| label.chars().count() > 64) - { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new( - 40001, - "API key label must be 64 characters or fewer", - (), - 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.client_api_key_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "API key repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - - // label 是后台显示备注，不能影响 cpr_ key 的 hash、prefix 或启用状态。 - match repo.set_label(&key_id, payload.label).await { - Ok(Some(key)) => AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(ClientApiKeyData::from(key), request_id), - ) - .into_response(), - Ok(None) => AdminResponse::new( - StatusCode::NOT_FOUND, - AdminEnvelope::new(40401, "API key not found", (), request_id), - ) - .into_response(), - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to update API key label", (), request_id), - ) - .into_response(), - } -} - -pub async fn update_api_key_status( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Path(key_id): Path, - Json(payload): Json, -) -> 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 enabled = match parse_client_api_key_enabled_status(&payload.status) { - Ok(enabled) => enabled, - Err(message) => { - return AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, message, (), request_id), - ) - .into_response(); - } - }; - let Some(repo) = state.client_api_key_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "API key repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - - match repo.set_enabled(&key_id, enabled).await { - Ok(true) => AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - UpdateClientApiKeyStatusData { - id: key_id, - enabled, - }, - request_id, - ), - ) - .into_response(), - Ok(false) => AdminResponse::new( - StatusCode::NOT_FOUND, - AdminEnvelope::new(40401, "API key not found", (), request_id), - ) - .into_response(), - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to update API key status", (), request_id), - ) - .into_response(), - } -} - -pub async fn delete_api_key( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Path(key_id): Path, -) -> 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.client_api_key_repository() else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new( - 50001, - "API key repository is not initialized", - (), - request_id, - ), - ) - .into_response(); - }; - - match repo.delete(&key_id).await { - Ok(true) => AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(DeleteClientApiKeyData { deleted: true }, request_id), - ) - .into_response(), - Ok(false) => AdminResponse::new( - StatusCode::NOT_FOUND, - AdminEnvelope::new(40401, "API key not found", (), request_id), - ) - .into_response(), - Err(_) => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to delete API key", (), request_id), - ) - .into_response(), - } -} - -#[derive(Debug)] -struct AdminUserRow { - id: String, - password_hash: String, -} - -async fn load_first_admin(pool: &sqlx::SqlitePool) -> Result, sqlx::Error> { - let row = - sqlx::query("select id, password_hash from admin_users order by created_at asc limit 1") - .fetch_optional(pool) - .await?; - Ok(row.map(|row| AdminUserRow { - id: row.get("id"), - password_hash: row.get("password_hash"), - })) -} - -async fn create_admin_session( - pool: &sqlx::SqlitePool, - session_id: &str, - user_id: &str, - expires_at: chrono::DateTime, -) -> Result<(), sqlx::Error> { - let now = Utc::now().to_rfc3339(); - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", - ) - .bind(session_id) - .bind(user_id) - .bind(expires_at.to_rfc3339()) - .bind(now) - .execute(pool) - .await?; - Ok(()) -} - -fn admin_session_set_cookie(session_id: &str, ttl_minutes: u64) -> Option { - let max_age = ttl_minutes.checked_mul(60)?; - let cookie = format!( - "cpr_admin_session={session_id}; Path=/; HttpOnly; SameSite=Lax; Max-Age={max_age}" - ); - HeaderValue::from_str(&cookie).ok() -} - -async fn validate_admin_session( - pool: &sqlx::SqlitePool, - headers: &HeaderMap, -) -> Result { - let Some(session_id) = admin_session_cookie(headers) else { - return Ok(false); - }; - let now = Utc::now().to_rfc3339(); - let count: (i64,) = - sqlx::query_as("select count(*) from admin_sessions where id = ? and expires_at > ?") - .bind(session_id) - .bind(now) - .fetch_one(pool) - .await?; - Ok(count.0 > 0) -} - -fn admin_session_cookie(headers: &HeaderMap) -> Option<&str> { - admin_session_id(headers) -} - -async fn require_admin_session( - pool: &sqlx::SqlitePool, - headers: &HeaderMap, - request_id: &str, -) -> Result<(), Response> { - match validate_admin_session(pool, headers).await { - Ok(true) => Ok(()), - Ok(false) => Err(AdminResponse::new( - StatusCode::UNAUTHORIZED, - AdminEnvelope::new(40101, "Admin session required", (), request_id), - ) - .into_response()), - Err(_) => Err(AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to validate admin session", (), request_id), - ) - .into_response()), - } -} - -fn filter_health_check_accounts( - accounts: Vec, - ids: Option<&[String]>, -) -> Vec { - 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, "a.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 { - 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 { - 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 { - 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 { - 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 { - 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 { - 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 { - 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::().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 { - 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, - account_id: &str, - status: AccountStatus, -) -> Result<(), AccountRepositoryError> { - let updated = state - .account_pool() - .lock() - .await - .set_status(account_id, status); - if updated || status != AccountStatus::Active { - return Ok(()); - } - - if let Some(account) = repo.get(account_id).await? { - state - .account_pool() - .lock() - .await - .insert(pool_account_from_stored(account)); - } - Ok(()) -} - -fn pool_account_from_stored(account: StoredAccount) -> Account { - Account { - id: account.id, - email: account.email, - account_id: account.account_id, - user_id: account.user_id, - label: account.label, - plan_type: account.plan_type, - access_token: account.access_token.expose_secret().to_string(), - refresh_token: account - .refresh_token - .as_ref() - .map(|token| token.expose_secret().to_string()), - access_token_expires_at: account.access_token_expires_at, - status: account.status, - quota_limit_reached: false, - quota_cooldown_until: None, - cloudflare_cooldown_until: None, - added_at: account.added_at.to_rfc3339(), - last_used_at: None, - } -} - -#[derive(Debug, Clone)] -struct ManualAccountClaims { - account_id: String, - user_id: Option, - email: Option, - plan_type: Option, - expires_at: DateTime, -} - -fn admin_account_data_from_stored(account: StoredAccount) -> AdminAccountData { - AdminAccountData { - id: account.id, - email: account.email, - account_id: account.account_id, - user_id: account.user_id, - label: account.label, - plan_type: account.plan_type, - status: account_status_value(account.status).to_string(), - access_token_expires_at: account - .access_token_expires_at - .map(|value| value.to_rfc3339()), - added_at: account.added_at.to_rfc3339(), - updated_at: account.updated_at.to_rfc3339(), - } -} - -#[derive(Debug)] -enum ValidatedAccountImportError { - TokenRequired, - TokenRefresherUnavailable, - RefreshTransport, - RefreshRejected, - InvalidToken(&'static str), - Inspect, - NotFound, - Update, - Insert, - Load, -} - -async fn store_validated_account_import( - state: &AppState, - repo: &crate::accounts::repository::AccountRepository, - token: Option, - refresh_token: Option, -) -> Result { - let (access_token, refresh_token_update, new_account_refresh_token) = match ( - empty_to_none(token.map(normalize_bearer_token)), - empty_to_none(refresh_token), - ) { - (Some(token), refresh_token) => (token, refresh_token.clone(), refresh_token), - (None, Some(refresh_token)) => { - let Some(refresher) = state.token_refresher() else { - return Err(ValidatedAccountImportError::TokenRefresherUnavailable); - }; - let tokens = match refresher.refresh(&refresh_token).await { - Ok(tokens) => tokens, - Err(RefreshFailure::Transport) => { - return Err(ValidatedAccountImportError::RefreshTransport); - } - Err(_) => return Err(ValidatedAccountImportError::RefreshRejected), - }; - let access_token = normalize_bearer_token(tokens.access_token); - let rotated_refresh_token = empty_to_none(tokens.refresh_token); - let new_account_refresh_token = rotated_refresh_token.clone().or(Some(refresh_token)); - ( - access_token, - rotated_refresh_token, - new_account_refresh_token, - ) - } - (None, None) => return Err(ValidatedAccountImportError::TokenRequired), - }; - - // 手动和 CLI 导入都只信任 ChatGPT JWT claim；请求体里的展示字段不能参与账号身份判定。 - let claims = manual_account_claims(&access_token, Utc::now()) - .map_err(ValidatedAccountImportError::InvalidToken)?; - let existing = repo - .find_by_chatgpt_identity(&claims.account_id, claims.user_id.as_deref()) - .await - .map_err(|_| ValidatedAccountImportError::Inspect)?; - - let account_id = if let Some(existing) = existing { - let updated = repo - .update_from_claims( - &existing.id, - AccountClaimsUpdate { - email: claims.email.clone(), - account_id: Some(claims.account_id.clone()), - user_id: claims.user_id.clone(), - plan_type: claims.plan_type.clone(), - access_token: SecretString::new(access_token.into()), - refresh_token: refresh_token_update - .map(|token| SecretString::new(token.into())), - access_token_expires_at: Some(claims.expires_at), - status: AccountStatus::Active, - }, - ) - .await - .map_err(|_| ValidatedAccountImportError::Update)?; - if !updated { - return Err(ValidatedAccountImportError::NotFound); - } - existing.id - } else { - let id = normalized_account_id(None); - let account = NewAccount { - id: id.clone(), - email: claims.email.clone(), - account_id: Some(claims.account_id.clone()), - user_id: claims.user_id.clone(), - label: None, - plan_type: claims.plan_type.clone(), - access_token: SecretString::new(access_token.into()), - refresh_token: new_account_refresh_token.map(|token| SecretString::new(token.into())), - access_token_expires_at: Some(claims.expires_at), - status: AccountStatus::Active, - }; - repo.insert(account) - .await - .map_err(|_| ValidatedAccountImportError::Insert)?; - id - }; - - let stored = repo - .get(&account_id) - .await - .map_err(|_| ValidatedAccountImportError::Load)? - .ok_or(ValidatedAccountImportError::NotFound)?; - state - .account_pool() - .lock() - .await - .insert(pool_account_from_stored(stored.clone())); - Ok(stored) -} - -fn validated_account_import_error_response( - error: ValidatedAccountImportError, - request_id: String, -) -> Response { - match error { - ValidatedAccountImportError::TokenRequired => AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new( - 40001, - "Either token or refreshToken is required", - (), - request_id, - ), - ) - .into_response(), - ValidatedAccountImportError::TokenRefresherUnavailable => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Token refresher is not initialized", (), request_id), - ) - .into_response(), - ValidatedAccountImportError::RefreshTransport => AdminResponse::new( - StatusCode::BAD_GATEWAY, - AdminEnvelope::new(50201, "Refresh token exchange failed", (), request_id), - ) - .into_response(), - ValidatedAccountImportError::RefreshRejected => AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, "Refresh token exchange failed", (), request_id), - ) - .into_response(), - ValidatedAccountImportError::InvalidToken(message) => AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new(40001, message, (), request_id), - ) - .into_response(), - ValidatedAccountImportError::Inspect => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to inspect account", (), request_id), - ) - .into_response(), - ValidatedAccountImportError::NotFound => AdminResponse::new( - StatusCode::NOT_FOUND, - AdminEnvelope::new(40401, "Account not found", (), request_id), - ) - .into_response(), - ValidatedAccountImportError::Update => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to update account", (), request_id), - ) - .into_response(), - ValidatedAccountImportError::Insert => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to create account", (), request_id), - ) - .into_response(), - ValidatedAccountImportError::Load => AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Failed to load account", (), request_id), - ) - .into_response(), - } -} - -fn oauth_error_response(error: OAuthError, request_id: String, message: &'static str) -> Response { - let status = match error { - OAuthError::Rejected(_) => StatusCode::BAD_REQUEST, - OAuthError::AuthorizationPending | OAuthError::SlowDown | OAuthError::Transport => { - StatusCode::BAD_GATEWAY - } - }; - let body_code = if status == StatusCode::BAD_REQUEST { - 40001 - } else { - 50201 - }; - AdminResponse::new( - status, - AdminEnvelope::new(body_code, message, (), request_id), - ) - .into_response() -} - -fn request_host(headers: &HeaderMap, config: &AppConfig) -> String { - headers - .get(HOST) - .and_then(|value| value.to_str().ok()) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) - .unwrap_or_else(|| format!("localhost:{}", config.server.port)) -} - -fn parse_callback_url(callback_url: &str) -> Result<(String, String), ()> { - let callback_url = callback_url.trim(); - if callback_url.is_empty() { - return Err(()); - } - let url = Url::parse(callback_url).map_err(|_| ())?; - if url.query_pairs().any(|(key, _)| key == "error") { - return Err(()); - } - let code = url - .query_pairs() - .find_map(|(key, value)| (key == "code").then(|| value.into_owned())) - .filter(|value| !value.trim().is_empty()) - .ok_or(())?; - let state = url - .query_pairs() - .find_map(|(key, value)| (key == "state").then(|| value.into_owned())) - .filter(|value| !value.trim().is_empty()) - .ok_or(())?; - Ok((code, state)) -} - -async fn acquire_pkce_session( - state: &AppState, - oauth_state: &str, - request_id: &str, -) -> Result, Response> { - let sessions = state.oauth_sessions(); - let mut sessions = sessions.lock().await; - if let Some(session) = sessions.try_acquire(oauth_state) { - return Ok(Some(session)); - } - if sessions.is_completed_or_exchanging(oauth_state) { - return Ok(None); - } - Err(AdminResponse::new( - StatusCode::BAD_REQUEST, - AdminEnvelope::new( - 40001, - "Invalid or expired OAuth session", - (), - request_id.to_string(), - ), - ) - .into_response()) -} - -async fn exchange_pkce_session( - state: &AppState, - repo: &crate::accounts::repository::AccountRepository, - oauth_client: &dyn OAuthClient, - oauth_state: &str, - code: &str, - session: PkceSession, - request_id: &str, -) -> Result { - let tokens = match oauth_client - .exchange_code(code, &session.code_verifier, &session.redirect_uri) - .await - { - Ok(tokens) => tokens, - Err(error) => { - state.oauth_sessions().lock().await.release(oauth_state); - return Err(oauth_error_response( - error, - request_id.to_string(), - "Token exchange failed", - )); - } - }; - match store_validated_account_import( - state, - repo, - Some(tokens.access_token), - tokens.refresh_token, - ) - .await - { - Ok(_) => { - state.oauth_sessions().lock().await.complete(oauth_state); - Ok(session.return_host) - } - Err(error) => { - state.oauth_sessions().lock().await.complete(oauth_state); - Err(validated_account_import_error_response( - error, - request_id.to_string(), - )) - } - } -} - -fn redirect_to_host(return_host: &str, request_id: String) -> Response { - let location = format!("http://{return_host}/"); - let Ok(location) = HeaderValue::from_str(&location) else { - return AdminResponse::new( - StatusCode::INTERNAL_SERVER_ERROR, - AdminEnvelope::new(50001, "Invalid OAuth return host", (), request_id), - ) - .into_response(); - }; - let mut response = StatusCode::SEE_OTHER.into_response(); - response.headers_mut().insert(LOCATION, location); - response -} - -fn manual_account_claims( - token: &str, - now: DateTime, -) -> Result { - let payload = decode_jwt_payload(token).ok_or("Invalid JWT format")?; - let exp = payload - .get("exp") - .and_then(Value::as_i64) - .ok_or("Token is expired")?; - if now.timestamp() >= exp { - return Err("Token is expired"); - } - let expires_at = DateTime::::from_timestamp(exp, 0).ok_or("Invalid JWT exp claim")?; - let auth = payload - .get("https://api.openai.com/auth") - .and_then(Value::as_object) - .ok_or("Token missing chatgpt_account_id claim")?; - let account_id = - string_claim(auth, "chatgpt_account_id").ok_or("Token missing chatgpt_account_id claim")?; - let profile = payload - .get("https://api.openai.com/profile") - .and_then(Value::as_object); - let user_id = string_claim(auth, "chatgpt_user_id") - .or_else(|| profile.and_then(|profile| string_claim(profile, "chatgpt_user_id"))); - let plan_type = string_claim(auth, "chatgpt_plan_type") - .or_else(|| profile.and_then(|profile| string_claim(profile, "chatgpt_plan_type"))); - let email = profile.and_then(|profile| string_claim(profile, "email")); - - Ok(ManualAccountClaims { - account_id, - user_id, - email, - plan_type, - expires_at, - }) -} - -fn decode_jwt_payload(token: &str) -> Option> { - let mut parts = token.split('.'); - let _header = parts.next()?; - let payload = parts.next()?; - if payload.is_empty() { - return None; - } - // OpenAI access token 这里按 TS reference 只解码 payload，不做本地签名验签。 - let bytes = URL_SAFE_NO_PAD.decode(payload).ok()?; - serde_json::from_slice::(&bytes) - .ok()? - .as_object() - .cloned() -} - -fn string_claim(map: &Map, key: &str) -> Option { - map.get(key) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToString::to_string) -} - -#[derive(Debug)] -enum StoredImportAccount { - Imported, - Skipped, -} - -#[derive(Debug)] -enum StoreImportAccountError { - Inspect, - Invalid(String), - Insert, -} - -async fn store_import_account_entry( - state: &AppState, - repo: &crate::accounts::repository::AccountRepository, - entry: AccountImportEntry, -) -> Result { - let access_token = entry.token.as_deref().unwrap_or_default().trim(); - if access_token.is_empty() { - return Ok(StoredImportAccount::Skipped); - } - if entry - .label - .as_ref() - .is_some_and(|label| label.chars().count() > 64) - { - return Err(StoreImportAccountError::Invalid( - "Account label must be 64 characters or fewer".to_string(), - )); - } - - let id = normalized_account_id(entry.id); - match repo.exists(&id).await { - Ok(true) => return Ok(StoredImportAccount::Skipped), - Ok(false) => {} - Err(_) => return Err(StoreImportAccountError::Inspect), - } - - let status = - parse_import_status(entry.status.as_deref()).map_err(StoreImportAccountError::Invalid)?; - let email = empty_to_none(entry.email); - let account_id = empty_to_none(entry.account_id); - let user_id = empty_to_none(entry.user_id); - let label = empty_to_none(entry.label); - let plan_type = empty_to_none(entry.plan_type); - let refresh_token = empty_to_none(entry.refresh_token); - let access_token = access_token.to_string(); - let now = chrono::Utc::now().to_rfc3339(); - let pool_account = Account { - id: id.clone(), - email: email.clone(), - account_id: account_id.clone(), - user_id: user_id.clone(), - label: label.clone(), - plan_type: plan_type.clone(), - access_token: access_token.clone(), - refresh_token: refresh_token.clone(), - access_token_expires_at: None, - status, - quota_limit_reached: false, - quota_cooldown_until: None, - cloudflare_cooldown_until: None, - added_at: now.clone(), - last_used_at: None, - }; - let account = NewAccount { - id: id.clone(), - email: email.clone(), - account_id: account_id.clone(), - user_id: user_id.clone(), - label: label.clone(), - plan_type: plan_type.clone(), - access_token: SecretString::new(access_token.into()), - refresh_token: refresh_token.map(|token| SecretString::new(token.into())), - access_token_expires_at: None, - status, - }; - repo.insert(account) - .await - .map_err(|_| StoreImportAccountError::Insert)?; - state.account_pool().lock().await.insert(pool_account); - - Ok(StoredImportAccount::Imported) -} - -fn normalized_account_id(id: Option) -> String { - id.and_then(|id| empty_to_none(Some(id))) - .unwrap_or_else(|| format!("acct_{}", Uuid::new_v4().simple())) -} - -fn account_export_ids(value: Option<&str>) -> Vec { - value - .into_iter() - .flat_map(|ids| ids.split(',')) - .filter_map(|id| { - let id = id.trim(); - (!id.is_empty()).then(|| id.to_string()) - }) - .collect() -} - -fn parse_client_api_key_import_payload(payload: &Value) -> Vec { - let payload = payload - .get("data") - .filter(|data| data.get("apiKeys").is_some() || data.get("keys").is_some()) - .unwrap_or(payload); - - if let Some(keys) = payload.get("apiKeys").and_then(Value::as_array) { - return keys - .iter() - .filter_map(client_api_key_import_entry_from_value) - .collect(); - } - if let Some(keys) = payload.get("keys").and_then(Value::as_array) { - return keys - .iter() - .filter_map(client_api_key_import_entry_from_value) - .collect(); - } - if let Some(keys) = payload.as_array() { - return keys - .iter() - .filter_map(client_api_key_import_entry_from_value) - .collect(); - } - - client_api_key_import_entry_from_value(payload) - .into_iter() - .collect() -} - -fn client_api_key_import_entry_from_value(value: &Value) -> Option { - value.as_object()?; - let name = first_string(value, &[&["name"]])?; - // 安全边界：导入即使收到 plaintext/keyHash，也只按元数据轮换生成新的本地 cpr_ key。 - Some(ClientApiKeyImportEntry { - source_id: first_string(value, &[&["id"], &["sourceId"]]), - source_prefix: first_string(value, &[&["prefix"], &["sourcePrefix"]]), - name, - label: first_string(value, &[&["label"]]), - enabled: client_api_key_import_enabled(value), - }) -} - -fn client_api_key_import_enabled(value: &Value) -> bool { - if let Some(enabled) = value.get("enabled").and_then(Value::as_bool) { - return enabled; - } - !first_string(value, &[&["status"]]) - .unwrap_or_else(|| "active".to_string()) - .trim() - .eq_ignore_ascii_case("disabled") -} - -fn parse_account_export_format(value: Option<&str>) -> Result { - match value.unwrap_or("native") { - "" | "native" | "full" => Ok(AccountExportFormat::Native), - "sub2api" => Ok(AccountExportFormat::Sub2Api), - _ => Err("Unsupported account export format"), - } -} - -fn native_account_export(accounts: Vec) -> Value { - json!({ - "sourceFormat": "native", - "accounts": accounts.into_iter().map(native_export_account).collect::>(), - }) -} - -fn native_export_account(account: StoredAccount) -> Value { - json!({ - "id": account.id, - "email": account.email, - "accountId": account.account_id, - "userId": account.user_id, - "label": account.label, - "planType": account.plan_type, - "token": account.access_token.expose_secret(), - "refreshToken": account - .refresh_token - .as_ref() - .map(|token| token.expose_secret().to_string()), - "status": account_status_value(account.status), - "accessTokenExpiresAt": account.access_token_expires_at.map(|value| value.to_rfc3339()), - "addedAt": account.added_at.to_rfc3339(), - "updatedAt": account.updated_at.to_rfc3339(), - }) -} - -fn sub2api_account_export(accounts: Vec) -> Value { - json!({ - "exported_at": Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), - "proxies": [], - "accounts": accounts.into_iter().map(sub2api_export_account).collect::>(), - "type": "sub2api-data", - "version": 1, - }) -} - -fn sub2api_export_account(account: StoredAccount) -> Value { - let mut credentials = Map::new(); - credentials.insert( - "access_token".to_string(), - Value::String(account.access_token.expose_secret().to_string()), - ); - insert_optional_string( - &mut credentials, - "refresh_token", - account - .refresh_token - .as_ref() - .map(|token| token.expose_secret().to_string()), - ); - insert_optional_string(&mut credentials, "email", account.email.clone()); - insert_optional_string(&mut credentials, "chatgpt_account_id", account.account_id); - insert_optional_string(&mut credentials, "chatgpt_user_id", account.user_id); - insert_optional_string(&mut credentials, "plan_type", account.plan_type); - insert_optional_string( - &mut credentials, - "expires_at", - account - .access_token_expires_at - .map(|value| value.to_rfc3339()), - ); - - json!({ - "name": account - .label - .as_deref() - .filter(|label| !label.trim().is_empty()) - .map(str::to_string) - .or_else(|| account.email.clone()) - .unwrap_or_else(|| account.id.clone()), - "platform": "openai", - "type": "oauth", - "credentials": credentials, - "concurrency": 0, - "priority": 0, - }) -} - -fn insert_optional_string(map: &mut Map, key: &str, value: Option) { - if let Some(value) = value.filter(|value| !value.is_empty()) { - map.insert(key.to_string(), Value::String(value)); - } -} - -fn distinct_active_plan_accounts(accounts: Vec) -> Vec<(String, Account)> { - let mut by_plan = std::collections::BTreeMap::new(); - for account in accounts { - if account.status != AccountStatus::Active { - continue; - } - let plan_type = account - .plan_type - .clone() - .unwrap_or_else(|| "unknown".to_string()); - // 同一 plan 的模型列表一致，刷新时只需要取一个账号，避免重复打上游。 - by_plan.entry(plan_type).or_insert(account); - } - by_plan.into_iter().collect() -} - -fn parse_account_import_payload(payload: &Value) -> ParsedAccountImportPayload { - if let Some(accounts) = parse_sub2api_oauth_payload(payload) { - return ParsedAccountImportPayload { - accounts, - source_format: AccountImportFormat::Sub2Api, - }; - } - - parse_native_account_payload(payload) -} - -fn parse_sub2api_oauth_payload(payload: &Value) -> Option> { - let accounts = payload.get("accounts")?.as_array()?; - let looks_like_sub2api = string_at(payload, &["type"]).as_deref() == Some("sub2api-data") - || payload.get("proxies").is_some() - || accounts - .iter() - .any(|account| account.get("credentials").is_some()); - if !looks_like_sub2api { - return None; - } - - Some( - accounts - .iter() - .filter_map(sub2api_oauth_account_entry) - .collect(), - ) -} - -fn sub2api_oauth_account_entry(account: &Value) -> Option { - let platform = string_at(account, &["platform"])?.to_ascii_lowercase(); - let account_type = string_at(account, &["type"])?.to_ascii_lowercase(); - if platform != "openai" || account_type != "oauth" { - return None; - } - let credentials = account.get("credentials")?; - let fallback_label = normalized_label(string_at(account, &["name"])); - let mut entry = account_entry_from_value(credentials, fallback_label); - if entry.token.is_none() && entry.refresh_token.is_none() { - return None; - } - if entry.email.is_none() { - entry.email = string_at(credentials, &["email"]); - } - if entry.account_id.is_none() { - entry.account_id = first_string( - credentials, - &[&["chatgpt_account_id"], &["account_id"], &["accountId"]], - ); - } - if entry.user_id.is_none() { - entry.user_id = first_string( - credentials, - &[&["chatgpt_user_id"], &["user_id"], &["userId"]], - ); - } - if entry.plan_type.is_none() { - entry.plan_type = first_string(credentials, &[&["plan_type"], &["planType"]]); - } - Some(entry) -} - -fn parse_native_account_payload(payload: &Value) -> ParsedAccountImportPayload { - if let Some(accounts) = payload.as_array() { - return ParsedAccountImportPayload { - accounts: accounts - .iter() - .filter_map(|account| { - let entry = account_entry_from_value(account, None); - (entry.token.is_some() || entry.refresh_token.is_some()).then_some(entry) - }) - .collect(), - source_format: AccountImportFormat::Native, - }; - } - - if let Some(accounts) = payload.get("accounts").and_then(Value::as_array) { - let source_format = if looks_like_sub2api_native_export(accounts) { - AccountImportFormat::Sub2Api - } else { - AccountImportFormat::Native - }; - return ParsedAccountImportPayload { - accounts: accounts - .iter() - .filter_map(|account| { - let entry = account_entry_from_value(account, None); - (entry.token.is_some() || entry.refresh_token.is_some()).then_some(entry) - }) - .collect(), - source_format, - }; - } - - let entry = account_entry_from_value(payload, None); - let accounts = if entry.token.is_some() || entry.refresh_token.is_some() { - vec![entry] - } else { - Vec::new() - }; - ParsedAccountImportPayload { - accounts, - source_format: AccountImportFormat::Native, - } -} - -fn looks_like_sub2api_native_export(accounts: &[Value]) -> bool { - accounts.iter().any(|account| { - // sub2api 兼容导出会携带代理/配额运行态字段；这里只用于格式识别，代理数据不进入 Rust 服务。 - account.get("proxyApiKey").is_some() - || account.get("cachedQuota").is_some() - || account.get("quotaVerifyRequired").is_some() - }) -} - -fn account_entry_from_value(value: &Value, fallback_label: Option) -> AccountImportEntry { - let token = first_string( - value, - &[ - &["token"], - &["accessToken"], - &["access_token"], - &["tokens", "accessToken"], - &["tokens", "access_token"], - &["credentials", "token"], - &["credentials", "accessToken"], - &["credentials", "access_token"], - ], - ) - .map(normalize_bearer_token); - let refresh_token = first_string( - value, - &[ - &["refreshToken"], - &["refresh_token"], - &["tokens", "refreshToken"], - &["tokens", "refresh_token"], - &["credentials", "refreshToken"], - &["credentials", "refresh_token"], - ], - ); - - AccountImportEntry { - id: first_string(value, &[&["id"]]), - email: first_string(value, &[&["email"], &["credentials", "email"]]), - account_id: first_string( - value, - &[ - &["accountId"], - &["account_id"], - &["chatgpt_account_id"], - &["credentials", "accountId"], - &["credentials", "account_id"], - &["credentials", "chatgpt_account_id"], - ], - ), - user_id: first_string( - value, - &[ - &["userId"], - &["user_id"], - &["chatgpt_user_id"], - &["credentials", "userId"], - &["credentials", "user_id"], - &["credentials", "chatgpt_user_id"], - ], - ), - label: label_from_value(value).or(fallback_label), - plan_type: first_string( - value, - &[ - &["planType"], - &["plan_type"], - &["credentials", "planType"], - &["credentials", "plan_type"], - ], - ), - token, - refresh_token, - status: first_string(value, &[&["status"]]), - } -} - -fn label_from_value(value: &Value) -> Option { - normalized_label(first_string( - value, - &[ - &["label"], - &["name"], - &["account_name"], - &["accountName"], - &["account_note"], - &["accountNote"], - &["note"], - ], - )) -} - -fn normalized_label(value: Option) -> Option { - value.map(|label| label.chars().take(64).collect()) -} - -fn first_string(value: &Value, paths: &[&[&str]]) -> Option { - paths.iter().find_map(|path| string_at(value, path)) -} - -fn string_at(value: &Value, path: &[&str]) -> Option { - let mut current = value; - for key in path { - current = current.get(*key)?; - } - current - .as_str() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToString::to_string) -} - -fn normalize_bearer_token(value: String) -> String { - let trimmed = value.trim(); - trimmed - .strip_prefix("Bearer ") - .or_else(|| trimmed.strip_prefix("bearer ")) - .unwrap_or(trimmed) - .trim() - .to_string() -} - -fn empty_to_none(value: Option) -> Option { - value - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) -} - -fn parse_import_status(status: Option<&str>) -> Result { - let normalized = status.unwrap_or("active").trim().to_ascii_lowercase(); - match normalized.as_str() { - "active" => Ok(AccountStatus::Active), - "expired" => Ok(AccountStatus::Expired), - "quota_exhausted" => Ok(AccountStatus::QuotaExhausted), - "refreshing" => Ok(AccountStatus::Refreshing), - "disabled" => Ok(AccountStatus::Disabled), - "banned" => Ok(AccountStatus::Banned), - other => Err(format!("Unsupported account status: {other}")), - } -} - -fn parse_admin_account_status(status: &str) -> Result { - match status.trim().to_ascii_lowercase().as_str() { - "active" => Ok(AccountStatus::Active), - "disabled" => Ok(AccountStatus::Disabled), - other => Err(format!("Unsupported account status: {other}")), - } -} - -fn parse_client_api_key_enabled_status(status: &str) -> Result { - match status.trim().to_ascii_lowercase().as_str() { - "active" => Ok(true), - "disabled" => Ok(false), - other => Err(format!("Unsupported API key status: {other}")), - } -} - -fn admin_cookie_header(value: &Value) -> Result { - if let Some(cookies) = value.as_str() { - let cookies = cookies.trim(); - if cookies.is_empty() { - return Err("cookies field is required"); - } - return Ok(cookies.to_string()); - } - let Some(object) = value.as_object() else { - return Err("cookies must be a string or object"); - }; - if object.is_empty() { - return Err("cookies field is required"); - } - let pairs = object - .iter() - .filter_map(|(name, value)| { - let value = value.as_str()?.trim(); - (!name.trim().is_empty() && !value.is_empty()) - .then(|| format!("{}={value}", name.trim())) - }) - .collect::>(); - if pairs.is_empty() { - Err("No valid cookies found") - } else { - Ok(pairs.join("; ")) - } -} - -fn account_status_value(status: AccountStatus) -> &'static str { - match status { - AccountStatus::Active => "active", - AccountStatus::Expired => "expired", - AccountStatus::QuotaExhausted => "quota_exhausted", - AccountStatus::Refreshing => "refreshing", - AccountStatus::Disabled => "disabled", - AccountStatus::Banned => "banned", - } -} diff --git a/src/http/admin/accounts.rs b/src/http/admin/accounts.rs new file mode 100644 index 000000000..3e4972e9b --- /dev/null +++ b/src/http/admin/accounts.rs @@ -0,0 +1,1627 @@ +use axum::{ + body::Bytes, + extract::{Path, Query, State}, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, Json, +}; +use chrono::{SecondsFormat, Utc}; +use secrecy::ExposeSecret; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Map, Value}; + +use crate::{ + app::state::AppState, + codex::accounts::service::{ + AccountImportEntry, AccountProbeResult, AccountQuotaError, AccountServiceError, + HealthCheckError, RefreshAccountError, StoreImportAccountError, + ValidatedAccountImportError, + }, + codex::accounts::{ + model::AccountStatus, + repository::{StoredAccount, StoredAccountMetadata}, + }, + codex::oauth::{default_codex_home, read_cli_auth_from_home}, + http::middleware::RequestId, + utils::{ + json::{first_string, string_at}, + pagination::{clamp_limit, Page}, + }, +}; + +use super::{require_admin_session, AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse}; + +fn account_service_error(error: AccountServiceError, request_id: &str) -> AdminError { + match error { + AccountServiceError::RepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account repository is not initialized", + request_id, + ), + AccountServiceError::UsageRepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account usage repository is not initialized", + request_id, + ), + AccountServiceError::CookieRepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Cookie repository is not initialized", + request_id, + ), + AccountServiceError::AccountNotFound => AdminError::new( + StatusCode::NOT_FOUND, + 40401, + "Account not found", + request_id, + ), + AccountServiceError::List => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to list accounts", + request_id, + ), + AccountServiceError::Export => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to export accounts", + request_id, + ), + AccountServiceError::Inspect => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to inspect account", + request_id, + ), + AccountServiceError::ResetUsage => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to reset account usage", + request_id, + ), + AccountServiceError::UpdateLabel => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to update account label", + request_id, + ), + AccountServiceError::UpdateStatus => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to update account status", + request_id, + ), + AccountServiceError::SyncStatus => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to sync account status", + request_id, + ), + AccountServiceError::Delete => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to delete account", + request_id, + ), + AccountServiceError::LoadCookies => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to load account cookies", + request_id, + ), + AccountServiceError::StoreCookies => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to store account cookies", + request_id, + ), + AccountServiceError::DeleteCookies => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to delete account cookies", + request_id, + ), + AccountServiceError::NoValidCookies => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "No valid cookies found", + request_id, + ), + AccountServiceError::EmptyIds => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Account ids are required", + request_id, + ), + AccountServiceError::InvalidStatus(message) => { + AdminError::new(StatusCode::BAD_REQUEST, 40001, message, request_id) + } + AccountServiceError::LabelTooLong => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Account label must be 64 characters or fewer", + request_id, + ), + } +} + +fn refresh_account_error(error: RefreshAccountError, request_id: &str) -> AdminError { + match error { + RefreshAccountError::RepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account repository is not initialized", + request_id, + ), + RefreshAccountError::Load => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to load account", + request_id, + ), + RefreshAccountError::NotFound => AdminError::new( + StatusCode::NOT_FOUND, + 40401, + "Account not found", + request_id, + ), + RefreshAccountError::TokenRefresherUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Token refresher is not initialized", + request_id, + ), + RefreshAccountError::StoreRefreshed => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to store refreshed account", + request_id, + ), + } +} + +pub(super) enum AccountQuotaHttpError { + Standard(AdminError), + UpstreamFetch { error: String, request_id: String }, +} + +impl From for AccountQuotaHttpError { + fn from(error: AdminError) -> Self { + Self::Standard(error) + } +} + +impl IntoResponse for AccountQuotaHttpError { + fn into_response(self) -> axum::response::Response { + match self { + Self::Standard(error) => error.into_response(), + Self::UpstreamFetch { error, request_id } => AdminResponse::new( + StatusCode::BAD_GATEWAY, + AdminEnvelope::new( + 50201, + "Failed to fetch quota from Codex API", + json!({ "error": error }), + request_id, + ), + ) + .into_response(), + } + } +} + +fn account_quota_error(error: AccountQuotaError, request_id: &str) -> AccountQuotaHttpError { + match error { + AccountQuotaError::RepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account repository is not initialized", + request_id, + ) + .into(), + AccountQuotaError::Load => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to load account", + request_id, + ) + .into(), + AccountQuotaError::NotFound => AdminError::new( + StatusCode::NOT_FOUND, + 40401, + "Account not found", + request_id, + ) + .into(), + AccountQuotaError::Inactive(status) => AdminError::new( + StatusCode::CONFLICT, + 40901, + format!( + "Account is {}, cannot query quota", + account_status_value(status) + ), + request_id, + ) + .into(), + AccountQuotaError::StoreQuota => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to store account quota", + request_id, + ) + .into(), + AccountQuotaError::Fetch(error) => AccountQuotaHttpError::UpstreamFetch { + error, + request_id: request_id.to_string(), + }, + } +} + +fn health_check_error(error: HealthCheckError, request_id: &str) -> AdminError { + match error { + HealthCheckError::RepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account repository is not initialized", + request_id, + ), + HealthCheckError::List => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to list accounts", + request_id, + ), + } +} + +fn account_probe_data_from_service(result: AccountProbeResult) -> AccountProbeData { + AccountProbeData { + id: result.id, + email: result.email, + previous_status: account_status_value(result.previous_status).to_string(), + result: result.outcome.as_str().to_string(), + status: result + .status + .map(account_status_value) + .map(ToString::to_string), + error: result.error, + duration_ms: result.duration_ms, + } +} + +pub async fn accounts( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Query(query): Query, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let limit = clamp_limit(query.limit.unwrap_or(50)); + let page = state + .services + .accounts + .list(query.cursor, limit) + .await + .map_err(|error| account_service_error(error, &request_id))?; + let Page { items, next_cursor } = page; + let page = Page { + items: items.into_iter().map(AdminAccountData::from).collect(), + next_cursor, + }; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminPageEnvelope::ok(page, limit, request_id), + )) +} + +pub async fn export_accounts( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Query(query): Query, +) -> Result { + let request_id = request_id.as_str().to_string(); + let format = match parse_account_export_format(query.format.as_deref()) { + Ok(format) => format, + Err(message) => { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + message, + request_id, + )); + } + }; + require_admin_session(&state, &headers, &request_id).await?; + + let ids = account_export_ids(query.ids.as_deref()); + let accounts = state + .services + .accounts + .export(ids) + .await + .map_err(|error| account_service_error(error, &request_id))?; + + // 账号导出会返回可重新导入的 OAuth token；只允许 admin session 访问，不写入日志。 + let data = match format { + AccountExportFormat::Native => native_account_export(accounts), + AccountExportFormat::Sub2Api => sub2api_account_export(accounts), + }; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(data, request_id), + )) +} + +pub async fn create_account( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let stored = state + .services + .accounts + .import_validated(payload.token, payload.refresh_token) + .await + .map_err(|error| validated_account_import_error(error, &request_id))?; + + // 手动添加账号的响应只返回可展示元数据，OAuth token 永不回显。 + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(admin_account_data_from_stored(stored), request_id), + )) +} + +pub async fn health_check_accounts( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + payload: Option>, +) -> Result { + 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 Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Account ids must not be empty", + request_id, + )); + } + if payload + .stagger_ms + .is_some_and(|value| !(500..=30_000).contains(&value)) + { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "staggerMs must be between 500 and 30000", + request_id, + )); + } + if payload + .concurrency + .is_some_and(|value| !(1..=10).contains(&value)) + { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "concurrency must be between 1 and 10", + request_id, + )); + } + require_admin_session(&state, &headers, &request_id).await?; + + let concurrency = usize::from(payload.concurrency.unwrap_or(2)); + let stagger_ms = payload.stagger_ms.unwrap_or(3_000); + let results = state + .services + .accounts + .health_check_accounts(payload.ids, concurrency, stagger_ms, &request_id) + .await + .map_err(|error| health_check_error(error, &request_id))? + .into_iter() + .map(account_probe_data_from_service) + .collect::>(); + 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(), + }; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(HealthCheckData { summary, results }, request_id), + )) +} + +pub async fn refresh_account( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(account_id): Path, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let result = state + .services + .accounts + .refresh_account(&account_id) + .await + .map_err(|error| refresh_account_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(account_probe_data_from_service(result), request_id), + )) +} + +pub async fn reset_account_usage( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(account_id): Path, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + if !state + .services + .accounts + .reset_usage(&account_id) + .await + .map_err(|error| account_service_error(error, &request_id))? + { + return Err(AdminError::new( + StatusCode::NOT_FOUND, + 40401, + "Account not found", + request_id, + )); + } + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + ResetAccountUsageData { + id: account_id, + reset: true, + }, + request_id, + ), + )) +} + +pub(super) async fn account_quota( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(account_id): Path, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let result = state + .services + .accounts + .account_quota(&account_id, &request_id) + .await + .map_err(|error| account_quota_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AccountQuotaData { + quota: result.quota, + raw: result.raw, + }, + request_id, + ), + )) +} + +pub async fn update_account_label( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(account_id): Path, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let label = payload.label; + let updated = state + .services + .accounts + .update_label(&account_id, label.clone()) + .await + .map_err(|error| account_service_error(error, &request_id))?; + + if !updated { + return Err(AdminError::new( + StatusCode::NOT_FOUND, + 40401, + "Account not found", + request_id, + )); + } + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + UpdateAccountLabelData { + id: account_id, + label, + }, + request_id, + ), + )) +} + +pub async fn update_account_status( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(account_id): Path, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let updated = state + .services + .accounts + .update_status(&account_id, &payload.status) + .await + .map_err(|error| account_service_error(error, &request_id))? + .ok_or_else(|| { + AdminError::new( + StatusCode::NOT_FOUND, + 40401, + "Account not found", + request_id.as_str(), + ) + })?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + UpdateAccountStatusData { + id: updated.id, + status: account_status_value(updated.status).to_string(), + }, + request_id, + ), + )) +} + +pub async fn delete_account( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(account_id): Path, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let deleted = state + .services + .accounts + .delete(&account_id) + .await + .map_err(|error| account_service_error(error, &request_id))?; + + if !deleted { + return Err(AdminError::new( + StatusCode::NOT_FOUND, + 40401, + "Account not found", + request_id, + )); + } + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(DeleteAccountData { deleted: true }, request_id), + )) +} + +pub async fn batch_delete_accounts( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + if payload.ids.is_empty() { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Account ids are required", + request_id, + )); + } + require_admin_session(&state, &headers, &request_id).await?; + + let deleted = state + .services + .accounts + .batch_delete(payload.ids) + .await + .map_err(|error| account_service_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + BatchDeleteAccountsData { + deleted: deleted.deleted, + not_found: deleted.not_found, + }, + request_id, + ), + )) +} + +pub async fn batch_update_account_status( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + if payload.ids.is_empty() { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Account ids are required", + request_id, + )); + } + parse_admin_account_status(&payload.status).map_err(|message| { + AdminError::new(StatusCode::BAD_REQUEST, 40001, message, request_id.as_str()) + })?; + require_admin_session(&state, &headers, &request_id).await?; + + let updated = state + .services + .accounts + .batch_update_status(payload.ids, &payload.status) + .await + .map_err(|error| account_service_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + BatchUpdateAccountStatusData { + updated: updated.updated, + not_found: updated.not_found, + }, + request_id, + ), + )) +} + +pub async fn get_account_cookies( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(account_id): Path, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let cookies = state + .services + .accounts + .get_cookies(&account_id) + .await + .map_err(|error| account_service_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountCookiesData { cookies }, request_id), + )) +} + +pub async fn set_account_cookies( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(account_id): Path, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + let cookie_header = match admin_cookie_header(&payload.cookies) { + Ok(cookie_header) => cookie_header, + Err(message) => { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + message, + request_id, + )); + } + }; + require_admin_session(&state, &headers, &request_id).await?; + + let cookies = state + .services + .accounts + .set_cookies(&account_id, &cookie_header) + .await + .map_err(|error| account_service_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountCookiesData { cookies }, request_id), + )) +} + +pub async fn delete_account_cookies( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(account_id): Path, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + state + .services + .accounts + .delete_cookies(&account_id) + .await + .map_err(|error| account_service_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(DeleteAccountCookiesData { deleted: true }, request_id), + )) +} + +pub async fn import_accounts( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + if !state.services.accounts.has_repository() { + return Err(account_service_error( + AccountServiceError::RepositoryUnavailable, + &request_id, + )); + } + let parsed = parse_account_import_payload(&payload); + if parsed.accounts.is_empty() { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "No importable accounts found", + request_id, + )); + } + + let counts = state + .services + .accounts + .import_entries(parsed.accounts) + .await + .map_err(|error| store_import_account_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AccountImportData { + imported: counts.imported, + skipped: counts.skipped, + source_format: parsed.source_format.as_str().to_string(), + }, + request_id, + ), + )) +} + +pub async fn import_cli_auth( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + body: Bytes, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let payload = if body.is_empty() { + ImportCliAuthRequest::default() + } else { + match serde_json::from_slice::(&body) { + Ok(payload) => payload, + Err(_) => { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Invalid CLI import request", + request_id, + )); + } + } + }; + let codex_home = match empty_to_none(payload.codex_home) { + Some(path) => std::path::PathBuf::from(path), + None => match default_codex_home() { + Ok(path) => path, + Err(error) => { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + error.to_string(), + request_id, + )); + } + }, + }; + let cli_auth = match read_cli_auth_from_home(&codex_home) { + Ok(auth) => auth, + Err(error) => { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + error.to_string(), + request_id, + )); + } + }; + let _stored = state + .services + .accounts + .import_validated( + Some(cli_auth.access_token().to_string()), + cli_auth.refresh_token().map(str::to_string), + ) + .await + .map_err(|error| validated_account_import_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AccountImportData { + imported: 1, + skipped: 0, + source_format: AccountImportFormat::CodexCli.as_str().to_string(), + }, + request_id, + ), + )) +} +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountsQuery { + pub cursor: Option, + pub limit: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountExportQuery { + pub ids: Option, + pub format: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AccountImportFormat { + Native, + Sub2Api, + CodexCli, +} + +impl AccountImportFormat { + fn as_str(self) -> &'static str { + match self { + Self::Native => "native", + Self::Sub2Api => "sub2api", + Self::CodexCli => "codex_cli", + } + } +} + +#[derive(Debug, Clone)] +struct ParsedAccountImportPayload { + accounts: Vec, + source_format: AccountImportFormat, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountImportData { + pub imported: u32, + pub skipped: u32, + pub source_format: String, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CreateAccountRequest { + pub token: Option, + pub refresh_token: Option, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ImportCliAuthRequest { + pub codex_home: Option, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct HealthCheckRequest { + pub ids: Option>, + #[serde(alias = "stagger_ms")] + pub stagger_ms: Option, + pub concurrency: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct HealthCheckData { + pub summary: HealthCheckSummary, + pub results: Vec, +} + +#[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, + pub previous_status: String, + pub result: String, + pub status: Option, + pub error: Option, + pub duration_ms: Option, +} + +#[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, + Sub2Api, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateAccountLabelRequest { + pub label: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateAccountLabelData { + pub id: String, + pub label: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateAccountStatusRequest { + pub status: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateAccountStatusData { + pub id: String, + pub status: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DeleteAccountData { + pub deleted: bool, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchDeleteAccountsRequest { + pub ids: Vec, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchDeleteAccountsData { + pub deleted: u32, + pub not_found: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchUpdateAccountStatusRequest { + pub ids: Vec, + pub status: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchUpdateAccountStatusData { + pub updated: u32, + pub not_found: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SetAccountCookiesRequest { + pub cookies: Value, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountCookiesData { + pub cookies: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DeleteAccountCookiesData { + pub deleted: bool, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAccountData { + pub id: String, + pub email: Option, + pub account_id: Option, + pub user_id: Option, + pub label: Option, + pub plan_type: Option, + pub status: String, + pub access_token_expires_at: Option, + pub added_at: String, + pub updated_at: String, +} + +impl From for AdminAccountData { + fn from(account: StoredAccountMetadata) -> Self { + Self { + id: account.id, + email: account.email, + account_id: account.account_id, + user_id: account.user_id, + label: account.label, + plan_type: account.plan_type, + status: account_status_value(account.status).to_string(), + access_token_expires_at: account + .access_token_expires_at + .map(|value| value.to_rfc3339()), + added_at: account.added_at.to_rfc3339(), + updated_at: account.updated_at.to_rfc3339(), + } + } +} + +fn admin_account_data_from_stored(account: StoredAccount) -> AdminAccountData { + AdminAccountData { + id: account.id, + email: account.email, + account_id: account.account_id, + user_id: account.user_id, + label: account.label, + plan_type: account.plan_type, + status: account_status_value(account.status).to_string(), + access_token_expires_at: account + .access_token_expires_at + .map(|value| value.to_rfc3339()), + added_at: account.added_at.to_rfc3339(), + updated_at: account.updated_at.to_rfc3339(), + } +} + +pub(super) fn validated_account_import_error( + error: ValidatedAccountImportError, + request_id: &str, +) -> AdminError { + match error { + ValidatedAccountImportError::RepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account repository is not initialized", + request_id, + ), + ValidatedAccountImportError::TokenRequired => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Either token or refreshToken is required", + request_id, + ), + ValidatedAccountImportError::TokenRefresherUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Token refresher is not initialized", + request_id, + ), + ValidatedAccountImportError::RefreshTransport => AdminError::new( + StatusCode::BAD_GATEWAY, + 50201, + "Refresh token exchange failed", + request_id, + ), + ValidatedAccountImportError::RefreshRejected => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Refresh token exchange failed", + request_id, + ), + ValidatedAccountImportError::InvalidToken(message) => { + AdminError::new(StatusCode::BAD_REQUEST, 40001, message, request_id) + } + ValidatedAccountImportError::Inspect => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to inspect account", + request_id, + ), + ValidatedAccountImportError::NotFound => AdminError::new( + StatusCode::NOT_FOUND, + 40401, + "Account not found", + request_id, + ), + ValidatedAccountImportError::Update => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to update account", + request_id, + ), + ValidatedAccountImportError::Insert => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to create account", + request_id, + ), + ValidatedAccountImportError::Load => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to load account", + request_id, + ), + } +} + +fn store_import_account_error(error: StoreImportAccountError, request_id: &str) -> AdminError { + match error { + StoreImportAccountError::RepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account repository is not initialized", + request_id, + ), + StoreImportAccountError::Inspect => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to inspect account", + request_id, + ), + StoreImportAccountError::Invalid(message) => { + AdminError::new(StatusCode::BAD_REQUEST, 40001, message, request_id) + } + StoreImportAccountError::Insert => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to import account", + request_id, + ), + } +} + +pub(super) fn account_export_ids(value: Option<&str>) -> Vec { + value + .into_iter() + .flat_map(|ids| ids.split(',')) + .filter_map(|id| { + let id = id.trim(); + (!id.is_empty()).then(|| id.to_string()) + }) + .collect() +} + +fn parse_account_export_format(value: Option<&str>) -> Result { + match value.unwrap_or("native") { + "" | "native" | "full" => Ok(AccountExportFormat::Native), + "sub2api" => Ok(AccountExportFormat::Sub2Api), + _ => Err("Unsupported account export format"), + } +} + +fn native_account_export(accounts: Vec) -> Value { + json!({ + "sourceFormat": "native", + "accounts": accounts.into_iter().map(native_export_account).collect::>(), + }) +} + +fn native_export_account(account: StoredAccount) -> Value { + json!({ + "id": account.id, + "email": account.email, + "accountId": account.account_id, + "userId": account.user_id, + "label": account.label, + "planType": account.plan_type, + "token": account.access_token.expose_secret(), + "refreshToken": account + .refresh_token + .as_ref() + .map(|token| token.expose_secret().to_string()), + "status": account_status_value(account.status), + "accessTokenExpiresAt": account.access_token_expires_at.map(|value| value.to_rfc3339()), + "addedAt": account.added_at.to_rfc3339(), + "updatedAt": account.updated_at.to_rfc3339(), + }) +} + +fn sub2api_account_export(accounts: Vec) -> Value { + json!({ + "exported_at": Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + "proxies": [], + "accounts": accounts.into_iter().map(sub2api_export_account).collect::>(), + "type": "sub2api-data", + "version": 1, + }) +} + +fn sub2api_export_account(account: StoredAccount) -> Value { + let mut credentials = Map::new(); + credentials.insert( + "access_token".to_string(), + Value::String(account.access_token.expose_secret().to_string()), + ); + insert_optional_string( + &mut credentials, + "refresh_token", + account + .refresh_token + .as_ref() + .map(|token| token.expose_secret().to_string()), + ); + insert_optional_string(&mut credentials, "email", account.email.clone()); + insert_optional_string(&mut credentials, "chatgpt_account_id", account.account_id); + insert_optional_string(&mut credentials, "chatgpt_user_id", account.user_id); + insert_optional_string(&mut credentials, "plan_type", account.plan_type); + insert_optional_string( + &mut credentials, + "expires_at", + account + .access_token_expires_at + .map(|value| value.to_rfc3339()), + ); + + json!({ + "name": account + .label + .as_deref() + .filter(|label| !label.trim().is_empty()) + .map(str::to_string) + .or_else(|| account.email.clone()) + .unwrap_or_else(|| account.id.clone()), + "platform": "openai", + "type": "oauth", + "credentials": credentials, + "concurrency": 0, + "priority": 0, + }) +} + +fn insert_optional_string(map: &mut Map, key: &str, value: Option) { + if let Some(value) = value.filter(|value| !value.is_empty()) { + map.insert(key.to_string(), Value::String(value)); + } +} + +fn parse_account_import_payload(payload: &Value) -> ParsedAccountImportPayload { + if let Some(accounts) = parse_sub2api_oauth_payload(payload) { + return ParsedAccountImportPayload { + accounts, + source_format: AccountImportFormat::Sub2Api, + }; + } + + parse_native_account_payload(payload) +} + +fn parse_sub2api_oauth_payload(payload: &Value) -> Option> { + let accounts = payload.get("accounts")?.as_array()?; + let looks_like_sub2api = string_at(payload, &["type"]).as_deref() == Some("sub2api-data") + || payload.get("proxies").is_some() + || accounts + .iter() + .any(|account| account.get("credentials").is_some()); + if !looks_like_sub2api { + return None; + } + + Some( + accounts + .iter() + .filter_map(sub2api_oauth_account_entry) + .collect(), + ) +} + +fn sub2api_oauth_account_entry(account: &Value) -> Option { + let platform = string_at(account, &["platform"])?.to_ascii_lowercase(); + let account_type = string_at(account, &["type"])?.to_ascii_lowercase(); + if platform != "openai" || account_type != "oauth" { + return None; + } + let credentials = account.get("credentials")?; + let fallback_label = normalized_label(string_at(account, &["name"])); + let mut entry = account_entry_from_value(credentials, fallback_label); + if entry.token.is_none() && entry.refresh_token.is_none() { + return None; + } + if entry.email.is_none() { + entry.email = string_at(credentials, &["email"]); + } + if entry.account_id.is_none() { + entry.account_id = first_string( + credentials, + &[&["chatgpt_account_id"], &["account_id"], &["accountId"]], + ); + } + if entry.user_id.is_none() { + entry.user_id = first_string( + credentials, + &[&["chatgpt_user_id"], &["user_id"], &["userId"]], + ); + } + if entry.plan_type.is_none() { + entry.plan_type = first_string(credentials, &[&["plan_type"], &["planType"]]); + } + Some(entry) +} + +fn parse_native_account_payload(payload: &Value) -> ParsedAccountImportPayload { + if let Some(accounts) = payload.as_array() { + return ParsedAccountImportPayload { + accounts: accounts + .iter() + .filter_map(|account| { + let entry = account_entry_from_value(account, None); + (entry.token.is_some() || entry.refresh_token.is_some()).then_some(entry) + }) + .collect(), + source_format: AccountImportFormat::Native, + }; + } + + if let Some(accounts) = payload.get("accounts").and_then(Value::as_array) { + let source_format = if looks_like_sub2api_native_export(accounts) { + AccountImportFormat::Sub2Api + } else { + AccountImportFormat::Native + }; + return ParsedAccountImportPayload { + accounts: accounts + .iter() + .filter_map(|account| { + let entry = account_entry_from_value(account, None); + (entry.token.is_some() || entry.refresh_token.is_some()).then_some(entry) + }) + .collect(), + source_format, + }; + } + + let entry = account_entry_from_value(payload, None); + let accounts = if entry.token.is_some() || entry.refresh_token.is_some() { + vec![entry] + } else { + Vec::new() + }; + ParsedAccountImportPayload { + accounts, + source_format: AccountImportFormat::Native, + } +} + +fn looks_like_sub2api_native_export(accounts: &[Value]) -> bool { + accounts.iter().any(|account| { + // sub2api 兼容导出会携带代理/配额运行态字段；这里只用于格式识别，代理数据不进入 Rust 服务。 + account.get("proxyApiKey").is_some() + || account.get("cachedQuota").is_some() + || account.get("quotaVerifyRequired").is_some() + }) +} + +fn account_entry_from_value(value: &Value, fallback_label: Option) -> AccountImportEntry { + let token = first_string( + value, + &[ + &["token"], + &["accessToken"], + &["access_token"], + &["tokens", "accessToken"], + &["tokens", "access_token"], + &["credentials", "token"], + &["credentials", "accessToken"], + &["credentials", "access_token"], + ], + ) + .map(normalize_bearer_token); + let refresh_token = first_string( + value, + &[ + &["refreshToken"], + &["refresh_token"], + &["tokens", "refreshToken"], + &["tokens", "refresh_token"], + &["credentials", "refreshToken"], + &["credentials", "refresh_token"], + ], + ); + + AccountImportEntry { + id: first_string(value, &[&["id"]]), + email: first_string(value, &[&["email"], &["credentials", "email"]]), + account_id: first_string( + value, + &[ + &["accountId"], + &["account_id"], + &["chatgpt_account_id"], + &["credentials", "accountId"], + &["credentials", "account_id"], + &["credentials", "chatgpt_account_id"], + ], + ), + user_id: first_string( + value, + &[ + &["userId"], + &["user_id"], + &["chatgpt_user_id"], + &["credentials", "userId"], + &["credentials", "user_id"], + &["credentials", "chatgpt_user_id"], + ], + ), + label: label_from_value(value).or(fallback_label), + plan_type: first_string( + value, + &[ + &["planType"], + &["plan_type"], + &["credentials", "planType"], + &["credentials", "plan_type"], + ], + ), + token, + refresh_token, + status: first_string(value, &[&["status"]]), + } +} + +fn label_from_value(value: &Value) -> Option { + normalized_label(first_string( + value, + &[ + &["label"], + &["name"], + &["account_name"], + &["accountName"], + &["account_note"], + &["accountNote"], + &["note"], + ], + )) +} + +fn normalized_label(value: Option) -> Option { + value.map(|label| label.chars().take(64).collect()) +} + +fn normalize_bearer_token(value: String) -> String { + let trimmed = value.trim(); + trimmed + .strip_prefix("Bearer ") + .or_else(|| trimmed.strip_prefix("bearer ")) + .unwrap_or(trimmed) + .trim() + .to_string() +} + +fn empty_to_none(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +fn parse_admin_account_status(status: &str) -> Result { + match status.trim().to_ascii_lowercase().as_str() { + "active" => Ok(AccountStatus::Active), + "disabled" => Ok(AccountStatus::Disabled), + other => Err(format!("Unsupported account status: {other}")), + } +} + +fn admin_cookie_header(value: &Value) -> Result { + if let Some(cookies) = value.as_str() { + let cookies = cookies.trim(); + if cookies.is_empty() { + return Err("cookies field is required"); + } + return Ok(cookies.to_string()); + } + let Some(object) = value.as_object() else { + return Err("cookies must be a string or object"); + }; + if object.is_empty() { + return Err("cookies field is required"); + } + let pairs = object + .iter() + .filter_map(|(name, value)| { + let value = value.as_str()?.trim(); + (!name.trim().is_empty() && !value.is_empty()) + .then(|| format!("{}={value}", name.trim())) + }) + .collect::>(); + if pairs.is_empty() { + Err("No valid cookies found") + } else { + Ok(pairs.join("; ")) + } +} + +pub(super) fn account_status_value(status: AccountStatus) -> &'static str { + match status { + AccountStatus::Active => "active", + AccountStatus::Expired => "expired", + AccountStatus::QuotaExhausted => "quota_exhausted", + AccountStatus::Refreshing => "refreshing", + AccountStatus::Disabled => "disabled", + AccountStatus::Banned => "banned", + } +} diff --git a/src/http/admin/api_keys.rs b/src/http/admin/api_keys.rs new file mode 100644 index 000000000..8a8716fe1 --- /dev/null +++ b/src/http/admin/api_keys.rs @@ -0,0 +1,520 @@ +use axum::{ + extract::{Path, Query, State}, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, Json, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::{ + app::state::AppState, + auth::api_key_repository::StoredClientApiKey, + http::middleware::RequestId, + service::api_key::{ApiKeyServiceError, ImportedClientApiKey}, + utils::pagination::{clamp_limit, Page}, +}; + +use super::{ + account_export_ids, require_admin_session, AdminEnvelope, AdminError, AdminPageEnvelope, + AdminResponse, +}; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ApiKeysQuery { + pub cursor: Option, + pub limit: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ApiKeyExportQuery { + pub ids: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CreateApiKeyRequest { + pub name: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ClientApiKeyData { + pub id: String, + pub name: String, + pub label: Option, + pub prefix: String, + pub enabled: bool, + pub created_at: String, + pub last_used_at: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct CreatedClientApiKeyData { + pub id: String, + pub name: String, + pub label: Option, + pub prefix: String, + pub enabled: bool, + pub created_at: String, + pub last_used_at: Option, + pub plaintext: String, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateClientApiKeyStatusRequest { + pub status: String, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateClientApiKeyLabelRequest { + pub label: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateClientApiKeyStatusData { + pub id: String, + pub enabled: bool, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchDeleteClientApiKeysRequest { + pub ids: Vec, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchDeleteClientApiKeysData { + pub deleted: u32, + pub not_found: Vec, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DeleteClientApiKeyData { + pub deleted: bool, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ClientApiKeyExportData { + pub source_format: &'static str, + pub rotation_required: bool, + pub api_keys: Vec, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ClientApiKeyExportEntry { + pub id: String, + pub name: String, + pub label: Option, + pub prefix: String, + pub enabled: bool, + pub created_at: String, + pub last_used_at: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ClientApiKeyImportData { + pub imported: u32, + pub skipped: u32, + pub rotated: bool, + pub keys: Vec, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ImportedClientApiKeyData { + pub source_id: Option, + pub source_prefix: Option, + pub id: String, + pub name: String, + pub label: Option, + pub prefix: String, + pub enabled: bool, + pub created_at: String, + pub last_used_at: Option, + pub plaintext: String, +} + +impl From for ClientApiKeyData { + fn from(key: StoredClientApiKey) -> Self { + Self { + id: key.id, + name: key.name, + label: key.label, + prefix: key.prefix, + enabled: key.enabled, + created_at: key.created_at, + last_used_at: key.last_used_at, + } + } +} + +impl From for ClientApiKeyExportEntry { + fn from(key: StoredClientApiKey) -> Self { + Self { + id: key.id, + name: key.name, + label: key.label, + prefix: key.prefix, + enabled: key.enabled, + created_at: key.created_at, + last_used_at: key.last_used_at, + } + } +} + +impl CreatedClientApiKeyData { + fn new(key: StoredClientApiKey, plaintext: String) -> Self { + Self { + id: key.id, + name: key.name, + label: key.label, + prefix: key.prefix, + enabled: key.enabled, + created_at: key.created_at, + last_used_at: key.last_used_at, + plaintext, + } + } +} + +impl ImportedClientApiKeyData { + fn new(imported: ImportedClientApiKey) -> Self { + let ImportedClientApiKey { + key, + plaintext, + source_id, + source_prefix, + } = imported; + Self { + source_id, + source_prefix, + id: key.id, + name: key.name, + label: key.label, + prefix: key.prefix, + enabled: key.enabled, + created_at: key.created_at, + last_used_at: key.last_used_at, + plaintext, + } + } +} + +fn api_key_service_error(error: ApiKeyServiceError, request_id: String) -> AdminError { + match error { + ApiKeyServiceError::RepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "API key repository is not initialized", + request_id, + ), + ApiKeyServiceError::HasherUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "API key hasher is not initialized", + request_id, + ), + ApiKeyServiceError::List => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to list API keys", + request_id, + ), + ApiKeyServiceError::Export => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to export API keys", + request_id, + ), + ApiKeyServiceError::Import => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to import API key", + request_id, + ), + ApiKeyServiceError::Create => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to create API key", + request_id, + ), + ApiKeyServiceError::Delete => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to delete API key", + request_id, + ), + ApiKeyServiceError::UpdateLabel => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to update API key label", + request_id, + ), + ApiKeyServiceError::UpdateStatus => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to update API key status", + request_id, + ), + ApiKeyServiceError::Verify => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to verify API key", + request_id, + ), + ApiKeyServiceError::InvalidStatus(message) => { + AdminError::new(StatusCode::BAD_REQUEST, 40001, message, request_id) + } + ApiKeyServiceError::EmptyName => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "API key name is required", + request_id, + ), + ApiKeyServiceError::EmptyIds => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "API key ids are required", + request_id, + ), + ApiKeyServiceError::LabelTooLong => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "API key label must be 64 characters or fewer", + request_id, + ), + ApiKeyServiceError::NoImportableKeys => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "No importable API keys found", + request_id, + ), + } +} + +fn api_key_not_found(request_id: String) -> AdminError { + AdminError::new( + StatusCode::NOT_FOUND, + 40401, + "API key not found", + request_id, + ) +} + +pub async fn api_keys( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Query(query): Query, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + let limit = clamp_limit(query.limit.unwrap_or(50)); + match state.services.api_keys.list(query.cursor, limit).await { + Ok(page) => { + let Page { items, next_cursor } = page; + let page = Page { + items: items.into_iter().map(ClientApiKeyData::from).collect(), + next_cursor, + }; + Ok(AdminResponse::new( + StatusCode::OK, + AdminPageEnvelope::ok(page, limit, request_id), + )) + } + Err(error) => Err(api_key_service_error(error, request_id)), + } +} + +pub async fn export_api_keys( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Query(query): Query, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + let ids = account_export_ids(query.ids.as_deref()); + let keys = match state.services.api_keys.export(ids).await { + Ok(keys) => keys, + Err(error) => return Err(api_key_service_error(error, request_id)), + }; + + // 安全边界：本地 cpr_ key 只导出可展示元数据，绝不导出 plaintext、key_hash 或 pepper。 + let data = ClientApiKeyExportData { + source_format: "rustLocalClientApiKeys", + rotation_required: true, + api_keys: keys + .into_iter() + .map(ClientApiKeyExportEntry::from) + .collect(), + }; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(data, request_id), + )) +} + +pub async fn import_api_keys( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + let imported = match state.services.api_keys.import(&payload).await { + Ok(imported) => imported, + Err(error) => return Err(api_key_service_error(error, request_id)), + }; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + ClientApiKeyImportData { + imported: imported.imported, + skipped: imported.skipped, + rotated: true, + keys: imported + .keys + .into_iter() + .map(ImportedClientApiKeyData::new) + .collect(), + }, + request_id, + ), + )) +} + +pub async fn create_api_key( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + match state.services.api_keys.create(&payload.name).await { + Ok(created) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + CreatedClientApiKeyData::new(created.key, created.plaintext), + request_id, + ), + )), + Err(error) => Err(api_key_service_error(error, request_id)), + } +} + +pub async fn batch_delete_api_keys( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + let deleted = match state.services.api_keys.batch_delete(payload.ids).await { + Ok(deleted) => deleted, + Err(error) => return Err(api_key_service_error(error, request_id)), + }; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + BatchDeleteClientApiKeysData { + deleted: deleted.deleted, + not_found: deleted.not_found, + }, + request_id, + ), + )) +} + +pub async fn update_api_key_label( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(key_id): Path, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + // label 是后台显示备注，不能影响 cpr_ key 的 hash、prefix 或启用状态。 + match state + .services + .api_keys + .update_label(&key_id, payload.label) + .await + { + Ok(Some(key)) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(ClientApiKeyData::from(key), request_id), + )), + Ok(None) => Err(api_key_not_found(request_id)), + Err(error) => Err(api_key_service_error(error, request_id)), + } +} + +pub async fn update_api_key_status( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(key_id): Path, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + match state + .services + .api_keys + .update_status(key_id, &payload.status) + .await + { + Ok(Some(updated)) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + UpdateClientApiKeyStatusData { + id: updated.id, + enabled: updated.enabled, + }, + request_id, + ), + )), + Ok(None) => Err(api_key_not_found(request_id)), + Err(error) => Err(api_key_service_error(error, request_id)), + } +} + +pub async fn delete_api_key( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Path(key_id): Path, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + match state.services.api_keys.delete(&key_id).await { + Ok(true) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(DeleteClientApiKeyData { deleted: true }, request_id), + )), + Ok(false) => Err(api_key_not_found(request_id)), + Err(error) => Err(api_key_service_error(error, request_id)), + } +} diff --git a/src/http/admin/auth.rs b/src/http/admin/auth.rs new file mode 100644 index 000000000..9eb0b009e --- /dev/null +++ b/src/http/admin/auth.rs @@ -0,0 +1,660 @@ +use axum::{ + extract::{Path, Query, State}, + http::{ + header::{HeaderValue, HOST, LOCATION, SET_COOKIE}, + HeaderMap, StatusCode, + }, + response::{IntoResponse, Response}, + Extension, Json, +}; +use reqwest::Url; +use serde::{Deserialize, Serialize}; + +use crate::{ + app::state::AppState, + codex::oauth::{DeviceCode, OAuthError}, + config::AppConfig, + http::auth::admin_session_id, + http::middleware::RequestId, + service::admin_auth::{ + AdminAuthOAuthError, AdminAuthPkceExchange, + AdminAuthPoolSummary as ServiceAdminAuthPoolSummary, AdminAuthServiceError, + AdminAuthStatus as ServiceAdminAuthStatus, AdminAuthUser as ServiceAdminAuthUser, + AdminSessionValidationError, + }, +}; + +use super::{ + account_status_value, validated_account_import_error, AdminEnvelope, AdminError, AdminResponse, +}; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct LoginRequest { + pub password: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct LoginData { + pub expires_at: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthStatusData { + pub authenticated: bool, + pub user: Option, + pub pool: AdminAuthPoolSummaryData, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthUserData { + pub id: String, + pub email: Option, + pub account_id: Option, + pub user_id: Option, + pub label: Option, + pub plan_type: Option, + pub status: String, + pub access_token_expires_at: Option, +} + +#[derive(Debug, Default, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthPoolSummaryData { + pub total: usize, + pub active: usize, + pub expired: usize, + pub quota_exhausted: usize, + pub refreshing: usize, + pub disabled: usize, + pub banned: usize, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthLogoutData { + pub success: bool, + pub deleted: u64, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthDeviceLoginData { + pub user_code: String, + pub verification_uri: String, + pub verification_uri_complete: String, + pub device_code: String, + pub expires_in: u64, + pub interval: u64, +} + +impl From for AdminAuthDeviceLoginData { + fn from(device: DeviceCode) -> Self { + Self { + user_code: device.user_code, + verification_uri: device.verification_uri, + verification_uri_complete: device.verification_uri_complete, + device_code: device.device_code, + expires_in: device.expires_in, + interval: device.interval, + } + } +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthDevicePollData { + pub success: bool, + pub pending: bool, + pub code: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthLoginStartData { + pub auth_url: String, + pub state: String, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthCodeRelayRequest { + pub callback_url: String, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthCallbackQuery { + pub code: Option, + pub state: Option, + pub error: Option, + pub error_description: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthCodeRelayData { + pub success: bool, +} + +fn auth_status_data(status: ServiceAdminAuthStatus) -> AdminAuthStatusData { + AdminAuthStatusData { + authenticated: status.authenticated, + user: status.user.map(auth_user_data), + pool: auth_pool_summary_data(status.pool), + } +} + +fn auth_user_data(user: ServiceAdminAuthUser) -> AdminAuthUserData { + AdminAuthUserData { + id: user.id, + email: user.email, + account_id: user.account_id, + user_id: user.user_id, + label: user.label, + plan_type: user.plan_type, + status: account_status_value(user.status).to_string(), + access_token_expires_at: user.access_token_expires_at.map(|value| value.to_rfc3339()), + } +} + +fn auth_pool_summary_data(summary: ServiceAdminAuthPoolSummary) -> AdminAuthPoolSummaryData { + AdminAuthPoolSummaryData { + total: summary.total, + active: summary.active, + expired: summary.expired, + quota_exhausted: summary.quota_exhausted, + refreshing: summary.refreshing, + disabled: summary.disabled, + banned: summary.banned, + } +} + +fn admin_session_set_cookie(session_id: &str, ttl_minutes: u64) -> Option { + let max_age = ttl_minutes.checked_mul(60)?; + let cookie = format!( + "cpr_admin_session={session_id}; Path=/; HttpOnly; SameSite=Lax; Max-Age={max_age}" + ); + HeaderValue::from_str(&cookie).ok() +} + +pub(super) async fn require_admin_session( + state: &AppState, + headers: &HeaderMap, + request_id: &str, +) -> Result<(), AdminError> { + match state + .services + .admin_auth + .validate_session(admin_session_id(headers)) + .await + { + Ok(true) => Ok(()), + Ok(false) => Err(AdminError::new( + StatusCode::UNAUTHORIZED, + 40101, + "Admin session required", + request_id, + )), + Err(AdminSessionValidationError::DatabaseUnavailable) => Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Database is not initialized", + request_id, + )), + Err(AdminSessionValidationError::ValidateSession) => Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to validate admin session", + request_id, + )), + } +} + +fn oauth_error(error: OAuthError, request_id: &str, message: &'static str) -> AdminError { + let status = match error { + OAuthError::Rejected(_) => StatusCode::BAD_REQUEST, + OAuthError::AuthorizationPending | OAuthError::SlowDown | OAuthError::Transport => { + StatusCode::BAD_GATEWAY + } + }; + let body_code = if status == StatusCode::BAD_REQUEST { + 40001 + } else { + 50201 + }; + AdminError::new(status, body_code, message, request_id) +} + +fn admin_oauth_error(error: AdminAuthOAuthError, request_id: &str) -> AdminError { + match error { + AdminAuthOAuthError::OAuthClientUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "OAuth client is not initialized", + request_id, + ), + AdminAuthOAuthError::AccountRepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account repository is not initialized", + request_id, + ), + AdminAuthOAuthError::DeviceCodeRequired => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Device code is required", + request_id, + ), + AdminAuthOAuthError::InvalidOrExpiredSession => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Invalid or expired OAuth session", + request_id, + ), + AdminAuthOAuthError::DeviceCodeRequest(error) => { + oauth_error(error, request_id, "Device code request failed") + } + AdminAuthOAuthError::DeviceAuthorization(error) => { + oauth_error(error, request_id, "Device authorization failed") + } + AdminAuthOAuthError::TokenExchange(error) => { + oauth_error(error, request_id, "Token exchange failed") + } + AdminAuthOAuthError::Import(error) => validated_account_import_error(error, request_id), + } +} + +fn request_host(headers: &HeaderMap, config: &AppConfig) -> String { + headers + .get(HOST) + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + .unwrap_or_else(|| format!("localhost:{}", config.server.port)) +} + +fn parse_callback_url(callback_url: &str) -> Result<(String, String), ()> { + let callback_url = callback_url.trim(); + if callback_url.is_empty() { + return Err(()); + } + let url = Url::parse(callback_url).map_err(|_| ())?; + if url.query_pairs().any(|(key, _)| key == "error") { + return Err(()); + } + let code = url + .query_pairs() + .find_map(|(key, value)| (key == "code").then(|| value.into_owned())) + .filter(|value| !value.trim().is_empty()) + .ok_or(())?; + let state = url + .query_pairs() + .find_map(|(key, value)| (key == "state").then(|| value.into_owned())) + .filter(|value| !value.trim().is_empty()) + .ok_or(())?; + Ok((code, state)) +} + +fn redirect_to_host(return_host: &str, request_id: &str) -> Result { + let location = format!("http://{return_host}/"); + let Ok(location) = HeaderValue::from_str(&location) else { + return Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Invalid OAuth return host", + request_id, + )); + }; + let mut response = StatusCode::SEE_OTHER.into_response(); + response.headers_mut().insert(LOCATION, location); + Ok(response) +} + +fn login_error(error: AdminAuthServiceError, request_id: &str) -> AdminError { + match error { + AdminAuthServiceError::DatabaseUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Database is not initialized", + request_id, + ), + AdminAuthServiceError::AdminPasswordInvalid => AdminError::new( + StatusCode::UNAUTHORIZED, + 40102, + "Admin password invalid", + request_id, + ), + AdminAuthServiceError::LoadAdminUser => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to load admin user", + request_id, + ), + AdminAuthServiceError::VerifyAdminPassword => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to verify admin password", + request_id, + ), + AdminAuthServiceError::InvalidSessionTtl => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Admin session ttl is invalid", + request_id, + ), + AdminAuthServiceError::CreateSession => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to create admin session", + request_id, + ), + AdminAuthServiceError::AccountRepositoryUnavailable + | AdminAuthServiceError::InspectAccountAuthStatus + | AdminAuthServiceError::ClearAccounts => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to create admin session", + request_id, + ), + } +} + +fn auth_status_error(error: AdminAuthServiceError, request_id: &str) -> AdminError { + match error { + AdminAuthServiceError::AccountRepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account repository is not initialized", + request_id, + ), + AdminAuthServiceError::InspectAccountAuthStatus => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to inspect account auth status", + request_id, + ), + AdminAuthServiceError::DatabaseUnavailable + | AdminAuthServiceError::LoadAdminUser + | AdminAuthServiceError::AdminPasswordInvalid + | AdminAuthServiceError::VerifyAdminPassword + | AdminAuthServiceError::InvalidSessionTtl + | AdminAuthServiceError::CreateSession + | AdminAuthServiceError::ClearAccounts => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to inspect account auth status", + request_id, + ), + } +} + +fn logout_error(error: AdminAuthServiceError, request_id: &str) -> AdminError { + match error { + AdminAuthServiceError::AccountRepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account repository is not initialized", + request_id, + ), + AdminAuthServiceError::ClearAccounts => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to clear accounts", + request_id, + ), + AdminAuthServiceError::DatabaseUnavailable + | AdminAuthServiceError::LoadAdminUser + | AdminAuthServiceError::AdminPasswordInvalid + | AdminAuthServiceError::VerifyAdminPassword + | AdminAuthServiceError::InvalidSessionTtl + | AdminAuthServiceError::CreateSession + | AdminAuthServiceError::InspectAccountAuthStatus => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to clear accounts", + request_id, + ), + } +} + +pub async fn login( + State(state): State, + Extension(request_id): Extension, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + let login = state + .services + .admin_auth + .login(&payload.password) + .await + .map_err(|error| login_error(error, &request_id))?; + + let Some(cookie) = admin_session_set_cookie(&login.session_id, login.ttl_minutes) else { + return Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to create admin session cookie", + request_id, + )); + }; + let mut response = AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + LoginData { + expires_at: login.expires_at.to_rfc3339(), + }, + request_id, + ), + ) + .into_response(); + response.headers_mut().insert(SET_COOKIE, cookie); + Ok(response) +} + +pub async fn auth_status( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let status = state + .services + .admin_auth + .status() + .await + .map_err(|error| auth_status_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(auth_status_data(status), request_id), + )) +} + +pub async fn auth_logout( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let logout = state + .services + .admin_auth + .logout() + .await + .map_err(|error| logout_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AdminAuthLogoutData { + success: true, + deleted: logout.deleted, + }, + request_id, + ), + )) +} + +pub async fn auth_device_login( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let device = state + .services + .admin_auth + .request_device_code() + .await + .map_err(|error| admin_oauth_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminAuthDeviceLoginData::from(device), request_id), + )) +} + +pub async fn auth_device_poll( + State(state): State, + Path(device_code): Path, + Extension(request_id): Extension, + headers: HeaderMap, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let result = state + .services + .admin_auth + .poll_device_token(&device_code) + .await + .map_err(|error| admin_oauth_error(error, &request_id))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AdminAuthDevicePollData { + success: result.success, + pending: result.pending, + code: result.code, + }, + request_id, + ), + )) +} + +pub async fn auth_login_start( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let return_host = request_host(&headers, state.config()); + let login = state + .services + .admin_auth + .start_pkce_login(&return_host) + .await; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AdminAuthLoginStartData { + auth_url: login.auth_url, + state: login.state, + }, + request_id, + ), + )) +} + +pub async fn auth_code_relay( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Json(payload): Json, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let Ok((code, oauth_state)) = parse_callback_url(&payload.callback_url) else { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Invalid OAuth callback URL", + request_id, + )); + }; + let exchange = state + .services + .admin_auth + .exchange_pkce_code(&oauth_state, &code) + .await + .map_err(|error| admin_oauth_error(error, &request_id))?; + + match exchange { + AdminAuthPkceExchange::Imported { .. } | AdminAuthPkceExchange::AlreadyCompleted => { + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminAuthCodeRelayData { success: true }, request_id), + )) + } + } +} + +pub async fn auth_callback( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Query(query): Query, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + if query.error.is_some() { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + query + .error_description + .unwrap_or_else(|| "OAuth error".to_string()), + request_id, + )); + } + let (Some(code), Some(oauth_state)) = (query.code.as_deref(), query.state.as_deref()) else { + return Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Missing OAuth code or state", + request_id, + )); + }; + let exchange = state + .services + .admin_auth + .exchange_pkce_code(oauth_state, code) + .await + .map_err(|error| admin_oauth_error(error, &request_id))?; + + match exchange { + AdminAuthPkceExchange::Imported { return_host } => { + redirect_to_host(&return_host, &request_id) + } + AdminAuthPkceExchange::AlreadyCompleted => { + redirect_to_host(&request_host(&headers, state.config()), &request_id) + } + } +} diff --git a/src/http/admin/logs.rs b/src/http/admin/logs.rs new file mode 100644 index 000000000..3d1bf2ba4 --- /dev/null +++ b/src/http/admin/logs.rs @@ -0,0 +1,57 @@ +use axum::{ + extract::{Query, State}, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, +}; +use serde::Deserialize; + +use crate::{ + app::state::AppState, http::middleware::RequestId, service::log::LogServiceError, + utils::pagination::clamp_limit, +}; + +use super::{require_admin_session, AdminError, AdminPageEnvelope, AdminResponse}; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct LogsQuery { + pub cursor: Option, + pub limit: Option, +} + +pub async fn logs( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Query(query): Query, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let limit = clamp_limit(query.limit.unwrap_or(50)); + match state.services.logs.list(query.cursor, limit).await { + Ok(page) => Ok(AdminResponse::new( + StatusCode::OK, + AdminPageEnvelope::ok(page, limit, request_id), + )), + Err(error) => Err(log_service_error(error, request_id)), + } +} + +fn log_service_error(error: LogServiceError, request_id: String) -> AdminError { + match error { + LogServiceError::RepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Event log repository is not initialized", + request_id, + ), + LogServiceError::List => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to list event logs", + request_id, + ), + } +} diff --git a/src/http/admin/mod.rs b/src/http/admin/mod.rs new file mode 100644 index 000000000..83e8646db --- /dev/null +++ b/src/http/admin/mod.rs @@ -0,0 +1,19 @@ +pub mod accounts; +pub mod api_keys; +pub mod auth; +pub mod logs; +pub mod models; +pub mod response; +pub mod router; +pub mod settings; +pub mod usage; + +use self::{ + accounts::{account_export_ids, account_status_value, validated_account_import_error}, + auth::require_admin_session, +}; + +pub use self::{ + response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse, PageMeta}, + router::router, +}; diff --git a/src/http/admin/models.rs b/src/http/admin/models.rs new file mode 100644 index 000000000..2c4e570ab --- /dev/null +++ b/src/http/admin/models.rs @@ -0,0 +1,114 @@ +use axum::{ + extract::State, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, +}; +use serde::Serialize; + +use crate::{ + app::state::AppState, + codex::models::service::{ModelRefreshResult, ModelServiceError}, + http::middleware::RequestId, +}; + +use super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct RefreshModelsData { + pub refreshed_plans: usize, + pub model_count: usize, + pub failed_plans: usize, +} + +impl From for RefreshModelsData { + fn from(result: ModelRefreshResult) -> Self { + Self { + refreshed_plans: result.refreshed_plans, + model_count: result.model_count, + failed_plans: result.failed_plans, + } + } +} + +pub async fn refresh_models( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + match state + .services + .models + .refresh_backend_models(&request_id) + .await + { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(RefreshModelsData::from(result), request_id), + )), + Err(error) => model_service_error_response(error, request_id), + } +} + +fn model_service_error_response( + error: ModelServiceError, + request_id: String, +) -> Result>, AdminError> { + match error { + ModelServiceError::AccountRepositoryUnavailable => Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account repository is not initialized", + request_id, + )), + ModelServiceError::ModelRepositoryUnavailable => Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Model repository is not initialized", + request_id, + )), + ModelServiceError::ListAccounts => Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to list accounts", + request_id, + )), + ModelServiceError::NoAccounts => Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "No accounts available for model refresh", + request_id, + )), + ModelServiceError::BuildClient => Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to build Codex client", + request_id, + )), + ModelServiceError::StoreSnapshot => Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to store model snapshot", + request_id, + )), + ModelServiceError::AllPlansFailed(result) => Ok(AdminResponse::new( + StatusCode::BAD_GATEWAY, + AdminEnvelope::new( + 50201, + "Failed to refresh backend models", + RefreshModelsData::from(result), + request_id, + ), + )), + ModelServiceError::LoadSnapshots => Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to load model snapshots", + request_id, + )), + } +} diff --git a/src/http/admin/response.rs b/src/http/admin/response.rs new file mode 100644 index 000000000..4da2ad6e6 --- /dev/null +++ b/src/http/admin/response.rs @@ -0,0 +1,133 @@ +use axum::{ + http::StatusCode, + response::{IntoResponse, Response}, + Json, +}; +use serde::Serialize; + +use crate::utils::pagination::Page; + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminEnvelope { + pub code: u32, + pub message: String, + pub data: T, + pub request_id: String, +} + +impl AdminEnvelope { + pub fn new( + code: u32, + message: impl Into, + data: T, + request_id: impl Into, + ) -> Self { + // Body codes drive frontend branches; HTTP status remains the transport truth. + Self { + code, + message: message.into(), + data, + request_id: request_id.into(), + } + } + + pub fn ok(data: T, request_id: impl Into) -> Self { + Self::new(200, "OK", data, request_id) + } +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct PageMeta { + pub limit: u32, + pub next_cursor: Option, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminPageEnvelope { + pub code: u32, + pub message: String, + pub data: Vec, + pub page: PageMeta, + pub request_id: String, +} + +impl AdminPageEnvelope { + pub fn new( + code: u32, + message: impl Into, + page: Page, + limit: u32, + request_id: impl Into, + ) -> Self { + let Page { items, next_cursor } = page; + Self { + code, + message: message.into(), + data: items, + page: PageMeta { limit, next_cursor }, + request_id: request_id.into(), + } + } + + pub fn ok(page: Page, limit: u32, request_id: impl Into) -> Self { + Self::new(200, "OK", page, limit, request_id) + } +} + +#[derive(Debug, Clone)] +pub struct AdminResponse { + pub status: StatusCode, + pub body: T, +} + +impl AdminResponse { + pub fn new(status: StatusCode, body: T) -> Self { + Self { status, body } + } +} + +#[derive(Debug, Clone)] +pub struct AdminError { + pub status: StatusCode, + pub code: u32, + pub message: String, + pub request_id: String, +} + +impl AdminError { + pub fn new( + status: StatusCode, + code: u32, + message: impl Into, + request_id: impl Into, + ) -> Self { + Self { + status, + code, + message: message.into(), + request_id: request_id.into(), + } + } +} + +impl IntoResponse for AdminError { + fn into_response(self) -> Response { + AdminResponse::new( + self.status, + AdminEnvelope::new(self.code, self.message, (), self.request_id), + ) + .into_response() + } +} + +impl IntoResponse for AdminResponse +where + T: Serialize, +{ + fn into_response(self) -> Response { + (self.status, Json(self.body)).into_response() + } +} diff --git a/src/app.rs b/src/http/admin/router.rs similarity index 63% rename from src/app.rs rename to src/http/admin/router.rs index 523353925..bc91bfaee 100644 --- a/src/app.rs +++ b/src/http/admin/router.rs @@ -1,43 +1,34 @@ use axum::{ - middleware::from_fn, routing::{delete, get, patch, post}, Router, }; -use tower_http::trace::TraceLayer; -use crate::{ - http::{ - admin::{ - account_quota, accounts, api_keys, auth_callback, auth_code_relay, auth_device_login, - auth_device_poll, auth_login_start, auth_logout, auth_status, 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, - export_api_keys, get_account_cookies, health_check_accounts, import_accounts, - import_api_keys, 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, - middleware::attach_request_id, - v1::{ - chat_completions, debug_models, model_catalog, model_detail, model_info, models, - responses, - }, +use crate::app::state::AppState; + +use super::{ + accounts::{ + account_quota, accounts, batch_delete_accounts, batch_update_account_status, + create_account, delete_account, delete_account_cookies, export_accounts, + get_account_cookies, health_check_accounts, import_accounts, import_cli_auth, + refresh_account, reset_account_usage, set_account_cookies, update_account_label, + update_account_status, }, - state::AppState, + api_keys::{ + api_keys, batch_delete_api_keys, create_api_key, delete_api_key, export_api_keys, + import_api_keys, update_api_key_label, update_api_key_status, + }, + auth::{ + auth_callback, auth_code_relay, auth_device_login, auth_device_poll, auth_login_start, + auth_logout, auth_status, login, + }, + logs::logs, + models::refresh_models, + settings::settings, + usage::{usage_stats, usage_stats_summary}, }; -pub fn build_router(state: AppState) -> Router { +pub fn router() -> Router { Router::new() - .route("/health", get(health)) - .route("/v1/responses", post(responses)) - .route("/v1/chat/completions", post(chat_completions)) - .route("/v1/models", get(models)) - .route("/v1/models/catalog", get(model_catalog)) - .route("/v1/models/{model_id}", get(model_detail)) - .route("/v1/models/{model_id}/info", get(model_info)) - .route("/debug/models", get(debug_models)) .route("/admin/login", post(login)) .route("/admin/auth/status", get(auth_status)) .route("/admin/auth/logout", post(auth_logout)) @@ -103,7 +94,4 @@ pub fn build_router(state: AppState) -> Router { "/admin/api-keys/{key_id}/status", patch(update_api_key_status), ) - .with_state(state) - .layer(TraceLayer::new_for_http()) - .layer(from_fn(attach_request_id)) } diff --git a/src/http/admin/settings.rs b/src/http/admin/settings.rs new file mode 100644 index 000000000..dd4cd6ec9 --- /dev/null +++ b/src/http/admin/settings.rs @@ -0,0 +1,77 @@ +use axum::{ + extract::State, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, +}; +use serde::Serialize; + +use crate::{ + app::state::AppState, + config::{AppConfig, QuotaWarningThresholds}, + http::middleware::RequestId, +}; + +use super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminSettingsData { + pub default_model: String, + pub default_reasoning_effort: Option, + pub service_tier: Option, + pub model_aliases: std::collections::BTreeMap, + pub refresh_enabled: bool, + pub refresh_margin_seconds: u64, + pub refresh_concurrency: u32, + pub max_concurrent_per_account: usize, + pub request_interval_ms: u64, + pub rotation_strategy: String, + pub tier_priority: Vec, + pub quota_refresh_interval_minutes: u64, + pub quota_warning_thresholds: QuotaWarningThresholds, + pub quota_skip_exhausted: bool, + pub logs_enabled: bool, + pub logs_capacity: u32, + pub logs_capture_body: bool, + pub usage_history_retention_days: Option, +} + +impl AdminSettingsData { + fn from_config(config: &AppConfig) -> Self { + Self { + default_model: config.model.default_model.clone(), + default_reasoning_effort: config.model.default_reasoning_effort.clone(), + service_tier: config.model.service_tier.clone(), + model_aliases: config.model.aliases.clone(), + refresh_enabled: config.auth.refresh_enabled, + refresh_margin_seconds: config.auth.refresh_margin_seconds, + refresh_concurrency: config.auth.refresh_concurrency, + max_concurrent_per_account: config.auth.max_concurrent_per_account, + request_interval_ms: config.auth.request_interval_ms, + rotation_strategy: config.auth.rotation_strategy.clone(), + tier_priority: config.auth.tier_priority.clone(), + quota_refresh_interval_minutes: config.quota.refresh_interval_minutes, + quota_warning_thresholds: config.quota.warning_thresholds.clone(), + quota_skip_exhausted: config.quota.skip_exhausted, + logs_enabled: config.logging.enabled, + logs_capacity: config.logging.capacity, + logs_capture_body: config.logging.capture_body, + usage_history_retention_days: config.usage_stats.history_retention_days, + } + } +} + +pub async fn settings( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminSettingsData::from_config(state.config()), request_id), + )) +} diff --git a/src/http/admin/usage.rs b/src/http/admin/usage.rs new file mode 100644 index 000000000..f56333b24 --- /dev/null +++ b/src/http/admin/usage.rs @@ -0,0 +1,158 @@ +use axum::{ + extract::{Query, State}, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, +}; +use serde::{Deserialize, Serialize}; + +use crate::{ + app::state::AppState, + codex::accounts::repository::{ + AccountRepositoryError, AccountUsageListRecord, AccountUsageSummary, + }, + http::middleware::RequestId, + service::usage::UsageServiceError, + utils::pagination::{clamp_limit, Page}, +}; + +use super::{require_admin_session, AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse}; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UsageStatsQuery { + pub cursor: Option, + pub limit: Option, +} + +#[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, + } + } +} + +pub async fn usage_stats( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + Query(query): Query, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + let limit = clamp_limit(query.limit.unwrap_or(50)); + match state.services.usage.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, + }; + Ok(AdminResponse::new( + StatusCode::OK, + AdminPageEnvelope::ok(page, limit, request_id), + )) + } + Err(error) => Err(usage_list_error(error, request_id)), + } +} + +fn usage_list_error(error: UsageServiceError, request_id: String) -> AdminError { + match error { + UsageServiceError::RepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account usage repository is not initialized", + request_id, + ), + UsageServiceError::Repository(AccountRepositoryError::InvalidCursor) => { + AdminError::new(StatusCode::BAD_REQUEST, 40002, "Invalid cursor", request_id) + } + UsageServiceError::Repository(_) => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to list usage stats", + request_id, + ), + } +} + +pub async fn usage_stats_summary( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + match state.services.usage.summary().await { + Ok(summary) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminUsageStatsSummaryData::from(summary), request_id), + )), + Err(error) => Err(usage_summary_error(error, request_id)), + } +} + +fn usage_summary_error(error: UsageServiceError, request_id: String) -> AdminError { + match error { + UsageServiceError::RepositoryUnavailable => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Account usage repository is not initialized", + request_id, + ), + UsageServiceError::Repository(_) => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to summarize usage stats", + request_id, + ), + } +} diff --git a/src/http/v1.rs b/src/http/v1.rs deleted file mode 100644 index 109eac1f1..000000000 --- a/src/http/v1.rs +++ /dev/null @@ -1,1998 +0,0 @@ -use std::{sync::Arc, time::Instant}; - -use axum::{ - body::{Body, Bytes}, - extract::{Path, State}, - http::{ - header::{CACHE_CONTROL, CONTENT_TYPE}, - HeaderMap, StatusCode, - }, - response::{IntoResponse, Response}, - Extension, Json, -}; -use chrono::{Duration, 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, AccountStatus}, - pool::{AccountAcquireRequest, AccountPool}, - repository::{TokenUpdate, UsageDelta}, - }, - auth::refresh::RefreshFailure, - codex::{ - client::{ - build_reqwest_client, CodexBackendClient, CodexBackendStream, CodexClientError, - CodexRequestContext, - }, - sse::parse_sse_events, - types::CodexResponsesRequest, - usage::{extract_sse_usage, TokenUsage}, - websocket::{transport_for_request, CodexTransport}, - }, - fingerprint::model::Fingerprint, - http::{auth::client_api_key, middleware::RequestId}, - logs::event::{EventLevel, EventLog}, - models::catalog::ModelCatalog, - state::AppState, - translation::{ - codex_to_openai::{ - chat_completion_from_codex_sse, chat_completion_stream_from_codex_sse, openai_error, - }, - openai_to_codex::{translate_chat_to_codex, ChatCompletionRequest}, - }, -}; - -const MODEL_CREATED_TIMESTAMP: i64 = 1_700_000_000; -const DEFAULT_RATE_LIMIT_BACKOFF_SECONDS: u64 = 60; -const MAX_RATE_LIMIT_BACKOFF_SECONDS: u64 = 86_400 * 7; -const CLOUDFLARE_CHALLENGE_COOLDOWN_SECONDS: u64 = 10; - -#[derive(Deserialize)] -struct ResponsesBody { - model: Option, - input: Option>, - instructions: Option, - reasoning: Option, - tools: Option>, - service_tier: Option, - tool_choice: Option, - parallel_tool_calls: Option, - text: Option, - prompt_cache_key: Option, - include: Option>, - client_metadata: Option, - previous_response_id: Option, - stream: Option, - #[serde(rename = "turnState")] - turn_state: Option, - #[serde(rename = "turnMetadata")] - turn_metadata: Option, - #[serde(rename = "betaFeatures")] - beta_features: Option, - #[serde(rename = "includeTimingMetrics")] - include_timing_metrics: Option, - #[serde(rename = "codexWindowId")] - codex_window_id: Option, - #[serde(rename = "parentThreadId")] - parent_thread_id: Option, - #[serde(alias = "useWebSocket")] - use_websocket: Option, -} - -pub async fn responses( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - body: Bytes, -) -> Response { - let started_at = Instant::now(); - if !authorize_client_api_key(&state, &headers).await { - return missing_client_api_key_response().into_response(); - } - - let default_model = state.config().model.default_model.clone(); - let body = serde_json::from_slice::(&body) - .unwrap_or_else(|_| default_body(default_model.clone())); - let client_stream = body.stream.unwrap_or(true); - let requested_model = body.model.clone().unwrap_or(default_model); - let catalog = model_catalog_for_state(&state).await; - if !catalog.is_recognized_model_name(&requested_model) { - return ( - StatusCode::NOT_FOUND, - Json(openai_error("Model not found", "model_not_found")), - ) - .into_response(); - } - let parsed_model = catalog.parse_model_name(&requested_model); - let mut codex_request = CodexResponsesRequest::new_http_sse( - parsed_model.model_id.clone(), - body.instructions.unwrap_or_default(), - body.input.unwrap_or_default(), - ); - codex_request.reasoning = body.reasoning; - codex_request.tools = body.tools; - codex_request.tool_choice = body.tool_choice; - codex_request.parallel_tool_calls = body.parallel_tool_calls; - codex_request.text = body.text; - codex_request.prompt_cache_key = body.prompt_cache_key; - codex_request.include = body.include; - codex_request.client_metadata = body.client_metadata; - codex_request.previous_response_id = body.previous_response_id; - codex_request.reasoning = responses_reasoning( - codex_request.reasoning.take(), - parsed_model.reasoning_effort.as_deref(), - state.config().model.default_reasoning_effort.as_deref(), - ); - codex_request.service_tier = body - .service_tier - .or(parsed_model.service_tier) - .or_else(|| state.config().model.service_tier.clone()) - .map(normalize_service_tier_for_upstream); - ensure_reasoning_include(&mut codex_request); - codex_request.turn_state = body - .turn_state - .or_else(|| header_string(&headers, "x-codex-turn-state")); - codex_request.turn_metadata = body - .turn_metadata - .or_else(|| header_string(&headers, "x-codex-turn-metadata")); - codex_request.beta_features = body - .beta_features - .or_else(|| header_string(&headers, "x-codex-beta-features")); - codex_request.include_timing_metrics = body - .include_timing_metrics - .or_else(|| header_string(&headers, "x-responsesapi-include-timing-metrics")); - codex_request.version = header_string(&headers, "version"); - codex_request.codex_window_id = body - .codex_window_id - .or_else(|| header_string(&headers, "x-codex-window-id")); - codex_request.parent_thread_id = body - .parent_thread_id - .or_else(|| header_string(&headers, "x-codex-parent-thread-id")); - codex_request.use_websocket = body.use_websocket.unwrap_or(false); - let acquired = { - state - .account_pool() - .lock() - .await - .acquire_with(AccountAcquireRequest::new(&codex_request.model, Utc::now())) - }; - let Some(acquired) = acquired else { - return no_available_accounts_response().into_response(); - }; - let mut account = acquired.account; - let mut log_context = V1LogContext::new( - request_id.as_str(), - &account.id, - &codex_request.model, - client_stream, - started_at, - ); - - if client_stream { - return responses_stream(state, codex_request, account, log_context).await; - } - - let mut excluded_account_ids = Vec::new(); - let response = loop { - 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); - - match response { - Ok(response) => break response, - Err(error) => { - if let Some(retry) = classify_upstream_account_retry(&error) { - if codex_request.previous_response_id.is_some() { - // previous_response_id 的历史由上游账号持有，换账号会静默丢失会话上下文。 - apply_upstream_account_retry(&state, &account, retry).await; - log_v1_response( - &state, - &log_context, - retry.status(), - EventLevel::Warn, - "v1 responses websocket history request kept on original account", - websocket_history_retry_metadata(retry, false), - ) - .await; - } else { - let fallback = apply_upstream_retry_and_acquire_fallback( - &state, - &account, - retry, - &codex_request.model, - &mut excluded_account_ids, - ) - .await; - log_v1_response( - &state, - &log_context, - retry.status(), - EventLevel::Warn, - "v1 responses upstream retrying with fallback account", - retry.metadata(false), - ) - .await; - if let Some(fallback) = fallback { - account = fallback; - log_context = log_context.with_account(&account.id); - continue; - } - } - } - let error_response = codex_client_error_response(error); - log_v1_response( - &state, - &log_context, - error_response.0, - EventLevel::Error, - "v1 responses upstream request failed", - json!({"stream": false}), - ) - .await; - return error_response.into_response(); - } - } - }; - if persist_upstream_cookies(&state, &account.id, &response.set_cookie_headers) - .await - .is_err() - { - log_v1_response( - &state, - &log_context, - StatusCode::INTERNAL_SERVER_ERROR, - EventLevel::Error, - "v1 responses cookie persistence failed", - json!({"stream": false, "cookieStoreError": true}), - ) - .await; - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(openai_error( - "Failed to persist upstream cookies", - "cookie_store_error", - )), - ) - .into_response(); - } - if let Some(usage) = response.usage { - if record_usage(&state, &account.id, usage).await.is_err() { - log_v1_response( - &state, - &log_context, - StatusCode::INTERNAL_SERVER_ERROR, - EventLevel::Error, - "v1 responses usage persistence failed", - json!({"stream": false, "usage": usage, "usageStoreError": true}), - ) - .await; - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(openai_error( - "Failed to record account usage", - "usage_store_error", - )), - ) - .into_response(); - } - } - - match completed_response_json(&response.body) { - Ok(CollectedResponse::Completed(body)) => { - log_v1_response( - &state, - &log_context, - StatusCode::OK, - EventLevel::Info, - "v1 responses completed", - json!({"stream": false, "usage": response.usage}), - ) - .await; - (StatusCode::OK, Json(body)).into_response() - } - Ok(CollectedResponse::Failed(failure)) => { - log_v1_response( - &state, - &log_context, - StatusCode::BAD_GATEWAY, - EventLevel::Error, - "v1 responses upstream SSE failed", - failure.metadata(false), - ) - .await; - ( - StatusCode::BAD_GATEWAY, - Json(openai_error( - &failure.openai_error_message(), - "upstream_error", - )), - ) - .into_response() - } - Ok(CollectedResponse::MissingCompleted) => { - log_v1_response( - &state, - &log_context, - StatusCode::BAD_GATEWAY, - EventLevel::Warn, - "v1 responses completed event missing", - json!({"stream": false, "usage": response.usage}), - ) - .await; - ( - StatusCode::BAD_GATEWAY, - Json(openai_error( - "Codex response did not include response.completed", - "empty_upstream_response", - )), - ) - .into_response() - } - Err(error) => { - log_v1_response( - &state, - &log_context, - StatusCode::BAD_GATEWAY, - EventLevel::Warn, - "v1 responses invalid SSE response", - json!({"stream": false, "sseParseError": error.to_string()}), - ) - .await; - ( - StatusCode::BAD_GATEWAY, - Json(openai_error( - "Invalid Codex SSE response", - "invalid_upstream_sse", - )), - ) - .into_response() - } - } -} - -pub async fn chat_completions( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - body: Bytes, -) -> Response { - let started_at = Instant::now(); - if !authorize_client_api_key(&state, &headers).await { - return missing_client_api_key_response().into_response(); - } - - let Ok(chat_request) = serde_json::from_slice::(&body) else { - return ( - StatusCode::BAD_REQUEST, - Json(openai_error( - "Invalid chat completion request", - "invalid_request", - )), - ) - .into_response(); - }; - let client_stream = chat_request.stream; - let requested_model = chat_request.model.clone(); - let catalog = model_catalog_for_state(&state).await; - if !catalog.is_recognized_model_name(&requested_model) { - return model_not_found_response().into_response(); - } - let parsed_model = catalog.parse_model_name(&requested_model); - let display_model = ModelCatalog::build_display_model_name(&parsed_model); - let mut codex_request = match translate_chat_to_codex(chat_request) { - Ok(request) => request, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(openai_error( - "Invalid chat completion request", - "invalid_request", - )), - ) - .into_response(); - } - }; - codex_request.model = parsed_model.model_id.clone(); - if codex_request.reasoning.is_none() { - let effort = parsed_model - .reasoning_effort - .clone() - .or_else(|| state.config().model.default_reasoning_effort.clone()); - if let Some(effort) = effort { - codex_request.reasoning = Some(json!({"effort": effort, "summary": "auto"})); - } - } - if codex_request.service_tier.is_none() { - codex_request.service_tier = parsed_model - .service_tier - .clone() - .or_else(|| state.config().model.service_tier.clone()); - } - codex_request.service_tier = codex_request - .service_tier - .map(normalize_service_tier_for_upstream); - let include_reasoning = codex_request.reasoning.is_some(); - - let acquired = { - state - .account_pool() - .lock() - .await - .acquire_with(AccountAcquireRequest::new(&codex_request.model, Utc::now())) - }; - let Some(acquired) = acquired else { - return no_available_accounts_response().into_response(); - }; - let mut account = acquired.account; - let mut log_context = V1LogContext::new( - request_id.as_str(), - &account.id, - &codex_request.model, - client_stream, - started_at, - ); - - let mut excluded_account_ids = Vec::new(); - let response = loop { - 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); - - match response { - Ok(response) => break response, - Err(error) => { - if let Some(retry) = classify_upstream_account_retry(&error) { - let fallback = apply_upstream_retry_and_acquire_fallback( - &state, - &account, - retry, - &codex_request.model, - &mut excluded_account_ids, - ) - .await; - log_v1_response( - &state, - &log_context, - retry.status(), - EventLevel::Warn, - "v1 chat completions upstream retrying with fallback account", - retry.metadata(client_stream), - ) - .await; - if let Some(fallback) = fallback { - account = fallback; - log_context = log_context.with_account(&account.id); - continue; - } - } - let error_response = codex_client_error_response(error); - log_v1_response( - &state, - &log_context, - error_response.0, - EventLevel::Error, - "v1 chat completions upstream request failed", - json!({"stream": client_stream}), - ) - .await; - return error_response.into_response(); - } - } - }; - if persist_upstream_cookies(&state, &account.id, &response.set_cookie_headers) - .await - .is_err() - { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(openai_error( - "Failed to persist upstream cookies", - "cookie_store_error", - )), - ) - .into_response(); - } - if let Some(usage) = response.usage { - if record_usage(&state, &account.id, usage).await.is_err() { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(openai_error( - "Failed to record account usage", - "usage_store_error", - )), - ) - .into_response(); - } - } - - if client_stream { - match chat_completion_stream_from_codex_sse( - &response.body, - &display_model, - include_reasoning, - ) { - Ok(Some(body)) => Response::builder() - .status(StatusCode::OK) - .header(CONTENT_TYPE, "text/event-stream") - .header(CACHE_CONTROL, "no-cache") - .body(Body::from(body)) - .unwrap_or_else(|_| { - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(openai_error( - "Failed to build stream response", - "stream_response_error", - )), - ) - .into_response() - }), - Ok(None) => ( - StatusCode::BAD_GATEWAY, - Json(openai_error( - "Codex response did not include response.completed", - "empty_upstream_response", - )), - ) - .into_response(), - Err(_) => ( - StatusCode::BAD_GATEWAY, - Json(openai_error( - "Invalid Codex SSE response", - "invalid_upstream_sse", - )), - ) - .into_response(), - } - } else { - match chat_completion_from_codex_sse(&response.body, &display_model, include_reasoning) { - Ok(Some(body)) => (StatusCode::OK, Json(body)).into_response(), - Ok(None) => ( - StatusCode::BAD_GATEWAY, - Json(openai_error( - "Codex response did not include response.completed", - "empty_upstream_response", - )), - ) - .into_response(), - Err(_) => ( - StatusCode::BAD_GATEWAY, - Json(openai_error( - "Invalid Codex SSE response", - "invalid_upstream_sse", - )), - ) - .into_response(), - } - } -} - -fn no_available_accounts_response() -> (StatusCode, Json) { - ( - StatusCode::SERVICE_UNAVAILABLE, - Json(openai_error( - "No available Codex accounts", - "no_available_accounts", - )), - ) -} - -async fn model_catalog_for_state(state: &AppState) -> ModelCatalog { - let Some(repo) = state.model_snapshot_repository() else { - return ModelCatalog::from_config(&state.config().model); - }; - match repo.list_plan_snapshots().await { - Ok(snapshots) if !snapshots.is_empty() => { - ModelCatalog::from_config_and_snapshots(&state.config().model, &snapshots) - } - Ok(_) => ModelCatalog::from_config(&state.config().model), - Err(error) => { - // 模型缓存只是加速和扩展 catalog；读取失败时回退静态配置，避免影响核心请求链路。 - tracing::warn!(?error, "failed to load cached model snapshots"); - ModelCatalog::from_config(&state.config().model) - } - } -} - -fn default_body(default_model: String) -> ResponsesBody { - ResponsesBody { - model: Some(default_model), - input: Some(Vec::new()), - instructions: Some(String::new()), - reasoning: None, - tools: None, - service_tier: None, - tool_choice: None, - parallel_tool_calls: None, - text: None, - prompt_cache_key: None, - include: None, - client_metadata: None, - previous_response_id: None, - stream: None, - turn_state: None, - turn_metadata: None, - beta_features: None, - include_timing_metrics: None, - codex_window_id: None, - parent_thread_id: None, - use_websocket: None, - } -} - -fn responses_reasoning( - client_reasoning: Option, - suffix_effort: Option<&str>, - default_effort: Option<&str>, -) -> Option { - let effort = client_reasoning - .as_ref() - .and_then(|reasoning| reasoning.get("effort")) - .and_then(Value::as_str) - .or(suffix_effort) - .or(default_effort); - let summary = client_reasoning - .as_ref() - .and_then(|reasoning| reasoning.get("summary")) - .and_then(Value::as_str) - .unwrap_or("auto"); - if effort.is_none() && client_reasoning.is_none() { - return None; - } - let mut reasoning = json!({"summary": summary}); - if let Some(effort) = effort { - reasoning["effort"] = Value::String(effort.to_string()); - } - Some(reasoning) -} - -fn normalize_service_tier_for_upstream(service_tier: String) -> String { - if service_tier == "fast" { - "priority".to_string() - } else { - service_tier - } -} - -fn ensure_reasoning_include(request: &mut CodexResponsesRequest) { - if request.reasoning.is_none() { - return; - } - let include = request.include.get_or_insert_with(Vec::new); - if include - .iter() - .any(|item| item == "reasoning.encrypted_content") - { - return; - } - include.push("reasoning.encrypted_content".to_string()); -} - -fn header_string(headers: &HeaderMap, name: &str) -> Option { - headers - .get(name) - .and_then(|value| value.to_str().ok()) - .filter(|value| !value.trim().is_empty()) - .map(ToString::to_string) -} - -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, - retry_after_seconds, - }) if status == StatusCode::UNAUTHORIZED => { - let Some(refreshed) = - refresh_account_after_unauthorized(state, request, account, request_id).await - else { - return Err(CodexClientError::Upstream { - status, - body, - retry_after_seconds, - }); - }; - send_codex_request(state, request, &refreshed, request_id).await - } - result => result, - } -} - -async fn send_codex_request( - state: &AppState, - request: &CodexResponsesRequest, - account: &Account, - request_id: &str, -) -> Result { - let request_domain = request_domain(&state.config().api.base_url); - let cookie_header = match (state.cookie_repository(), request_domain.as_deref()) { - (Some(repo), Some(domain)) => repo.cookie_header(&account.id, domain).await.ok().flatten(), - _ => None, - }; - let client = CodexBackendClient::new( - build_reqwest_client(state.config().tls.force_http11)?, - state.config().api.base_url.clone(), - Fingerprint::default_codex_desktop(), - ); - client - .create_response( - request, - CodexRequestContext { - access_token: &account.access_token, - account_id: account.account_id.as_deref(), - request_id, - turn_state: request.turn_state.as_deref(), - turn_metadata: request.turn_metadata.as_deref(), - beta_features: request.beta_features.as_deref(), - include_timing_metrics: request.include_timing_metrics.as_deref(), - version: request.version.as_deref(), - codex_window_id: request.codex_window_id.as_deref(), - parent_thread_id: request.parent_thread_id.as_deref(), - cookie_header: cookie_header.as_deref(), - }, - ) - .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, - retry_after_seconds, - }) if status == StatusCode::UNAUTHORIZED => { - let Some(refreshed) = - refresh_account_after_unauthorized(state, request, account, request_id).await - else { - return Err(CodexClientError::Upstream { - status, - body, - retry_after_seconds, - }); - }; - send_codex_stream_request(state, request, &refreshed, request_id).await - } - result => result, - } -} - -async fn send_codex_stream_request( - state: &AppState, - request: &CodexResponsesRequest, - account: &Account, - request_id: &str, -) -> Result { - let request_domain = request_domain(&state.config().api.base_url); - let cookie_header = match (state.cookie_repository(), request_domain.as_deref()) { - (Some(repo), Some(domain)) => repo.cookie_header(&account.id, domain).await.ok().flatten(), - _ => None, - }; - let client = CodexBackendClient::new( - build_reqwest_client(state.config().tls.force_http11)?, - state.config().api.base_url.clone(), - Fingerprint::default_codex_desktop(), - ); - client - .stream_response( - request, - CodexRequestContext { - access_token: &account.access_token, - account_id: account.account_id.as_deref(), - request_id, - turn_state: request.turn_state.as_deref(), - turn_metadata: request.turn_metadata.as_deref(), - beta_features: request.beta_features.as_deref(), - include_timing_metrics: request.include_timing_metrics.as_deref(), - version: request.version.as_deref(), - codex_window_id: request.codex_window_id.as_deref(), - parent_thread_id: request.parent_thread_id.as_deref(), - cookie_header: cookie_header.as_deref(), - }, - ) - .await -} - -#[derive(Debug, Clone, Copy)] -enum UpstreamAccountRetry { - RateLimited { retry_after_seconds: u64 }, - QuotaExhausted, - CloudflareChallenge { cooldown_seconds: u64 }, - Banned, -} - -impl UpstreamAccountRetry { - fn status(self) -> StatusCode { - match self { - Self::RateLimited { .. } => StatusCode::TOO_MANY_REQUESTS, - Self::QuotaExhausted => StatusCode::PAYMENT_REQUIRED, - Self::CloudflareChallenge { .. } => StatusCode::FORBIDDEN, - Self::Banned => StatusCode::FORBIDDEN, - } - } - - fn metadata(self, stream: bool) -> Value { - match self { - Self::RateLimited { - retry_after_seconds, - } => json!({ - "stream": stream, - "retry": true, - "reason": "rateLimited", - "retryAfterSeconds": retry_after_seconds, - }), - Self::QuotaExhausted => json!({ - "stream": stream, - "retry": true, - "reason": "quotaExhausted", - }), - Self::CloudflareChallenge { cooldown_seconds } => json!({ - "stream": stream, - "retry": true, - "reason": "cloudflareChallenge", - "cooldownSeconds": cooldown_seconds, - }), - Self::Banned => json!({ - "stream": stream, - "retry": true, - "reason": "banned", - }), - } - } -} - -fn classify_upstream_account_retry(error: &CodexClientError) -> Option { - match error { - CodexClientError::Upstream { - status, - retry_after_seconds, - .. - } if *status == StatusCode::TOO_MANY_REQUESTS => Some(UpstreamAccountRetry::RateLimited { - retry_after_seconds: retry_after_seconds - .unwrap_or(DEFAULT_RATE_LIMIT_BACKOFF_SECONDS) - .min(MAX_RATE_LIMIT_BACKOFF_SECONDS), - }), - CodexClientError::Upstream { status, .. } if *status == StatusCode::PAYMENT_REQUIRED => { - Some(UpstreamAccountRetry::QuotaExhausted) - } - CodexClientError::Upstream { status, body, .. } if *status == StatusCode::FORBIDDEN => { - if is_cloudflare_challenge(body) { - Some(UpstreamAccountRetry::CloudflareChallenge { - cooldown_seconds: CLOUDFLARE_CHALLENGE_COOLDOWN_SECONDS, - }) - } else { - Some(UpstreamAccountRetry::Banned) - } - } - _ => None, - } -} - -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") -} - -async fn apply_upstream_retry_and_acquire_fallback( - state: &AppState, - account: &Account, - retry: UpstreamAccountRetry, - model: &str, - excluded_account_ids: &mut Vec, -) -> Option { - apply_upstream_account_retry(state, account, retry).await; - excluded_account_ids.push(account.id.clone()); - state - .account_pool() - .lock() - .await - .acquire_with( - AccountAcquireRequest::new(model, Utc::now()) - .with_exclude_account_ids(excluded_account_ids.iter().cloned()), - ) - .map(|fallback| fallback.account) -} - -async fn apply_upstream_account_retry( - state: &AppState, - account: &Account, - retry: UpstreamAccountRetry, -) { - match retry { - UpstreamAccountRetry::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); - if record_request_attempt(state, &account.id).await.is_err() { - tracing::warn!( - account_id = %account.id, - "failed to record rate-limited account attempt" - ); - } - } - UpstreamAccountRetry::QuotaExhausted => { - set_account_status(state, account, AccountStatus::QuotaExhausted).await; - } - UpstreamAccountRetry::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); - } - UpstreamAccountRetry::Banned => { - set_account_status(state, account, AccountStatus::Banned).await; - } - } -} - -async fn set_account_status(state: &AppState, account: &Account, status: AccountStatus) { - if let Some(repo) = state.account_repository() { - if repo.set_status(&account.id, status).await.is_err() { - tracing::warn!( - account_id = %account.id, - "failed to persist upstream account status" - ); - } - } - state - .account_pool() - .lock() - .await - .set_status(&account.id, status); -} - -async fn responses_stream( - state: AppState, - request: CodexResponsesRequest, - mut account: Account, - mut log_context: V1LogContext, -) -> Response { - if transport_for_request(&request) == CodexTransport::WebSocketRequired { - return responses_websocket_stream(state, request, account, log_context).await; - } - - let mut excluded_account_ids = Vec::new(); - let stream_response = loop { - let stream_response = send_codex_stream_request_with_refresh_retry( - &state, - &request, - &account, - log_context.request_id.as_str(), - ) - .await; - match stream_response { - Ok(response) => break response, - Err(error) => { - state.account_pool().lock().await.release(&account.id); - if let Some(retry) = classify_upstream_account_retry(&error) { - let fallback = apply_upstream_retry_and_acquire_fallback( - &state, - &account, - retry, - &request.model, - &mut excluded_account_ids, - ) - .await; - log_v1_response( - &state, - &log_context, - retry.status(), - EventLevel::Warn, - "v1 responses stream upstream retrying with fallback account", - retry.metadata(true), - ) - .await; - if let Some(fallback) = fallback { - account = fallback; - log_context = log_context.with_account(&account.id); - continue; - } - } - let error_response = codex_client_error_response(error); - log_v1_response( - &state, - &log_context, - error_response.0, - EventLevel::Error, - "v1 responses stream upstream request failed", - json!({"stream": true}), - ) - .await; - return error_response.into_response(); - } - } - }; - - if persist_upstream_cookies(&state, &account.id, &stream_response.set_cookie_headers) - .await - .is_err() - { - state.account_pool().lock().await.release(&account.id); - log_v1_response( - &state, - &log_context, - StatusCode::INTERNAL_SERVER_ERROR, - EventLevel::Error, - "v1 responses stream cookie persistence failed", - json!({"stream": true, "cookieStoreError": true}), - ) - .await; - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(openai_error( - "Failed to persist upstream cookies", - "cookie_store_error", - )), - ) - .into_response(); - } - - let upstream = Box::pin(stream_response.response.bytes_stream()); - let audit = StreamAudit::new(state, log_context, account.id); - let body_stream = stream::unfold(Some((upstream, Vec::new(), audit)), |state| async move { - let (mut upstream, mut collected, mut audit) = state?; - match upstream.next().await { - Some(Ok(chunk)) => { - collected.extend_from_slice(&chunk); - Some((Ok(chunk), Some((upstream, collected, audit)))) - } - Some(Err(error)) => { - audit.log_transport_error(&error).await; - Some((Err(error), None)) - } - None => { - audit.complete(&collected).await; - None - } - } - }); - - Response::builder() - .status(StatusCode::OK) - .header(CONTENT_TYPE, "text/event-stream") - .header(CACHE_CONTROL, "no-cache") - .body(Body::from_stream(body_stream)) - .unwrap_or_else(|_| { - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(openai_error( - "Failed to build stream response", - "stream_response_error", - )), - ) - .into_response() - }) -} - -async fn responses_websocket_stream( - state: AppState, - request: CodexResponsesRequest, - mut account: Account, - mut log_context: V1LogContext, -) -> Response { - let mut excluded_account_ids = Vec::new(); - let response = loop { - let response = send_codex_request_with_refresh_retry( - &state, - &request, - &account, - log_context.request_id.as_str(), - ) - .await; - state.account_pool().lock().await.release(&account.id); - - match response { - Ok(response) => break response, - Err(error) => { - if let Some(retry) = classify_upstream_account_retry(&error) { - if request.previous_response_id.is_some() { - // previous_response_id 的历史由上游账号持有，换账号会静默丢失会话上下文。 - apply_upstream_account_retry(&state, &account, retry).await; - log_v1_response( - &state, - &log_context, - retry.status(), - EventLevel::Warn, - "v1 responses websocket history request kept on original account", - websocket_history_retry_metadata(retry, true), - ) - .await; - } else { - let fallback = apply_upstream_retry_and_acquire_fallback( - &state, - &account, - retry, - &request.model, - &mut excluded_account_ids, - ) - .await; - log_v1_response( - &state, - &log_context, - retry.status(), - EventLevel::Warn, - "v1 responses websocket upstream retrying with fallback account", - retry.metadata(true), - ) - .await; - if let Some(fallback) = fallback { - account = fallback; - log_context = log_context.with_account(&account.id); - continue; - } - } - } - let error_response = codex_client_error_response(error); - log_v1_response( - &state, - &log_context, - error_response.0, - EventLevel::Error, - "v1 responses websocket stream upstream request failed", - json!({"stream": true, "transport": "websocket"}), - ) - .await; - return error_response.into_response(); - } - } - }; - - if persist_upstream_cookies(&state, &account.id, &response.set_cookie_headers) - .await - .is_err() - { - log_v1_response( - &state, - &log_context, - StatusCode::INTERNAL_SERVER_ERROR, - EventLevel::Error, - "v1 responses websocket stream cookie persistence failed", - json!({"stream": true, "transport": "websocket", "cookieStoreError": true}), - ) - .await; - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(openai_error( - "Failed to persist upstream cookies", - "cookie_store_error", - )), - ) - .into_response(); - } - - let rate_limit_headers = response.rate_limit_headers.clone(); - let mut level = EventLevel::Info; - let mut message = "v1 responses websocket stream completed"; - let mut metadata = json!({ - "stream": true, - "transport": "websocket", - "usage": response.usage, - "rateLimitHeaders": rate_limit_headers.clone(), - }); - if let Some(usage) = response.usage { - if record_usage(&state, &account.id, usage).await.is_err() { - level = EventLevel::Warn; - message = "v1 responses websocket stream completed with usage store error"; - metadata = json!({ - "stream": true, - "transport": "websocket", - "usage": usage, - "rateLimitHeaders": rate_limit_headers.clone(), - "usageStoreError": true, - }); - } - } - log_v1_response( - &state, - &log_context, - StatusCode::OK, - level, - message, - metadata, - ) - .await; - - Response::builder() - .status(StatusCode::OK) - .header(CONTENT_TYPE, "text/event-stream") - .header(CACHE_CONTROL, "no-cache") - .body(Body::from(response.body)) - .unwrap_or_else(|_| { - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(openai_error( - "Failed to build stream response", - "stream_response_error", - )), - ) - .into_response() - }) -} - -fn websocket_history_retry_metadata(retry: UpstreamAccountRetry, stream: bool) -> Value { - match retry { - UpstreamAccountRetry::RateLimited { - retry_after_seconds, - } => json!({ - "stream": stream, - "transport": "websocket", - "retry": false, - "reason": "rateLimited", - "retryAfterSeconds": retry_after_seconds, - "accountAffinity": "previousResponseId", - }), - UpstreamAccountRetry::QuotaExhausted => json!({ - "stream": stream, - "transport": "websocket", - "retry": false, - "reason": "quotaExhausted", - "accountAffinity": "previousResponseId", - }), - UpstreamAccountRetry::CloudflareChallenge { cooldown_seconds } => json!({ - "stream": stream, - "transport": "websocket", - "retry": false, - "reason": "cloudflareChallenge", - "cooldownSeconds": cooldown_seconds, - "accountAffinity": "previousResponseId", - }), - UpstreamAccountRetry::Banned => json!({ - "stream": stream, - "transport": "websocket", - "retry": false, - "reason": "banned", - "accountAffinity": "previousResponseId", - }), - } -} - -async fn refresh_account_after_unauthorized( - state: &AppState, - request: &CodexResponsesRequest, - account: &Account, - request_id: &str, -) -> 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, request_id, &request.model).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, - request_id: &str, - model: &str, -) { - let status = status_for_refresh_failure(failure); - if let Some(status) = status { - 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); - } - log_account_refresh_failure(state, account, failure, status, request_id, model).await; -} - -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 log_account_refresh_failure( - state: &AppState, - account: &Account, - failure: RefreshFailure, - status: Option, - request_id: &str, - model: &str, -) { - let Some(repo) = state.event_logs() else { - return; - }; - let mut event = EventLog::new( - "account.refresh", - EventLevel::Warn, - "account refresh failed after upstream 401", - ); - event.request_id = Some(request_id.to_string()); - event.account_id = Some(account.id.clone()); - event.route = Some("/v1/responses".to_string()); - event.model = Some(model.to_string()); - event.status_code = Some(i64::from(StatusCode::UNAUTHORIZED.as_u16())); - event.metadata = json!({ - "trigger": "upstream_401", - "failure": refresh_failure_value(failure), - "accountStatus": status.map(account_status_value), - }); - if let Err(error) = repo.insert(event).await { - tracing::warn!(?error, "failed to insert account refresh event log"); - } -} - -fn refresh_failure_value(failure: RefreshFailure) -> &'static str { - match failure { - RefreshFailure::InvalidGrant => "invalidGrant", - RefreshFailure::QuotaExhausted => "quotaExhausted", - RefreshFailure::Banned => "banned", - RefreshFailure::Disabled => "disabled", - RefreshFailure::Transport => "transport", - } -} - -fn account_status_value(status: AccountStatus) -> &'static str { - match status { - AccountStatus::Active => "active", - AccountStatus::Expired => "expired", - AccountStatus::QuotaExhausted => "quotaExhausted", - AccountStatus::Refreshing => "refreshing", - AccountStatus::Disabled => "disabled", - AccountStatus::Banned => "banned", - } -} - -async fn persist_upstream_cookies( - state: &AppState, - account_id: &str, - set_cookie_headers: &[String], -) -> Result<(), ()> { - let Some(cookie_repo) = state.cookie_repository() else { - return Ok(()); - }; - for cookie in set_cookie_headers { - cookie_repo - .capture_set_cookie(account_id, cookie) - .await - .map_err(|_| ())?; - } - Ok(()) -} - -async fn record_usage(state: &AppState, account_id: &str, usage: TokenUsage) -> Result<(), ()> { - let Some(repo) = state.account_repository() else { - return Ok(()); - }; - repo.record_usage( - account_id, - UsageDelta { - input_tokens: u64_to_i64_saturating(usage.input_tokens), - output_tokens: u64_to_i64_saturating(usage.output_tokens), - cached_tokens: u64_to_i64_saturating(usage.cached_tokens), - }, - ) - .await - .map_err(|_| ()) -} - -async fn record_request_attempt(state: &AppState, account_id: &str) -> Result<(), ()> { - let Some(repo) = state.account_repository() else { - return Ok(()); - }; - repo.record_usage( - account_id, - UsageDelta { - input_tokens: 0, - output_tokens: 0, - cached_tokens: 0, - }, - ) - .await - .map_err(|_| ()) -} - -#[derive(Clone)] -struct V1LogContext { - request_id: String, - account_id: String, - model: String, - stream: bool, - started_at: Instant, -} - -impl V1LogContext { - fn new( - request_id: &str, - account_id: &str, - model: &str, - stream: bool, - started_at: Instant, - ) -> Self { - Self { - request_id: request_id.to_string(), - account_id: account_id.to_string(), - model: model.to_string(), - stream, - started_at, - } - } - - fn latency_ms(&self) -> i64 { - self.started_at.elapsed().as_millis().min(i64::MAX as u128) as i64 - } - - fn with_account(&self, account_id: &str) -> Self { - Self { - request_id: self.request_id.clone(), - account_id: account_id.to_string(), - model: self.model.clone(), - stream: self.stream, - started_at: self.started_at, - } - } -} - -struct StreamAudit { - state: AppState, - context: V1LogContext, - account_slot: AccountSlotGuard, -} - -impl StreamAudit { - fn new(state: AppState, context: V1LogContext, account_id: String) -> Self { - let account_slot = AccountSlotGuard::new(state.account_pool(), account_id); - Self { - state, - context, - account_slot, - } - } - - async fn complete(&mut self, body: &[u8]) { - let body = String::from_utf8_lossy(body); - let mut status = StatusCode::OK; - let mut level = EventLevel::Info; - let mut message = "v1 responses stream completed"; - let mut metadata = match extract_sse_usage(&body) { - Ok(usage) => { - if let Some(usage) = usage { - if record_usage(&self.state, &self.context.account_id, usage) - .await - .is_err() - { - level = EventLevel::Warn; - message = "v1 responses stream completed with usage store error"; - json!({"stream": true, "usage": usage, "usageStoreError": true}) - } else { - json!({"stream": true, "usage": usage}) - } - } else { - json!({"stream": true, "usage": null}) - } - } - Err(error) => { - level = EventLevel::Warn; - message = "v1 responses stream completed with invalid SSE usage"; - json!({"stream": true, "sseParseError": error.to_string()}) - } - }; - match responses_sse_failure(&body) { - Ok(Some(failure)) => { - // SSE 响应头已发出，HTTP 状态不能回滚；用终止事件透传给客户端，并在审计里标记上游失败。 - status = StatusCode::BAD_GATEWAY; - level = EventLevel::Error; - message = "v1 responses stream upstream SSE failed"; - failure.extend_metadata(&mut metadata); - } - Ok(None) => {} - Err(error) => { - level = EventLevel::Warn; - message = "v1 responses stream completed with invalid SSE failure metadata"; - metadata = json!({"stream": true, "sseParseError": error.to_string()}); - } - } - ensure_stream_metadata(&mut metadata, true); - log_v1_response(&self.state, &self.context, status, level, message, metadata).await; - self.account_slot.release().await; - } - - async fn log_transport_error(&mut self, error: &reqwest::Error) { - log_v1_response( - &self.state, - &self.context, - StatusCode::BAD_GATEWAY, - EventLevel::Error, - "v1 responses stream transport failed", - json!({"stream": true, "transportError": error.to_string()}), - ) - .await; - self.account_slot.release().await; - } -} - -struct AccountSlotGuard { - pool: Arc>, - account_id: String, - released: bool, -} - -impl AccountSlotGuard { - fn new(pool: Arc>, account_id: String) -> Self { - Self { - pool, - account_id, - released: false, - } - } - - async fn release(&mut self) { - if self.released { - return; - } - self.pool.lock().await.release(&self.account_id); - self.released = true; - } -} - -impl Drop for AccountSlotGuard { - fn drop(&mut self) { - if self.released { - return; - } - let pool = self.pool.clone(); - let account_id = self.account_id.clone(); - let Ok(handle) = tokio::runtime::Handle::try_current() else { - return; - }; - handle.spawn(async move { - pool.lock().await.release(&account_id); - }); - } -} - -async fn log_v1_response( - state: &AppState, - context: &V1LogContext, - status: StatusCode, - level: EventLevel, - message: &str, - mut metadata: Value, -) { - let Some(repo) = state.event_logs() else { - return; - }; - ensure_stream_metadata(&mut metadata, context.stream); - let mut event = EventLog::new("v1.response", level, message); - event.request_id = Some(context.request_id.clone()); - event.account_id = Some(context.account_id.clone()); - event.route = Some("/v1/responses".to_string()); - event.model = Some(context.model.clone()); - event.status_code = Some(i64::from(status.as_u16())); - event.latency_ms = Some(context.latency_ms()); - event.metadata = metadata; - if let Err(error) = repo.insert(event).await { - tracing::warn!(?error, "failed to insert v1 response event log"); - } -} - -fn ensure_stream_metadata(metadata: &mut Value, stream_value: bool) { - let Some(object) = metadata.as_object_mut() else { - *metadata = json!({"stream": stream_value}); - return; - }; - object - .entry("stream".to_string()) - .or_insert_with(|| json!(stream_value)); -} - -enum CollectedResponse { - Completed(Value), - Failed(ResponsesSseFailure), - MissingCompleted, -} - -#[derive(Debug, Clone)] -struct ResponsesSseFailure { - event: String, - message: String, - upstream_code: Option, -} - -impl ResponsesSseFailure { - fn from_event(event: &str, value: &Value) -> Self { - Self { - event: event.to_string(), - message: failure_message(value).unwrap_or_else(|| "Codex upstream SSE failed".into()), - upstream_code: failure_code(value), - } - } - - fn openai_error_message(&self) -> String { - format!("Codex upstream SSE failed: {}", self.message) - } - - fn metadata(&self, stream: bool) -> Value { - let mut metadata = json!({"stream": stream}); - self.extend_metadata(&mut metadata); - metadata - } - - fn extend_metadata(&self, metadata: &mut Value) { - let Some(object) = metadata.as_object_mut() else { - *metadata = self.metadata(true); - return; - }; - object.insert("failureEvent".to_string(), json!(self.event)); - object.insert("failureMessage".to_string(), json!(self.message)); - if let Some(code) = &self.upstream_code { - object.insert("upstreamCode".to_string(), json!(code)); - } - } -} - -fn completed_response_json(body: &str) -> Result { - let events = parse_sse_events(body)?; - let mut output_text = String::new(); - let mut output_items = Vec::new(); - let mut completed_response = None; - let mut failed_response = None; - for event in events { - let Ok(value) = serde_json::from_str::(&event.data) else { - continue; - }; - match event.event.as_deref() { - Some("response.output_text.delta") => { - if let Some(delta) = value.get("delta").and_then(Value::as_str) { - output_text.push_str(delta); - } - } - Some("response.output_item.done") => { - if let Some(item) = value.get("item") { - output_items.push(item.clone()); - } - } - Some("response.completed") => { - if let Some(response) = value.get("response") { - completed_response = Some(response.clone()); - } - } - Some(event_name @ ("error" | "response.failed")) if failed_response.is_none() => { - failed_response = Some(ResponsesSseFailure::from_event(event_name, &value)); - } - _ => {} - } - } - if let Some(failure) = failed_response { - return Ok(CollectedResponse::Failed(failure)); - } - let Some(mut response) = completed_response else { - return Ok(CollectedResponse::MissingCompleted); - }; - ensure_completed_response_output(&mut response, &output_items, &output_text); - sync_output_text_from_output(&mut response); - Ok(CollectedResponse::Completed(response)) -} - -fn responses_sse_failure( - body: &str, -) -> Result, crate::codex::sse::SseError> { - for event in parse_sse_events(body)? { - let Some(event_name @ ("error" | "response.failed")) = event.event.as_deref() else { - continue; - }; - let Ok(value) = serde_json::from_str::(&event.data) else { - continue; - }; - return Ok(Some(ResponsesSseFailure::from_event(event_name, &value))); - } - Ok(None) -} - -fn failure_message(value: &Value) -> Option { - value - .pointer("/response/error/message") - .or_else(|| value.pointer("/error/message")) - .or_else(|| value.get("message")) - .and_then(Value::as_str) - .filter(|message| !message.trim().is_empty()) - .map(ToString::to_string) -} - -fn failure_code(value: &Value) -> Option { - value - .pointer("/response/error/code") - .or_else(|| value.pointer("/error/code")) - .or_else(|| value.get("code")) - .and_then(Value::as_str) - .filter(|code| !code.trim().is_empty()) - .map(ToString::to_string) -} - -fn ensure_completed_response_output( - response: &mut Value, - output_items: &[Value], - output_text: &str, -) { - let output_is_empty = response - .get("output") - .and_then(Value::as_array) - .is_none_or(Vec::is_empty); - if !output_is_empty { - return; - } - - if !output_items.is_empty() { - response["output"] = Value::Array(output_items.to_vec()); - return; - } - if output_text.is_empty() { - return; - } - - // 原版 passthrough 会用 done item 或文本 delta 回填 completed.output，避免非流式客户端拿到空正文。 - response["output"] = json!([{ - "type": "message", - "status": "completed", - "role": "assistant", - "content": [{ - "type": "output_text", - "text": output_text, - "annotations": [] - }] - }]); -} - -fn sync_output_text_from_output(response: &mut Value) { - let output_text_is_empty = response - .get("output_text") - .and_then(Value::as_str) - .is_none_or(str::is_empty); - if !output_text_is_empty { - return; - } - let Some(items) = response.get("output").and_then(Value::as_array) else { - return; - }; - let texts = items - .iter() - .filter_map(output_text_from_item) - .collect::>(); - if texts.is_empty() { - return; - } - response["output_text"] = Value::String(texts.join("\n\n")); -} - -fn output_text_from_item(item: &Value) -> Option { - let content = item.get("content")?.as_array()?; - let text = content - .iter() - .filter_map(|part| { - let part_type = part.get("type")?.as_str()?; - if part_type != "output_text" && part_type != "text" { - return None; - } - part.get("text")?.as_str() - }) - .collect::(); - (!text.is_empty()).then_some(text) -} - -fn codex_client_error_response(error: CodexClientError) -> (StatusCode, Json) { - match error { - CodexClientError::UnsupportedTransport(_) => ( - StatusCode::BAD_REQUEST, - Json(openai_error( - "previous_response_id requires Codex WebSocket transport", - "websocket_required", - )), - ), - CodexClientError::Upstream { status, body, .. } => ( - status, - Json(openai_error( - &format!( - "Codex upstream error: {}", - body.chars().take(300).collect::() - ), - "upstream_error", - )), - ), - _ => ( - StatusCode::BAD_GATEWAY, - Json(openai_error( - "Codex upstream request failed", - "upstream_error", - )), - ), - } -} - -fn request_domain(base_url: &str) -> Option { - Url::parse(base_url) - .ok() - .and_then(|url| url.host_str().map(ToString::to_string)) -} - -fn u64_to_i64_saturating(value: u64) -> i64 { - value.min(i64::MAX as u64) as i64 -} - -pub async fn models(State(state): State, headers: HeaderMap) -> impl IntoResponse { - if !authorize_client_api_key(&state, &headers).await { - return missing_client_api_key_response(); - } - - let catalog = model_catalog_for_state(&state).await; - let data = catalog - .models() - .iter() - .map(|model| openai_model_json(&model.id)) - .collect::>(); - ( - StatusCode::OK, - Json(json!({ - "object": "list", - "data": data - })), - ) -} - -pub async fn model_catalog(State(state): State, headers: HeaderMap) -> impl IntoResponse { - if !authorize_client_api_key(&state, &headers).await { - return missing_client_api_key_response(); - } - - let catalog = model_catalog_for_state(&state).await; - (StatusCode::OK, Json(json!(catalog.models()))) -} - -pub async fn model_detail( - State(state): State, - headers: HeaderMap, - Path(model_id): Path, -) -> impl IntoResponse { - if !authorize_client_api_key(&state, &headers).await { - return missing_client_api_key_response(); - } - - let catalog = model_catalog_for_state(&state).await; - if catalog.model_info(&model_id).is_none() { - return model_not_found_response(); - } - (StatusCode::OK, Json(openai_model_json(&model_id))) -} - -pub async fn model_info( - State(state): State, - headers: HeaderMap, - Path(model_id): Path, -) -> impl IntoResponse { - if !authorize_client_api_key(&state, &headers).await { - return missing_client_api_key_response(); - } - - let catalog = model_catalog_for_state(&state).await; - let Some(info) = catalog.model_info(&model_id) else { - return model_not_found_response(); - }; - (StatusCode::OK, Json(json!(info))) -} - -pub async fn debug_models(State(state): State, headers: HeaderMap) -> impl IntoResponse { - if !authorize_client_api_key(&state, &headers).await { - return missing_client_api_key_response(); - } - - let catalog = model_catalog_for_state(&state).await; - (StatusCode::OK, Json(json!(catalog.debug()))) -} - -async fn authorize_client_api_key(state: &AppState, headers: &HeaderMap) -> bool { - let Some(api_key) = client_api_key(headers) else { - return false; - }; - let Some(repo) = state.client_api_key_repository() else { - return false; - }; - let Some(hasher) = state.api_key_hasher().cloned() else { - return false; - }; - repo.verify_and_touch(api_key.as_str(), &hasher) - .await - .unwrap_or(false) -} - -fn missing_client_api_key_response() -> (StatusCode, Json) { - ( - StatusCode::UNAUTHORIZED, - Json(openai_error("Missing client API key", "invalid_api_key")), - ) -} - -fn model_not_found_response() -> (StatusCode, Json) { - ( - StatusCode::NOT_FOUND, - Json(openai_error("Model not found", "model_not_found")), - ) -} - -fn openai_model_json(id: &str) -> Value { - json!({ - "id": id, - "object": "model", - "created": MODEL_CREATED_TIMESTAMP, - "owned_by": "openai" - }) -} diff --git a/src/http/v1/auth.rs b/src/http/v1/auth.rs new file mode 100644 index 000000000..68cad6138 --- /dev/null +++ b/src/http/v1/auth.rs @@ -0,0 +1,15 @@ +use axum::http::HeaderMap; + +use crate::{app::state::AppState, http::auth::client_api_key}; + +pub(super) async fn authorize_client_api_key(state: &AppState, headers: &HeaderMap) -> bool { + let Some(api_key) = client_api_key(headers) else { + return false; + }; + state + .services + .api_keys + .verify(api_key.as_str()) + .await + .unwrap_or(false) +} diff --git a/src/http/v1/chat.rs b/src/http/v1/chat.rs new file mode 100644 index 000000000..a59cbd00d --- /dev/null +++ b/src/http/v1/chat.rs @@ -0,0 +1,44 @@ +use axum::{ + body::Bytes, + extract::State, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + Extension, Json, +}; +use std::time::Instant; + +use crate::{ + app::state::AppState, + codex::protocol::{codex_to_openai::openai_error, openai_to_codex::ChatCompletionRequest}, + http::middleware::RequestId, +}; + +use super::{auth::authorize_client_api_key, errors::missing_client_api_key_response}; + +pub async fn chat_completions( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + body: Bytes, +) -> Response { + let started_at = Instant::now(); + if !authorize_client_api_key(&state, &headers).await { + return missing_client_api_key_response().into_response(); + } + + let Ok(chat_request) = serde_json::from_slice::(&body) else { + return ( + StatusCode::BAD_REQUEST, + Json(openai_error( + "Invalid chat completion request", + "invalid_request", + )), + ) + .into_response(); + }; + state + .services + .chat + .handle(request_id.as_str(), chat_request, started_at) + .await +} diff --git a/src/http/v1/errors.rs b/src/http/v1/errors.rs new file mode 100644 index 000000000..304669c20 --- /dev/null +++ b/src/http/v1/errors.rs @@ -0,0 +1,49 @@ +use axum::{http::StatusCode, Json}; +use serde_json::Value; + +use crate::{ + codex::protocol::codex_to_openai::openai_error, codex::transport::client::CodexClientError, +}; + +pub(crate) fn codex_client_error_response(error: CodexClientError) -> (StatusCode, Json) { + match error { + CodexClientError::UnsupportedTransport(_) => ( + StatusCode::BAD_REQUEST, + Json(openai_error( + "previous_response_id requires Codex WebSocket transport", + "websocket_required", + )), + ), + CodexClientError::Upstream { status, body, .. } => ( + status, + Json(openai_error( + &format!( + "Codex upstream error: {}", + body.chars().take(300).collect::() + ), + "upstream_error", + )), + ), + _ => ( + StatusCode::BAD_GATEWAY, + Json(openai_error( + "Codex upstream request failed", + "upstream_error", + )), + ), + } +} + +pub(super) fn missing_client_api_key_response() -> (StatusCode, Json) { + ( + StatusCode::UNAUTHORIZED, + Json(openai_error("Missing client API key", "invalid_api_key")), + ) +} + +pub(crate) fn model_not_found_response() -> (StatusCode, Json) { + ( + StatusCode::NOT_FOUND, + Json(openai_error("Model not found", "model_not_found")), + ) +} diff --git a/src/http/v1/mod.rs b/src/http/v1/mod.rs new file mode 100644 index 000000000..c6b32ace8 --- /dev/null +++ b/src/http/v1/mod.rs @@ -0,0 +1,8 @@ +pub mod auth; +pub mod chat; +pub mod errors; +pub mod models; +pub mod responses; +pub mod router; + +pub use self::router::router; diff --git a/src/http/v1/models.rs b/src/http/v1/models.rs new file mode 100644 index 000000000..2f27621f3 --- /dev/null +++ b/src/http/v1/models.rs @@ -0,0 +1,99 @@ +use axum::{ + extract::{Path, State}, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Json, +}; +use serde_json::{json, Value}; + +use crate::{app::state::AppState, codex::models::catalog::ModelCatalog}; + +use super::{ + auth::authorize_client_api_key, + errors::{missing_client_api_key_response, model_not_found_response}, +}; + +const MODEL_CREATED_TIMESTAMP: i64 = 1_700_000_000; + +pub async fn models(State(state): State, headers: HeaderMap) -> impl IntoResponse { + if !authorize_client_api_key(&state, &headers).await { + return missing_client_api_key_response(); + } + + let catalog = model_catalog_for_state(&state).await; + let data = catalog + .models() + .iter() + .map(|model| openai_model_json(&model.id)) + .collect::>(); + ( + StatusCode::OK, + Json(json!({ + "object": "list", + "data": data + })), + ) +} + +pub async fn model_catalog(State(state): State, headers: HeaderMap) -> impl IntoResponse { + if !authorize_client_api_key(&state, &headers).await { + return missing_client_api_key_response(); + } + + let catalog = model_catalog_for_state(&state).await; + (StatusCode::OK, Json(json!(catalog.models()))) +} + +pub async fn model_detail( + State(state): State, + headers: HeaderMap, + Path(model_id): Path, +) -> impl IntoResponse { + if !authorize_client_api_key(&state, &headers).await { + return missing_client_api_key_response(); + } + + let catalog = model_catalog_for_state(&state).await; + if catalog.model_info(&model_id).is_none() { + return model_not_found_response(); + } + (StatusCode::OK, Json(openai_model_json(&model_id))) +} + +pub async fn model_info( + State(state): State, + headers: HeaderMap, + Path(model_id): Path, +) -> impl IntoResponse { + if !authorize_client_api_key(&state, &headers).await { + return missing_client_api_key_response(); + } + + let catalog = model_catalog_for_state(&state).await; + let Some(info) = catalog.model_info(&model_id) else { + return model_not_found_response(); + }; + (StatusCode::OK, Json(json!(info))) +} + +pub async fn debug_models(State(state): State, headers: HeaderMap) -> impl IntoResponse { + if !authorize_client_api_key(&state, &headers).await { + return missing_client_api_key_response(); + } + + let catalog = model_catalog_for_state(&state).await; + (StatusCode::OK, Json(json!(catalog.debug()))) +} + +pub(super) async fn model_catalog_for_state(state: &AppState) -> ModelCatalog { + state.services.models.catalog().await +} + +fn openai_model_json(id: &str) -> Value { + json!({ + "id": id, + "object": "model", + "created": MODEL_CREATED_TIMESTAMP, + "owned_by": "openai" + }) +} diff --git a/src/http/v1/responses.rs b/src/http/v1/responses.rs new file mode 100644 index 000000000..0a1df6675 --- /dev/null +++ b/src/http/v1/responses.rs @@ -0,0 +1,31 @@ +use std::time::Instant; + +use axum::{ + body::Bytes, + extract::State, + http::HeaderMap, + response::{IntoResponse, Response}, + Extension, +}; + +use crate::{app::state::AppState, http::middleware::RequestId}; + +use super::{auth::authorize_client_api_key, errors::missing_client_api_key_response}; + +pub async fn responses( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + body: Bytes, +) -> Response { + let started_at = Instant::now(); + if !authorize_client_api_key(&state, &headers).await { + return missing_client_api_key_response().into_response(); + } + + state + .services + .responses + .handle(request_id.as_str(), headers, body, started_at) + .await +} diff --git a/src/http/v1/router.rs b/src/http/v1/router.rs new file mode 100644 index 000000000..2365c0eeb --- /dev/null +++ b/src/http/v1/router.rs @@ -0,0 +1,23 @@ +use axum::{ + routing::{get, post}, + Router, +}; + +use crate::app::state::AppState; + +use super::{ + chat::chat_completions, + models::{debug_models, model_catalog, model_detail, model_info, models}, + responses::responses, +}; + +pub fn router() -> Router { + Router::new() + .route("/v1/responses", post(responses)) + .route("/v1/chat/completions", post(chat_completions)) + .route("/v1/models", get(models)) + .route("/v1/models/catalog", get(model_catalog)) + .route("/v1/models/{model_id}", get(model_detail)) + .route("/v1/models/{model_id}/info", get(model_info)) + .route("/debug/models", get(debug_models)) +} diff --git a/src/lib.rs b/src/lib.rs index ed5363673..3ebe73d6e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,16 +1,11 @@ -pub mod accounts; pub mod app; pub mod auth; pub mod codex; + pub mod config; -pub mod cookies; -pub mod crypto; -pub mod error; -pub mod fingerprint; +pub mod utils; + pub mod http; pub mod logs; -pub mod models; -pub mod pagination; -pub mod state; +pub mod service; pub mod storage; -pub mod translation; diff --git a/src/logs/repository.rs b/src/logs/repository.rs index af7fd9e6c..7bc25399a 100644 --- a/src/logs/repository.rs +++ b/src/logs/repository.rs @@ -2,7 +2,7 @@ use sqlx::{Row, SqlitePool}; use crate::{ logs::event::{EventLevel, EventLog}, - pagination::{clamp_limit, decode_cursor, encode_cursor, Page}, + utils::pagination::{clamp_limit, decode_cursor, encode_cursor, Page}, }; #[derive(Clone)] diff --git a/src/main.rs b/src/main.rs index 28f9d4ecd..9e93cfa56 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,7 @@ use codex_proxy_rs::{ - app::build_router, - auth::{api_key::ApiKeyHasher, oauth::OpenAiOAuthRefresher}, + app::{bootstrap::build_state, build_router}, config::AppConfig, - crypto::SecretBox, logs::rotation::{init_tracing, RotationConfig}, - state::AppState, - storage::db::connect_sqlite, }; #[tokio::main] @@ -19,23 +15,7 @@ async fn main() -> anyhow::Result<()> { ))?; let host = config.server.host.clone(); 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 oauth_client = 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_api_key_hasher_and_oauth_client( - config, - pool, - secret_box, - api_key_hasher, - oauth_client, - ); - let restored_accounts = state.reload_account_pool_from_repository().await?; + let (state, restored_accounts) = build_state(config).await?; tracing::info!(restored_accounts, "account pool restored from sqlite"); let app = build_router(state); let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?; diff --git a/src/service/admin_auth.rs b/src/service/admin_auth.rs new file mode 100644 index 000000000..38b0559c7 --- /dev/null +++ b/src/service/admin_auth.rs @@ -0,0 +1,432 @@ +use std::sync::Arc; + +use chrono::{DateTime, Duration, Utc}; +use sqlx::{Row, SqlitePool}; +use tokio::sync::Mutex; +use uuid::Uuid; + +use crate::{ + auth::admin_session::verify_admin_password, + codex::accounts::service::{AccountService, ValidatedAccountImportError}, + codex::accounts::{ + model::AccountStatus, + pool::AccountPool, + repository::{AccountRepository, StoredAccountMetadata}, + }, + codex::oauth::{ + DeviceCode, OAuthClient, OAuthConfig, OAuthError, PkceSession, PkceSessionStore, + }, + config::AppConfig, +}; + +#[derive(Clone)] +pub struct AdminAuthService { + config: AppConfig, + db: Option, + account_repository: Option, + account_pool: Arc>, + oauth_client: Option>, + oauth_sessions: Arc>, + accounts: AccountService, +} + +#[derive(Debug)] +pub enum AdminAuthServiceError { + DatabaseUnavailable, + LoadAdminUser, + AdminPasswordInvalid, + VerifyAdminPassword, + InvalidSessionTtl, + CreateSession, + AccountRepositoryUnavailable, + InspectAccountAuthStatus, + ClearAccounts, +} + +#[derive(Debug)] +pub enum AdminSessionValidationError { + DatabaseUnavailable, + ValidateSession, +} + +#[derive(Debug)] +pub enum AdminAuthOAuthError { + OAuthClientUnavailable, + AccountRepositoryUnavailable, + DeviceCodeRequired, + InvalidOrExpiredSession, + DeviceCodeRequest(OAuthError), + DeviceAuthorization(OAuthError), + TokenExchange(OAuthError), + Import(ValidatedAccountImportError), +} + +#[derive(Debug)] +pub struct AdminLogin { + pub session_id: String, + pub expires_at: DateTime, + pub ttl_minutes: u64, +} + +#[derive(Debug)] +pub struct AdminAuthStatus { + pub authenticated: bool, + pub user: Option, + pub pool: AdminAuthPoolSummary, +} + +#[derive(Debug)] +pub struct AdminAuthUser { + pub id: String, + pub email: Option, + pub account_id: Option, + pub user_id: Option, + pub label: Option, + pub plan_type: Option, + pub status: AccountStatus, + pub access_token_expires_at: Option>, +} + +#[derive(Debug, Default)] +pub struct AdminAuthPoolSummary { + pub total: usize, + pub active: usize, + pub expired: usize, + pub quota_exhausted: usize, + pub refreshing: usize, + pub disabled: usize, + pub banned: usize, +} + +#[derive(Debug)] +pub struct AdminLogout { + pub deleted: u64, +} + +#[derive(Debug)] +pub struct AdminAuthLoginStart { + pub auth_url: String, + pub state: String, +} + +#[derive(Debug)] +pub struct AdminAuthDevicePoll { + pub success: bool, + pub pending: bool, + pub code: Option, +} + +#[derive(Debug)] +pub enum AdminAuthPkceExchange { + Imported { return_host: String }, + AlreadyCompleted, +} + +#[derive(Debug)] +struct AdminUserRow { + id: String, + password_hash: String, +} + +impl AdminAuthService { + pub fn new( + config: AppConfig, + db: Option, + account_repository: Option, + account_pool: Arc>, + oauth_client: Option>, + oauth_sessions: Arc>, + accounts: AccountService, + ) -> Self { + Self { + config, + db, + account_repository, + account_pool, + oauth_client, + oauth_sessions, + accounts, + } + } + + pub async fn login(&self, password: &str) -> Result { + let pool = self.pool()?; + let admin = load_first_admin(pool) + .await + .map_err(|_| AdminAuthServiceError::LoadAdminUser)? + .ok_or(AdminAuthServiceError::AdminPasswordInvalid)?; + + match verify_admin_password(password, &admin.password_hash) { + Ok(true) => {} + Ok(false) => return Err(AdminAuthServiceError::AdminPasswordInvalid), + Err(_) => return Err(AdminAuthServiceError::VerifyAdminPassword), + } + + let ttl_minutes = self.config.admin.session_ttl_minutes; + let ttl_minutes_i64 = + i64::try_from(ttl_minutes).map_err(|_| AdminAuthServiceError::InvalidSessionTtl)?; + let expires_at = Utc::now() + Duration::minutes(ttl_minutes_i64); + let session_id = format!("sess_{}", Uuid::new_v4().simple()); + create_admin_session(pool, &session_id, &admin.id, expires_at) + .await + .map_err(|_| AdminAuthServiceError::CreateSession)?; + + Ok(AdminLogin { + session_id, + expires_at, + ttl_minutes, + }) + } + + pub async fn validate_session( + &self, + session_id: Option<&str>, + ) -> Result { + let Some(session_id) = session_id else { + return Ok(false); + }; + let pool = self + .pool() + .map_err(|_| AdminSessionValidationError::DatabaseUnavailable)?; + validate_admin_session(pool, session_id) + .await + .map_err(|_| AdminSessionValidationError::ValidateSession) + } + + pub async fn status(&self) -> Result { + let repo = self.account_repository()?; + let accounts = repo + .list_all_metadata() + .await + .map_err(|_| AdminAuthServiceError::InspectAccountAuthStatus)?; + let pool = account_auth_pool_summary(&accounts); + let user = accounts.first().map(account_auth_user); + Ok(AdminAuthStatus { + authenticated: pool.total > 0, + user, + pool, + }) + } + + pub async fn logout(&self) -> Result { + let repo = self.account_repository()?; + let deleted = repo + .delete_all() + .await + .map_err(|_| AdminAuthServiceError::ClearAccounts)?; + self.account_pool.lock().await.clear(); + Ok(AdminLogout { deleted }) + } + + pub async fn request_device_code(&self) -> Result { + self.oauth_client()? + .request_device_code() + .await + .map_err(AdminAuthOAuthError::DeviceCodeRequest) + } + + pub async fn poll_device_token( + &self, + device_code: &str, + ) -> Result { + let oauth_client = self.oauth_client()?; + if !self.accounts.has_repository() { + return Err(AdminAuthOAuthError::AccountRepositoryUnavailable); + } + let device_code = device_code.trim(); + if device_code.is_empty() { + return Err(AdminAuthOAuthError::DeviceCodeRequired); + } + + match oauth_client.poll_device_token(device_code).await { + Ok(tokens) => { + self.accounts + .import_validated(Some(tokens.access_token), tokens.refresh_token) + .await + .map_err(AdminAuthOAuthError::Import)?; + Ok(AdminAuthDevicePoll { + success: true, + pending: false, + code: None, + }) + } + Err(error) => { + if let Some(code) = error.pending_code() { + return Ok(AdminAuthDevicePoll { + success: false, + pending: true, + code: Some(code.to_string()), + }); + } + Err(AdminAuthOAuthError::DeviceAuthorization(error)) + } + } + } + + pub async fn start_pkce_login(&self, return_host: &str) -> AdminAuthLoginStart { + let login = self + .oauth_sessions + .lock() + .await + .start_login(return_host, &OAuthConfig::codex_default()); + AdminAuthLoginStart { + auth_url: login.auth_url, + state: login.state, + } + } + + pub async fn exchange_pkce_code( + &self, + oauth_state: &str, + code: &str, + ) -> Result { + if !self.accounts.has_repository() { + return Err(AdminAuthOAuthError::AccountRepositoryUnavailable); + } + let oauth_client = self.oauth_client()?; + let Some(session) = self.acquire_pkce_session(oauth_state).await? else { + return Ok(AdminAuthPkceExchange::AlreadyCompleted); + }; + self.exchange_pkce_session(oauth_client.as_ref(), oauth_state, code, session) + .await + } + + fn pool(&self) -> Result<&SqlitePool, AdminAuthServiceError> { + self.db + .as_ref() + .ok_or(AdminAuthServiceError::DatabaseUnavailable) + } + + fn account_repository(&self) -> Result<&AccountRepository, AdminAuthServiceError> { + self.account_repository + .as_ref() + .ok_or(AdminAuthServiceError::AccountRepositoryUnavailable) + } + + fn oauth_client(&self) -> Result, AdminAuthOAuthError> { + self.oauth_client + .clone() + .ok_or(AdminAuthOAuthError::OAuthClientUnavailable) + } + + async fn acquire_pkce_session( + &self, + oauth_state: &str, + ) -> Result, AdminAuthOAuthError> { + let mut sessions = self.oauth_sessions.lock().await; + if let Some(session) = sessions.try_acquire(oauth_state) { + return Ok(Some(session)); + } + if sessions.is_completed_or_exchanging(oauth_state) { + return Ok(None); + } + Err(AdminAuthOAuthError::InvalidOrExpiredSession) + } + + async fn exchange_pkce_session( + &self, + oauth_client: &dyn OAuthClient, + oauth_state: &str, + code: &str, + session: PkceSession, + ) -> Result { + let tokens = match oauth_client + .exchange_code(code, &session.code_verifier, &session.redirect_uri) + .await + { + Ok(tokens) => tokens, + Err(error) => { + self.oauth_sessions.lock().await.release(oauth_state); + return Err(AdminAuthOAuthError::TokenExchange(error)); + } + }; + match self + .accounts + .import_validated(Some(tokens.access_token), tokens.refresh_token) + .await + { + Ok(_) => { + self.oauth_sessions.lock().await.complete(oauth_state); + Ok(AdminAuthPkceExchange::Imported { + return_host: session.return_host, + }) + } + Err(error) => { + self.oauth_sessions.lock().await.complete(oauth_state); + Err(AdminAuthOAuthError::Import(error)) + } + } + } +} + +async fn load_first_admin(pool: &SqlitePool) -> Result, sqlx::Error> { + let row = + sqlx::query("select id, password_hash from admin_users order by created_at asc limit 1") + .fetch_optional(pool) + .await?; + Ok(row.map(|row| AdminUserRow { + id: row.get("id"), + password_hash: row.get("password_hash"), + })) +} + +async fn validate_admin_session(pool: &SqlitePool, session_id: &str) -> Result { + let now = Utc::now().to_rfc3339(); + let count: (i64,) = + sqlx::query_as("select count(*) from admin_sessions where id = ? and expires_at > ?") + .bind(session_id) + .bind(now) + .fetch_one(pool) + .await?; + Ok(count.0 > 0) +} + +async fn create_admin_session( + pool: &SqlitePool, + session_id: &str, + user_id: &str, + expires_at: DateTime, +) -> Result<(), sqlx::Error> { + let now = Utc::now().to_rfc3339(); + sqlx::query( + "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", + ) + .bind(session_id) + .bind(user_id) + .bind(expires_at.to_rfc3339()) + .bind(now) + .execute(pool) + .await?; + Ok(()) +} + +fn account_auth_user(account: &StoredAccountMetadata) -> AdminAuthUser { + AdminAuthUser { + id: account.id.clone(), + email: account.email.clone(), + account_id: account.account_id.clone(), + user_id: account.user_id.clone(), + label: account.label.clone(), + plan_type: account.plan_type.clone(), + status: account.status, + access_token_expires_at: account.access_token_expires_at, + } +} + +fn account_auth_pool_summary(accounts: &[StoredAccountMetadata]) -> AdminAuthPoolSummary { + let mut summary = AdminAuthPoolSummary { + total: accounts.len(), + ..AdminAuthPoolSummary::default() + }; + for account in accounts { + match account.status { + AccountStatus::Active => summary.active += 1, + AccountStatus::Expired => summary.expired += 1, + AccountStatus::QuotaExhausted => summary.quota_exhausted += 1, + AccountStatus::Refreshing => summary.refreshing += 1, + AccountStatus::Disabled => summary.disabled += 1, + AccountStatus::Banned => summary.banned += 1, + } + } + summary +} diff --git a/src/service/api_key.rs b/src/service/api_key.rs new file mode 100644 index 000000000..280752677 --- /dev/null +++ b/src/service/api_key.rs @@ -0,0 +1,334 @@ +use serde_json::Value; + +use crate::{ + auth::{ + api_key::ApiKeyHasher, + api_key_repository::{ClientApiKeyRepository, StoredClientApiKey}, + }, + utils::{json::first_string, pagination::Page}, +}; + +#[derive(Clone)] +pub struct ApiKeyService { + repository: Option, + hasher: Option, +} + +#[derive(Debug)] +pub enum ApiKeyServiceError { + RepositoryUnavailable, + HasherUnavailable, + List, + Export, + Import, + Create, + Delete, + UpdateLabel, + UpdateStatus, + Verify, + InvalidStatus(String), + EmptyName, + EmptyIds, + LabelTooLong, + NoImportableKeys, +} + +#[derive(Debug)] +pub struct CreatedClientApiKey { + pub key: StoredClientApiKey, + pub plaintext: String, +} + +#[derive(Debug)] +pub struct ImportedClientApiKey { + pub source_id: Option, + pub source_prefix: Option, + pub key: StoredClientApiKey, + pub plaintext: String, +} + +#[derive(Debug)] +pub struct ImportedClientApiKeys { + pub imported: u32, + pub skipped: u32, + pub keys: Vec, +} + +#[derive(Debug)] +pub struct BatchDeleteClientApiKeys { + pub deleted: u32, + pub not_found: Vec, +} + +#[derive(Debug)] +pub struct UpdateClientApiKeyStatus { + pub id: String, + pub enabled: bool, +} + +#[derive(Debug, Clone)] +struct ClientApiKeyImportEntry { + source_id: Option, + source_prefix: Option, + name: String, + label: Option, + enabled: bool, +} + +impl ApiKeyService { + pub fn new(repository: Option, hasher: Option) -> Self { + Self { repository, hasher } + } + + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> Result, ApiKeyServiceError> { + self.repository()? + .list(cursor, limit) + .await + .map_err(|_| ApiKeyServiceError::List) + } + + pub async fn export( + &self, + ids: Vec, + ) -> Result, ApiKeyServiceError> { + let repo = self.repository()?; + if ids.is_empty() { + return repo + .list_all() + .await + .map_err(|_| ApiKeyServiceError::Export); + } + + let mut keys = Vec::with_capacity(ids.len()); + for id in ids { + match repo.get(&id).await { + Ok(Some(key)) => keys.push(key), + Ok(None) => {} + Err(_) => return Err(ApiKeyServiceError::Export), + } + } + Ok(keys) + } + + pub async fn import( + &self, + payload: &Value, + ) -> Result { + let repo = self.repository()?; + let hasher = self.hasher()?; + let entries = parse_client_api_key_import_payload(payload); + if entries.is_empty() { + return Err(ApiKeyServiceError::NoImportableKeys); + } + + let mut imported = 0u32; + let mut skipped = 0u32; + let mut keys = Vec::with_capacity(entries.len()); + for entry in entries { + let name = entry.name.trim(); + if name.is_empty() { + skipped += 1; + continue; + } + if entry + .label + .as_ref() + .is_some_and(|label| label.chars().count() > 64) + { + return Err(ApiKeyServiceError::LabelTooLong); + } + + let generated = hasher.generate_client_api_key(name); + let plaintext = generated.plaintext.clone(); + let source_id = entry.source_id; + let source_prefix = entry.source_prefix; + match repo + .insert_generated_with_metadata( + name, + entry.label.as_deref(), + entry.enabled, + &generated, + ) + .await + { + Ok(key) => { + imported += 1; + keys.push(ImportedClientApiKey { + source_id, + source_prefix, + key, + plaintext, + }); + } + Err(_) => return Err(ApiKeyServiceError::Import), + } + } + + Ok(ImportedClientApiKeys { + imported, + skipped, + keys, + }) + } + + pub async fn create(&self, name: &str) -> Result { + let name = name.trim(); + if name.is_empty() { + return Err(ApiKeyServiceError::EmptyName); + } + let repo = self.repository()?; + let hasher = self.hasher()?; + let generated = hasher.generate_client_api_key(name); + let plaintext = generated.plaintext.clone(); + repo.insert_generated(name, &generated) + .await + .map(|key| CreatedClientApiKey { key, plaintext }) + .map_err(|_| ApiKeyServiceError::Create) + } + + pub async fn batch_delete( + &self, + ids: Vec, + ) -> Result { + if ids.is_empty() { + return Err(ApiKeyServiceError::EmptyIds); + } + let repo = self.repository()?; + let mut deleted = 0u32; + let mut not_found = Vec::new(); + for key_id in ids { + match repo.delete(&key_id).await { + Ok(true) => deleted += 1, + Ok(false) => not_found.push(key_id), + Err(_) => return Err(ApiKeyServiceError::Delete), + } + } + Ok(BatchDeleteClientApiKeys { deleted, not_found }) + } + + pub async fn update_label( + &self, + key_id: &str, + label: Option, + ) -> Result, ApiKeyServiceError> { + if label + .as_ref() + .is_some_and(|label| label.chars().count() > 64) + { + return Err(ApiKeyServiceError::LabelTooLong); + } + self.repository()? + .set_label(key_id, label) + .await + .map_err(|_| ApiKeyServiceError::UpdateLabel) + } + + pub async fn update_status( + &self, + key_id: String, + status: &str, + ) -> Result, ApiKeyServiceError> { + let enabled = parse_client_api_key_enabled_status(status)?; + match self.repository()?.set_enabled(&key_id, enabled).await { + Ok(true) => Ok(Some(UpdateClientApiKeyStatus { + id: key_id, + enabled, + })), + Ok(false) => Ok(None), + Err(_) => Err(ApiKeyServiceError::UpdateStatus), + } + } + + pub async fn delete(&self, key_id: &str) -> Result { + self.repository()? + .delete(key_id) + .await + .map_err(|_| ApiKeyServiceError::Delete) + } + + pub async fn verify(&self, plaintext: &str) -> Result { + let repo = self.repository()?; + let hasher = self.hasher()?; + repo.verify_and_touch(plaintext, hasher) + .await + .map_err(|_| ApiKeyServiceError::Verify) + } + + fn repository(&self) -> Result<&ClientApiKeyRepository, ApiKeyServiceError> { + self.repository + .as_ref() + .ok_or(ApiKeyServiceError::RepositoryUnavailable) + } + + fn hasher(&self) -> Result<&ApiKeyHasher, ApiKeyServiceError> { + self.hasher + .as_ref() + .ok_or(ApiKeyServiceError::HasherUnavailable) + } +} + +fn parse_client_api_key_import_payload(payload: &Value) -> Vec { + let payload = payload + .get("data") + .filter(|data| data.get("apiKeys").is_some() || data.get("keys").is_some()) + .unwrap_or(payload); + + if let Some(keys) = payload.get("apiKeys").and_then(Value::as_array) { + return keys + .iter() + .filter_map(client_api_key_import_entry_from_value) + .collect(); + } + if let Some(keys) = payload.get("keys").and_then(Value::as_array) { + return keys + .iter() + .filter_map(client_api_key_import_entry_from_value) + .collect(); + } + if let Some(keys) = payload.as_array() { + return keys + .iter() + .filter_map(client_api_key_import_entry_from_value) + .collect(); + } + + client_api_key_import_entry_from_value(payload) + .into_iter() + .collect() +} + +fn client_api_key_import_entry_from_value(value: &Value) -> Option { + value.as_object()?; + let name = first_string(value, &[&["name"]])?; + Some(ClientApiKeyImportEntry { + source_id: first_string(value, &[&["id"], &["sourceId"]]), + source_prefix: first_string(value, &[&["prefix"], &["sourcePrefix"]]), + name, + label: first_string(value, &[&["label"]]), + enabled: client_api_key_import_enabled(value), + }) +} + +fn client_api_key_import_enabled(value: &Value) -> bool { + if let Some(enabled) = value.get("enabled").and_then(Value::as_bool) { + return enabled; + } + !first_string(value, &[&["status"]]) + .unwrap_or_else(|| "active".to_string()) + .trim() + .eq_ignore_ascii_case("disabled") +} + +fn parse_client_api_key_enabled_status(status: &str) -> Result { + match status.trim().to_ascii_lowercase().as_str() { + "active" => Ok(true), + "disabled" => Ok(false), + other => Err(ApiKeyServiceError::InvalidStatus(format!( + "Unsupported API key status: {other}" + ))), + } +} diff --git a/src/service/chat.rs b/src/service/chat.rs new file mode 100644 index 000000000..e78e16f10 --- /dev/null +++ b/src/service/chat.rs @@ -0,0 +1,255 @@ +use std::time::Instant; + +use axum::{ + body::Body, + http::{ + header::{CACHE_CONTROL, CONTENT_TYPE}, + StatusCode, + }, + response::{IntoResponse, Response}, + Json, +}; +use serde_json::json; + +use crate::{ + codex::models::catalog::ModelCatalog, + codex::models::service::ModelService, + codex::protocol::{ + codex_to_openai::{ + chat_completion_from_codex_sse, chat_completion_stream_from_codex_sse, openai_error, + }, + openai_to_codex::{translate_chat_to_codex, ChatCompletionRequest}, + }, + codex::upstream::{ + classify_upstream_account_retry, no_available_accounts_response, + normalize_service_tier_for_upstream, CodexRequestLogContext, CodexUpstreamService, + }, + config::ModelConfig, + http::v1::errors::{codex_client_error_response, model_not_found_response}, + logs::event::EventLevel, +}; + +#[derive(Clone)] +pub struct ChatService { + model_config: ModelConfig, + models: ModelService, + upstream: CodexUpstreamService, +} + +impl ChatService { + pub(crate) fn new( + model_config: ModelConfig, + models: ModelService, + upstream: CodexUpstreamService, + ) -> Self { + Self { + model_config, + models, + upstream, + } + } + + pub async fn handle( + &self, + request_id: &str, + chat_request: ChatCompletionRequest, + started_at: Instant, + ) -> Response { + let client_stream = chat_request.stream; + let requested_model = chat_request.model.clone(); + let catalog = self.models.catalog().await; + if !catalog.is_recognized_model_name(&requested_model) { + return model_not_found_response().into_response(); + } + let parsed_model = catalog.parse_model_name(&requested_model); + let display_model = ModelCatalog::build_display_model_name(&parsed_model); + let mut codex_request = match translate_chat_to_codex(chat_request) { + Ok(request) => request, + Err(_) => { + return ( + StatusCode::BAD_REQUEST, + Json(openai_error( + "Invalid chat completion request", + "invalid_request", + )), + ) + .into_response(); + } + }; + codex_request.model = parsed_model.model_id.clone(); + if codex_request.reasoning.is_none() { + let effort = parsed_model + .reasoning_effort + .clone() + .or_else(|| self.model_config.default_reasoning_effort.clone()); + if let Some(effort) = effort { + codex_request.reasoning = Some(json!({"effort": effort, "summary": "auto"})); + } + } + if codex_request.service_tier.is_none() { + codex_request.service_tier = parsed_model + .service_tier + .clone() + .or_else(|| self.model_config.service_tier.clone()); + } + codex_request.service_tier = codex_request + .service_tier + .map(normalize_service_tier_for_upstream); + let include_reasoning = codex_request.reasoning.is_some(); + + let Some(mut account) = self.upstream.acquire_account(&codex_request.model).await else { + return no_available_accounts_response().into_response(); + }; + let mut log_context = CodexRequestLogContext::new( + request_id, + &account.id, + &codex_request.model, + client_stream, + started_at, + ); + + let mut excluded_account_ids = Vec::new(); + let response = loop { + let response = self + .upstream + .send_codex_request_with_refresh_retry(&codex_request, &account, request_id) + .await; + self.upstream.release_account(&account.id).await; + + match response { + Ok(response) => break response, + Err(error) => { + if let Some(retry) = classify_upstream_account_retry(&error) { + let fallback = self + .upstream + .apply_retry_and_acquire_fallback( + &account, + retry, + &codex_request.model, + &mut excluded_account_ids, + ) + .await; + self.upstream + .log_response( + &log_context, + retry.status(), + EventLevel::Warn, + "v1 chat completions upstream retrying with fallback account", + retry.metadata(client_stream), + ) + .await; + if let Some(fallback) = fallback { + account = fallback; + log_context = log_context.with_account(&account.id); + continue; + } + } + let error_response = codex_client_error_response(error); + self.upstream + .log_response( + &log_context, + error_response.0, + EventLevel::Error, + "v1 chat completions upstream request failed", + json!({"stream": client_stream}), + ) + .await; + return error_response.into_response(); + } + } + }; + if self + .upstream + .persist_cookies(&account.id, &response.set_cookie_headers) + .await + .is_err() + { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(openai_error( + "Failed to persist upstream cookies", + "cookie_store_error", + )), + ) + .into_response(); + } + if let Some(usage) = response.usage { + if self + .upstream + .record_usage(&account.id, usage) + .await + .is_err() + { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(openai_error( + "Failed to record account usage", + "usage_store_error", + )), + ) + .into_response(); + } + } + + if client_stream { + match chat_completion_stream_from_codex_sse( + &response.body, + &display_model, + include_reasoning, + ) { + Ok(Some(body)) => Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream") + .header(CACHE_CONTROL, "no-cache") + .body(Body::from(body)) + .unwrap_or_else(|_| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(openai_error( + "Failed to build stream response", + "stream_response_error", + )), + ) + .into_response() + }), + Ok(None) => ( + StatusCode::BAD_GATEWAY, + Json(openai_error( + "Codex response did not include response.completed", + "empty_upstream_response", + )), + ) + .into_response(), + Err(_) => ( + StatusCode::BAD_GATEWAY, + Json(openai_error( + "Invalid Codex SSE response", + "invalid_upstream_sse", + )), + ) + .into_response(), + } + } else { + match chat_completion_from_codex_sse(&response.body, &display_model, include_reasoning) + { + Ok(Some(body)) => (StatusCode::OK, Json(body)).into_response(), + Ok(None) => ( + StatusCode::BAD_GATEWAY, + Json(openai_error( + "Codex response did not include response.completed", + "empty_upstream_response", + )), + ) + .into_response(), + Err(_) => ( + StatusCode::BAD_GATEWAY, + Json(openai_error( + "Invalid Codex SSE response", + "invalid_upstream_sse", + )), + ) + .into_response(), + } + } + } +} diff --git a/src/service/log.rs b/src/service/log.rs new file mode 100644 index 000000000..d14ac2a54 --- /dev/null +++ b/src/service/log.rs @@ -0,0 +1,38 @@ +use crate::{ + logs::{event::EventLog, repository::EventLogRepository}, + utils::pagination::Page, +}; + +#[derive(Clone)] +pub struct LogService { + repository: Option, +} + +#[derive(Debug)] +pub enum LogServiceError { + RepositoryUnavailable, + List, +} + +impl LogService { + pub fn new(repository: Option) -> Self { + Self { repository } + } + + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> Result, LogServiceError> { + self.repository()? + .list(cursor, limit) + .await + .map_err(|_| LogServiceError::List) + } + + fn repository(&self) -> Result<&EventLogRepository, LogServiceError> { + self.repository + .as_ref() + .ok_or(LogServiceError::RepositoryUnavailable) + } +} diff --git a/src/service/mod.rs b/src/service/mod.rs new file mode 100644 index 000000000..f0fe75900 --- /dev/null +++ b/src/service/mod.rs @@ -0,0 +1,6 @@ +pub mod admin_auth; +pub mod api_key; +pub mod chat; +pub mod log; +pub mod responses; +pub mod usage; diff --git a/src/service/responses.rs b/src/service/responses.rs new file mode 100644 index 000000000..7627a78b0 --- /dev/null +++ b/src/service/responses.rs @@ -0,0 +1,423 @@ +use std::time::Instant; + +use axum::{ + body::Bytes, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + Json, +}; +use serde::Deserialize; +use serde_json::{json, Value}; + +use crate::{ + codex::models::service::ModelService, + codex::protocol::codex_to_openai::openai_error, + codex::transport::types::CodexResponsesRequest, + codex::upstream::{ + classify_upstream_account_retry, completed_response_json, no_available_accounts_response, + normalize_service_tier_for_upstream, websocket_history_retry_metadata, + CodexRequestLogContext, CodexUpstreamService, CollectedResponse, + }, + config::ModelConfig, + http::v1::errors::codex_client_error_response, + logs::event::EventLevel, +}; + +#[derive(Clone)] +pub struct ResponsesService { + model_config: ModelConfig, + models: ModelService, + upstream: CodexUpstreamService, +} + +impl ResponsesService { + pub(crate) fn new( + model_config: ModelConfig, + models: ModelService, + upstream: CodexUpstreamService, + ) -> Self { + Self { + model_config, + models, + upstream, + } + } + + pub async fn handle( + &self, + request_id: &str, + headers: HeaderMap, + body: Bytes, + started_at: Instant, + ) -> Response { + let default_model = self.model_config.default_model.clone(); + let body = serde_json::from_slice::(&body) + .unwrap_or_else(|_| default_body(default_model.clone())); + let client_stream = body.stream.unwrap_or(true); + let requested_model = body.model.clone().unwrap_or(default_model); + let catalog = self.models.catalog().await; + if !catalog.is_recognized_model_name(&requested_model) { + return ( + StatusCode::NOT_FOUND, + Json(openai_error("Model not found", "model_not_found")), + ) + .into_response(); + } + let parsed_model = catalog.parse_model_name(&requested_model); + let mut codex_request = CodexResponsesRequest::new_http_sse( + parsed_model.model_id.clone(), + body.instructions.unwrap_or_default(), + body.input.unwrap_or_default(), + ); + codex_request.reasoning = body.reasoning; + codex_request.tools = body.tools; + codex_request.tool_choice = body.tool_choice; + codex_request.parallel_tool_calls = body.parallel_tool_calls; + codex_request.text = body.text; + codex_request.prompt_cache_key = body.prompt_cache_key; + codex_request.include = body.include; + codex_request.client_metadata = body.client_metadata; + codex_request.previous_response_id = body.previous_response_id; + codex_request.reasoning = responses_reasoning( + codex_request.reasoning.take(), + parsed_model.reasoning_effort.as_deref(), + self.model_config.default_reasoning_effort.as_deref(), + ); + codex_request.service_tier = body + .service_tier + .or(parsed_model.service_tier) + .or_else(|| self.model_config.service_tier.clone()) + .map(normalize_service_tier_for_upstream); + ensure_reasoning_include(&mut codex_request); + codex_request.turn_state = body + .turn_state + .or_else(|| header_string(&headers, "x-codex-turn-state")); + codex_request.turn_metadata = body + .turn_metadata + .or_else(|| header_string(&headers, "x-codex-turn-metadata")); + codex_request.beta_features = body + .beta_features + .or_else(|| header_string(&headers, "x-codex-beta-features")); + codex_request.include_timing_metrics = body + .include_timing_metrics + .or_else(|| header_string(&headers, "x-responsesapi-include-timing-metrics")); + codex_request.version = header_string(&headers, "version"); + codex_request.codex_window_id = body + .codex_window_id + .or_else(|| header_string(&headers, "x-codex-window-id")); + codex_request.parent_thread_id = body + .parent_thread_id + .or_else(|| header_string(&headers, "x-codex-parent-thread-id")); + codex_request.use_websocket = body.use_websocket.unwrap_or(false); + let Some(mut account) = self.upstream.acquire_account(&codex_request.model).await else { + return no_available_accounts_response().into_response(); + }; + let mut log_context = CodexRequestLogContext::new( + request_id, + &account.id, + &codex_request.model, + client_stream, + started_at, + ); + + if client_stream { + return self + .upstream + .responses_stream(codex_request, account, log_context) + .await; + } + + let mut excluded_account_ids = Vec::new(); + let response = loop { + let response = self + .upstream + .send_codex_request_with_refresh_retry(&codex_request, &account, request_id) + .await; + self.upstream.release_account(&account.id).await; + + match response { + Ok(response) => break response, + Err(error) => { + if let Some(retry) = classify_upstream_account_retry(&error) { + if codex_request.previous_response_id.is_some() { + // previous_response_id history is account-affine upstream. + self.upstream.apply_account_retry(&account, retry).await; + self.upstream + .log_response( + &log_context, + retry.status(), + EventLevel::Warn, + "v1 responses websocket history request kept on original account", + websocket_history_retry_metadata(retry, false), + ) + .await; + } else { + let fallback = self + .upstream + .apply_retry_and_acquire_fallback( + &account, + retry, + &codex_request.model, + &mut excluded_account_ids, + ) + .await; + self.upstream + .log_response( + &log_context, + retry.status(), + EventLevel::Warn, + "v1 responses upstream retrying with fallback account", + retry.metadata(false), + ) + .await; + if let Some(fallback) = fallback { + account = fallback; + log_context = log_context.with_account(&account.id); + continue; + } + } + } + let error_response = codex_client_error_response(error); + self.upstream + .log_response( + &log_context, + error_response.0, + EventLevel::Error, + "v1 responses upstream request failed", + json!({"stream": false}), + ) + .await; + return error_response.into_response(); + } + } + }; + if self + .upstream + .persist_cookies(&account.id, &response.set_cookie_headers) + .await + .is_err() + { + self.upstream + .log_response( + &log_context, + StatusCode::INTERNAL_SERVER_ERROR, + EventLevel::Error, + "v1 responses cookie persistence failed", + json!({"stream": false, "cookieStoreError": true}), + ) + .await; + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(openai_error( + "Failed to persist upstream cookies", + "cookie_store_error", + )), + ) + .into_response(); + } + if let Some(usage) = response.usage { + if self + .upstream + .record_usage(&account.id, usage) + .await + .is_err() + { + self.upstream + .log_response( + &log_context, + StatusCode::INTERNAL_SERVER_ERROR, + EventLevel::Error, + "v1 responses usage persistence failed", + json!({"stream": false, "usage": usage, "usageStoreError": true}), + ) + .await; + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(openai_error( + "Failed to record account usage", + "usage_store_error", + )), + ) + .into_response(); + } + } + + match completed_response_json(&response.body) { + Ok(CollectedResponse::Completed(body)) => { + self.upstream + .log_response( + &log_context, + StatusCode::OK, + EventLevel::Info, + "v1 responses completed", + json!({"stream": false, "usage": response.usage}), + ) + .await; + (StatusCode::OK, Json(body)).into_response() + } + Ok(CollectedResponse::Failed(failure)) => { + self.upstream + .log_response( + &log_context, + StatusCode::BAD_GATEWAY, + EventLevel::Error, + "v1 responses upstream SSE failed", + failure.metadata(false), + ) + .await; + ( + StatusCode::BAD_GATEWAY, + Json(openai_error( + &failure.openai_error_message(), + "upstream_error", + )), + ) + .into_response() + } + Ok(CollectedResponse::MissingCompleted) => { + self.upstream + .log_response( + &log_context, + StatusCode::BAD_GATEWAY, + EventLevel::Warn, + "v1 responses completed event missing", + json!({"stream": false, "usage": response.usage}), + ) + .await; + ( + StatusCode::BAD_GATEWAY, + Json(openai_error( + "Codex response did not include response.completed", + "empty_upstream_response", + )), + ) + .into_response() + } + Err(error) => { + self.upstream + .log_response( + &log_context, + StatusCode::BAD_GATEWAY, + EventLevel::Warn, + "v1 responses invalid SSE response", + json!({"stream": false, "sseParseError": error.to_string()}), + ) + .await; + ( + StatusCode::BAD_GATEWAY, + Json(openai_error( + "Invalid Codex SSE response", + "invalid_upstream_sse", + )), + ) + .into_response() + } + } + } +} + +#[derive(Deserialize)] +struct ResponsesBody { + model: Option, + input: Option>, + instructions: Option, + reasoning: Option, + tools: Option>, + service_tier: Option, + tool_choice: Option, + parallel_tool_calls: Option, + text: Option, + prompt_cache_key: Option, + include: Option>, + client_metadata: Option, + previous_response_id: Option, + stream: Option, + #[serde(rename = "turnState")] + turn_state: Option, + #[serde(rename = "turnMetadata")] + turn_metadata: Option, + #[serde(rename = "betaFeatures")] + beta_features: Option, + #[serde(rename = "includeTimingMetrics")] + include_timing_metrics: Option, + #[serde(rename = "codexWindowId")] + codex_window_id: Option, + #[serde(rename = "parentThreadId")] + parent_thread_id: Option, + #[serde(alias = "useWebSocket")] + use_websocket: Option, +} + +fn default_body(default_model: String) -> ResponsesBody { + ResponsesBody { + model: Some(default_model), + input: Some(Vec::new()), + instructions: Some(String::new()), + reasoning: None, + tools: None, + service_tier: None, + tool_choice: None, + parallel_tool_calls: None, + text: None, + prompt_cache_key: None, + include: None, + client_metadata: None, + previous_response_id: None, + stream: None, + turn_state: None, + turn_metadata: None, + beta_features: None, + include_timing_metrics: None, + codex_window_id: None, + parent_thread_id: None, + use_websocket: None, + } +} + +fn responses_reasoning( + client_reasoning: Option, + suffix_effort: Option<&str>, + default_effort: Option<&str>, +) -> Option { + let effort = client_reasoning + .as_ref() + .and_then(|reasoning| reasoning.get("effort")) + .and_then(Value::as_str) + .or(suffix_effort) + .or(default_effort); + let summary = client_reasoning + .as_ref() + .and_then(|reasoning| reasoning.get("summary")) + .and_then(Value::as_str) + .unwrap_or("auto"); + if effort.is_none() && client_reasoning.is_none() { + return None; + } + let mut reasoning = json!({"summary": summary}); + if let Some(effort) = effort { + reasoning["effort"] = Value::String(effort.to_string()); + } + Some(reasoning) +} + +fn ensure_reasoning_include(request: &mut CodexResponsesRequest) { + if request.reasoning.is_none() { + return; + } + let include = request.include.get_or_insert_with(Vec::new); + if include + .iter() + .any(|item| item == "reasoning.encrypted_content") + { + return; + } + include.push("reasoning.encrypted_content".to_string()); +} + +fn header_string(headers: &HeaderMap, name: &str) -> Option { + headers + .get(name) + .and_then(|value| value.to_str().ok()) + .filter(|value| !value.trim().is_empty()) + .map(ToString::to_string) +} diff --git a/src/service/usage.rs b/src/service/usage.rs new file mode 100644 index 000000000..4edaa1ba0 --- /dev/null +++ b/src/service/usage.rs @@ -0,0 +1,47 @@ +use crate::{ + codex::accounts::repository::{ + AccountRepositoryError, AccountUsageListRecord, AccountUsageRepository, AccountUsageSummary, + }, + utils::pagination::Page, +}; + +#[derive(Clone)] +pub struct UsageService { + repository: Option, +} + +#[derive(Debug)] +pub enum UsageServiceError { + RepositoryUnavailable, + Repository(AccountRepositoryError), +} + +impl UsageService { + pub fn new(repository: Option) -> Self { + Self { repository } + } + + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> Result, UsageServiceError> { + self.repository()? + .list(cursor, limit) + .await + .map_err(UsageServiceError::Repository) + } + + pub async fn summary(&self) -> Result { + self.repository()? + .summary() + .await + .map_err(UsageServiceError::Repository) + } + + fn repository(&self) -> Result<&AccountUsageRepository, UsageServiceError> { + self.repository + .as_ref() + .ok_or(UsageServiceError::RepositoryUnavailable) + } +} diff --git a/src/state.rs b/src/state.rs deleted file mode 100644 index b0c5e9517..000000000 --- a/src/state.rs +++ /dev/null @@ -1,257 +0,0 @@ -use std::sync::Arc; - -use sqlx::SqlitePool; -use tokio::sync::Mutex; - -use crate::accounts::{ - pool::{AccountPool, AccountPoolOptions, RotationStrategy}, - repository::{AccountRepository, AccountRepositoryResult, AccountUsageRepository}, -}; -use crate::auth::{ - api_key::ApiKeyHasher, - api_key_repository::ClientApiKeyRepository, - oauth::{OAuthClient, PkceSessionStore}, - refresh::TokenRefresher, -}; -use crate::config::AppConfig; -use crate::cookies::repository::CookieRepository; -use crate::crypto::SecretBox; -use crate::logs::repository::EventLogRepository; -use crate::models::repository::ModelSnapshotRepository; - -#[derive(Clone)] -pub struct AppState { - pub services: Arc, -} - -pub struct AppServices { - pub config: AppConfig, - pub db: Option, - pub event_logs: Option, - pub secret_box: Option, - pub api_key_hasher: Option, - pub token_refresher: Option>, - pub oauth_client: Option>, - pub oauth_sessions: Arc>, - pub account_pool: Arc>, -} - -impl AppState { - pub fn new(config: AppConfig) -> Self { - let account_pool = account_pool_from_config(&config); - Self { - services: Arc::new(AppServices { - config, - db: None, - event_logs: None, - secret_box: None, - api_key_hasher: None, - token_refresher: None, - oauth_client: None, - oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())), - account_pool, - }), - } - } - - pub fn with_pool(config: AppConfig, pool: SqlitePool) -> Self { - 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: None, - api_key_hasher: None, - token_refresher: None, - oauth_client: None, - oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())), - account_pool, - }), - } - } - - pub fn with_pool_and_secret_box( - config: AppConfig, - pool: SqlitePool, - secret_box: SecretBox, - ) -> Self { - 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: None, - token_refresher: None, - oauth_client: None, - oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())), - account_pool, - }), - } - } - - pub fn with_pool_secret_and_api_key_hasher( - config: AppConfig, - pool: SqlitePool, - secret_box: SecretBox, - api_key_hasher: ApiKeyHasher, - ) -> Self { - 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: None, - oauth_client: None, - oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())), - 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)), - oauth_client: None, - oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())), - account_pool, - }), - } - } - - pub fn with_pool_secret_api_key_hasher_and_oauth_client( - config: AppConfig, - pool: SqlitePool, - secret_box: SecretBox, - api_key_hasher: ApiKeyHasher, - oauth_client: C, - ) -> Self - where - C: OAuthClient + TokenRefresher, - { - let account_pool = account_pool_from_config(&config); - let oauth_client = Arc::new(oauth_client); - let token_refresher: Arc = oauth_client.clone(); - let oauth_client: Arc = oauth_client; - 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(token_refresher), - oauth_client: Some(oauth_client), - oauth_sessions: Arc::new(Mutex::new(PkceSessionStore::default())), - account_pool, - }), - } - } - - pub fn config(&self) -> &AppConfig { - &self.services.config - } - - pub fn db(&self) -> Option<&SqlitePool> { - self.services.db.as_ref() - } - - pub fn event_logs(&self) -> Option<&EventLogRepository> { - self.services.event_logs.as_ref() - } - - pub fn account_repository(&self) -> Option { - Some(AccountRepository::new( - self.db()?.clone(), - self.services.secret_box.clone()?, - )) - } - - 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(), - self.services.secret_box.clone()?, - )) - } - - pub fn account_pool(&self) -> Arc> { - self.services.account_pool.clone() - } - - pub fn client_api_key_repository(&self) -> Option { - Some(ClientApiKeyRepository::new(self.db()?.clone())) - } - - pub fn model_snapshot_repository(&self) -> Option { - Some(ModelSnapshotRepository::new(self.db()?.clone())) - } - - pub fn api_key_hasher(&self) -> Option<&ApiKeyHasher> { - self.services.api_key_hasher.as_ref() - } - - pub fn token_refresher(&self) -> Option> { - self.services.token_refresher.clone() - } - - pub fn oauth_client(&self) -> Option> { - self.services.oauth_client.clone() - } - - pub fn oauth_sessions(&self) -> Arc> { - self.services.oauth_sessions.clone() - } - - pub async fn reload_account_pool_from_repository(&self) -> AccountRepositoryResult { - let Some(repo) = self.account_repository() else { - return Ok(0); - }; - let accounts = repo.list_pool_accounts().await?; - let restored = accounts.len(); - let mut pool = self.services.account_pool.lock().await; - for account in accounts { - pool.insert(account); - } - Ok(restored) - } -} - -fn account_pool_from_config(config: &AppConfig) -> Arc> { - Arc::new(Mutex::new(AccountPool::with_options(AccountPoolOptions { - max_concurrent_per_account: config.auth.max_concurrent_per_account, - stale_slot_ttl: chrono::Duration::minutes(5), - rotation_strategy: match config.auth.rotation_strategy.as_str() { - "round_robin" => RotationStrategy::RoundRobin, - "sticky" => RotationStrategy::Sticky, - _ => RotationStrategy::LeastUsed, - }, - skip_quota_limited: config.quota.skip_exhausted, - tier_priority: config.auth.tier_priority.clone(), - model_plan_allowlist: std::collections::BTreeMap::new(), - }))) -} diff --git a/src/crypto.rs b/src/utils/crypto.rs similarity index 100% rename from src/crypto.rs rename to src/utils/crypto.rs diff --git a/src/utils/json.rs b/src/utils/json.rs new file mode 100644 index 000000000..0f6519517 --- /dev/null +++ b/src/utils/json.rs @@ -0,0 +1,17 @@ +use serde_json::Value; + +pub fn first_string(value: &Value, paths: &[&[&str]]) -> Option { + paths.iter().find_map(|path| string_at(value, path)) +} + +pub fn string_at(value: &Value, path: &[&str]) -> Option { + let mut current = value; + for segment in path { + current = current.get(*segment)?; + } + current + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs new file mode 100644 index 000000000..15e187381 --- /dev/null +++ b/src/utils/mod.rs @@ -0,0 +1,3 @@ +pub mod crypto; +pub mod json; +pub mod pagination; diff --git a/src/pagination.rs b/src/utils/pagination.rs similarity index 100% rename from src/pagination.rs rename to src/utils/pagination.rs diff --git a/tests/account_pool_scheduling_test.rs b/tests/account_pool_scheduling_test.rs index 876d05283..5863267b5 100644 --- a/tests/account_pool_scheduling_test.rs +++ b/tests/account_pool_scheduling_test.rs @@ -1,7 +1,7 @@ use std::collections::BTreeMap; use chrono::{Duration, TimeZone, Utc}; -use codex_proxy_rs::accounts::{ +use codex_proxy_rs::codex::accounts::{ model::{Account, AccountStatus}, pool::{AccountAcquireRequest, AccountPool, AccountPoolOptions, RotationStrategy}, }; diff --git a/tests/account_pool_startup_test.rs b/tests/account_pool_startup_test.rs index 9d0225069..d410c0f4c 100644 --- a/tests/account_pool_startup_test.rs +++ b/tests/account_pool_startup_test.rs @@ -2,7 +2,8 @@ use chrono::Utc; use secrecy::SecretString; use codex_proxy_rs::{ - accounts::{ + app::state::AppState, + codex::accounts::{ model::AccountStatus, repository::{NewAccount, UsageDelta}, }, @@ -11,9 +12,8 @@ use codex_proxy_rs::{ QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, }, - crypto::SecretBox, - state::AppState, storage::db::connect_sqlite, + utils::crypto::SecretBox, }; fn test_config(database_url: String) -> AppConfig { @@ -80,7 +80,7 @@ async fn app_state_should_restore_account_pool_from_sqlite_accounts() { let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); let secret_box = SecretBox::new([31u8; 32]); - let repo = codex_proxy_rs::accounts::repository::AccountRepository::new( + let repo = codex_proxy_rs::codex::accounts::repository::AccountRepository::new( pool.clone(), secret_box.clone(), ); @@ -114,10 +114,10 @@ async fn app_state_should_restore_account_pool_from_sqlite_accounts() { assert_eq!(restored, 1); let acquired = state - .account_pool() - .lock() + .services + .accounts + .acquire_runtime_account("gpt-5.5") .await - .acquire("gpt-5.5") .unwrap(); assert_eq!(acquired.id, "acct_restored"); assert_eq!(acquired.access_token, "access-secret"); diff --git a/tests/account_refresh_test.rs b/tests/account_refresh_test.rs index 61cb94f39..0bdc46e05 100644 --- a/tests/account_refresh_test.rs +++ b/tests/account_refresh_test.rs @@ -11,12 +11,9 @@ use chrono::{Duration as ChronoDuration, Utc}; use tokio::time::sleep; use codex_proxy_rs::{ - accounts::model::{Account, AccountStatus}, - auth::{ - refresh::{ - RefreshFailure, RefreshPolicy, RefreshScheduler, RefreshTrigger, TokenRefresher, - }, - token::TokenPair, + codex::accounts::model::{Account, AccountStatus}, + codex::oauth::{ + RefreshFailure, RefreshPolicy, RefreshScheduler, RefreshTrigger, TokenPair, TokenRefresher, }, }; diff --git a/tests/account_repository_test.rs b/tests/account_repository_test.rs index be3426f08..ccc96b94b 100644 --- a/tests/account_repository_test.rs +++ b/tests/account_repository_test.rs @@ -2,12 +2,12 @@ use chrono::{Duration, Utc}; use secrecy::{ExposeSecret, SecretString}; use codex_proxy_rs::{ - accounts::{ + codex::accounts::{ model::AccountStatus, repository::{AccountRepository, NewAccount, TokenUpdate, UsageDelta}, }, - crypto::SecretBox, storage::db::connect_sqlite, + utils::crypto::SecretBox, }; #[tokio::test] diff --git a/tests/admin_accounts_cookies_quota_test.rs b/tests/admin_accounts_cookies_quota_test.rs new file mode 100644 index 000000000..978da7992 --- /dev/null +++ b/tests/admin_accounts_cookies_quota_test.rs @@ -0,0 +1,440 @@ +use std::sync::Arc; + +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use serde_json::json; +use tokio::sync::Mutex; +use tower::ServiceExt; +use wiremock::{ + matchers::{header, method, path}, + Mock, MockServer, ResponseTemplate, +}; + +use codex_proxy_rs::codex::oauth::{RefreshFailure, TokenPair}; + +mod common; + +use common::{ + admin_accounts::{ + admin_accounts_test_app, admin_accounts_test_app_with_base_url, + admin_accounts_test_app_with_refresher, import_test_account, FailingTokenRefresher, + StaticTokenRefresher, + }, + response_json, +}; + +#[tokio::test] +async fn admin_account_cookies_should_set_get_and_delete_encrypted_cookie_header() { + let (app, _state, pool, _dir) = + admin_accounts_test_app("admin-account-cookies.sqlite", 22).await; + import_test_account(&app, "session_1", "acct_cookies").await; + + let set_response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/acct_cookies/cookies") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + r#"{"cookies":"cf_clearance=clear-secret; __cf_bm=bm-secret"}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(set_response.status(), StatusCode::OK); + let body = response_json(set_response).await; + assert_eq!( + body["data"]["cookies"], + "__cf_bm=bm-secret; cf_clearance=clear-secret" + ); + + let stored = sqlx::query_as::<_, (String, String)>( + "select name, value_cipher from account_cookies where account_id = ? order by name asc", + ) + .bind("acct_cookies") + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!(stored.len(), 2); + assert!(stored.iter().all(|(_, cipher)| cipher.starts_with("v1:"))); + assert!(stored + .iter() + .all(|(_, cipher)| !cipher.contains("clear-secret") && !cipher.contains("bm-secret"))); + + let get_response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/accounts/acct_cookies/cookies") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(get_response.status(), StatusCode::OK); + let body = response_json(get_response).await; + assert_eq!( + body["data"]["cookies"], + "__cf_bm=bm-secret; cf_clearance=clear-secret" + ); + + let delete_response = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri("/admin/accounts/acct_cookies/cookies") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(delete_response.status(), StatusCode::OK); + + let get_empty = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/accounts/acct_cookies/cookies") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(get_empty.status(), StatusCode::OK); + let body = response_json(get_empty).await; + assert!(body["data"]["cookies"].is_null()); +} + +#[tokio::test] +async fn admin_account_cookies_should_require_existing_account_and_non_empty_cookies() { + let (app, _state, _pool, _dir) = + admin_accounts_test_app("admin-account-cookies-invalid.sqlite", 23).await; + + let missing = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/accounts/missing/cookies") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(missing.status(), StatusCode::NOT_FOUND); + + import_test_account(&app, "session_1", "acct_cookie_invalid").await; + let invalid = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/acct_cookie_invalid/cookies") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"cookies":""}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(invalid.status(), StatusCode::BAD_REQUEST); +} + +#[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 + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .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 + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .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) = 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 + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .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"); +} diff --git a/tests/admin_accounts_import_export_test.rs b/tests/admin_accounts_import_export_test.rs new file mode 100644 index 000000000..fd7c9631c --- /dev/null +++ b/tests/admin_accounts_import_export_test.rs @@ -0,0 +1,887 @@ +use std::{fs, sync::Arc}; + +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use secrecy::ExposeSecret; +use serde_json::json; +use sqlx::Row; +use tokio::sync::Mutex; +use tower::ServiceExt; + +use codex_proxy_rs::{ + app::build_router, app::state::AppState, codex::accounts::repository::AccountRepository, + codex::oauth::TokenPair, storage::db::connect_sqlite, utils::crypto::SecretBox, +}; + +mod common; + +use common::{ + admin_accounts::{ + admin_accounts_test_app, admin_accounts_test_app_with_refresher, import_test_account, + post_admin_account, test_config, test_jwt, StaticTokenRefresher, + }, + response_json, seed_admin_session, +}; + +#[tokio::test] +async fn admin_accounts_import_should_require_admin_session_cookie() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-accounts.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let app = build_router(AppState::with_pool_and_secret_box( + test_config(url), + pool, + SecretBox::new([11u8; 32]), + )); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/import") + .header("content-type", "application/json") + .header("x-request-id", "req_accounts") + .body(Body::from(r#"{"accounts":[]}"#)) + .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_accounts"); +} + +#[tokio::test] +async fn admin_accounts_import_should_store_tokens_encrypted_and_list_sanitized_accounts() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-accounts.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + let app = build_router(AppState::with_pool_and_secret_box( + test_config(url), + pool.clone(), + SecretBox::new([12u8; 32]), + )); + let import_body = json!({ + "accounts": [{ + "id": "acct_imported", + "email": "user@example.com", + "accountId": "chatgpt-account", + "userId": "chatgpt-user", + "label": "primary", + "planType": "plus", + "token": "access-secret", + "refreshToken": "refresh-secret", + "status": "active" + }] + }); + + 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") + .header("x-request-id", "req_accounts") + .body(Body::from(import_body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(import_response.status(), StatusCode::OK); + let body = response_json(import_response).await; + assert_eq!(body["data"]["imported"], 1); + assert_eq!(body["data"]["skipped"], 0); + let stored: (String, String) = sqlx::query_as( + "select access_token_cipher, refresh_token_cipher from accounts where id = ?", + ) + .bind("acct_imported") + .fetch_one(&pool) + .await + .unwrap(); + assert!(stored.0.starts_with("v1:")); + assert!(!stored.0.contains("access-secret")); + assert!(stored.1.starts_with("v1:")); + assert!(!stored.1.contains("refresh-secret")); + + let list_response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/accounts?limit=10") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_accounts_list") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(list_response.status(), StatusCode::OK); + let body = response_json(list_response).await; + assert_eq!(body["data"][0]["id"], "acct_imported"); + assert_eq!(body["data"][0]["email"], "user@example.com"); + assert!(body["data"][0].get("token").is_none()); + assert!(body["data"][0].get("refreshToken").is_none()); + assert_eq!(body["page"]["limit"], 10); +} + +#[tokio::test] +async fn admin_accounts_import_should_accept_sub2api_oauth_export_and_mark_format() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-sub2api.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + let app = build_router(AppState::with_pool_and_secret_box( + test_config(url), + pool.clone(), + SecretBox::new([14u8; 32]), + )); + let import_body = json!({ + "type": "sub2api-data", + "version": 1, + "proxies": [], + "accounts": [ + { + "name": "Sub2API Team", + "platform": "openai", + "type": "oauth", + "credentials": { + "access_token": "Bearer sub2api-access-secret", + "refresh_token": "rt_sub2api", + "email": "team@example.com", + "chatgpt_account_id": "chatgpt-account", + "chatgpt_user_id": "chatgpt-user", + "plan_type": "team" + }, + "concurrency": 0, + "priority": 0 + }, + { + "name": "Other Provider", + "platform": "anthropic", + "type": "oauth", + "credentials": { + "access_token": "ignored-secret" + } + } + ] + }); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/import") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_sub2api") + .body(Body::from(import_body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["sourceFormat"], "sub2api"); + assert_eq!(body["data"]["imported"], 1); + assert_eq!(body["data"]["skipped"], 0); + let stored: (String, String, String, String, String) = + sqlx::query_as("select email, account_id, user_id, label, plan_type from accounts") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(stored.0, "team@example.com"); + assert_eq!(stored.1, "chatgpt-account"); + assert_eq!(stored.2, "chatgpt-user"); + assert_eq!(stored.3, "Sub2API Team"); + assert_eq!(stored.4, "team"); +} + +#[tokio::test] +async fn admin_accounts_import_should_accept_sub2api_native_account_export_without_proxy_data() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-sub2api-native.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + let app = build_router(AppState::with_pool_and_secret_box( + test_config(url), + pool, + SecretBox::new([15u8; 32]), + )); + let import_body = json!({ + "accounts": [{ + "id": "acct_sub2api_native", + "token": "native-access-secret", + "refreshToken": "native-refresh-secret", + "email": "native@example.com", + "accountId": "native-account", + "userId": "native-user", + "label": "Native Sub2API", + "planType": "plus", + "proxyApiKey": "ignored-proxy-secret", + "status": "active", + "usage": { + "request_count": 1, + "input_tokens": 0, + "output_tokens": 0, + "cached_tokens": 0 + }, + "cachedQuota": null, + "quotaVerifyRequired": false + }] + }); + + 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") + .header("x-request-id", "req_sub2api_native") + .body(Body::from(import_body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(import_response.status(), StatusCode::OK); + let body = response_json(import_response).await; + assert_eq!(body["data"]["sourceFormat"], "sub2api"); + assert_eq!(body["data"]["imported"], 1); + + let list_response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/accounts?limit=10") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_sub2api_native_list") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let list = response_json(list_response).await; + assert_eq!(list["data"][0]["id"], "acct_sub2api_native"); + assert_eq!(list["data"][0]["planType"], "plus"); + assert!(list["data"][0].get("proxyApiKey").is_none()); + assert!(list["data"][0].get("usage").is_none()); +} + +#[tokio::test] +async fn admin_accounts_export_should_return_native_accounts_with_tokens_and_filter_ids() { + let (app, _state, _pool, _dir) = + admin_accounts_test_app("admin-account-export-native.sqlite", 24).await; + import_test_account(&app, "session_1", "acct_export_a").await; + import_test_account(&app, "session_1", "acct_export_b").await; + + let response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/accounts/export?ids=acct_export_a") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["sourceFormat"], "native"); + assert_eq!(body["data"]["accounts"].as_array().unwrap().len(), 1); + assert_eq!(body["data"]["accounts"][0]["id"], "acct_export_a"); + assert_eq!(body["data"]["accounts"][0]["token"], "access-acct_export_a"); + assert_eq!( + body["data"]["accounts"][0]["refreshToken"], + "refresh-acct_export_a" + ); + assert!(body["data"]["accounts"][0].get("proxyApiKey").is_none()); + + let invalid = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/accounts/export?format=proxy") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(invalid.status(), StatusCode::BAD_REQUEST); +} + +#[tokio::test] +async fn admin_accounts_export_should_return_sub2api_openai_oauth_payload_without_proxy_fields() { + let (app, _state, _pool, _dir) = + admin_accounts_test_app("admin-account-export-sub2api.sqlite", 25).await; + import_test_account(&app, "session_1", "acct_export_sub2api").await; + + let label_response = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri("/admin/accounts/acct_export_sub2api/label") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"label":"Sub2API Export"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(label_response.status(), StatusCode::OK); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/accounts/export?format=sub2api") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["type"], "sub2api-data"); + assert_eq!(body["data"]["version"], 1); + assert_eq!(body["data"]["proxies"], json!([])); + assert_eq!(body["data"]["accounts"].as_array().unwrap().len(), 1); + assert_eq!(body["data"]["accounts"][0]["name"], "Sub2API Export"); + assert_eq!(body["data"]["accounts"][0]["platform"], "openai"); + assert_eq!(body["data"]["accounts"][0]["type"], "oauth"); + assert_eq!( + body["data"]["accounts"][0]["credentials"]["access_token"], + "access-acct_export_sub2api" + ); + assert_eq!( + body["data"]["accounts"][0]["credentials"]["refresh_token"], + "refresh-acct_export_sub2api" + ); + assert_eq!( + body["data"]["accounts"][0]["credentials"]["email"], + "acct_export_sub2api@example.com" + ); + assert_eq!( + body["data"]["accounts"][0]["credentials"]["plan_type"], + "plus" + ); + assert!(body["data"]["accounts"][0].get("proxy").is_none()); + assert!(body["data"]["accounts"][0].get("proxyUrl").is_none()); +} + +#[tokio::test] +async fn manual_add_should_derive_claims_ignore_metadata_and_sync_pool() { + let (app, state, pool, _dir) = + admin_accounts_test_app("admin-account-manual-add.sqlite", 26).await; + let token = test_jwt( + Some("jwt-account"), + Some("jwt-user"), + Some("jwt@example.com"), + Some("team"), + 3600, + ); + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({ + "id": "caller-id", + "email": "caller@example.com", + "accountId": "caller-account", + "userId": "caller-user", + "label": "Caller Label", + "planType": "caller-plan", + "token": token, + "refreshToken": "manual-refresh-secret", + "status": "disabled" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + let response_id = body["data"]["id"].as_str().unwrap(); + assert_ne!(response_id, "caller-id"); + assert_eq!(body["data"]["email"], "jwt@example.com"); + assert_eq!(body["data"]["accountId"], "jwt-account"); + assert_eq!(body["data"]["userId"], "jwt-user"); + assert_eq!(body["data"]["planType"], "team"); + assert!(body["data"]["label"].is_null()); + assert_eq!(body["data"]["status"], "active"); + assert!(body["data"].get("token").is_none()); + assert!(body["data"].get("refreshToken").is_none()); + + let stored = sqlx::query( + "select id, email, account_id, user_id, label, plan_type, status, access_token_cipher, refresh_token_cipher from accounts", + ) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(stored.get::("id"), response_id); + assert_eq!( + stored.get::, _>("email").as_deref(), + Some("jwt@example.com") + ); + assert_eq!( + stored.get::, _>("account_id").as_deref(), + Some("jwt-account") + ); + assert_eq!( + stored.get::, _>("user_id").as_deref(), + Some("jwt-user") + ); + assert_eq!(stored.get::, _>("label"), None); + assert_eq!( + stored.get::, _>("plan_type").as_deref(), + Some("team") + ); + assert_eq!(stored.get::("status"), "active"); + let access_token_cipher = stored.get::("access_token_cipher"); + assert!(access_token_cipher.starts_with("v1:")); + assert!(!access_token_cipher.contains(&token)); + let refresh_token_cipher = stored.get::("refresh_token_cipher"); + assert!(refresh_token_cipher.starts_with("v1:")); + assert!(!refresh_token_cipher.contains("manual-refresh-secret")); + + let acquired = state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .unwrap(); + assert_eq!(acquired.id, response_id); + assert_eq!(acquired.email.as_deref(), Some("jwt@example.com")); + assert_eq!(acquired.account_id.as_deref(), Some("jwt-account")); + assert_eq!(acquired.user_id.as_deref(), Some("jwt-user")); + assert_eq!(acquired.plan_type.as_deref(), Some("team")); +} + +#[tokio::test] +async fn admin_account_manual_add_should_reject_missing_invalid_expired_or_unbound_tokens() { + let (app, _state, _pool, _dir) = + admin_accounts_test_app("admin-account-manual-add-invalid.sqlite", 27).await; + + let missing_tokens = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(missing_tokens.status(), StatusCode::BAD_REQUEST); + + let invalid = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({ + "token": "not-a-jwt" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(invalid.status(), StatusCode::BAD_REQUEST); + + let expired = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({ + "token": test_jwt( + Some("expired-account"), + Some("expired-user"), + Some("expired@example.com"), + Some("plus"), + -3600, + ) + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(expired.status(), StatusCode::BAD_REQUEST); + + let missing_account_claim = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({ + "token": test_jwt( + None, + Some("claimless-user"), + Some("claimless@example.com"), + Some("free"), + 3600, + ) + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(missing_account_claim.status(), StatusCode::BAD_REQUEST); +} + +#[tokio::test] +async fn admin_account_manual_add_should_exchange_refresh_token_only_and_rotate_refresh_token() { + let token = test_jwt( + Some("rt-account"), + Some("rt-user"), + Some("rt@example.com"), + Some("plus"), + 3600, + ); + let calls = Arc::new(Mutex::new(Vec::new())); + let refresher = StaticTokenRefresher { + result: Ok(TokenPair { + access_token: token.clone(), + refresh_token: Some("rotated-refresh".to_string()), + }), + calls: calls.clone(), + }; + let (app, state, pool, _dir) = admin_accounts_test_app_with_refresher( + "admin-account-manual-refresh-only.sqlite", + 28, + refresher, + ) + .await; + + let response = post_admin_account( + &app, + json!({ + "refreshToken": "initial-refresh", + "email": "caller@example.com", + "planType": "caller-plan" + }), + ) + .await; + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["email"], "rt@example.com"); + assert_eq!(body["data"]["accountId"], "rt-account"); + assert!(body["data"].get("token").is_none()); + assert!(body["data"].get("refreshToken").is_none()); + assert_eq!(*calls.lock().await, vec!["initial-refresh".to_string()]); + + let repo = AccountRepository::new(pool, SecretBox::new([28u8; 32])); + let stored = repo + .get(body["data"]["id"].as_str().unwrap()) + .await + .unwrap() + .unwrap(); + assert_eq!(stored.access_token.expose_secret(), &token); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "rotated-refresh" + ); + let acquired = state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .unwrap(); + assert_eq!(acquired.refresh_token.as_deref(), Some("rotated-refresh")); +} + +#[tokio::test] +async fn admin_account_manual_add_should_preserve_input_refresh_token_when_exchange_omits_rotation() +{ + let token = test_jwt( + Some("rt-preserve-account"), + Some("rt-preserve-user"), + Some("preserve@example.com"), + Some("free"), + 3600, + ); + let refresher = StaticTokenRefresher { + result: Ok(TokenPair { + access_token: token, + refresh_token: None, + }), + calls: Arc::new(Mutex::new(Vec::new())), + }; + let (app, _state, pool, _dir) = admin_accounts_test_app_with_refresher( + "admin-account-manual-refresh-preserve.sqlite", + 29, + refresher, + ) + .await; + + let response = post_admin_account(&app, json!({ "refreshToken": "preserved-refresh" })).await; + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + let repo = AccountRepository::new(pool, SecretBox::new([29u8; 32])); + let stored = repo + .get(body["data"]["id"].as_str().unwrap()) + .await + .unwrap() + .unwrap(); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "preserved-refresh" + ); +} + +#[tokio::test] +async fn admin_account_manual_add_should_update_existing_account_for_same_account_and_user() { + let (app, state, pool, _dir) = + admin_accounts_test_app("admin-account-manual-update-existing.sqlite", 30).await; + let first_token = test_jwt( + Some("team-account"), + Some("same-user"), + Some("first@example.com"), + Some("free"), + 3600, + ); + let first_response = post_admin_account( + &app, + json!({ + "token": first_token, + "refreshToken": "first-refresh" + }), + ) + .await; + assert_eq!(first_response.status(), StatusCode::OK); + let first_body = response_json(first_response).await; + let first_id = first_body["data"]["id"].as_str().unwrap().to_string(); + + let second_token = test_jwt( + Some("team-account"), + Some("same-user"), + Some("second@example.com"), + Some("team"), + 7200, + ); + let second_response = post_admin_account(&app, json!({ "token": second_token })).await; + + assert_eq!(second_response.status(), StatusCode::OK); + let second_body = response_json(second_response).await; + assert_eq!(second_body["data"]["id"], first_id); + assert_eq!(second_body["data"]["email"], "second@example.com"); + assert_eq!(second_body["data"]["planType"], "team"); + + let count: (i64,) = sqlx::query_as("select count(*) from accounts") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(count.0, 1); + let repo = AccountRepository::new(pool, SecretBox::new([30u8; 32])); + let stored = repo.get(&first_id).await.unwrap().unwrap(); + assert_eq!(stored.access_token.expose_secret(), &second_token); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "first-refresh" + ); + assert_eq!(stored.email.as_deref(), Some("second@example.com")); + assert_eq!(stored.plan_type.as_deref(), Some("team")); + + let acquired = state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .unwrap(); + assert_eq!(acquired.id, first_id); + assert_eq!(acquired.access_token, second_token); + assert_eq!(acquired.refresh_token.as_deref(), Some("first-refresh")); +} + +#[tokio::test] +async fn manual_add_should_preserve_existing_refresh_when_refresh_only_omits_rotation() { + let refreshed_token = test_jwt( + Some("rt-existing-account"), + Some("rt-existing-user"), + Some("second@example.com"), + Some("team"), + 7200, + ); + let refresher = StaticTokenRefresher { + result: Ok(TokenPair { + access_token: refreshed_token.clone(), + refresh_token: None, + }), + calls: Arc::new(Mutex::new(Vec::new())), + }; + let (app, _state, pool, _dir) = admin_accounts_test_app_with_refresher( + "admin-account-manual-refresh-existing-preserve.sqlite", + 31, + refresher, + ) + .await; + let first_token = test_jwt( + Some("rt-existing-account"), + Some("rt-existing-user"), + Some("first@example.com"), + Some("free"), + 3600, + ); + let first_response = post_admin_account( + &app, + json!({ + "token": first_token, + "refreshToken": "old-refresh" + }), + ) + .await; + assert_eq!(first_response.status(), StatusCode::OK); + let first_body = response_json(first_response).await; + let account_id = first_body["data"]["id"].as_str().unwrap().to_string(); + + let second_response = + post_admin_account(&app, json!({ "refreshToken": "incoming-refresh" })).await; + + assert_eq!(second_response.status(), StatusCode::OK); + let repo = AccountRepository::new(pool, SecretBox::new([31u8; 32])); + let stored = repo.get(&account_id).await.unwrap().unwrap(); + assert_eq!(stored.access_token.expose_secret(), &refreshed_token); + assert_eq!(stored.refresh_token.unwrap().expose_secret(), "old-refresh"); + assert_eq!(stored.email.as_deref(), Some("second@example.com")); + assert_eq!(stored.plan_type.as_deref(), Some("team")); +} + +#[tokio::test] +async fn admin_accounts_import_cli_should_read_codex_auth_file_store_encrypted_and_sync_pool() { + let (app, state, pool, dir) = + admin_accounts_test_app("admin-account-import-cli.sqlite", 28).await; + let codex_home = dir.path().join("codex-home"); + fs::create_dir_all(&codex_home).unwrap(); + let token = test_jwt( + Some("cli-account"), + Some("cli-user"), + Some("cli@example.com"), + Some("plus"), + 3600, + ); + fs::write( + codex_home.join("auth.json"), + json!({ + "access_token": token, + "refresh_token": "cli-refresh-secret" + }) + .to_string(), + ) + .unwrap(); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/import-cli") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_import_cli") + .body(Body::from( + json!({ + "codexHome": codex_home.display().to_string(), + "id": "caller-id", + "email": "caller@example.com", + "accountId": "caller-account", + "userId": "caller-user", + "label": "caller-label", + "planType": "caller-plan" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["sourceFormat"], "codex_cli"); + assert_eq!(body["data"]["imported"], 1); + assert!(body["data"].get("token").is_none()); + assert!(body["data"].get("refreshToken").is_none()); + + let stored: ( + String, + String, + String, + String, + Option, + String, + String, + String, + ) = sqlx::query_as( + "select id, email, account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher from accounts limit 1", + ) + .fetch_one(&pool) + .await + .unwrap(); + assert_ne!(stored.0, "caller-id"); + assert_eq!(stored.1, "cli@example.com"); + assert_eq!(stored.2, "cli-account"); + assert_eq!(stored.3, "cli-user"); + assert_eq!(stored.4, None); + assert_eq!(stored.5, "plus"); + assert!(stored.6.starts_with("v1:")); + assert!(!stored.6.contains(&token)); + assert!(stored.7.starts_with("v1:")); + assert!(!stored.7.contains("cli-refresh-secret")); + + let acquired = state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .unwrap(); + assert_eq!(acquired.id, stored.0); +} diff --git a/tests/admin_accounts_list_test.rs b/tests/admin_accounts_list_test.rs new file mode 100644 index 000000000..27688b527 --- /dev/null +++ b/tests/admin_accounts_list_test.rs @@ -0,0 +1,58 @@ +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use chrono::Utc; +use tower::ServiceExt; + +use codex_proxy_rs::{ + app::build_router, app::state::AppState, storage::db::connect_sqlite, utils::crypto::SecretBox, +}; + +mod common; + +use common::{admin_accounts::test_config, response_json, seed_admin_session}; + +#[tokio::test] +async fn admin_accounts_list_should_not_decrypt_account_tokens() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-accounts.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + let now = Utc::now().to_rfc3339(); + sqlx::query( + "insert into accounts (id, email, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)", + ) + .bind("acct_corrupt") + .bind("user@example.com") + .bind("not-a-secret-box-cipher") + .bind("active") + .bind(&now) + .bind(&now) + .execute(&pool) + .await + .unwrap(); + let app = build_router(AppState::with_pool_and_secret_box( + test_config(url), + pool, + SecretBox::new([13u8; 32]), + )); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/accounts?limit=10") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_accounts_list") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"][0]["id"], "acct_corrupt"); +} diff --git a/tests/admin_accounts_mutation_test.rs b/tests/admin_accounts_mutation_test.rs new file mode 100644 index 000000000..2bf28e689 --- /dev/null +++ b/tests/admin_accounts_mutation_test.rs @@ -0,0 +1,332 @@ +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use serde_json::json; +use tower::ServiceExt; + +mod common; + +use common::{ + admin_accounts::{admin_accounts_test_app, import_test_account}, + response_json, +}; + +#[tokio::test] +async fn admin_account_label_should_update_and_clear_label() { + let (app, state, pool, _dir) = admin_accounts_test_app("admin-account-label.sqlite", 16).await; + import_test_account(&app, "session_1", "acct_label").await; + + let set_response = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri("/admin/accounts/acct_label/label") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_label") + .body(Body::from(r#"{"label":"Team Alpha"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(set_response.status(), StatusCode::OK); + let body = response_json(set_response).await; + assert_eq!(body["data"]["label"], "Team Alpha"); + + let clear_response = app + .oneshot( + Request::builder() + .method("PATCH") + .uri("/admin/accounts/acct_label/label") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"label":null}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(clear_response.status(), StatusCode::OK); + let body = response_json(clear_response).await; + assert!(body["data"]["label"].is_null()); + + let stored: (Option,) = sqlx::query_as("select label from accounts where id = ?") + .bind("acct_label") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(stored.0, None); + assert!(state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .is_some()); +} + +#[tokio::test] +async fn admin_account_label_should_reject_too_long_or_missing_account() { + let (app, _state, _pool, _dir) = + admin_accounts_test_app("admin-account-label-invalid.sqlite", 17).await; + import_test_account(&app, "session_1", "acct_label_invalid").await; + let long_label = "x".repeat(65); + + let too_long = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri("/admin/accounts/acct_label_invalid/label") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(json!({ "label": long_label }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(too_long.status(), StatusCode::BAD_REQUEST); + + let missing = app + .oneshot( + Request::builder() + .method("PATCH") + .uri("/admin/accounts/missing/label") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"label":"Team Alpha"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(missing.status(), StatusCode::NOT_FOUND); +} + +#[tokio::test] +async fn admin_account_status_should_update_database_and_runtime_pool() { + let (app, state, pool, _dir) = admin_accounts_test_app("admin-account-status.sqlite", 18).await; + import_test_account(&app, "session_1", "acct_status").await; + + let response = app + .oneshot( + Request::builder() + .method("PATCH") + .uri("/admin/accounts/acct_status/status") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_status") + .body(Body::from(r#"{"status":"disabled"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["status"], "disabled"); + + let stored: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_status") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(stored.0, "disabled"); + assert!(state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .is_none()); +} + +#[tokio::test] +async fn admin_account_delete_should_remove_database_row_and_runtime_pool_entry() { + let (app, state, pool, _dir) = admin_accounts_test_app("admin-account-delete.sqlite", 19).await; + import_test_account(&app, "session_1", "acct_delete").await; + + let response = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri("/admin/accounts/acct_delete") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_delete") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["deleted"], true); + + let row_count: (i64,) = sqlx::query_as("select count(*) from accounts where id = ?") + .bind("acct_delete") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(row_count.0, 0); + assert!(state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .is_none()); + + let missing = app + .oneshot( + Request::builder() + .method("DELETE") + .uri("/admin/accounts/acct_delete") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(missing.status(), StatusCode::NOT_FOUND); +} + +#[tokio::test] +async fn admin_accounts_batch_delete_should_delete_found_accounts_and_report_missing_ids() { + let (app, state, pool, _dir) = + admin_accounts_test_app("admin-accounts-batch-delete.sqlite", 20).await; + import_test_account(&app, "session_1", "acct_batch_delete_a").await; + import_test_account(&app, "session_1", "acct_batch_delete_b").await; + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/batch-delete") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({ + "ids": ["acct_batch_delete_a", "ghost", "acct_batch_delete_b"] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["deleted"], 2); + assert_eq!(body["data"]["notFound"], json!(["ghost"])); + + let row_count: (i64,) = sqlx::query_as("select count(*) from accounts") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(row_count.0, 0); + assert!(state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .is_none()); + + let empty = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/batch-delete") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"ids":[]}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(empty.status(), StatusCode::BAD_REQUEST); +} + +#[tokio::test] +async fn admin_accounts_batch_status_should_update_found_accounts_and_reject_invalid_status() { + let (app, state, pool, _dir) = + admin_accounts_test_app("admin-accounts-batch-status.sqlite", 21).await; + import_test_account(&app, "session_1", "acct_batch_status_a").await; + import_test_account(&app, "session_1", "acct_batch_status_b").await; + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/batch-status") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({ + "ids": ["acct_batch_status_a", "ghost"], + "status": "disabled" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["updated"], 1); + assert_eq!(body["data"]["notFound"], json!(["ghost"])); + + let statuses = + sqlx::query_as::<_, (String, String)>("select id, status from accounts order by id asc") + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!( + statuses, + vec![ + ("acct_batch_status_a".to_string(), "disabled".to_string()), + ("acct_batch_status_b".to_string(), "active".to_string()) + ] + ); + let acquired = state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .unwrap(); + assert_eq!(acquired.id, "acct_batch_status_b"); + + let invalid = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/batch-status") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({ + "ids": ["acct_batch_status_a"], + "status": "expired" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(invalid.status(), StatusCode::BAD_REQUEST); + + let empty = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/batch-status") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"ids":[],"status":"active"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(empty.status(), StatusCode::BAD_REQUEST); +} diff --git a/tests/admin_accounts_oauth_test.rs b/tests/admin_accounts_oauth_test.rs new file mode 100644 index 000000000..c69e4b499 --- /dev/null +++ b/tests/admin_accounts_oauth_test.rs @@ -0,0 +1,606 @@ +use std::sync::Arc; + +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use secrecy::ExposeSecret; +use serde_json::json; +use tokio::sync::Mutex; +use tower::ServiceExt; + +use codex_proxy_rs::{ + codex::accounts::repository::AccountRepository, + codex::oauth::{DeviceCode, OAuthError, TokenPair}, + utils::crypto::SecretBox, +}; + +mod common; + +use common::{ + admin_accounts::{ + admin_accounts_test_app, admin_accounts_test_app_with_oauth_client, import_test_account, + test_jwt, StaticOAuthClient, + }, + response_json, +}; + +#[tokio::test] +async fn admin_auth_status_should_require_admin_session_cookie() { + let (app, _state, _pool, _dir) = + admin_accounts_test_app("admin-auth-status-auth.sqlite", 33).await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/auth/status") + .header("x-request-id", "req_auth_status_no_session") + .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_auth_status_no_session"); +} + +#[tokio::test] +async fn admin_auth_status_should_report_empty_account_pool() { + let (app, _state, _pool, _dir) = + admin_accounts_test_app("admin-auth-status-empty.sqlite", 34).await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/auth/status") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_auth_status_empty") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["authenticated"], false); + assert!(body["data"]["user"].is_null()); + assert_eq!(body["data"]["pool"]["total"], 0); + assert_eq!(body["data"]["pool"]["active"], 0); +} + +#[tokio::test] +async fn admin_auth_status_should_report_user_and_pool_summary_without_secrets() { + let (app, _state, pool, _dir) = + admin_accounts_test_app("admin-auth-status-summary.sqlite", 35).await; + import_test_account(&app, "session_1", "acct_auth_status").await; + import_test_account(&app, "session_1", "acct_auth_disabled").await; + sqlx::query("update accounts set status = 'disabled' where id = ?") + .bind("acct_auth_disabled") + .execute(&pool) + .await + .unwrap(); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/auth/status") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_auth_status_summary") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["authenticated"], true); + assert_eq!(body["data"]["pool"]["total"], 2); + assert_eq!(body["data"]["pool"]["active"], 1); + assert_eq!(body["data"]["pool"]["disabled"], 1); + assert!(body["data"]["user"].get("accessToken").is_none()); + assert!(body["data"]["user"].get("token").is_none()); + assert!(body["data"]["user"].get("refreshToken").is_none()); +} + +#[tokio::test] +async fn admin_auth_logout_should_clear_accounts_and_runtime_pool() { + let (app, state, pool, _dir) = admin_accounts_test_app("admin-auth-logout.sqlite", 36).await; + import_test_account(&app, "session_1", "acct_auth_logout").await; + sqlx::query( + "insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens) values (?, 2, 3, 4, 1)", + ) + .bind("acct_auth_logout") + .execute(&pool) + .await + .unwrap(); + let set_cookies = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/acct_auth_logout/cookies") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"cookies":"cf_clearance=secret"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(set_cookies.status(), StatusCode::OK); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/auth/logout") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_auth_logout") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["success"], true); + assert_eq!(body["data"]["deleted"], 1); + let account_count: (i64,) = sqlx::query_as("select count(*) from accounts") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(account_count.0, 0); + let usage_count: (i64,) = sqlx::query_as("select count(*) from account_usage") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(usage_count.0, 0); + let cookie_count: (i64,) = sqlx::query_as("select count(*) from account_cookies") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(cookie_count.0, 0); + assert!(state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .is_none()); +} + +#[tokio::test] +async fn admin_auth_device_login_should_require_admin_session_cookie() { + let (app, _state, _pool, _dir) = + admin_accounts_test_app("admin-auth-device-login-auth.sqlite", 37).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/auth/device-login") + .header("x-request-id", "req_device_login_no_session") + .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_device_login_no_session"); +} + +#[tokio::test] +async fn admin_auth_device_login_should_return_openai_device_code() { + let client = StaticOAuthClient { + device_response: Ok(DeviceCode { + device_code: "device-secret".to_string(), + user_code: "ABCD-EFGH".to_string(), + verification_uri: "https://auth.openai.com/activate".to_string(), + verification_uri_complete: "https://auth.openai.com/activate?user_code=ABCD-EFGH" + .to_string(), + expires_in: 900, + interval: 5, + }), + poll_response: Err(OAuthError::AuthorizationPending), + exchange_response: Err(OAuthError::Rejected("not used".to_string())), + poll_calls: Arc::new(Mutex::new(Vec::new())), + exchange_calls: Arc::new(Mutex::new(Vec::new())), + }; + let (app, _state, _pool, _dir) = + admin_accounts_test_app_with_oauth_client("admin-auth-device-login.sqlite", 38, client) + .await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/auth/device-login") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_device_login") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["userCode"], "ABCD-EFGH"); + assert_eq!( + body["data"]["verificationUri"], + "https://auth.openai.com/activate" + ); + assert_eq!(body["data"]["deviceCode"], "device-secret"); + assert_eq!(body["data"]["expiresIn"], 900); + assert_eq!(body["data"]["interval"], 5); +} + +#[tokio::test] +async fn admin_auth_device_poll_should_return_pending_without_importing_account() { + let poll_calls = Arc::new(Mutex::new(Vec::new())); + let client = StaticOAuthClient { + device_response: Err(OAuthError::Rejected("not used".to_string())), + poll_response: Err(OAuthError::SlowDown), + exchange_response: Err(OAuthError::Rejected("not used".to_string())), + poll_calls: poll_calls.clone(), + exchange_calls: Arc::new(Mutex::new(Vec::new())), + }; + let (app, _state, pool, _dir) = + admin_accounts_test_app_with_oauth_client("admin-auth-device-pending.sqlite", 39, client) + .await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/auth/device-poll/device-secret") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_device_pending") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["pending"], true); + assert_eq!(body["data"]["success"], false); + assert_eq!(body["data"]["code"], "slow_down"); + assert_eq!(*poll_calls.lock().await, vec!["device-secret".to_string()]); + let account_count: (i64,) = sqlx::query_as("select count(*) from accounts") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(account_count.0, 0); +} + +#[tokio::test] +async fn admin_auth_device_poll_should_import_tokens_without_returning_secrets() { + let token = test_jwt( + Some("device-account"), + Some("device-user"), + Some("device@example.com"), + Some("plus"), + 3600, + ); + let poll_calls = Arc::new(Mutex::new(Vec::new())); + let client = StaticOAuthClient { + device_response: Err(OAuthError::Rejected("not used".to_string())), + poll_response: Ok(TokenPair { + access_token: token.clone(), + refresh_token: Some("device-refresh-secret".to_string()), + }), + exchange_response: Err(OAuthError::Rejected("not used".to_string())), + poll_calls: poll_calls.clone(), + exchange_calls: Arc::new(Mutex::new(Vec::new())), + }; + let (app, state, pool, _dir) = + admin_accounts_test_app_with_oauth_client("admin-auth-device-success.sqlite", 40, client) + .await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/auth/device-poll/device-success") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_device_success") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["success"], true); + assert_eq!(body["data"]["pending"], false); + assert!(body["data"].get("accessToken").is_none()); + assert!(body["data"].get("refreshToken").is_none()); + assert_eq!(*poll_calls.lock().await, vec!["device-success".to_string()]); + + let stored = AccountRepository::new(pool.clone(), SecretBox::new([40; 32])) + .find_by_chatgpt_identity("device-account", Some("device-user")) + .await + .unwrap() + .unwrap(); + assert_eq!(stored.email.as_deref(), Some("device@example.com")); + assert_eq!(stored.account_id.as_deref(), Some("device-account")); + assert!(!stored.access_token.expose_secret().is_empty()); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "device-refresh-secret" + ); + let raw: (String, String) = sqlx::query_as( + "select access_token_cipher, refresh_token_cipher from accounts where account_id = ?", + ) + .bind("device-account") + .fetch_one(&pool) + .await + .unwrap(); + assert!(!raw.0.contains(&token)); + assert!(!raw.1.contains("device-refresh-secret")); + let acquired = state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .unwrap(); + assert_eq!(acquired.email.as_deref(), Some("device@example.com")); +} + +#[tokio::test] +async fn admin_auth_login_start_should_require_admin_session_cookie() { + let (app, _state, _pool, _dir) = + admin_accounts_test_app("admin-auth-login-start-auth.sqlite", 41).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/auth/login-start") + .header("x-request-id", "req_login_start_no_session") + .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_login_start_no_session"); +} + +#[tokio::test] +async fn admin_auth_login_start_should_return_pkce_auth_url_and_state() { + let (app, _state, _pool, _dir) = + admin_accounts_test_app("admin-auth-login-start.sqlite", 42).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/auth/login-start") + .header("cookie", "cpr_admin_session=session_1") + .header("host", "127.0.0.1:8080") + .header("x-request-id", "req_login_start") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + let state = body["data"]["state"].as_str().unwrap(); + let auth_url = body["data"]["authUrl"].as_str().unwrap(); + assert_eq!(state.len(), 32); + assert!(auth_url.starts_with("https://auth.openai.com/oauth/authorize?")); + assert!(auth_url.contains("response_type=code")); + assert!(auth_url.contains("client_id=app_EMoamEEZ73f0CkXaXp7hrann")); + assert!(auth_url.contains("redirect_uri=http%3A%2F%2Flocalhost%3A1455%2Fauth%2Fcallback")); + assert!(auth_url.contains("scope=openid%20profile%20email%20offline_access")); + assert!(auth_url.contains("code_challenge_method=S256")); + assert!(auth_url.contains("originator=codex_cli_rs")); + assert!(auth_url.contains(&format!("state={state}"))); +} + +#[tokio::test] +async fn admin_auth_code_relay_should_exchange_code_and_import_account() { + let token = test_jwt( + Some("pkce-account"), + Some("pkce-user"), + Some("pkce@example.com"), + Some("plus"), + 3600, + ); + let exchange_calls = Arc::new(Mutex::new(Vec::new())); + let client = StaticOAuthClient { + device_response: Err(OAuthError::Rejected("not used".to_string())), + poll_response: Err(OAuthError::AuthorizationPending), + exchange_response: Ok(TokenPair { + access_token: token.clone(), + refresh_token: Some("pkce-refresh-secret".to_string()), + }), + poll_calls: Arc::new(Mutex::new(Vec::new())), + exchange_calls: exchange_calls.clone(), + }; + let (app, state, pool, _dir) = + admin_accounts_test_app_with_oauth_client("admin-auth-code-relay.sqlite", 43, client).await; + let start = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/auth/login-start") + .header("cookie", "cpr_admin_session=session_1") + .header("host", "127.0.0.1:8080") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let start_body = response_json(start).await; + let state_value = start_body["data"]["state"].as_str().unwrap(); + let callback_url = + format!("http://localhost:1455/auth/callback?code=oauth-code&state={state_value}"); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/auth/code-relay") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_code_relay") + .body(Body::from( + json!({ "callbackUrl": callback_url }).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["success"], true); + assert!(body["data"].get("accessToken").is_none()); + assert!(body["data"].get("refreshToken").is_none()); + let calls = exchange_calls.lock().await; + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].code, "oauth-code"); + assert!(!calls[0].code_verifier.is_empty()); + assert_eq!(calls[0].redirect_uri, "http://localhost:1455/auth/callback"); + drop(calls); + let stored = AccountRepository::new(pool.clone(), SecretBox::new([43; 32])) + .find_by_chatgpt_identity("pkce-account", Some("pkce-user")) + .await + .unwrap() + .unwrap(); + assert_eq!(stored.email.as_deref(), Some("pkce@example.com")); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "pkce-refresh-secret" + ); + let acquired = state + .services + .accounts + .acquire_runtime_account("gpt-5.5") + .await + .unwrap(); + assert_eq!(acquired.account_id.as_deref(), Some("pkce-account")); +} + +#[tokio::test] +async fn admin_auth_code_relay_should_reject_invalid_callback_url() { + let client = StaticOAuthClient { + device_response: Err(OAuthError::Rejected("not used".to_string())), + poll_response: Err(OAuthError::AuthorizationPending), + exchange_response: Err(OAuthError::Rejected("not used".to_string())), + poll_calls: Arc::new(Mutex::new(Vec::new())), + exchange_calls: Arc::new(Mutex::new(Vec::new())), + }; + let (app, _state, _pool, _dir) = admin_accounts_test_app_with_oauth_client( + "admin-auth-code-relay-invalid.sqlite", + 44, + client, + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/auth/code-relay") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_code_relay_invalid") + .body(Body::from( + json!({ "callbackUrl": "not a url" }).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body = response_json(response).await; + assert_eq!(body["code"], 40001); +} + +#[tokio::test] +async fn admin_auth_callback_should_exchange_code_and_redirect_to_return_host() { + let token = test_jwt( + Some("callback-account"), + Some("callback-user"), + Some("callback@example.com"), + Some("plus"), + 3600, + ); + let exchange_calls = Arc::new(Mutex::new(Vec::new())); + let client = StaticOAuthClient { + device_response: Err(OAuthError::Rejected("not used".to_string())), + poll_response: Err(OAuthError::AuthorizationPending), + exchange_response: Ok(TokenPair { + access_token: token, + refresh_token: Some("callback-refresh-secret".to_string()), + }), + poll_calls: Arc::new(Mutex::new(Vec::new())), + exchange_calls: exchange_calls.clone(), + }; + let (app, _state, pool, _dir) = + admin_accounts_test_app_with_oauth_client("admin-auth-callback.sqlite", 45, client).await; + let start = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/auth/login-start") + .header("cookie", "cpr_admin_session=session_1") + .header("host", "codex.local:1455") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let start_body = response_json(start).await; + let state_value = start_body["data"]["state"].as_str().unwrap(); + let callback_path = format!("/auth/callback?code=callback-code&state={state_value}"); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri(callback_path) + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::SEE_OTHER); + assert_eq!( + response.headers().get("location").unwrap(), + "http://codex.local:1455/" + ); + assert_eq!(exchange_calls.lock().await[0].code, "callback-code"); + let stored = AccountRepository::new(pool.clone(), SecretBox::new([45; 32])) + .find_by_chatgpt_identity("callback-account", Some("callback-user")) + .await + .unwrap() + .unwrap(); + assert_eq!(stored.email.as_deref(), Some("callback@example.com")); +} diff --git a/tests/admin_accounts_route_test.rs b/tests/admin_accounts_route_test.rs deleted file mode 100644 index d47a3a708..000000000 --- a/tests/admin_accounts_route_test.rs +++ /dev/null @@ -1,2588 +0,0 @@ -use std::{collections::BTreeMap, fs, sync::Arc}; - -use async_trait::async_trait; -use axum::{ - body::{to_bytes, Body}, - http::{Request, StatusCode}, - Router, -}; -use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; -use chrono::Utc; -use secrecy::ExposeSecret; -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, - app::build_router, - auth::{ - api_key::ApiKeyHasher, - oauth::{DeviceCode, OAuthClient, OAuthError}, - refresh::{RefreshFailure, TokenRefresher}, - token::TokenPair, - }, - config::{ - AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, - QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, - }, - crypto::SecretBox, - state::AppState, - storage::db::connect_sqlite, -}; - -#[derive(Clone)] -struct StaticTokenRefresher { - result: Result, - calls: Arc>>, -} - -#[derive(Clone)] -struct FailingTokenRefresher { - failure: RefreshFailure, -} - -#[derive(Clone)] -struct StaticOAuthClient { - device_response: Result, - poll_response: Result, - exchange_response: Result, - poll_calls: Arc>>, - exchange_calls: Arc>>, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct ExchangeCall { - code: String, - code_verifier: String, - redirect_uri: String, -} - -#[async_trait] -impl TokenRefresher for StaticTokenRefresher { - async fn refresh(&self, refresh_token: &str) -> Result { - 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 { - Err(self.failure) - } -} - -#[async_trait] -impl TokenRefresher for StaticOAuthClient { - async fn refresh(&self, _refresh_token: &str) -> Result { - Err(RefreshFailure::Transport) - } -} - -#[async_trait] -impl OAuthClient for StaticOAuthClient { - async fn exchange_code( - &self, - code: &str, - code_verifier: &str, - redirect_uri: &str, - ) -> Result { - self.exchange_calls.lock().await.push(ExchangeCall { - code: code.to_string(), - code_verifier: code_verifier.to_string(), - redirect_uri: redirect_uri.to_string(), - }); - self.exchange_response.clone() - } - - async fn request_device_code(&self) -> Result { - self.device_response.clone() - } - - async fn poll_device_token(&self, device_code: &str) -> Result { - self.poll_calls.lock().await.push(device_code.to_string()); - self.poll_response.clone() - } -} - -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, - }, - } -} - -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(); - let db = dir.path().join("admin-accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let app = build_router(AppState::with_pool_and_secret_box( - test_config(url), - pool, - SecretBox::new([11u8; 32]), - )); - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/import") - .header("content-type", "application/json") - .header("x-request-id", "req_accounts") - .body(Body::from(r#"{"accounts":[]}"#)) - .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_accounts"); -} - -#[tokio::test] -async fn admin_auth_status_should_require_admin_session_cookie() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-auth-status-auth.sqlite", 33).await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/auth/status") - .header("x-request-id", "req_auth_status_no_session") - .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_auth_status_no_session"); -} - -#[tokio::test] -async fn admin_auth_status_should_report_empty_account_pool() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-auth-status-empty.sqlite", 34).await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/auth/status") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_auth_status_empty") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["authenticated"], false); - assert!(body["data"]["user"].is_null()); - assert_eq!(body["data"]["pool"]["total"], 0); - assert_eq!(body["data"]["pool"]["active"], 0); -} - -#[tokio::test] -async fn admin_auth_status_should_report_user_and_pool_summary_without_secrets() { - let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-auth-status-summary.sqlite", 35).await; - import_test_account(&app, "session_1", "acct_auth_status").await; - import_test_account(&app, "session_1", "acct_auth_disabled").await; - sqlx::query("update accounts set status = 'disabled' where id = ?") - .bind("acct_auth_disabled") - .execute(&pool) - .await - .unwrap(); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/auth/status") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_auth_status_summary") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["authenticated"], true); - assert_eq!(body["data"]["pool"]["total"], 2); - assert_eq!(body["data"]["pool"]["active"], 1); - assert_eq!(body["data"]["pool"]["disabled"], 1); - assert!(body["data"]["user"].get("accessToken").is_none()); - assert!(body["data"]["user"].get("token").is_none()); - assert!(body["data"]["user"].get("refreshToken").is_none()); -} - -#[tokio::test] -async fn admin_auth_logout_should_clear_accounts_and_runtime_pool() { - let (app, state, pool, _dir) = admin_accounts_test_app("admin-auth-logout.sqlite", 36).await; - import_test_account(&app, "session_1", "acct_auth_logout").await; - sqlx::query( - "insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens) values (?, 2, 3, 4, 1)", - ) - .bind("acct_auth_logout") - .execute(&pool) - .await - .unwrap(); - let set_cookies = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/acct_auth_logout/cookies") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"cookies":"cf_clearance=secret"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(set_cookies.status(), StatusCode::OK); - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/auth/logout") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_auth_logout") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["success"], true); - assert_eq!(body["data"]["deleted"], 1); - let account_count: (i64,) = sqlx::query_as("select count(*) from accounts") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(account_count.0, 0); - let usage_count: (i64,) = sqlx::query_as("select count(*) from account_usage") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(usage_count.0, 0); - let cookie_count: (i64,) = sqlx::query_as("select count(*) from account_cookies") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(cookie_count.0, 0); - assert!(state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .is_none()); -} - -#[tokio::test] -async fn admin_auth_device_login_should_require_admin_session_cookie() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-auth-device-login-auth.sqlite", 37).await; - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/auth/device-login") - .header("x-request-id", "req_device_login_no_session") - .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_device_login_no_session"); -} - -#[tokio::test] -async fn admin_auth_device_login_should_return_openai_device_code() { - let client = StaticOAuthClient { - device_response: Ok(DeviceCode { - device_code: "device-secret".to_string(), - user_code: "ABCD-EFGH".to_string(), - verification_uri: "https://auth.openai.com/activate".to_string(), - verification_uri_complete: "https://auth.openai.com/activate?user_code=ABCD-EFGH" - .to_string(), - expires_in: 900, - interval: 5, - }), - poll_response: Err(OAuthError::AuthorizationPending), - exchange_response: Err(OAuthError::Rejected("not used".to_string())), - poll_calls: Arc::new(Mutex::new(Vec::new())), - exchange_calls: Arc::new(Mutex::new(Vec::new())), - }; - let (app, _state, _pool, _dir) = - admin_accounts_test_app_with_oauth_client("admin-auth-device-login.sqlite", 38, client) - .await; - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/auth/device-login") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_device_login") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["userCode"], "ABCD-EFGH"); - assert_eq!( - body["data"]["verificationUri"], - "https://auth.openai.com/activate" - ); - assert_eq!(body["data"]["deviceCode"], "device-secret"); - assert_eq!(body["data"]["expiresIn"], 900); - assert_eq!(body["data"]["interval"], 5); -} - -#[tokio::test] -async fn admin_auth_device_poll_should_return_pending_without_importing_account() { - let poll_calls = Arc::new(Mutex::new(Vec::new())); - let client = StaticOAuthClient { - device_response: Err(OAuthError::Rejected("not used".to_string())), - poll_response: Err(OAuthError::SlowDown), - exchange_response: Err(OAuthError::Rejected("not used".to_string())), - poll_calls: poll_calls.clone(), - exchange_calls: Arc::new(Mutex::new(Vec::new())), - }; - let (app, _state, pool, _dir) = - admin_accounts_test_app_with_oauth_client("admin-auth-device-pending.sqlite", 39, client) - .await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/auth/device-poll/device-secret") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_device_pending") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["pending"], true); - assert_eq!(body["data"]["success"], false); - assert_eq!(body["data"]["code"], "slow_down"); - assert_eq!(*poll_calls.lock().await, vec!["device-secret".to_string()]); - let account_count: (i64,) = sqlx::query_as("select count(*) from accounts") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(account_count.0, 0); -} - -#[tokio::test] -async fn admin_auth_device_poll_should_import_tokens_without_returning_secrets() { - let token = test_jwt( - Some("device-account"), - Some("device-user"), - Some("device@example.com"), - Some("plus"), - 3600, - ); - let poll_calls = Arc::new(Mutex::new(Vec::new())); - let client = StaticOAuthClient { - device_response: Err(OAuthError::Rejected("not used".to_string())), - poll_response: Ok(TokenPair { - access_token: token.clone(), - refresh_token: Some("device-refresh-secret".to_string()), - }), - exchange_response: Err(OAuthError::Rejected("not used".to_string())), - poll_calls: poll_calls.clone(), - exchange_calls: Arc::new(Mutex::new(Vec::new())), - }; - let (app, state, pool, _dir) = - admin_accounts_test_app_with_oauth_client("admin-auth-device-success.sqlite", 40, client) - .await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/auth/device-poll/device-success") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_device_success") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["success"], true); - assert_eq!(body["data"]["pending"], false); - assert!(body["data"].get("accessToken").is_none()); - assert!(body["data"].get("refreshToken").is_none()); - assert_eq!(*poll_calls.lock().await, vec!["device-success".to_string()]); - - let stored = AccountRepository::new(pool.clone(), SecretBox::new([40; 32])) - .find_by_chatgpt_identity("device-account", Some("device-user")) - .await - .unwrap() - .unwrap(); - assert_eq!(stored.email.as_deref(), Some("device@example.com")); - assert_eq!(stored.account_id.as_deref(), Some("device-account")); - assert!(!stored.access_token.expose_secret().is_empty()); - assert_eq!( - stored.refresh_token.unwrap().expose_secret(), - "device-refresh-secret" - ); - let raw: (String, String) = sqlx::query_as( - "select access_token_cipher, refresh_token_cipher from accounts where account_id = ?", - ) - .bind("device-account") - .fetch_one(&pool) - .await - .unwrap(); - assert!(!raw.0.contains(&token)); - assert!(!raw.1.contains("device-refresh-secret")); - let acquired = state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .unwrap(); - assert_eq!(acquired.email.as_deref(), Some("device@example.com")); -} - -#[tokio::test] -async fn admin_auth_login_start_should_require_admin_session_cookie() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-auth-login-start-auth.sqlite", 41).await; - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/auth/login-start") - .header("x-request-id", "req_login_start_no_session") - .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_login_start_no_session"); -} - -#[tokio::test] -async fn admin_auth_login_start_should_return_pkce_auth_url_and_state() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-auth-login-start.sqlite", 42).await; - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/auth/login-start") - .header("cookie", "cpr_admin_session=session_1") - .header("host", "127.0.0.1:8080") - .header("x-request-id", "req_login_start") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - let state = body["data"]["state"].as_str().unwrap(); - let auth_url = body["data"]["authUrl"].as_str().unwrap(); - assert_eq!(state.len(), 32); - assert!(auth_url.starts_with("https://auth.openai.com/oauth/authorize?")); - assert!(auth_url.contains("response_type=code")); - assert!(auth_url.contains("client_id=app_EMoamEEZ73f0CkXaXp7hrann")); - assert!(auth_url.contains("redirect_uri=http%3A%2F%2Flocalhost%3A1455%2Fauth%2Fcallback")); - assert!(auth_url.contains("scope=openid%20profile%20email%20offline_access")); - assert!(auth_url.contains("code_challenge_method=S256")); - assert!(auth_url.contains("originator=codex_cli_rs")); - assert!(auth_url.contains(&format!("state={state}"))); -} - -#[tokio::test] -async fn admin_auth_code_relay_should_exchange_code_and_import_account() { - let token = test_jwt( - Some("pkce-account"), - Some("pkce-user"), - Some("pkce@example.com"), - Some("plus"), - 3600, - ); - let exchange_calls = Arc::new(Mutex::new(Vec::new())); - let client = StaticOAuthClient { - device_response: Err(OAuthError::Rejected("not used".to_string())), - poll_response: Err(OAuthError::AuthorizationPending), - exchange_response: Ok(TokenPair { - access_token: token.clone(), - refresh_token: Some("pkce-refresh-secret".to_string()), - }), - poll_calls: Arc::new(Mutex::new(Vec::new())), - exchange_calls: exchange_calls.clone(), - }; - let (app, state, pool, _dir) = - admin_accounts_test_app_with_oauth_client("admin-auth-code-relay.sqlite", 43, client).await; - let start = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/auth/login-start") - .header("cookie", "cpr_admin_session=session_1") - .header("host", "127.0.0.1:8080") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let start_body = response_json(start).await; - let state_value = start_body["data"]["state"].as_str().unwrap(); - let callback_url = - format!("http://localhost:1455/auth/callback?code=oauth-code&state={state_value}"); - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/auth/code-relay") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_code_relay") - .body(Body::from( - json!({ "callbackUrl": callback_url }).to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["success"], true); - assert!(body["data"].get("accessToken").is_none()); - assert!(body["data"].get("refreshToken").is_none()); - let calls = exchange_calls.lock().await; - assert_eq!(calls.len(), 1); - assert_eq!(calls[0].code, "oauth-code"); - assert!(!calls[0].code_verifier.is_empty()); - assert_eq!(calls[0].redirect_uri, "http://localhost:1455/auth/callback"); - drop(calls); - let stored = AccountRepository::new(pool.clone(), SecretBox::new([43; 32])) - .find_by_chatgpt_identity("pkce-account", Some("pkce-user")) - .await - .unwrap() - .unwrap(); - assert_eq!(stored.email.as_deref(), Some("pkce@example.com")); - assert_eq!( - stored.refresh_token.unwrap().expose_secret(), - "pkce-refresh-secret" - ); - let acquired = state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .unwrap(); - assert_eq!(acquired.account_id.as_deref(), Some("pkce-account")); -} - -#[tokio::test] -async fn admin_auth_code_relay_should_reject_invalid_callback_url() { - let client = StaticOAuthClient { - device_response: Err(OAuthError::Rejected("not used".to_string())), - poll_response: Err(OAuthError::AuthorizationPending), - exchange_response: Err(OAuthError::Rejected("not used".to_string())), - poll_calls: Arc::new(Mutex::new(Vec::new())), - exchange_calls: Arc::new(Mutex::new(Vec::new())), - }; - let (app, _state, _pool, _dir) = admin_accounts_test_app_with_oauth_client( - "admin-auth-code-relay-invalid.sqlite", - 44, - client, - ) - .await; - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/auth/code-relay") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_code_relay_invalid") - .body(Body::from( - json!({ "callbackUrl": "not a url" }).to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_REQUEST); - let body = response_json(response).await; - assert_eq!(body["code"], 40001); -} - -#[tokio::test] -async fn admin_auth_callback_should_exchange_code_and_redirect_to_return_host() { - let token = test_jwt( - Some("callback-account"), - Some("callback-user"), - Some("callback@example.com"), - Some("plus"), - 3600, - ); - let exchange_calls = Arc::new(Mutex::new(Vec::new())); - let client = StaticOAuthClient { - device_response: Err(OAuthError::Rejected("not used".to_string())), - poll_response: Err(OAuthError::AuthorizationPending), - exchange_response: Ok(TokenPair { - access_token: token, - refresh_token: Some("callback-refresh-secret".to_string()), - }), - poll_calls: Arc::new(Mutex::new(Vec::new())), - exchange_calls: exchange_calls.clone(), - }; - let (app, _state, pool, _dir) = - admin_accounts_test_app_with_oauth_client("admin-auth-callback.sqlite", 45, client).await; - let start = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/auth/login-start") - .header("cookie", "cpr_admin_session=session_1") - .header("host", "codex.local:1455") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let start_body = response_json(start).await; - let state_value = start_body["data"]["state"].as_str().unwrap(); - let callback_path = format!("/auth/callback?code=callback-code&state={state_value}"); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri(callback_path) - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::SEE_OTHER); - assert_eq!( - response.headers().get("location").unwrap(), - "http://codex.local:1455/" - ); - assert_eq!(exchange_calls.lock().await[0].code, "callback-code"); - let stored = AccountRepository::new(pool.clone(), SecretBox::new([45; 32])) - .find_by_chatgpt_identity("callback-account", Some("callback-user")) - .await - .unwrap() - .unwrap(); - assert_eq!(stored.email.as_deref(), Some("callback@example.com")); -} - -#[tokio::test] -async fn admin_accounts_import_should_store_tokens_encrypted_and_list_sanitized_accounts() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let app = build_router(AppState::with_pool_and_secret_box( - test_config(url), - pool.clone(), - SecretBox::new([12u8; 32]), - )); - let import_body = json!({ - "accounts": [{ - "id": "acct_imported", - "email": "user@example.com", - "accountId": "chatgpt-account", - "userId": "chatgpt-user", - "label": "primary", - "planType": "plus", - "token": "access-secret", - "refreshToken": "refresh-secret", - "status": "active" - }] - }); - - 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") - .header("x-request-id", "req_accounts") - .body(Body::from(import_body.to_string())) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(import_response.status(), StatusCode::OK); - let body = response_json(import_response).await; - assert_eq!(body["data"]["imported"], 1); - assert_eq!(body["data"]["skipped"], 0); - let stored: (String, String) = sqlx::query_as( - "select access_token_cipher, refresh_token_cipher from accounts where id = ?", - ) - .bind("acct_imported") - .fetch_one(&pool) - .await - .unwrap(); - assert!(stored.0.starts_with("v1:")); - assert!(!stored.0.contains("access-secret")); - assert!(stored.1.starts_with("v1:")); - assert!(!stored.1.contains("refresh-secret")); - - let list_response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/accounts?limit=10") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_accounts_list") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(list_response.status(), StatusCode::OK); - let body = response_json(list_response).await; - assert_eq!(body["data"][0]["id"], "acct_imported"); - assert_eq!(body["data"][0]["email"], "user@example.com"); - assert!(body["data"][0].get("token").is_none()); - assert!(body["data"][0].get("refreshToken").is_none()); - assert_eq!(body["page"]["limit"], 10); -} - -#[tokio::test] -async fn admin_accounts_import_should_accept_sub2api_oauth_export_and_mark_format() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-sub2api.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let app = build_router(AppState::with_pool_and_secret_box( - test_config(url), - pool.clone(), - SecretBox::new([14u8; 32]), - )); - let import_body = json!({ - "type": "sub2api-data", - "version": 1, - "proxies": [], - "accounts": [ - { - "name": "Sub2API Team", - "platform": "openai", - "type": "oauth", - "credentials": { - "access_token": "Bearer sub2api-access-secret", - "refresh_token": "rt_sub2api", - "email": "team@example.com", - "chatgpt_account_id": "chatgpt-account", - "chatgpt_user_id": "chatgpt-user", - "plan_type": "team" - }, - "concurrency": 0, - "priority": 0 - }, - { - "name": "Other Provider", - "platform": "anthropic", - "type": "oauth", - "credentials": { - "access_token": "ignored-secret" - } - } - ] - }); - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/import") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_sub2api") - .body(Body::from(import_body.to_string())) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["sourceFormat"], "sub2api"); - assert_eq!(body["data"]["imported"], 1); - assert_eq!(body["data"]["skipped"], 0); - let stored: (String, String, String, String, String) = - sqlx::query_as("select email, account_id, user_id, label, plan_type from accounts") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(stored.0, "team@example.com"); - assert_eq!(stored.1, "chatgpt-account"); - assert_eq!(stored.2, "chatgpt-user"); - assert_eq!(stored.3, "Sub2API Team"); - assert_eq!(stored.4, "team"); -} - -#[tokio::test] -async fn admin_accounts_import_should_accept_sub2api_native_account_export_without_proxy_data() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-sub2api-native.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let app = build_router(AppState::with_pool_and_secret_box( - test_config(url), - pool, - SecretBox::new([15u8; 32]), - )); - let import_body = json!({ - "accounts": [{ - "id": "acct_sub2api_native", - "token": "native-access-secret", - "refreshToken": "native-refresh-secret", - "email": "native@example.com", - "accountId": "native-account", - "userId": "native-user", - "label": "Native Sub2API", - "planType": "plus", - "proxyApiKey": "ignored-proxy-secret", - "status": "active", - "usage": { - "request_count": 1, - "input_tokens": 0, - "output_tokens": 0, - "cached_tokens": 0 - }, - "cachedQuota": null, - "quotaVerifyRequired": false - }] - }); - - 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") - .header("x-request-id", "req_sub2api_native") - .body(Body::from(import_body.to_string())) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(import_response.status(), StatusCode::OK); - let body = response_json(import_response).await; - assert_eq!(body["data"]["sourceFormat"], "sub2api"); - assert_eq!(body["data"]["imported"], 1); - - let list_response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/accounts?limit=10") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_sub2api_native_list") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let list = response_json(list_response).await; - assert_eq!(list["data"][0]["id"], "acct_sub2api_native"); - assert_eq!(list["data"][0]["planType"], "plus"); - assert!(list["data"][0].get("proxyApiKey").is_none()); - assert!(list["data"][0].get("usage").is_none()); -} - -#[tokio::test] -async fn admin_accounts_list_should_not_decrypt_account_tokens() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let now = Utc::now().to_rfc3339(); - sqlx::query( - "insert into accounts (id, email, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)", - ) - .bind("acct_corrupt") - .bind("user@example.com") - .bind("not-a-secret-box-cipher") - .bind("active") - .bind(&now) - .bind(&now) - .execute(&pool) - .await - .unwrap(); - let app = build_router(AppState::with_pool_and_secret_box( - test_config(url), - pool, - SecretBox::new([13u8; 32]), - )); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/accounts?limit=10") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_accounts_list") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"][0]["id"], "acct_corrupt"); -} - -#[tokio::test] -async fn admin_account_label_should_update_and_clear_label() { - let (app, state, pool, _dir) = admin_accounts_test_app("admin-account-label.sqlite", 16).await; - import_test_account(&app, "session_1", "acct_label").await; - - let set_response = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri("/admin/accounts/acct_label/label") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_label") - .body(Body::from(r#"{"label":"Team Alpha"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(set_response.status(), StatusCode::OK); - let body = response_json(set_response).await; - assert_eq!(body["data"]["label"], "Team Alpha"); - - let clear_response = app - .oneshot( - Request::builder() - .method("PATCH") - .uri("/admin/accounts/acct_label/label") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"label":null}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(clear_response.status(), StatusCode::OK); - let body = response_json(clear_response).await; - assert!(body["data"]["label"].is_null()); - - let stored: (Option,) = sqlx::query_as("select label from accounts where id = ?") - .bind("acct_label") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(stored.0, None); - assert!(state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .is_some()); -} - -#[tokio::test] -async fn admin_account_label_should_reject_too_long_or_missing_account() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-account-label-invalid.sqlite", 17).await; - import_test_account(&app, "session_1", "acct_label_invalid").await; - let long_label = "x".repeat(65); - - let too_long = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri("/admin/accounts/acct_label_invalid/label") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(json!({ "label": long_label }).to_string())) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(too_long.status(), StatusCode::BAD_REQUEST); - - let missing = app - .oneshot( - Request::builder() - .method("PATCH") - .uri("/admin/accounts/missing/label") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"label":"Team Alpha"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(missing.status(), StatusCode::NOT_FOUND); -} - -#[tokio::test] -async fn admin_account_status_should_update_database_and_runtime_pool() { - let (app, state, pool, _dir) = admin_accounts_test_app("admin-account-status.sqlite", 18).await; - import_test_account(&app, "session_1", "acct_status").await; - - let response = app - .oneshot( - Request::builder() - .method("PATCH") - .uri("/admin/accounts/acct_status/status") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_status") - .body(Body::from(r#"{"status":"disabled"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["status"], "disabled"); - - let stored: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_status") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(stored.0, "disabled"); - assert!(state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .is_none()); -} - -#[tokio::test] -async fn admin_account_delete_should_remove_database_row_and_runtime_pool_entry() { - let (app, state, pool, _dir) = admin_accounts_test_app("admin-account-delete.sqlite", 19).await; - import_test_account(&app, "session_1", "acct_delete").await; - - let response = app - .clone() - .oneshot( - Request::builder() - .method("DELETE") - .uri("/admin/accounts/acct_delete") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_delete") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["deleted"], true); - - let row_count: (i64,) = sqlx::query_as("select count(*) from accounts where id = ?") - .bind("acct_delete") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(row_count.0, 0); - assert!(state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .is_none()); - - let missing = app - .oneshot( - Request::builder() - .method("DELETE") - .uri("/admin/accounts/acct_delete") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(missing.status(), StatusCode::NOT_FOUND); -} - -#[tokio::test] -async fn admin_accounts_batch_delete_should_delete_found_accounts_and_report_missing_ids() { - let (app, state, pool, _dir) = - admin_accounts_test_app("admin-accounts-batch-delete.sqlite", 20).await; - import_test_account(&app, "session_1", "acct_batch_delete_a").await; - import_test_account(&app, "session_1", "acct_batch_delete_b").await; - - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/batch-delete") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - json!({ - "ids": ["acct_batch_delete_a", "ghost", "acct_batch_delete_b"] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["deleted"], 2); - assert_eq!(body["data"]["notFound"], json!(["ghost"])); - - let row_count: (i64,) = sqlx::query_as("select count(*) from accounts") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(row_count.0, 0); - assert!(state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .is_none()); - - let empty = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/batch-delete") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"ids":[]}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(empty.status(), StatusCode::BAD_REQUEST); -} - -#[tokio::test] -async fn admin_accounts_batch_status_should_update_found_accounts_and_reject_invalid_status() { - let (app, state, pool, _dir) = - admin_accounts_test_app("admin-accounts-batch-status.sqlite", 21).await; - import_test_account(&app, "session_1", "acct_batch_status_a").await; - import_test_account(&app, "session_1", "acct_batch_status_b").await; - - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/batch-status") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - json!({ - "ids": ["acct_batch_status_a", "ghost"], - "status": "disabled" - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["updated"], 1); - assert_eq!(body["data"]["notFound"], json!(["ghost"])); - - let statuses = - sqlx::query_as::<_, (String, String)>("select id, status from accounts order by id asc") - .fetch_all(&pool) - .await - .unwrap(); - assert_eq!( - statuses, - vec![ - ("acct_batch_status_a".to_string(), "disabled".to_string()), - ("acct_batch_status_b".to_string(), "active".to_string()) - ] - ); - let acquired = state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .unwrap(); - assert_eq!(acquired.id, "acct_batch_status_b"); - - let invalid = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/batch-status") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - json!({ - "ids": ["acct_batch_status_a"], - "status": "expired" - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(invalid.status(), StatusCode::BAD_REQUEST); - - let empty = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/batch-status") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"ids":[],"status":"active"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(empty.status(), StatusCode::BAD_REQUEST); -} - -#[tokio::test] -async fn admin_account_cookies_should_set_get_and_delete_encrypted_cookie_header() { - let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-account-cookies.sqlite", 22).await; - import_test_account(&app, "session_1", "acct_cookies").await; - - let set_response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/acct_cookies/cookies") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - r#"{"cookies":"cf_clearance=clear-secret; __cf_bm=bm-secret"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(set_response.status(), StatusCode::OK); - let body = response_json(set_response).await; - assert_eq!( - body["data"]["cookies"], - "__cf_bm=bm-secret; cf_clearance=clear-secret" - ); - - let stored = sqlx::query_as::<_, (String, String)>( - "select name, value_cipher from account_cookies where account_id = ? order by name asc", - ) - .bind("acct_cookies") - .fetch_all(&pool) - .await - .unwrap(); - assert_eq!(stored.len(), 2); - assert!(stored.iter().all(|(_, cipher)| cipher.starts_with("v1:"))); - assert!(stored - .iter() - .all(|(_, cipher)| !cipher.contains("clear-secret") && !cipher.contains("bm-secret"))); - - let get_response = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/accounts/acct_cookies/cookies") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(get_response.status(), StatusCode::OK); - let body = response_json(get_response).await; - assert_eq!( - body["data"]["cookies"], - "__cf_bm=bm-secret; cf_clearance=clear-secret" - ); - - let delete_response = app - .clone() - .oneshot( - Request::builder() - .method("DELETE") - .uri("/admin/accounts/acct_cookies/cookies") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(delete_response.status(), StatusCode::OK); - - let get_empty = app - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/accounts/acct_cookies/cookies") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(get_empty.status(), StatusCode::OK); - let body = response_json(get_empty).await; - assert!(body["data"]["cookies"].is_null()); -} - -#[tokio::test] -async fn admin_account_cookies_should_require_existing_account_and_non_empty_cookies() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-account-cookies-invalid.sqlite", 23).await; - - let missing = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/accounts/missing/cookies") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(missing.status(), StatusCode::NOT_FOUND); - - import_test_account(&app, "session_1", "acct_cookie_invalid").await; - let invalid = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/acct_cookie_invalid/cookies") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"cookies":""}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(invalid.status(), StatusCode::BAD_REQUEST); -} - -#[tokio::test] -async fn admin_accounts_export_should_return_native_accounts_with_tokens_and_filter_ids() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-account-export-native.sqlite", 24).await; - import_test_account(&app, "session_1", "acct_export_a").await; - import_test_account(&app, "session_1", "acct_export_b").await; - - let response = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/accounts/export?ids=acct_export_a") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["sourceFormat"], "native"); - assert_eq!(body["data"]["accounts"].as_array().unwrap().len(), 1); - assert_eq!(body["data"]["accounts"][0]["id"], "acct_export_a"); - assert_eq!(body["data"]["accounts"][0]["token"], "access-acct_export_a"); - assert_eq!( - body["data"]["accounts"][0]["refreshToken"], - "refresh-acct_export_a" - ); - assert!(body["data"]["accounts"][0].get("proxyApiKey").is_none()); - - let invalid = app - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/accounts/export?format=proxy") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(invalid.status(), StatusCode::BAD_REQUEST); -} - -#[tokio::test] -async fn admin_accounts_export_should_return_sub2api_openai_oauth_payload_without_proxy_fields() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-account-export-sub2api.sqlite", 25).await; - import_test_account(&app, "session_1", "acct_export_sub2api").await; - - let label_response = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri("/admin/accounts/acct_export_sub2api/label") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"label":"Sub2API Export"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(label_response.status(), StatusCode::OK); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/admin/accounts/export?format=sub2api") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["type"], "sub2api-data"); - assert_eq!(body["data"]["version"], 1); - assert_eq!(body["data"]["proxies"], json!([])); - assert_eq!(body["data"]["accounts"].as_array().unwrap().len(), 1); - assert_eq!(body["data"]["accounts"][0]["name"], "Sub2API Export"); - assert_eq!(body["data"]["accounts"][0]["platform"], "openai"); - assert_eq!(body["data"]["accounts"][0]["type"], "oauth"); - assert_eq!( - body["data"]["accounts"][0]["credentials"]["access_token"], - "access-acct_export_sub2api" - ); - assert_eq!( - body["data"]["accounts"][0]["credentials"]["refresh_token"], - "refresh-acct_export_sub2api" - ); - assert_eq!( - body["data"]["accounts"][0]["credentials"]["email"], - "acct_export_sub2api@example.com" - ); - assert_eq!( - body["data"]["accounts"][0]["credentials"]["plan_type"], - "plus" - ); - assert!(body["data"]["accounts"][0].get("proxy").is_none()); - assert!(body["data"]["accounts"][0].get("proxyUrl").is_none()); -} - -#[tokio::test] -async fn manual_add_should_derive_claims_ignore_metadata_and_sync_pool() { - let (app, state, pool, _dir) = - admin_accounts_test_app("admin-account-manual-add.sqlite", 26).await; - let token = test_jwt( - Some("jwt-account"), - Some("jwt-user"), - Some("jwt@example.com"), - Some("team"), - 3600, - ); - - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - json!({ - "id": "caller-id", - "email": "caller@example.com", - "accountId": "caller-account", - "userId": "caller-user", - "label": "Caller Label", - "planType": "caller-plan", - "token": token, - "refreshToken": "manual-refresh-secret", - "status": "disabled" - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - let response_id = body["data"]["id"].as_str().unwrap(); - assert_ne!(response_id, "caller-id"); - assert_eq!(body["data"]["email"], "jwt@example.com"); - assert_eq!(body["data"]["accountId"], "jwt-account"); - assert_eq!(body["data"]["userId"], "jwt-user"); - assert_eq!(body["data"]["planType"], "team"); - assert!(body["data"]["label"].is_null()); - assert_eq!(body["data"]["status"], "active"); - assert!(body["data"].get("token").is_none()); - assert!(body["data"].get("refreshToken").is_none()); - - let stored = sqlx::query( - "select id, email, account_id, user_id, label, plan_type, status, access_token_cipher, refresh_token_cipher from accounts", - ) - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(stored.get::("id"), response_id); - assert_eq!( - stored.get::, _>("email").as_deref(), - Some("jwt@example.com") - ); - assert_eq!( - stored.get::, _>("account_id").as_deref(), - Some("jwt-account") - ); - assert_eq!( - stored.get::, _>("user_id").as_deref(), - Some("jwt-user") - ); - assert_eq!(stored.get::, _>("label"), None); - assert_eq!( - stored.get::, _>("plan_type").as_deref(), - Some("team") - ); - assert_eq!(stored.get::("status"), "active"); - let access_token_cipher = stored.get::("access_token_cipher"); - assert!(access_token_cipher.starts_with("v1:")); - assert!(!access_token_cipher.contains(&token)); - let refresh_token_cipher = stored.get::("refresh_token_cipher"); - assert!(refresh_token_cipher.starts_with("v1:")); - assert!(!refresh_token_cipher.contains("manual-refresh-secret")); - - let acquired = state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .unwrap(); - assert_eq!(acquired.id, response_id); - assert_eq!(acquired.email.as_deref(), Some("jwt@example.com")); - assert_eq!(acquired.account_id.as_deref(), Some("jwt-account")); - assert_eq!(acquired.user_id.as_deref(), Some("jwt-user")); - assert_eq!(acquired.plan_type.as_deref(), Some("team")); -} - -#[tokio::test] -async fn admin_account_manual_add_should_reject_missing_invalid_expired_or_unbound_tokens() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-account-manual-add-invalid.sqlite", 27).await; - - let missing_tokens = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(missing_tokens.status(), StatusCode::BAD_REQUEST); - - let invalid = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - json!({ - "token": "not-a-jwt" - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(invalid.status(), StatusCode::BAD_REQUEST); - - let expired = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - json!({ - "token": test_jwt( - Some("expired-account"), - Some("expired-user"), - Some("expired@example.com"), - Some("plus"), - -3600, - ) - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(expired.status(), StatusCode::BAD_REQUEST); - - let missing_account_claim = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - json!({ - "token": test_jwt( - None, - Some("claimless-user"), - Some("claimless@example.com"), - Some("free"), - 3600, - ) - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(missing_account_claim.status(), StatusCode::BAD_REQUEST); -} - -#[tokio::test] -async fn admin_account_manual_add_should_exchange_refresh_token_only_and_rotate_refresh_token() { - let token = test_jwt( - Some("rt-account"), - Some("rt-user"), - Some("rt@example.com"), - Some("plus"), - 3600, - ); - let calls = Arc::new(Mutex::new(Vec::new())); - let refresher = StaticTokenRefresher { - result: Ok(TokenPair { - access_token: token.clone(), - refresh_token: Some("rotated-refresh".to_string()), - }), - calls: calls.clone(), - }; - let (app, state, pool, _dir) = admin_accounts_test_app_with_refresher( - "admin-account-manual-refresh-only.sqlite", - 28, - refresher, - ) - .await; - - let response = post_admin_account( - &app, - json!({ - "refreshToken": "initial-refresh", - "email": "caller@example.com", - "planType": "caller-plan" - }), - ) - .await; - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["email"], "rt@example.com"); - assert_eq!(body["data"]["accountId"], "rt-account"); - assert!(body["data"].get("token").is_none()); - assert!(body["data"].get("refreshToken").is_none()); - assert_eq!(*calls.lock().await, vec!["initial-refresh".to_string()]); - - let repo = AccountRepository::new(pool, SecretBox::new([28u8; 32])); - let stored = repo - .get(body["data"]["id"].as_str().unwrap()) - .await - .unwrap() - .unwrap(); - assert_eq!(stored.access_token.expose_secret(), &token); - assert_eq!( - stored.refresh_token.unwrap().expose_secret(), - "rotated-refresh" - ); - let acquired = state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .unwrap(); - assert_eq!(acquired.refresh_token.as_deref(), Some("rotated-refresh")); -} - -#[tokio::test] -async fn admin_account_manual_add_should_preserve_input_refresh_token_when_exchange_omits_rotation() -{ - let token = test_jwt( - Some("rt-preserve-account"), - Some("rt-preserve-user"), - Some("preserve@example.com"), - Some("free"), - 3600, - ); - let refresher = StaticTokenRefresher { - result: Ok(TokenPair { - access_token: token, - refresh_token: None, - }), - calls: Arc::new(Mutex::new(Vec::new())), - }; - let (app, _state, pool, _dir) = admin_accounts_test_app_with_refresher( - "admin-account-manual-refresh-preserve.sqlite", - 29, - refresher, - ) - .await; - - let response = post_admin_account(&app, json!({ "refreshToken": "preserved-refresh" })).await; - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - let repo = AccountRepository::new(pool, SecretBox::new([29u8; 32])); - let stored = repo - .get(body["data"]["id"].as_str().unwrap()) - .await - .unwrap() - .unwrap(); - assert_eq!( - stored.refresh_token.unwrap().expose_secret(), - "preserved-refresh" - ); -} - -#[tokio::test] -async fn admin_account_manual_add_should_update_existing_account_for_same_account_and_user() { - let (app, state, pool, _dir) = - admin_accounts_test_app("admin-account-manual-update-existing.sqlite", 30).await; - let first_token = test_jwt( - Some("team-account"), - Some("same-user"), - Some("first@example.com"), - Some("free"), - 3600, - ); - let first_response = post_admin_account( - &app, - json!({ - "token": first_token, - "refreshToken": "first-refresh" - }), - ) - .await; - assert_eq!(first_response.status(), StatusCode::OK); - let first_body = response_json(first_response).await; - let first_id = first_body["data"]["id"].as_str().unwrap().to_string(); - - let second_token = test_jwt( - Some("team-account"), - Some("same-user"), - Some("second@example.com"), - Some("team"), - 7200, - ); - let second_response = post_admin_account(&app, json!({ "token": second_token })).await; - - assert_eq!(second_response.status(), StatusCode::OK); - let second_body = response_json(second_response).await; - assert_eq!(second_body["data"]["id"], first_id); - assert_eq!(second_body["data"]["email"], "second@example.com"); - assert_eq!(second_body["data"]["planType"], "team"); - - let count: (i64,) = sqlx::query_as("select count(*) from accounts") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(count.0, 1); - let repo = AccountRepository::new(pool, SecretBox::new([30u8; 32])); - let stored = repo.get(&first_id).await.unwrap().unwrap(); - assert_eq!(stored.access_token.expose_secret(), &second_token); - assert_eq!( - stored.refresh_token.unwrap().expose_secret(), - "first-refresh" - ); - assert_eq!(stored.email.as_deref(), Some("second@example.com")); - assert_eq!(stored.plan_type.as_deref(), Some("team")); - - let acquired = state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .unwrap(); - assert_eq!(acquired.id, first_id); - assert_eq!(acquired.access_token, second_token); - assert_eq!(acquired.refresh_token.as_deref(), Some("first-refresh")); -} - -#[tokio::test] -async fn manual_add_should_preserve_existing_refresh_when_refresh_only_omits_rotation() { - let refreshed_token = test_jwt( - Some("rt-existing-account"), - Some("rt-existing-user"), - Some("second@example.com"), - Some("team"), - 7200, - ); - let refresher = StaticTokenRefresher { - result: Ok(TokenPair { - access_token: refreshed_token.clone(), - refresh_token: None, - }), - calls: Arc::new(Mutex::new(Vec::new())), - }; - let (app, _state, pool, _dir) = admin_accounts_test_app_with_refresher( - "admin-account-manual-refresh-existing-preserve.sqlite", - 31, - refresher, - ) - .await; - let first_token = test_jwt( - Some("rt-existing-account"), - Some("rt-existing-user"), - Some("first@example.com"), - Some("free"), - 3600, - ); - let first_response = post_admin_account( - &app, - json!({ - "token": first_token, - "refreshToken": "old-refresh" - }), - ) - .await; - assert_eq!(first_response.status(), StatusCode::OK); - let first_body = response_json(first_response).await; - let account_id = first_body["data"]["id"].as_str().unwrap().to_string(); - - let second_response = - post_admin_account(&app, json!({ "refreshToken": "incoming-refresh" })).await; - - assert_eq!(second_response.status(), StatusCode::OK); - let repo = AccountRepository::new(pool, SecretBox::new([31u8; 32])); - let stored = repo.get(&account_id).await.unwrap().unwrap(); - assert_eq!(stored.access_token.expose_secret(), &refreshed_token); - assert_eq!(stored.refresh_token.unwrap().expose_secret(), "old-refresh"); - assert_eq!(stored.email.as_deref(), Some("second@example.com")); - assert_eq!(stored.plan_type.as_deref(), Some("team")); -} - -#[tokio::test] -async fn admin_accounts_import_cli_should_read_codex_auth_file_store_encrypted_and_sync_pool() { - let (app, state, pool, dir) = - admin_accounts_test_app("admin-account-import-cli.sqlite", 28).await; - let codex_home = dir.path().join("codex-home"); - fs::create_dir_all(&codex_home).unwrap(); - let token = test_jwt( - Some("cli-account"), - Some("cli-user"), - Some("cli@example.com"), - Some("plus"), - 3600, - ); - fs::write( - codex_home.join("auth.json"), - json!({ - "access_token": token, - "refresh_token": "cli-refresh-secret" - }) - .to_string(), - ) - .unwrap(); - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/import-cli") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_import_cli") - .body(Body::from( - json!({ - "codexHome": codex_home.display().to_string(), - "id": "caller-id", - "email": "caller@example.com", - "accountId": "caller-account", - "userId": "caller-user", - "label": "caller-label", - "planType": "caller-plan" - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["sourceFormat"], "codex_cli"); - assert_eq!(body["data"]["imported"], 1); - assert!(body["data"].get("token").is_none()); - assert!(body["data"].get("refreshToken").is_none()); - - let stored: ( - String, - String, - String, - String, - Option, - String, - String, - String, - ) = sqlx::query_as( - "select id, email, account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher from accounts limit 1", - ) - .fetch_one(&pool) - .await - .unwrap(); - assert_ne!(stored.0, "caller-id"); - assert_eq!(stored.1, "cli@example.com"); - assert_eq!(stored.2, "cli-account"); - assert_eq!(stored.3, "cli-user"); - assert_eq!(stored.4, None); - assert_eq!(stored.5, "plus"); - assert!(stored.6.starts_with("v1:")); - assert!(!stored.6.contains(&token)); - assert!(stored.7.starts_with("v1:")); - assert!(!stored.7.contains("cli-refresh-secret")); - - let acquired = state - .account_pool() - .lock() - .await - .acquire("gpt-5.5") - .unwrap(); - 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) = 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); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let state = AppState::with_pool_and_secret_box( - test_config_with_base_url(url, base_url), - pool.clone(), - SecretBox::new([key_byte; 32]), - ); - let app = build_router(state.clone()); - (app, state, pool, dir) -} - -async fn admin_accounts_test_app_with_refresher( - db_name: &str, - key_byte: u8, - token_refresher: C, -) -> (Router, AppState, sqlx::SqlitePool, tempfile::TempDir) -where - C: TokenRefresher, -{ - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let state = AppState::with_pool_secret_api_key_hasher_and_token_refresher( - test_config(url), - pool.clone(), - SecretBox::new([key_byte; 32]), - ApiKeyHasher::new([key_byte; 32]), - token_refresher, - ); - let app = build_router(state.clone()); - (app, state, pool, dir) -} - -async fn admin_accounts_test_app_with_oauth_client( - db_name: &str, - key_byte: u8, - oauth_client: C, -) -> (Router, AppState, sqlx::SqlitePool, tempfile::TempDir) -where - C: OAuthClient + TokenRefresher, -{ - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let state = AppState::with_pool_secret_api_key_hasher_and_oauth_client( - test_config(url), - pool.clone(), - SecretBox::new([key_byte; 32]), - ApiKeyHasher::new([key_byte; 32]), - oauth_client, - ); - let app = build_router(state.clone()); - (app, state, pool, dir) -} - -async fn post_admin_account(app: &Router, payload: Value) -> axum::response::Response { - app.clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(payload.to_string())) - .unwrap(), - ) - .await - .unwrap() -} - -async fn import_test_account(app: &Router, session_id: &str, account_id: &str) { - let import_body = json!({ - "accounts": [{ - "id": account_id, - "email": format!("{account_id}@example.com"), - "planType": "plus", - "token": format!("access-{account_id}"), - "refreshToken": format!("refresh-{account_id}"), - "status": "active" - }] - }); - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/admin/accounts/import") - .header("content-type", "application/json") - .header("cookie", format!("cpr_admin_session={session_id}")) - .body(Body::from(import_body.to_string())) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); -} - -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 response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} - -fn test_jwt( - account_id: Option<&str>, - user_id: Option<&str>, - email: Option<&str>, - plan_type: Option<&str>, - exp_offset_seconds: i64, -) -> String { - let mut auth = serde_json::Map::new(); - if let Some(account_id) = account_id { - auth.insert( - "chatgpt_account_id".to_string(), - Value::String(account_id.to_string()), - ); - } - if let Some(user_id) = user_id { - auth.insert( - "chatgpt_user_id".to_string(), - Value::String(user_id.to_string()), - ); - } - if let Some(plan_type) = plan_type { - auth.insert( - "chatgpt_plan_type".to_string(), - Value::String(plan_type.to_string()), - ); - } - let mut profile = serde_json::Map::new(); - if let Some(email) = email { - profile.insert("email".to_string(), Value::String(email.to_string())); - } - let payload = json!({ - "exp": Utc::now().timestamp() + exp_offset_seconds, - "https://api.openai.com/auth": Value::Object(auth), - "https://api.openai.com/profile": Value::Object(profile), - }); - let header = json!({ "alg": "none", "typ": "JWT" }); - format!("{}.{}.", jwt_part(&header), jwt_part(&payload),) -} - -fn jwt_part(value: &Value) -> String { - URL_SAFE_NO_PAD.encode(serde_json::to_vec(value).unwrap()) -} diff --git a/tests/admin_api_contract_test.rs b/tests/admin_api_contract_test.rs index 0ac33b990..c095b9102 100644 --- a/tests/admin_api_contract_test.rs +++ b/tests/admin_api_contract_test.rs @@ -1,12 +1,13 @@ use axum::{ + body::to_bytes, http::StatusCode, response::{IntoResponse, Response}, }; use serde_json::{json, Value}; use codex_proxy_rs::{ - http::admin::{AdminEnvelope, AdminPageEnvelope, AdminResponse}, - pagination::Page, + http::admin::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse}, + utils::pagination::Page, }; #[test] @@ -67,3 +68,27 @@ fn admin_error_body_uses_null_data() { assert_eq!(value["data"], Value::Null); } + +#[tokio::test] +async fn admin_error_into_response_uses_admin_envelope() { + let response: Response = AdminError::new( + StatusCode::UNAUTHORIZED, + 40101, + "Admin login required", + "req_1", + ) + .into_response(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let value: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!( + value, + json!({ + "code": 40101, + "message": "Admin login required", + "data": null, + "requestId": "req_1" + }) + ); +} diff --git a/tests/admin_api_keys_route_test.rs b/tests/admin_api_keys_route_test.rs index c19e7feaf..3e11a6b27 100644 --- a/tests/admin_api_keys_route_test.rs +++ b/tests/admin_api_keys_route_test.rs @@ -1,26 +1,29 @@ use std::collections::BTreeMap; use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; -use chrono::Utc; -use serde_json::{json, Value}; +use serde_json::json; use tower::ServiceExt; use codex_proxy_rs::{ app::build_router, + app::state::AppState, auth::api_key::ApiKeyHasher, config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, }, - crypto::SecretBox, - state::AppState, storage::db::connect_sqlite, + utils::crypto::SecretBox, }; +mod common; + +use common::{response_json, seed_admin_session}; + fn test_config(database_url: String) -> AppConfig { AppConfig { server: ServerConfig { @@ -628,32 +631,3 @@ async fn admin_api_keys_import_should_rotate_exported_metadata_and_return_new_pl .unwrap(); assert_eq!(new_accepted.status(), StatusCode::OK); } - -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 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/admin_login_route_test.rs b/tests/admin_login_route_test.rs index 5fc8e6339..e6b21c0e6 100644 --- a/tests/admin_login_route_test.rs +++ b/tests/admin_login_route_test.rs @@ -10,13 +10,13 @@ use tower::ServiceExt; use codex_proxy_rs::{ app::build_router, + app::state::AppState, auth::admin_session::hash_admin_password, config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, }, - state::AppState, storage::db::connect_sqlite, }; diff --git a/tests/admin_logs_route_test.rs b/tests/admin_logs_route_test.rs index 8f582f36b..31d1bfb07 100644 --- a/tests/admin_logs_route_test.rs +++ b/tests/admin_logs_route_test.rs @@ -1,15 +1,14 @@ use std::collections::BTreeMap; use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; -use chrono::Utc; -use serde_json::Value; use tower::ServiceExt; use codex_proxy_rs::{ app::build_router, + app::state::AppState, config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, @@ -19,10 +18,13 @@ use codex_proxy_rs::{ event::{EventLevel, EventLog}, repository::EventLogRepository, }, - state::AppState, storage::db::connect_sqlite, }; +mod common; + +use common::{response_json, seed_admin_session}; + fn test_config(database_url: String) -> AppConfig { AppConfig { server: ServerConfig { @@ -143,32 +145,3 @@ async fn admin_logs_reject_missing_admin_session_cookie() { assert_eq!(body["code"], 40101); assert_eq!(body["requestId"], "req_admin"); } - -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 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/admin_models_route_test.rs b/tests/admin_models_route_test.rs index 34df39858..a2f6308f8 100644 --- a/tests/admin_models_route_test.rs +++ b/tests/admin_models_route_test.rs @@ -1,12 +1,11 @@ use std::collections::BTreeMap; use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; -use chrono::Utc; use secrecy::SecretString; -use serde_json::{json, Value}; +use serde_json::json; use tower::ServiceExt; use wiremock::{ matchers::{method, path}, @@ -14,24 +13,27 @@ use wiremock::{ }; use codex_proxy_rs::{ - accounts::{ + app::build_router, + app::state::AppState, + auth::{api_key::ApiKeyHasher, api_key_repository::ClientApiKeyRepository}, + codex::accounts::{ model::{Account, AccountStatus}, - pool::AccountAcquireRequest, repository::{AccountRepository, NewAccount}, }, - app::build_router, - auth::{api_key::ApiKeyHasher, api_key_repository::ClientApiKeyRepository}, + codex::models::repository::ModelSnapshotRepository, config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, }, - crypto::SecretBox, - models::repository::ModelSnapshotRepository, - state::AppState, storage::db::connect_sqlite, + utils::crypto::SecretBox, }; +mod common; + +use common::{response_json, seed_admin_session}; + fn test_config(database_url: String, base_url: String) -> AppConfig { AppConfig { server: ServerConfig { @@ -161,13 +163,21 @@ async fn admin_refresh_models_should_fetch_snapshots_for_distinct_account_plans( secret_box, hasher, ); - { - let account_pool = state.account_pool(); - let mut account_pool = account_pool.lock().await; - account_pool.insert(pool_account("acct_aaa_free", "free")); - account_pool.insert(pool_account("acct_plus_a", "plus")); - account_pool.insert(pool_account("acct_team", "team")); - } + state + .services + .accounts + .insert_runtime_account(pool_account("acct_aaa_free", "free")) + .await; + state + .services + .accounts + .insert_runtime_account(pool_account("acct_plus_a", "plus")) + .await; + state + .services + .accounts + .insert_runtime_account(pool_account("acct_team", "team")) + .await; let app = build_router(state.clone()); let response = app @@ -212,12 +222,12 @@ async fn admin_refresh_models_should_fetch_snapshots_for_distinct_account_plans( assert_eq!(catalog[0]["source"], "backend"); let acquired = state - .account_pool() - .lock() + .services + .accounts + .acquire_runtime_account("gpt-refresh") .await - .acquire_with(AccountAcquireRequest::new("gpt-refresh", Utc::now())) .unwrap(); - assert_ne!(acquired.account.plan_type.as_deref(), Some("free")); + assert_ne!(acquired.plan_type.as_deref(), Some("free")); } async fn insert_account( @@ -247,32 +257,3 @@ fn pool_account(id: &str, plan_type: &str) -> Account { account.plan_type = Some(plan_type.to_string()); account } - -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 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/admin_settings_route_test.rs b/tests/admin_settings_route_test.rs index 3126518ac..cff10d9b4 100644 --- a/tests/admin_settings_route_test.rs +++ b/tests/admin_settings_route_test.rs @@ -1,24 +1,27 @@ use std::collections::BTreeMap; use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; -use chrono::Utc; -use serde_json::{json, Value}; +use serde_json::json; use tower::ServiceExt; use codex_proxy_rs::{ app::build_router, + app::state::AppState, config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, }, - state::AppState, storage::db::connect_sqlite, }; +mod common; + +use common::{response_json, seed_admin_session}; + fn test_config(database_url: String) -> AppConfig { let mut aliases = BTreeMap::new(); aliases.insert("codex-fast".to_string(), "gpt-5.5".to_string()); @@ -170,32 +173,3 @@ async fn admin_settings_should_return_only_in_scope_runtime_fields() { ); } } - -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 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/admin_usage_stats_route_test.rs b/tests/admin_usage_stats_route_test.rs index e20600fe4..7a2272653 100644 --- a/tests/admin_usage_stats_route_test.rs +++ b/tests/admin_usage_stats_route_test.rs @@ -1,24 +1,27 @@ use std::collections::BTreeMap; use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; use chrono::Utc; -use serde_json::Value; use tower::ServiceExt; use codex_proxy_rs::{ app::build_router, + app::state::AppState, config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, }, - state::AppState, storage::db::connect_sqlite, }; +mod common; + +use common::{response_json, seed_admin_session}; + fn test_config(database_url: String) -> AppConfig { AppConfig { server: ServerConfig { @@ -188,30 +191,6 @@ async fn admin_usage_stats_should_reject_missing_admin_session_cookie() { 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 [ @@ -257,8 +236,3 @@ async fn seed_usage(pool: &sqlx::SqlitePool) { .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/chat_completions_route_test.rs b/tests/chat_completions_route_test.rs index e13e5cf5f..88e9e7f64 100644 --- a/tests/chat_completions_route_test.rs +++ b/tests/chat_completions_route_test.rs @@ -1,10 +1,9 @@ use std::collections::BTreeMap; use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; -use chrono::Utc; use serde_json::{json, Value}; use tower::ServiceExt; use wiremock::{ @@ -14,17 +13,30 @@ use wiremock::{ use codex_proxy_rs::{ app::build_router, + app::state::AppState, auth::{api_key::ApiKeyHasher, api_key_repository::ClientApiKeyRepository}, config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, }, - crypto::SecretBox, - state::AppState, storage::db::connect_sqlite, + utils::crypto::SecretBox, }; +mod common; + +use common::{response_json, response_sse_data, seed_admin_session}; + +const CHAT_RETRY_AFTER_SUCCESS_SSE: &str = include_str!("fixtures/chat_retry_after_success.sse"); +const CHAT_NON_STREAM_TEXT_SSE: &str = include_str!("fixtures/chat_non_stream_text.sse"); +const CHAT_STREAM_TEXT_SSE: &str = include_str!("fixtures/chat_stream_text.sse"); +const CHAT_PARALLEL_TOOLS_SUCCESS_SSE: &str = + include_str!("fixtures/chat_parallel_tools_success.sse"); +const CHAT_TOOL_REASONING_COMPLETE_SSE: &str = + include_str!("fixtures/chat_tool_reasoning_complete.sse"); +const CHAT_TOOL_STREAM_SSE: &str = include_str!("fixtures/chat_tool_stream.sse"); + struct ImportedApp { app: axum::Router, client_api_key: String, @@ -166,14 +178,6 @@ async fn build_imported_app_with_accounts( #[tokio::test] async fn chat_completions_stream_should_retry_next_account_after_429_retry_after() { let server = MockServer::start().await; - let sse_body = concat!( - "event: response.output_text.delta\n", - "data: {\"delta\":\"ok\"}\n", - "\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_chat_after_429\",\"object\":\"response\",\"status\":\"completed\",\"usage\":{\"input_tokens\":4,\"output_tokens\":1}}}\n", - "\n", - ); Mock::given(method("POST")) .and(path("/codex/responses")) .and(header("authorization", "Bearer access-a")) @@ -193,7 +197,7 @@ async fn chat_completions_stream_should_retry_next_account_after_429_retry_after .respond_with( ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), + .set_body_string(CHAT_RETRY_AFTER_SUCCESS_SSE), ) .mount(&server) .await; @@ -257,14 +261,6 @@ async fn chat_completions_stream_should_retry_next_account_after_429_retry_after #[tokio::test] async fn chat_completions_translates_messages_and_returns_openai_response() { let server = MockServer::start().await; - let sse_body = concat!( - "event: response.output_text.delta\n", - "data: {\"delta\":\"hello\"}\n", - "\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_chat\",\"object\":\"response\",\"status\":\"completed\",\"usage\":{\"input_tokens\":9,\"output_tokens\":3}}}\n", - "\n", - ); Mock::given(method("POST")) .and(path("/codex/responses")) .and(header("authorization", "Bearer access-secret")) @@ -282,7 +278,7 @@ async fn chat_completions_translates_messages_and_returns_openai_response() { .respond_with( ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), + .set_body_string(CHAT_NON_STREAM_TEXT_SSE), ) .mount(&server) .await; @@ -330,17 +326,6 @@ async fn chat_completions_translates_messages_and_returns_openai_response() { #[tokio::test] async fn chat_completions_streams_openai_chunks_from_codex_sse() { let server = MockServer::start().await; - let sse_body = concat!( - "event: response.output_text.delta\n", - "data: {\"delta\":\"hel\"}\n", - "\n", - "event: response.output_text.delta\n", - "data: {\"delta\":\"lo\"}\n", - "\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_chat\",\"object\":\"response\",\"status\":\"completed\",\"usage\":{\"input_tokens\":9,\"output_tokens\":3,\"input_tokens_details\":{\"cached_tokens\":4}}}}\n", - "\n", - ); Mock::given(method("POST")) .and(path("/codex/responses")) .and(header("authorization", "Bearer access-secret")) @@ -358,7 +343,7 @@ async fn chat_completions_streams_openai_chunks_from_codex_sse() { .respond_with( ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), + .set_body_string(CHAT_STREAM_TEXT_SSE), ) .mount(&server) .await; @@ -417,14 +402,6 @@ async fn chat_completions_streams_openai_chunks_from_codex_sse() { #[tokio::test] async fn chat_completions_forwards_parallel_tools_and_model_suffix_options() { let server = MockServer::start().await; - let sse_body = concat!( - "event: response.output_text.delta\n", - "data: {\"delta\":\"ok\"}\n", - "\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_chat\",\"object\":\"response\",\"status\":\"completed\",\"usage\":{\"input_tokens\":2,\"output_tokens\":1}}}\n", - "\n", - ); Mock::given(method("POST")) .and(path("/codex/responses")) .and(header("authorization", "Bearer access-secret")) @@ -448,7 +425,7 @@ async fn chat_completions_forwards_parallel_tools_and_model_suffix_options() { .respond_with( ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), + .set_body_string(CHAT_PARALLEL_TOOLS_SUCCESS_SSE), ) .mount(&server) .await; @@ -488,26 +465,6 @@ async fn chat_completions_forwards_parallel_tools_and_model_suffix_options() { #[tokio::test] async fn chat_completions_collects_tool_calls_and_reasoning_content() { let server = MockServer::start().await; - let sse_body = concat!( - "event: response.reasoning_summary_text.delta\n", - "data: {\"delta\":\"thinking \"}\n", - "\n", - "event: response.output_item.added\n", - "data: {\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"item_1\",\"call_id\":\"call_abc\",\"name\":\"lookup\"}}\n", - "\n", - "event: response.function_call_arguments.delta\n", - "data: {\"item_id\":\"item_1\",\"delta\":\"{\\\"city\\\"\"}\n", - "\n", - "event: response.function_call_arguments.delta\n", - "data: {\"item_id\":\"item_1\",\"delta\":\":\\\"Paris\\\"}\"}\n", - "\n", - "event: response.function_call_arguments.done\n", - "data: {\"item_id\":\"item_1\",\"arguments\":\"{\\\"city\\\":\\\"Paris\\\"}\"}\n", - "\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_chat\",\"object\":\"response\",\"status\":\"completed\",\"usage\":{\"input_tokens\":5,\"output_tokens\":2,\"output_tokens_details\":{\"reasoning_tokens\":7}}}}\n", - "\n", - ); Mock::given(method("POST")) .and(path("/codex/responses")) .and(header("authorization", "Bearer access-secret")) @@ -529,7 +486,7 @@ async fn chat_completions_collects_tool_calls_and_reasoning_content() { .respond_with( ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), + .set_body_string(CHAT_TOOL_REASONING_COMPLETE_SSE), ) .mount(&server) .await; @@ -583,20 +540,6 @@ async fn chat_completions_collects_tool_calls_and_reasoning_content() { #[tokio::test] async fn chat_completions_streams_tool_call_chunks_from_codex_sse() { let server = MockServer::start().await; - let sse_body = concat!( - "event: response.output_item.added\n", - "data: {\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"item_1\",\"call_id\":\"call_abc\",\"name\":\"lookup\"}}\n", - "\n", - "event: response.function_call_arguments.delta\n", - "data: {\"item_id\":\"item_1\",\"delta\":\"{\\\"city\\\"\"}\n", - "\n", - "event: response.function_call_arguments.delta\n", - "data: {\"item_id\":\"item_1\",\"delta\":\":\\\"Paris\\\"}\"}\n", - "\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_chat\",\"object\":\"response\",\"status\":\"completed\",\"usage\":{\"input_tokens\":5,\"output_tokens\":2}}}\n", - "\n", - ); Mock::given(method("POST")) .and(path("/codex/responses")) .and(header("authorization", "Bearer access-secret")) @@ -614,7 +557,7 @@ async fn chat_completions_streams_tool_call_chunks_from_codex_sse() { .respond_with( ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), + .set_body_string(CHAT_TOOL_STREAM_SSE), ) .mount(&server) .await; @@ -670,41 +613,3 @@ async fn chat_completions_streams_tool_call_chunks_from_codex_sse() { ); assert_eq!(chunks[4]["choices"][0]["finish_reason"], "tool_calls"); } - -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 response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} - -async fn response_sse_data(response: axum::response::Response) -> Vec { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - String::from_utf8(bytes.to_vec()) - .unwrap() - .lines() - .filter_map(|line| line.strip_prefix("data: ").map(ToString::to_string)) - .collect() -} diff --git a/tests/cli_auth_import_test.rs b/tests/cli_auth_import_test.rs index dc6ae0246..24148a180 100644 --- a/tests/cli_auth_import_test.rs +++ b/tests/cli_auth_import_test.rs @@ -1,6 +1,6 @@ use std::fs; -use codex_proxy_rs::auth::cli_import::{parse_cli_auth_json, read_cli_auth_from_home}; +use codex_proxy_rs::codex::oauth::{parse_cli_auth_json, read_cli_auth_from_home}; #[test] fn parse_cli_auth_json_should_accept_codex_cli_tokens() { diff --git a/tests/codex_client_test.rs b/tests/codex_client_test.rs index 93720e33d..f735c4a93 100644 --- a/tests/codex_client_test.rs +++ b/tests/codex_client_test.rs @@ -5,12 +5,12 @@ use wiremock::{ }; use codex_proxy_rs::{ - codex::{ + codex::fingerprint::model::Fingerprint, + codex::transport::{ client::{build_reqwest_client, CodexBackendClient, CodexRequestContext}, types::CodexResponsesRequest, usage::TokenUsage, }, - fingerprint::model::Fingerprint, }; #[tokio::test] diff --git a/tests/codex_headers_test.rs b/tests/codex_headers_test.rs index 696c11367..862ef9900 100644 --- a/tests/codex_headers_test.rs +++ b/tests/codex_headers_test.rs @@ -1,6 +1,6 @@ use codex_proxy_rs::{ - codex::headers::build_codex_headers, - fingerprint::{model::Fingerprint, updater::parse_update_manifest}, + codex::fingerprint::{model::Fingerprint, updater::parse_update_manifest}, + codex::transport::headers::build_codex_headers, }; #[test] diff --git a/tests/codex_sse_usage_test.rs b/tests/codex_sse_usage_test.rs index 3fd31c8f7..0099820fc 100644 --- a/tests/codex_sse_usage_test.rs +++ b/tests/codex_sse_usage_test.rs @@ -1,6 +1,6 @@ use serde_json::json; -use codex_proxy_rs::codex::{ +use codex_proxy_rs::codex::transport::{ sse::{encode_sse_event, parse_sse_events}, usage::{extract_usage, TokenUsage}, }; diff --git a/tests/codex_websocket_test.rs b/tests/codex_websocket_test.rs index a1b467ad1..68f3669ab 100644 --- a/tests/codex_websocket_test.rs +++ b/tests/codex_websocket_test.rs @@ -15,7 +15,7 @@ use tokio_tungstenite::{ }, }; -use codex_proxy_rs::codex::{ +use codex_proxy_rs::codex::transport::{ client::{build_reqwest_client, CodexBackendClient, CodexClientError, CodexRequestContext}, types::CodexResponsesRequest, websocket::{ensure_http_sse_supported, transport_for_request, CodexTransport}, @@ -113,7 +113,7 @@ async fn previous_response_id_should_use_websocket_transport() { let client = CodexBackendClient::new( build_reqwest_client(false).unwrap(), format!("http://{addr}"), - codex_proxy_rs::fingerprint::model::Fingerprint::default_for_tests(), + codex_proxy_rs::codex::fingerprint::model::Fingerprint::default_for_tests(), ); let response = client @@ -174,7 +174,7 @@ async fn websocket_handshake_429_should_surface_as_upstream_error_before_body_is let client = CodexBackendClient::new( build_reqwest_client(false).unwrap(), format!("http://{addr}"), - codex_proxy_rs::fingerprint::model::Fingerprint::default_for_tests(), + codex_proxy_rs::codex::fingerprint::model::Fingerprint::default_for_tests(), ); let error = client @@ -266,7 +266,7 @@ async fn websocket_success_should_capture_handshake_headers_and_rate_limit_heade let client = CodexBackendClient::new( build_reqwest_client(false).unwrap(), format!("http://{addr}"), - codex_proxy_rs::fingerprint::model::Fingerprint::default_for_tests(), + codex_proxy_rs::codex::fingerprint::model::Fingerprint::default_for_tests(), ); let response = client diff --git a/tests/common/admin_accounts.rs b/tests/common/admin_accounts.rs new file mode 100644 index 000000000..175d5d9fd --- /dev/null +++ b/tests/common/admin_accounts.rs @@ -0,0 +1,334 @@ +use std::{collections::BTreeMap, sync::Arc}; + +use async_trait::async_trait; +use axum::{ + body::Body, + http::{Request, StatusCode}, + Router, +}; +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use chrono::Utc; +use serde_json::{json, Value}; +use tokio::sync::Mutex; +use tower::ServiceExt; + +use codex_proxy_rs::{ + app::build_router, + app::state::AppState, + auth::api_key::ApiKeyHasher, + codex::oauth::{ + DeviceCode, OAuthClient, OAuthError, RefreshFailure, TokenPair, TokenRefresher, + }, + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, + UsageStatsConfig, + }, + storage::db::connect_sqlite, + utils::crypto::SecretBox, +}; + +use super::seed_admin_session; + +#[derive(Clone)] +pub struct StaticTokenRefresher { + pub result: Result, + pub calls: Arc>>, +} + +#[derive(Clone)] +pub struct FailingTokenRefresher { + pub failure: RefreshFailure, +} + +#[derive(Clone)] +pub struct StaticOAuthClient { + pub device_response: Result, + pub poll_response: Result, + pub exchange_response: Result, + pub poll_calls: Arc>>, + pub exchange_calls: Arc>>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExchangeCall { + pub code: String, + pub code_verifier: String, + pub redirect_uri: String, +} + +#[async_trait] +impl TokenRefresher for StaticTokenRefresher { + async fn refresh(&self, refresh_token: &str) -> Result { + 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 { + Err(self.failure) + } +} + +#[async_trait] +impl TokenRefresher for StaticOAuthClient { + async fn refresh(&self, _refresh_token: &str) -> Result { + Err(RefreshFailure::Transport) + } +} + +#[async_trait] +impl OAuthClient for StaticOAuthClient { + async fn exchange_code( + &self, + code: &str, + code_verifier: &str, + redirect_uri: &str, + ) -> Result { + self.exchange_calls.lock().await.push(ExchangeCall { + code: code.to_string(), + code_verifier: code_verifier.to_string(), + redirect_uri: redirect_uri.to_string(), + }); + self.exchange_response.clone() + } + + async fn request_device_code(&self) -> Result { + self.device_response.clone() + } + + async fn poll_device_token(&self, device_code: &str) -> Result { + self.poll_calls.lock().await.push(device_code.to_string()); + self.poll_response.clone() + } +} + +pub 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, + }, + } +} + +pub fn test_config_with_base_url(database_url: String, base_url: String) -> AppConfig { + AppConfig { + api: ApiConfig { base_url }, + ..test_config(database_url) + } +} + +pub 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 +} + +pub 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); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + let state = AppState::with_pool_and_secret_box( + test_config_with_base_url(url, base_url), + pool.clone(), + SecretBox::new([key_byte; 32]), + ); + let app = build_router(state.clone()); + (app, state, pool, dir) +} + +pub async fn admin_accounts_test_app_with_refresher( + db_name: &str, + key_byte: u8, + token_refresher: C, +) -> (Router, AppState, sqlx::SqlitePool, tempfile::TempDir) +where + C: TokenRefresher, +{ + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join(db_name); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + let state = AppState::with_pool_secret_api_key_hasher_and_token_refresher( + test_config(url), + pool.clone(), + SecretBox::new([key_byte; 32]), + ApiKeyHasher::new([key_byte; 32]), + token_refresher, + ); + let app = build_router(state.clone()); + (app, state, pool, dir) +} + +pub async fn admin_accounts_test_app_with_oauth_client( + db_name: &str, + key_byte: u8, + oauth_client: C, +) -> (Router, AppState, sqlx::SqlitePool, tempfile::TempDir) +where + C: OAuthClient + TokenRefresher, +{ + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join(db_name); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + let state = AppState::with_pool_secret_api_key_hasher_and_oauth_client( + test_config(url), + pool.clone(), + SecretBox::new([key_byte; 32]), + ApiKeyHasher::new([key_byte; 32]), + oauth_client, + ); + let app = build_router(state.clone()); + (app, state, pool, dir) +} + +pub async fn post_admin_account(app: &Router, payload: Value) -> axum::response::Response { + app.clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(payload.to_string())) + .unwrap(), + ) + .await + .unwrap() +} + +pub async fn import_test_account(app: &Router, session_id: &str, account_id: &str) { + let import_body = json!({ + "accounts": [{ + "id": account_id, + "email": format!("{account_id}@example.com"), + "planType": "plus", + "token": format!("access-{account_id}"), + "refreshToken": format!("refresh-{account_id}"), + "status": "active" + }] + }); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/accounts/import") + .header("content-type", "application/json") + .header("cookie", format!("cpr_admin_session={session_id}")) + .body(Body::from(import_body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); +} + +pub fn test_jwt( + account_id: Option<&str>, + user_id: Option<&str>, + email: Option<&str>, + plan_type: Option<&str>, + exp_offset_seconds: i64, +) -> String { + let mut auth = serde_json::Map::new(); + if let Some(account_id) = account_id { + auth.insert( + "chatgpt_account_id".to_string(), + Value::String(account_id.to_string()), + ); + } + if let Some(user_id) = user_id { + auth.insert( + "chatgpt_user_id".to_string(), + Value::String(user_id.to_string()), + ); + } + if let Some(plan_type) = plan_type { + auth.insert( + "chatgpt_plan_type".to_string(), + Value::String(plan_type.to_string()), + ); + } + let mut profile = serde_json::Map::new(); + if let Some(email) = email { + profile.insert("email".to_string(), Value::String(email.to_string())); + } + let payload = json!({ + "exp": Utc::now().timestamp() + exp_offset_seconds, + "https://api.openai.com/auth": Value::Object(auth), + "https://api.openai.com/profile": Value::Object(profile), + }); + let header = json!({ "alg": "none", "typ": "JWT" }); + format!("{}.{}.", jwt_part(&header), jwt_part(&payload),) +} + +fn jwt_part(value: &Value) -> String { + URL_SAFE_NO_PAD.encode(serde_json::to_vec(value).unwrap()) +} diff --git a/tests/common/mod.rs b/tests/common/mod.rs new file mode 100644 index 000000000..5b00b3e88 --- /dev/null +++ b/tests/common/mod.rs @@ -0,0 +1,52 @@ +#![allow(dead_code)] +// Each integration test crate imports this module independently, so a helper can be +// used by one route test and intentionally unused by another. + +use axum::body::to_bytes; +use chrono::Utc; +use serde_json::Value; + +pub mod admin_accounts; +pub mod upstream; + +pub 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(); +} + +pub 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() +} + +pub async fn response_text(response: axum::response::Response) -> String { + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + String::from_utf8(bytes.to_vec()).unwrap() +} + +pub async fn response_sse_data(response: axum::response::Response) -> Vec { + response_text(response) + .await + .lines() + .filter_map(|line| line.strip_prefix("data: ").map(ToString::to_string)) + .collect() +} diff --git a/tests/common/upstream.rs b/tests/common/upstream.rs new file mode 100644 index 000000000..ccbe22cf2 --- /dev/null +++ b/tests/common/upstream.rs @@ -0,0 +1,315 @@ +use std::collections::BTreeMap; + +use async_trait::async_trait; +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use serde_json::{json, Value}; +use tower::ServiceExt; + +use codex_proxy_rs::{ + app::build_router, + app::state::AppState, + auth::{api_key::ApiKeyHasher, api_key_repository::ClientApiKeyRepository}, + codex::oauth::{RefreshFailure, TokenPair, TokenRefresher}, + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, + UsageStatsConfig, + }, + storage::db::connect_sqlite, + utils::crypto::SecretBox, +}; + +use super::seed_admin_session; + +pub struct ImportedApp { + pub app: axum::Router, + pub pool: sqlx::SqlitePool, + pub secret_box: SecretBox, + pub client_api_key: String, + _tempdir: tempfile::TempDir, +} + +#[derive(Clone)] +pub struct StaticTokenRefresher { + pub access_token: String, + pub refresh_token: Option, +} + +#[derive(Clone)] +pub struct FailingTokenRefresher { + pub failure: RefreshFailure, +} + +pub struct ImportAccount { + pub id: &'static str, + pub account_id: &'static str, + pub token: &'static str, + pub refresh_token: &'static str, +} + +#[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(), + }) + } +} + +#[async_trait] +impl TokenRefresher for FailingTokenRefresher { + async fn refresh(&self, _refresh_token: &str) -> Result { + Err(self.failure) + } +} + +pub fn test_config(database_url: String, base_url: String) -> AppConfig { + AppConfig { + server: ServerConfig { + host: "127.0.0.1".to_string(), + port: 0, + }, + api: ApiConfig { base_url }, + 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, + }, + } +} + +pub async fn build_imported_app(base_url: String) -> ImportedApp { + build_imported_app_with_accounts( + base_url, + &[ImportAccount { + id: "acct_imported", + account_id: "chatgpt-account", + token: "access-secret", + refresh_token: "refresh-secret", + }], + ) + .await +} + +pub async fn build_imported_app_with_accounts( + base_url: String, + accounts: &[ImportAccount], +) -> ImportedApp { + let tempdir = tempfile::tempdir().unwrap(); + let db = tempdir.path().join("v1-upstream.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([21u8; 32]); + let hasher = ApiKeyHasher::new([22u8; 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_and_api_key_hasher( + test_config(url, base_url), + pool.clone(), + secret_box.clone(), + hasher, + )); + let accounts = accounts + .iter() + .map(|account| { + json!({ + "id": account.id, + "accountId": account.account_id, + "token": account.token, + "refreshToken": account.refresh_token, + "status": "active" + }) + }) + .collect::>(); + + 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": accounts }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(import_response.status(), StatusCode::OK); + + ImportedApp { + app, + pool, + secret_box, + client_api_key: generated.plaintext, + _tempdir: tempdir, + } +} + +pub async fn build_imported_app_with_refresher( + base_url: String, + access_token: &str, +) -> ImportedApp { + build_imported_app_with_token_refresher( + base_url, + StaticTokenRefresher { + access_token: access_token.to_string(), + refresh_token: None, + }, + ) + .await +} + +pub async fn build_imported_app_with_token_refresher( + base_url: String, + token_refresher: C, +) -> ImportedApp +where + C: TokenRefresher, +{ + build_imported_app_with_accounts_and_token_refresher( + base_url, + &[ImportAccount { + id: "acct_imported", + account_id: "chatgpt-account", + token: "access-secret", + refresh_token: "refresh-secret", + }], + token_refresher, + ) + .await +} + +pub async fn build_imported_app_with_accounts_and_token_refresher( + base_url: String, + accounts: &[ImportAccount], + token_refresher: C, +) -> ImportedApp +where + C: TokenRefresher, +{ + 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, + token_refresher, + ), + ); + let accounts = accounts + .iter() + .map(|account| { + json!({ + "id": account.id, + "accountId": account.account_id, + "token": account.token, + "refreshToken": account.refresh_token, + "status": "active" + }) + }) + .collect::>(); + + 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": accounts }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(import_response.status(), StatusCode::OK); + + ImportedApp { + app, + pool, + secret_box, + client_api_key: generated.plaintext, + _tempdir: tempdir, + } +} + +pub async fn fetch_v1_event_log( + pool: &sqlx::SqlitePool, + request_id: &str, +) -> (String, String, String, i64, Value) { + let row: (String, String, String, i64, String) = sqlx::query_as( + "select account_id, route, model, status_code, metadata_json from event_logs where request_id = ? and kind = 'v1.response'", + ) + .bind(request_id) + .fetch_one(pool) + .await + .unwrap(); + ( + row.0, + row.1, + row.2, + row.3, + serde_json::from_str(&row.4).unwrap(), + ) +} diff --git a/tests/config_test.rs b/tests/config_test.rs index 5716ce8bd..0f69f007c 100644 --- a/tests/config_test.rs +++ b/tests/config_test.rs @@ -3,55 +3,8 @@ use std::fs; #[test] fn default_config_keeps_only_codex_backend() { - let yaml = r#" -server: - host: 127.0.0.1 - port: 8080 -api: - base_url: https://chatgpt.com/backend-api -model: - default: gpt-5.5 - default_reasoning_effort: null - default_service_tier: null - aliases: {} -auth: - refresh_margin_seconds: 300 - refresh_enabled: true - refresh_concurrency: 2 - max_concurrent_per_account: 3 - request_interval_ms: 50 - rotation_strategy: least_used - tier_priority: [] -quota: - refresh_interval_minutes: 5 - warning_thresholds: - primary: - - 80 - - 90 - secondary: - - 80 - - 90 - skip_exhausted: true -usage_stats: - history_retention_days: null -database: - url: sqlite://data/codex-proxy-rs.sqlite -security: - master_key_file: data/master.key - api_key_pepper_file: data/api-key-pepper.key -tls: - force_http11: false -admin: - session_ttl_minutes: 1440 -logging: - directory: logs - max_file_bytes: 10485760 - retention_days: 14 - enabled: false - capacity: 2000 - capture_body: false -"#; - let cfg: AppConfig = serde_yaml::from_str(yaml).unwrap(); + let cfg: AppConfig = serde_yaml::from_str(include_str!("../config.yaml")).unwrap(); + assert_eq!(cfg.server.host, "0.0.0.0"); assert_eq!(cfg.api.base_url, "https://chatgpt.com/backend-api"); assert_eq!(cfg.model.default_model, "gpt-5.5"); assert_eq!(cfg.auth.refresh_margin_seconds, 300); @@ -61,10 +14,10 @@ logging: } #[test] -fn config_loader_merges_default_local_and_cprs_env_overrides() { +fn config_loader_merges_config_and_local_yaml() { let dir = tempfile::tempdir().unwrap(); fs::write( - dir.path().join("default.yaml"), + dir.path().join("config.yaml"), r#" server: host: 127.0.0.1 @@ -126,23 +79,13 @@ logging: ) .unwrap(); - let cfg = AppConfig::load_from_dir_with_env( - dir.path(), - [ - ("CPRS_PORT", "19090"), - ("CPRS_DATABASE_URL", "sqlite://override.sqlite"), - ("CPRS_LOG_MAX_FILE_BYTES", "4096"), - ("CPRS_DEFAULT_MODEL", "gpt-5.5"), - ("CPRS_MAX_CONCURRENT_PER_ACCOUNT", "4"), - ], - ) - .unwrap(); + let cfg = AppConfig::load_from_dir(dir.path()).unwrap(); assert_eq!(cfg.server.host, "0.0.0.0"); - assert_eq!(cfg.server.port, 19090); - assert_eq!(cfg.database.url, "sqlite://override.sqlite"); + assert_eq!(cfg.server.port, 8080); + assert_eq!(cfg.database.url, "sqlite://data/codex-proxy-rs.sqlite"); assert_eq!(cfg.model.default_model, "gpt-5.5"); - assert_eq!(cfg.auth.max_concurrent_per_account, 4); + assert_eq!(cfg.auth.max_concurrent_per_account, 3); assert_eq!(cfg.logging.directory, "local-logs"); - assert_eq!(cfg.logging.max_file_bytes, 4096); + assert_eq!(cfg.logging.max_file_bytes, 10485760); } diff --git a/tests/cookie_store_test.rs b/tests/cookie_store_test.rs index f6ee6e6a3..49e9eeb81 100644 --- a/tests/cookie_store_test.rs +++ b/tests/cookie_store_test.rs @@ -1,9 +1,9 @@ use chrono::Utc; use codex_proxy_rs::{ - cookies::{jar::CookieJar, repository::CookieRepository}, - crypto::SecretBox, + codex::cookies::{jar::CookieJar, repository::CookieRepository}, storage::db::connect_sqlite, + utils::crypto::SecretBox, }; #[test] diff --git a/tests/crypto_test.rs b/tests/crypto_test.rs index 00451bb26..f8ea74b9c 100644 --- a/tests/crypto_test.rs +++ b/tests/crypto_test.rs @@ -1,4 +1,4 @@ -use codex_proxy_rs::crypto::SecretBox; +use codex_proxy_rs::utils::crypto::SecretBox; use secrecy::{ExposeSecret, SecretString}; #[test] diff --git a/tests/fingerprint_update_test.rs b/tests/fingerprint_update_test.rs index acb71e082..171399afd 100644 --- a/tests/fingerprint_update_test.rs +++ b/tests/fingerprint_update_test.rs @@ -1,9 +1,9 @@ use codex_proxy_rs::{ - codex::client::build_reqwest_client, - fingerprint::{ + codex::fingerprint::{ repository::FingerprintRepository, updater::{FingerprintUpdater, CODEX_DESKTOP_UPDATE_SOURCE}, }, + codex::transport::client::build_reqwest_client, storage::db::connect_sqlite, }; use wiremock::{ diff --git a/tests/fixtures/chat_non_stream_text.sse b/tests/fixtures/chat_non_stream_text.sse new file mode 100644 index 000000000..86edd1a19 --- /dev/null +++ b/tests/fixtures/chat_non_stream_text.sse @@ -0,0 +1,6 @@ +event: response.output_text.delta +data: {"delta":"hello"} + +event: response.completed +data: {"response":{"id":"resp_chat","object":"response","status":"completed","usage":{"input_tokens":9,"output_tokens":3}}} + diff --git a/tests/fixtures/chat_parallel_tools_success.sse b/tests/fixtures/chat_parallel_tools_success.sse new file mode 100644 index 000000000..3e00ab755 --- /dev/null +++ b/tests/fixtures/chat_parallel_tools_success.sse @@ -0,0 +1,6 @@ +event: response.output_text.delta +data: {"delta":"ok"} + +event: response.completed +data: {"response":{"id":"resp_chat","object":"response","status":"completed","usage":{"input_tokens":2,"output_tokens":1}}} + diff --git a/tests/fixtures/chat_retry_after_success.sse b/tests/fixtures/chat_retry_after_success.sse new file mode 100644 index 000000000..83cf7a8e2 --- /dev/null +++ b/tests/fixtures/chat_retry_after_success.sse @@ -0,0 +1,6 @@ +event: response.output_text.delta +data: {"delta":"ok"} + +event: response.completed +data: {"response":{"id":"resp_chat_after_429","object":"response","status":"completed","usage":{"input_tokens":4,"output_tokens":1}}} + diff --git a/tests/fixtures/chat_stream_text.sse b/tests/fixtures/chat_stream_text.sse new file mode 100644 index 000000000..49e83639c --- /dev/null +++ b/tests/fixtures/chat_stream_text.sse @@ -0,0 +1,9 @@ +event: response.output_text.delta +data: {"delta":"hel"} + +event: response.output_text.delta +data: {"delta":"lo"} + +event: response.completed +data: {"response":{"id":"resp_chat","object":"response","status":"completed","usage":{"input_tokens":9,"output_tokens":3,"input_tokens_details":{"cached_tokens":4}}}} + diff --git a/tests/fixtures/chat_tool_reasoning_complete.sse b/tests/fixtures/chat_tool_reasoning_complete.sse new file mode 100644 index 000000000..8e68557f7 --- /dev/null +++ b/tests/fixtures/chat_tool_reasoning_complete.sse @@ -0,0 +1,18 @@ +event: response.reasoning_summary_text.delta +data: {"delta":"thinking "} + +event: response.output_item.added +data: {"output_index":0,"item":{"type":"function_call","id":"item_1","call_id":"call_abc","name":"lookup"}} + +event: response.function_call_arguments.delta +data: {"item_id":"item_1","delta":"{\"city\""} + +event: response.function_call_arguments.delta +data: {"item_id":"item_1","delta":":\"Paris\"}"} + +event: response.function_call_arguments.done +data: {"item_id":"item_1","arguments":"{\"city\":\"Paris\"}"} + +event: response.completed +data: {"response":{"id":"resp_chat","object":"response","status":"completed","usage":{"input_tokens":5,"output_tokens":2,"output_tokens_details":{"reasoning_tokens":7}}}} + diff --git a/tests/fixtures/chat_tool_stream.sse b/tests/fixtures/chat_tool_stream.sse new file mode 100644 index 000000000..53c1453f6 --- /dev/null +++ b/tests/fixtures/chat_tool_stream.sse @@ -0,0 +1,12 @@ +event: response.output_item.added +data: {"output_index":0,"item":{"type":"function_call","id":"item_1","call_id":"call_abc","name":"lookup"}} + +event: response.function_call_arguments.delta +data: {"item_id":"item_1","delta":"{\"city\""} + +event: response.function_call_arguments.delta +data: {"item_id":"item_1","delta":":\"Paris\"}"} + +event: response.completed +data: {"response":{"id":"resp_chat","object":"response","status":"completed","usage":{"input_tokens":5,"output_tokens":2}}} + diff --git a/tests/fixtures/v1_responses_after_402.sse b/tests/fixtures/v1_responses_after_402.sse new file mode 100644 index 000000000..c6c86ce44 --- /dev/null +++ b/tests/fixtures/v1_responses_after_402.sse @@ -0,0 +1,3 @@ +event: response.completed +data: {"response":{"id":"resp_after_402","object":"response","usage":{"input_tokens":6,"output_tokens":3}}} + diff --git a/tests/fixtures/v1_responses_after_403.sse b/tests/fixtures/v1_responses_after_403.sse new file mode 100644 index 000000000..f8c04ef17 --- /dev/null +++ b/tests/fixtures/v1_responses_after_403.sse @@ -0,0 +1,3 @@ +event: response.completed +data: {"response":{"id":"resp_after_403","object":"response","usage":{"input_tokens":6,"output_tokens":3}}} + diff --git a/tests/fixtures/v1_responses_after_429.sse b/tests/fixtures/v1_responses_after_429.sse new file mode 100644 index 000000000..066f733cc --- /dev/null +++ b/tests/fixtures/v1_responses_after_429.sse @@ -0,0 +1,3 @@ +event: response.completed +data: {"response":{"id":"resp_after_429","object":"response","usage":{"input_tokens":5,"output_tokens":2}}} + diff --git a/tests/fixtures/v1_responses_after_cloudflare_403.sse b/tests/fixtures/v1_responses_after_cloudflare_403.sse new file mode 100644 index 000000000..eb456495f --- /dev/null +++ b/tests/fixtures/v1_responses_after_cloudflare_403.sse @@ -0,0 +1,3 @@ +event: response.completed +data: {"response":{"id":"resp_after_cf","object":"response","usage":{"input_tokens":6,"output_tokens":3}}} + diff --git a/tests/fixtures/v1_responses_after_refresh.sse b/tests/fixtures/v1_responses_after_refresh.sse new file mode 100644 index 000000000..bda9297ba --- /dev/null +++ b/tests/fixtures/v1_responses_after_refresh.sse @@ -0,0 +1,3 @@ +event: response.completed +data: {"response":{"id":"resp_after_refresh","object":"response","usage":{"input_tokens":2,"output_tokens":1}}} + diff --git a/tests/fixtures/v1_responses_completed_fields.sse b/tests/fixtures/v1_responses_completed_fields.sse new file mode 100644 index 000000000..d8fe8ddd2 --- /dev/null +++ b/tests/fixtures/v1_responses_completed_fields.sse @@ -0,0 +1,3 @@ +event: response.completed +data: {"response":{"id":"resp_fields","object":"response","usage":{"input_tokens":1,"output_tokens":1}}} + diff --git a/tests/fixtures/v1_responses_completed_reasoning_include.sse b/tests/fixtures/v1_responses_completed_reasoning_include.sse new file mode 100644 index 000000000..24aefade4 --- /dev/null +++ b/tests/fixtures/v1_responses_completed_reasoning_include.sse @@ -0,0 +1,3 @@ +event: response.completed +data: {"response":{"id":"resp_reasoning_include","object":"response","usage":{"input_tokens":1,"output_tokens":1}}} + diff --git a/tests/fixtures/v1_responses_completed_usage.sse b/tests/fixtures/v1_responses_completed_usage.sse new file mode 100644 index 000000000..3259297b9 --- /dev/null +++ b/tests/fixtures/v1_responses_completed_usage.sse @@ -0,0 +1,3 @@ +event: response.completed +data: {"response":{"id":"resp_1","object":"response","usage":{"input_tokens":7,"output_tokens":4,"input_tokens_details":{"cached_tokens":2}}}} + diff --git a/tests/fixtures/v1_responses_default_stream.sse b/tests/fixtures/v1_responses_default_stream.sse new file mode 100644 index 000000000..abe7f3308 --- /dev/null +++ b/tests/fixtures/v1_responses_default_stream.sse @@ -0,0 +1,6 @@ +event: response.output_text.delta +data: {"delta":"pong"} + +event: response.completed +data: {"response":{"id":"resp_default_stream","object":"response","usage":{"input_tokens":3,"output_tokens":5}}} + diff --git a/tests/fixtures/v1_responses_done_item_completed.sse b/tests/fixtures/v1_responses_done_item_completed.sse new file mode 100644 index 000000000..f1f2d9f66 --- /dev/null +++ b/tests/fixtures/v1_responses_done_item_completed.sse @@ -0,0 +1,6 @@ +event: response.output_item.done +data: {"item":{"type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"来自 done item"}]}} + +event: response.completed +data: {"response":{"id":"resp_item","object":"response","status":"completed","output":[],"usage":{"input_tokens":5,"output_tokens":3}}} + diff --git a/tests/fixtures/v1_responses_error_event.sse b/tests/fixtures/v1_responses_error_event.sse new file mode 100644 index 000000000..bb15949d0 --- /dev/null +++ b/tests/fixtures/v1_responses_error_event.sse @@ -0,0 +1,3 @@ +event: error +data: {"error":{"message":"model overloaded","type":"server_error","code":"server_error"}} + diff --git a/tests/fixtures/v1_responses_failed_event.sse b/tests/fixtures/v1_responses_failed_event.sse new file mode 100644 index 000000000..6d3f922fc --- /dev/null +++ b/tests/fixtures/v1_responses_failed_event.sse @@ -0,0 +1,3 @@ +event: response.failed +data: {"response":{"id":"resp_failed","object":"response","status":"failed","error":{"code":"rate_limit_exceeded","message":"quota reached"}}} + diff --git a/tests/fixtures/v1_responses_stream_after_429.sse b/tests/fixtures/v1_responses_stream_after_429.sse new file mode 100644 index 000000000..5f0ba21ba --- /dev/null +++ b/tests/fixtures/v1_responses_stream_after_429.sse @@ -0,0 +1,6 @@ +event: response.output_text.delta +data: {"delta":"ok"} + +event: response.completed +data: {"response":{"id":"resp_stream_after_429","object":"response","usage":{"input_tokens":4,"output_tokens":1}}} + diff --git a/tests/fixtures/v1_responses_stream_after_refresh.sse b/tests/fixtures/v1_responses_stream_after_refresh.sse new file mode 100644 index 000000000..77d2a77b9 --- /dev/null +++ b/tests/fixtures/v1_responses_stream_after_refresh.sse @@ -0,0 +1,6 @@ +event: response.output_text.delta +data: {"delta":"pong"} + +event: response.completed +data: {"response":{"id":"resp_stream_after_refresh","object":"response","usage":{"input_tokens":4,"output_tokens":2}}} + diff --git a/tests/fixtures/v1_responses_stream_error_event.sse b/tests/fixtures/v1_responses_stream_error_event.sse new file mode 100644 index 000000000..f150de072 --- /dev/null +++ b/tests/fixtures/v1_responses_stream_error_event.sse @@ -0,0 +1,3 @@ +event: error +data: {"error":{"message":"upstream stream failed","type":"server_error","code":"upstream_error"}} + diff --git a/tests/fixtures/v1_responses_stream_failed_event.sse b/tests/fixtures/v1_responses_stream_failed_event.sse new file mode 100644 index 000000000..40bbb861e --- /dev/null +++ b/tests/fixtures/v1_responses_stream_failed_event.sse @@ -0,0 +1,3 @@ +event: response.failed +data: {"response":{"id":"resp_stream_failed","object":"response","status":"failed","error":{"code":"server_error","message":"stream failed"}}} + diff --git a/tests/fixtures/v1_responses_stream_usage.sse b/tests/fixtures/v1_responses_stream_usage.sse new file mode 100644 index 000000000..5ca0ff17e --- /dev/null +++ b/tests/fixtures/v1_responses_stream_usage.sse @@ -0,0 +1,6 @@ +event: response.output_text.delta +data: {"delta":"pong"} + +event: response.completed +data: {"response":{"id":"resp_stream","object":"response","usage":{"input_tokens":3,"output_tokens":5,"input_tokens_details":{"cached_tokens":1}}}} + diff --git a/tests/fixtures/v1_responses_text_deltas_completed.sse b/tests/fixtures/v1_responses_text_deltas_completed.sse new file mode 100644 index 000000000..98ca02dc5 --- /dev/null +++ b/tests/fixtures/v1_responses_text_deltas_completed.sse @@ -0,0 +1,9 @@ +event: response.output_text.delta +data: {"delta":"你好，我是"} + +event: response.output_text.delta +data: {"delta":"一个中文助手。"} + +event: response.completed +data: {"response":{"id":"resp_text","object":"response","status":"completed","output":[],"usage":{"input_tokens":6,"output_tokens":4}}} + diff --git a/tests/fixtures/v1_responses_websocket_completed.json b/tests/fixtures/v1_responses_websocket_completed.json new file mode 100644 index 000000000..355cfee51 --- /dev/null +++ b/tests/fixtures/v1_responses_websocket_completed.json @@ -0,0 +1,11 @@ +{ + "type": "response.completed", + "response": { + "id": "resp_fixture", + "object": "response", + "usage": { + "input_tokens": 3, + "output_tokens": 1 + } + } +} diff --git a/tests/fixtures/v1_responses_websocket_first_account_limited.json b/tests/fixtures/v1_responses_websocket_first_account_limited.json new file mode 100644 index 000000000..e88fea407 --- /dev/null +++ b/tests/fixtures/v1_responses_websocket_first_account_limited.json @@ -0,0 +1 @@ +{"error":{"message":"first account limited"}} diff --git a/tests/fixtures/v1_responses_websocket_history_rate_limited.json b/tests/fixtures/v1_responses_websocket_history_rate_limited.json new file mode 100644 index 000000000..049498d77 --- /dev/null +++ b/tests/fixtures/v1_responses_websocket_history_rate_limited.json @@ -0,0 +1 @@ +{"error":{"message":"history account rate limited"}} diff --git a/tests/fixtures/v1_responses_websocket_rate_limited.json b/tests/fixtures/v1_responses_websocket_rate_limited.json new file mode 100644 index 000000000..711ca5a8e --- /dev/null +++ b/tests/fixtures/v1_responses_websocket_rate_limited.json @@ -0,0 +1 @@ +{"error":{"message":"rate limited"}} diff --git a/tests/fixtures/v1_responses_websocket_second_account_limited.json b/tests/fixtures/v1_responses_websocket_second_account_limited.json new file mode 100644 index 000000000..b3da3c897 --- /dev/null +++ b/tests/fixtures/v1_responses_websocket_second_account_limited.json @@ -0,0 +1 @@ +{"error":{"message":"second account limited"}} diff --git a/tests/fixtures/v1_responses_websocket_token_revoked.json b/tests/fixtures/v1_responses_websocket_token_revoked.json new file mode 100644 index 000000000..f0cda0b80 --- /dev/null +++ b/tests/fixtures/v1_responses_websocket_token_revoked.json @@ -0,0 +1 @@ +{"error":{"code":"token_revoked"}} diff --git a/tests/model_catalog_test.rs b/tests/model_catalog_test.rs index e2eb503e8..a04838004 100644 --- a/tests/model_catalog_test.rs +++ b/tests/model_catalog_test.rs @@ -1,14 +1,14 @@ use std::collections::BTreeMap; use codex_proxy_rs::{ - config::ModelConfig, - models::{ + codex::models::{ catalog::{ BackendModelEntry, BackendReasoningEffort, BackendTruncationPolicy, ModelCatalog, ModelPlanSnapshot, ParsedModelName, }, repository::ModelSnapshotRepository, }, + config::ModelConfig, storage::db::connect_sqlite, }; diff --git a/tests/oauth_refresh_test.rs b/tests/oauth_refresh_test.rs index 4d7e07513..2c765a540 100644 --- a/tests/oauth_refresh_test.rs +++ b/tests/oauth_refresh_test.rs @@ -3,9 +3,8 @@ use wiremock::{ Mock, MockServer, ResponseTemplate, }; -use codex_proxy_rs::auth::{ - oauth::{OAuthClient, OAuthConfig, OAuthError, OpenAiOAuthRefresher}, - refresh::TokenRefresher, +use codex_proxy_rs::codex::oauth::{ + OAuthClient, OAuthConfig, OAuthError, OpenAiOAuthRefresher, TokenRefresher, }; #[tokio::test] diff --git a/tests/refresh_scheduler_test.rs b/tests/refresh_scheduler_test.rs index 073d09f2c..a7a9c762d 100644 --- a/tests/refresh_scheduler_test.rs +++ b/tests/refresh_scheduler_test.rs @@ -1,4 +1,4 @@ -use codex_proxy_rs::accounts::{ +use codex_proxy_rs::codex::accounts::{ model::{Account, AccountStatus}, pool::AccountPool, }; diff --git a/tests/routes_chat_test.rs b/tests/routes_chat_test.rs index b58297ce3..52a9b551b 100644 --- a/tests/routes_chat_test.rs +++ b/tests/routes_chat_test.rs @@ -1,6 +1,6 @@ use serde_json::json; -use codex_proxy_rs::translation::openai_to_codex::{ +use codex_proxy_rs::codex::protocol::openai_to_codex::{ translate_chat_to_codex, ChatCompletionRequest, ChatMessage, }; diff --git a/tests/routes_responses_test.rs b/tests/routes_responses_test.rs index 585db58c1..c68dc09d4 100644 --- a/tests/routes_responses_test.rs +++ b/tests/routes_responses_test.rs @@ -10,19 +10,19 @@ use tower::ServiceExt; use codex_proxy_rs::{ app::build_router, + app::state::AppState, auth::{api_key::ApiKeyHasher, api_key_repository::ClientApiKeyRepository}, + codex::models::{ + catalog::{BackendModelEntry, ModelPlanSnapshot}, + repository::ModelSnapshotRepository, + }, config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, }, - crypto::SecretBox, - models::{ - catalog::{BackendModelEntry, ModelPlanSnapshot}, - repository::ModelSnapshotRepository, - }, - state::AppState, storage::db::connect_sqlite, + utils::crypto::SecretBox, }; fn test_config() -> AppConfig { diff --git a/tests/v1_responses_http_sse_test.rs b/tests/v1_responses_http_sse_test.rs new file mode 100644 index 000000000..040c089c5 --- /dev/null +++ b/tests/v1_responses_http_sse_test.rs @@ -0,0 +1,467 @@ +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use serde_json::json; +use tower::ServiceExt; +use wiremock::{ + matchers::{body_json, header, method, path}, + Mock, MockServer, ResponseTemplate, +}; + +use codex_proxy_rs::codex::cookies::repository::CookieRepository; + +mod common; + +use common::{ + response_json, response_text, + upstream::{build_imported_app, fetch_v1_event_log}, +}; + +const COMPLETED_USAGE_SSE: &str = include_str!("fixtures/v1_responses_completed_usage.sse"); +const COMPLETED_FIELDS_SSE: &str = include_str!("fixtures/v1_responses_completed_fields.sse"); +const COMPLETED_REASONING_INCLUDE_SSE: &str = + include_str!("fixtures/v1_responses_completed_reasoning_include.sse"); +const TEXT_DELTAS_COMPLETED_SSE: &str = + include_str!("fixtures/v1_responses_text_deltas_completed.sse"); +const DONE_ITEM_COMPLETED_SSE: &str = include_str!("fixtures/v1_responses_done_item_completed.sse"); +const STREAM_USAGE_SSE: &str = include_str!("fixtures/v1_responses_stream_usage.sse"); +const DEFAULT_STREAM_SSE: &str = include_str!("fixtures/v1_responses_default_stream.sse"); + +#[tokio::test] +async fn v1_responses_should_use_imported_account_and_record_usage() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .and(header("chatgpt-account-id", "chatgpt-account")) + .and(body_json(json!({ + "model": "gpt-5.5", + "instructions": "", + "input": [], + "stream": true, + "store": false + }))) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .insert_header( + "set-cookie", + "cf_clearance=new; Domain=.chatgpt.com; Path=/", + ) + .set_body_string(COMPLETED_USAGE_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).await; + + let response = imported + .app + .clone() + .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_non_stream") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["id"], "resp_1"); + assert_eq!(body["usage"]["input_tokens"], 7); + 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, 7, 4, 2)); + let cookie_header = CookieRepository::new(imported.pool.clone(), imported.secret_box) + .cookie_header("acct_imported", "chatgpt.com") + .await + .unwrap(); + assert_eq!(cookie_header.as_deref(), Some("cf_clearance=new")); + let event = fetch_v1_event_log(&imported.pool, "req_non_stream").await; + assert_eq!(event.0, "acct_imported"); + assert_eq!(event.1, "/v1/responses"); + assert_eq!(event.2, "gpt-5.5"); + assert_eq!(event.3, 200); + assert_eq!(event.4["stream"], false); + assert_eq!(event.4["usage"]["inputTokens"], 7); +} + +#[tokio::test] +async fn v1_responses_should_forward_parity_fields_and_context_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .and(header("chatgpt-account-id", "chatgpt-account")) + .and(header("x-codex-turn-state", "turn-body")) + .and(header("x-codex-turn-metadata", "meta-header")) + .and(header("x-codex-beta-features", "beta-header")) + .and(header("x-responsesapi-include-timing-metrics", "true")) + .and(header("version", "2026-06-12")) + .and(header("x-codex-window-id", "window-1")) + .and(header("x-codex-parent-thread-id", "parent-1")) + .and(body_json(json!({ + "model": "gpt-5.5", + "instructions": "", + "input": [], + "stream": true, + "store": false, + "reasoning": { + "effort": "high", + "summary": "auto" + }, + "service_tier": "priority", + "tool_choice": { + "type": "function", + "function": {"name": "lookup"} + }, + "parallel_tool_calls": true, + "text": { + "format": { + "type": "json_schema", + "name": "Answer", + "schema": {"type": "object"}, + "strict": true + } + }, + "prompt_cache_key": "pcache", + "include": ["reasoning.encrypted_content"], + "client_metadata": { + "safe": "yes" + } + }))) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(COMPLETED_FIELDS_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).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-codex-turn-state", "turn-header") + .header("x-codex-turn-metadata", "meta-header") + .header("x-codex-beta-features", "beta-header") + .header("x-responsesapi-include-timing-metrics", "true") + .header("version", "2026-06-12") + .header("x-codex-window-id", "window-1") + .header("x-codex-parent-thread-id", "parent-1") + .body(Body::from( + json!({ + "model": "gpt-5.5-fast", + "stream": false, + "input": [], + "reasoning": {"effort": "high"}, + "service_tier": "fast", + "tool_choice": { + "type": "function", + "function": {"name": "lookup"} + }, + "parallel_tool_calls": true, + "text": { + "format": { + "type": "json_schema", + "name": "Answer", + "schema": {"type": "object"}, + "strict": true + } + }, + "prompt_cache_key": "pcache", + "include": ["reasoning.encrypted_content"], + "client_metadata": { + "safe": "yes" + }, + "turnState": "turn-body" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["id"], "resp_fields"); +} + +#[tokio::test] +async fn v1_responses_should_include_encrypted_reasoning_by_default() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(body_json(json!({ + "model": "gpt-5.5", + "instructions": "", + "input": [], + "stream": true, + "store": false, + "reasoning": { + "effort": "high", + "summary": "auto" + }, + "include": ["reasoning.encrypted_content"] + }))) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(COMPLETED_REASONING_INCLUDE_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).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") + .body(Body::from( + r#"{"model":"gpt-5.5","stream":false,"input":[],"reasoning":{"effort":"high"}}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["id"], "resp_reasoning_include"); +} + +#[tokio::test] +async fn v1_responses_should_reconstruct_non_stream_output_text_from_sse_deltas() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(TEXT_DELTAS_COMPLETED_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).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") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["id"], "resp_text"); + assert_eq!( + body["output"][0]["content"][0]["text"], + "你好，我是一个中文助手。" + ); + assert_eq!(body["output_text"], "你好，我是一个中文助手。"); + assert_eq!(body["output"][0]["role"], "assistant"); +} + +#[tokio::test] +async fn v1_responses_should_use_done_output_items_when_completed_output_is_empty() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(DONE_ITEM_COMPLETED_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).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") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["id"], "resp_item"); + assert_eq!(body["output"][0]["content"][0]["text"], "来自 done item"); + assert_eq!(body["output_text"], "来自 done item"); +} + +#[tokio::test] +async fn v1_responses_should_passthrough_stream_and_record_usage_and_log() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(body_json(json!({ + "model": "gpt-5.5", + "instructions": "", + "input": [], + "stream": true, + "store": false + }))) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(STREAM_USAGE_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).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_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.output_text.delta")); + 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, 3, 5, 1)); + let event = fetch_v1_event_log(&imported.pool, "req_stream").await; + assert_eq!(event.0, "acct_imported"); + assert_eq!(event.1, "/v1/responses"); + assert_eq!(event.2, "gpt-5.5"); + assert_eq!(event.3, 200); + assert_eq!(event.4["stream"], true); + assert_eq!(event.4["usage"]["outputTokens"], 5); +} + +#[tokio::test] +async fn v1_responses_should_default_to_streaming_when_stream_is_omitted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(body_json(json!({ + "model": "gpt-5.5", + "instructions": "", + "input": [], + "stream": true, + "store": false + }))) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(DEFAULT_STREAM_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).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") + .body(Body::from(r#"{"model":"gpt-5.5","input":[]}"#)) + .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.output_text.delta")); + assert!(body.contains("event: response.completed")); +} diff --git a/tests/v1_responses_websocket_test.rs b/tests/v1_responses_websocket_test.rs new file mode 100644 index 000000000..42bf149a7 --- /dev/null +++ b/tests/v1_responses_websocket_test.rs @@ -0,0 +1,503 @@ +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use futures::{SinkExt, StreamExt}; +use secrecy::ExposeSecret; +use serde_json::{json, Value}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + sync::oneshot, +}; +use tokio_tungstenite::{ + accept_async, accept_hdr_async, + tungstenite::{ + handshake::server::{Request as WsRequest, Response as WsResponse}, + Message, + }, +}; +use tower::ServiceExt; + +use codex_proxy_rs::codex::accounts::repository::AccountRepository; + +mod common; + +use common::{ + response_text, + upstream::{ + build_imported_app, build_imported_app_with_accounts, + build_imported_app_with_accounts_and_token_refresher, ImportAccount, StaticTokenRefresher, + }, +}; + +const WEBSOCKET_COMPLETED_RESPONSE: &str = + include_str!("fixtures/v1_responses_websocket_completed.json"); +const WEBSOCKET_HISTORY_RATE_LIMITED: &str = + include_str!("fixtures/v1_responses_websocket_history_rate_limited.json"); +const WEBSOCKET_RATE_LIMITED: &str = + include_str!("fixtures/v1_responses_websocket_rate_limited.json"); +const WEBSOCKET_TOKEN_REVOKED: &str = + include_str!("fixtures/v1_responses_websocket_token_revoked.json"); +const WEBSOCKET_FIRST_ACCOUNT_LIMITED: &str = + include_str!("fixtures/v1_responses_websocket_first_account_limited.json"); +const WEBSOCKET_SECOND_ACCOUNT_LIMITED: &str = + include_str!("fixtures/v1_responses_websocket_second_account_limited.json"); + +#[tokio::test] +async fn v1_responses_should_use_websocket_for_previous_response_id_streaming() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (request_tx, request_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let message = websocket.next().await.unwrap().unwrap(); + let request = serde_json::from_str::(&message.into_text().unwrap()).unwrap(); + request_tx.send(request).unwrap(); + websocket + .send(Message::Text( + websocket_completed_response("resp_route_ws", 8, 5).into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let imported = build_imported_app(format!("http://{addr}")).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") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"previous_response_id":"resp_prev"}"#, + )) + .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")); + assert!(body.contains("\"id\":\"resp_route_ws\"")); + let request = request_rx.await.unwrap(); + assert_eq!(request["type"], "response.create"); + assert_eq!(request["previous_response_id"], "resp_prev"); + server.await.unwrap(); +} + +#[tokio::test] +async fn v1_responses_previous_response_id_websocket_429_should_not_retry_different_account() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + reject_next_websocket_upgrade( + &listener, + "Bearer access-a", + 429, + "Too Many Requests", + Some(77), + WEBSOCKET_HISTORY_RATE_LIMITED, + ) + .await; + }); + let imported = build_imported_app_with_accounts( + format!("http://{addr}"), + &[ + ImportAccount { + id: "acct_a", + account_id: "chatgpt-a", + token: "access-a", + refresh_token: "refresh-a", + }, + ImportAccount { + id: "acct_b", + account_id: "chatgpt-b", + token: "access-b", + refresh_token: "refresh-b", + }, + ], + ) + .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_ws_history_429") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"previous_response_id":"resp_prev"}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + server.await.unwrap(); + let account_b_usage = + sqlx::query_as::<_, (i64,)>("select count(*) from account_usage where account_id = ?") + .bind("acct_b") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(account_b_usage.0, 0); +} + +#[tokio::test] +async fn v1_responses_non_stream_previous_response_id_websocket_429_should_not_retry_different_account( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + reject_next_websocket_upgrade( + &listener, + "Bearer access-a", + 429, + "Too Many Requests", + Some(77), + WEBSOCKET_HISTORY_RATE_LIMITED, + ) + .await; + }); + let imported = build_imported_app_with_accounts( + format!("http://{addr}"), + &[ + ImportAccount { + id: "acct_a", + account_id: "chatgpt-a", + token: "access-a", + refresh_token: "refresh-a", + }, + ImportAccount { + id: "acct_b", + account_id: "chatgpt-b", + token: "access-b", + refresh_token: "refresh-b", + }, + ], + ) + .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_ws_history_429_non_stream") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false,"previous_response_id":"resp_prev"}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + server.await.unwrap(); + let account_b_usage = + sqlx::query_as::<_, (i64,)>("select count(*) from account_usage where account_id = ?") + .bind("acct_b") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(account_b_usage.0, 0); +} + +#[tokio::test] +async fn v1_responses_websocket_without_history_should_fallback_and_refresh_fallback_account() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + reject_next_websocket_upgrade( + &listener, + "Bearer access-a", + 429, + "Too Many Requests", + Some(30), + WEBSOCKET_RATE_LIMITED, + ) + .await; + reject_next_websocket_upgrade( + &listener, + "Bearer access-b", + 401, + "Unauthorized", + None, + WEBSOCKET_TOKEN_REVOKED, + ) + .await; + accept_successful_websocket_response(&listener, "Bearer refreshed-b", "resp_ws_fallback") + .await + }); + let imported = build_imported_app_with_accounts_and_token_refresher( + format!("http://{addr}"), + &[ + ImportAccount { + id: "acct_a", + account_id: "chatgpt-a", + token: "access-a", + refresh_token: "refresh-a", + }, + ImportAccount { + id: "acct_b", + account_id: "chatgpt-b", + token: "access-b", + refresh_token: "refresh-b", + }, + ], + StaticTokenRefresher { + access_token: "refreshed-b".to_string(), + refresh_token: None, + }, + ) + .await; + + let response = imported + .app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header( + "authorization", + format!("Bearer {}", imported.client_api_key), + ) + .header("content-type", "application/json") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":true,"use_websocket":true}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_text(response).await; + assert!(body.contains("\"id\":\"resp_ws_fallback\"")); + let websocket_request = server.await.unwrap(); + assert_eq!(websocket_request["type"], "response.create"); + assert!(websocket_request.get("previous_response_id").is_none()); + let repo = AccountRepository::new(imported.pool.clone(), imported.secret_box); + let account_b = repo.get("acct_b").await.unwrap().unwrap(); + assert_eq!(account_b.access_token.expose_secret(), "refreshed-b"); + assert_eq!( + account_b.refresh_token.unwrap().expose_secret(), + "refresh-b" + ); + let usage_a: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_a") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(usage_a, (1, 0, 0)); + let usage_b: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_b") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(usage_b, (1, 3, 1)); +} + +#[tokio::test] +async fn v1_responses_websocket_without_history_should_return_429_when_fallback_accounts_exhausted() +{ + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + reject_next_websocket_upgrade( + &listener, + "Bearer access-a", + 429, + "Too Many Requests", + Some(11), + WEBSOCKET_FIRST_ACCOUNT_LIMITED, + ) + .await; + reject_next_websocket_upgrade( + &listener, + "Bearer access-b", + 429, + "Too Many Requests", + Some(22), + WEBSOCKET_SECOND_ACCOUNT_LIMITED, + ) + .await; + }); + let imported = build_imported_app_with_accounts( + format!("http://{addr}"), + &[ + ImportAccount { + id: "acct_a", + account_id: "chatgpt-a", + token: "access-a", + refresh_token: "refresh-a", + }, + ImportAccount { + id: "acct_b", + account_id: "chatgpt-b", + token: "access-b", + refresh_token: "refresh-b", + }, + ], + ) + .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") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":true,"use_websocket":true}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + server.await.unwrap(); + let usage_a: (i64,) = + sqlx::query_as("select request_count from account_usage where account_id = ?") + .bind("acct_a") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(usage_a.0, 1); + let usage_b: (i64,) = + sqlx::query_as("select request_count from account_usage where account_id = ?") + .bind("acct_b") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(usage_b.0, 1); +} + +async fn reject_next_websocket_upgrade( + listener: &TcpListener, + expected_authorization: &str, + status: u16, + reason: &str, + retry_after_seconds: Option, + body: &str, +) { + let (mut stream, _) = listener.accept().await.unwrap(); + let request = read_http_upgrade_request(&mut stream).await; + assert!(request.starts_with("GET /codex/responses HTTP/1.1")); + assert!( + request.contains(&format!("authorization: {expected_authorization}")), + "unexpected websocket authorization header in request:\n{request}" + ); + let retry_after = retry_after_seconds + .map(|seconds| format!("retry-after: {seconds}\r\n")) + .unwrap_or_default(); + let response = format!( + "HTTP/1.1 {status} {reason}\r\n{retry_after}content-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + stream.write_all(response.as_bytes()).await.unwrap(); +} + +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn accept_successful_websocket_response( + listener: &TcpListener, + expected_authorization: &str, + response_id: &str, +) -> Value { + let (stream, _) = listener.accept().await.unwrap(); + let expected_authorization = expected_authorization.to_string(); + let mut websocket = + accept_hdr_async(stream, move |request: &WsRequest, response: WsResponse| { + assert_eq!( + request + .headers() + .get("authorization") + .and_then(|value| value.to_str().ok()), + Some(expected_authorization.as_str()) + ); + Ok(response) + }) + .await + .unwrap(); + let message = websocket.next().await.unwrap().unwrap(); + let request = serde_json::from_str::(&message.into_text().unwrap()).unwrap(); + websocket + .send(Message::Text( + websocket_completed_response(response_id, 3, 1).into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + request +} + +fn websocket_completed_response( + response_id: &str, + input_tokens: i64, + output_tokens: i64, +) -> String { + let mut value: Value = serde_json::from_str(WEBSOCKET_COMPLETED_RESPONSE).unwrap(); + value["response"]["id"] = Value::String(response_id.to_string()); + value["response"]["usage"]["input_tokens"] = json!(input_tokens); + value["response"]["usage"]["output_tokens"] = json!(output_tokens); + value.to_string() +} + +async fn read_http_upgrade_request(stream: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0u8; 1024]; + loop { + let read = stream.read(&mut buffer).await.unwrap(); + if read == 0 { + break; + } + request.extend_from_slice(&buffer[..read]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + String::from_utf8(request).unwrap() +} diff --git a/tests/v1_upstream_errors_test.rs b/tests/v1_upstream_errors_test.rs new file mode 100644 index 000000000..f694cb7af --- /dev/null +++ b/tests/v1_upstream_errors_test.rs @@ -0,0 +1,192 @@ +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use tower::ServiceExt; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +mod common; + +use common::{ + response_json, response_text, + upstream::{build_imported_app, fetch_v1_event_log}, +}; + +const ERROR_EVENT_SSE: &str = include_str!("fixtures/v1_responses_error_event.sse"); +const FAILED_EVENT_SSE: &str = include_str!("fixtures/v1_responses_failed_event.sse"); +const STREAM_ERROR_EVENT_SSE: &str = include_str!("fixtures/v1_responses_stream_error_event.sse"); +const STREAM_FAILED_EVENT_SSE: &str = include_str!("fixtures/v1_responses_stream_failed_event.sse"); + +#[tokio::test] +async fn v1_responses_non_stream_should_return_upstream_error_when_sse_error_event_arrives() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(ERROR_EVENT_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).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") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + let body = response_json(response).await; + assert_eq!(body["error"]["code"], "upstream_error"); + assert_eq!(body["error"]["type"], "server_error"); + assert!(body["error"]["message"] + .as_str() + .unwrap() + .contains("model overloaded")); +} + +#[tokio::test] +async fn v1_responses_non_stream_should_return_upstream_error_when_response_failed_arrives() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(FAILED_EVENT_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).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") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + let body = response_json(response).await; + assert_eq!(body["error"]["code"], "upstream_error"); + assert!(body["error"]["message"] + .as_str() + .unwrap() + .contains("quota reached")); +} + +#[tokio::test] +async fn v1_responses_stream_should_passthrough_error_event_and_log_failure() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(STREAM_ERROR_EVENT_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).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_stream_error_event") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":true}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_text(response).await; + assert!(body.contains("event: error")); + assert!(body.contains("upstream stream failed")); + let event = fetch_v1_event_log(&imported.pool, "req_stream_error_event").await; + assert_eq!(event.3, 502); + assert_eq!(event.4["failureEvent"], "error"); +} + +#[tokio::test] +async fn v1_responses_stream_should_passthrough_response_failed_and_log_failure() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(STREAM_FAILED_EVENT_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).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_stream_response_failed") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":true}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_text(response).await; + assert!(body.contains("event: response.failed")); + assert!(body.contains("resp_stream_failed")); + let event = fetch_v1_event_log(&imported.pool, "req_stream_response_failed").await; + assert_eq!(event.3, 502); + assert_eq!(event.4["failureEvent"], "response.failed"); +} diff --git a/tests/v1_upstream_fallback_test.rs b/tests/v1_upstream_fallback_test.rs new file mode 100644 index 000000000..c21d9ad02 --- /dev/null +++ b/tests/v1_upstream_fallback_test.rs @@ -0,0 +1,591 @@ +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use secrecy::ExposeSecret; +use serde_json::{json, Value}; +use tower::ServiceExt; +use wiremock::{ + matchers::{header, method, path}, + Mock, MockServer, ResponseTemplate, +}; + +use codex_proxy_rs::{ + codex::accounts::repository::AccountRepository, codex::oauth::RefreshFailure, +}; + +mod common; + +use common::{ + response_json, response_text, + upstream::{ + build_imported_app_with_accounts, build_imported_app_with_refresher, + build_imported_app_with_token_refresher, FailingTokenRefresher, ImportAccount, + }, +}; + +const AFTER_429_SSE: &str = include_str!("fixtures/v1_responses_after_429.sse"); +const STREAM_AFTER_429_SSE: &str = include_str!("fixtures/v1_responses_stream_after_429.sse"); +const AFTER_402_SSE: &str = include_str!("fixtures/v1_responses_after_402.sse"); +const AFTER_403_SSE: &str = include_str!("fixtures/v1_responses_after_403.sse"); +const AFTER_CLOUDFLARE_403_SSE: &str = + include_str!("fixtures/v1_responses_after_cloudflare_403.sse"); +const AFTER_REFRESH_SSE: &str = include_str!("fixtures/v1_responses_after_refresh.sse"); +const STREAM_AFTER_REFRESH_SSE: &str = + include_str!("fixtures/v1_responses_stream_after_refresh.sse"); + +#[tokio::test] +async fn v1_responses_should_retry_next_account_after_429_retry_after() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-a")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("retry-after", "120") + .set_body_json(json!({ + "error": { + "message": "rate limited", + "resets_in_seconds": 120 + } + })), + ) + .up_to_n_times(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-b")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(AFTER_429_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app_with_accounts( + server.uri(), + &[ + ImportAccount { + id: "acct_a", + account_id: "chatgpt-a", + token: "access-a", + refresh_token: "refresh-a", + }, + ImportAccount { + id: "acct_b", + account_id: "chatgpt-b", + token: "access-b", + refresh_token: "refresh-b", + }, + ], + ) + .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") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["id"], "resp_after_429"); + let usage_a: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_a") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(usage_a, (1, 0, 0)); + let usage_b: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_b") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(usage_b, (1, 5, 2)); +} + +#[tokio::test] +async fn v1_responses_stream_should_retry_next_account_after_429_retry_after() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-a")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("retry-after", "90") + .set_body_json(json!({ + "error": {"message": "rate limited"} + })), + ) + .up_to_n_times(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-b")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(STREAM_AFTER_429_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app_with_accounts( + server.uri(), + &[ + ImportAccount { + id: "acct_a", + account_id: "chatgpt-a", + token: "access-a", + refresh_token: "refresh-a", + }, + ImportAccount { + id: "acct_b", + account_id: "chatgpt-b", + token: "access-b", + refresh_token: "refresh-b", + }, + ], + ) + .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") + .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("resp_stream_after_429")); +} + +#[tokio::test] +async fn v1_responses_should_mark_quota_exhausted_after_402_and_retry_next_account() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-a")) + .respond_with(ResponseTemplate::new(402).set_body_json(json!({ + "error": {"message": "quota exhausted"} + }))) + .up_to_n_times(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-b")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(AFTER_402_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app_with_accounts( + server.uri(), + &[ + ImportAccount { + id: "acct_a", + account_id: "chatgpt-a", + token: "access-a", + refresh_token: "refresh-a", + }, + ImportAccount { + id: "acct_b", + account_id: "chatgpt-b", + token: "access-b", + refresh_token: "refresh-b", + }, + ], + ) + .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") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["id"], "resp_after_402"); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_a") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(account_status.0, "quota_exhausted"); +} + +#[tokio::test] +async fn v1_responses_should_mark_banned_after_403_and_retry_next_account() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-a")) + .respond_with(ResponseTemplate::new(403).set_body_json(json!({ + "error": {"message": "account banned"} + }))) + .up_to_n_times(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-b")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(AFTER_403_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app_with_accounts( + server.uri(), + &[ + ImportAccount { + id: "acct_a", + account_id: "chatgpt-a", + token: "access-a", + refresh_token: "refresh-a", + }, + ImportAccount { + id: "acct_b", + account_id: "chatgpt-b", + token: "access-b", + refresh_token: "refresh-b", + }, + ], + ) + .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") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["id"], "resp_after_403"); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_a") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(account_status.0, "banned"); +} + +#[tokio::test] +async fn v1_responses_should_cool_down_cloudflare_403_and_retry_next_account() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-a")) + .respond_with(ResponseTemplate::new(403).set_body_string( + "Just a moment...cf_chl challenge", + )) + .up_to_n_times(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-b")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(AFTER_CLOUDFLARE_403_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app_with_accounts( + server.uri(), + &[ + ImportAccount { + id: "acct_a", + account_id: "chatgpt-a", + token: "access-a", + refresh_token: "refresh-a", + }, + ImportAccount { + id: "acct_b", + account_id: "chatgpt-b", + token: "access-b", + refresh_token: "refresh-b", + }, + ], + ) + .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") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["id"], "resp_after_cf"); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_a") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(account_status.0, "active"); +} + +#[tokio::test] +async fn v1_responses_should_refresh_after_401_and_retry_non_stream() { + let server = MockServer::start().await; + 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(AFTER_REFRESH_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":[],"stream":false}"#, + )) + .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; + 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(STREAM_AFTER_REFRESH_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)); +} + +#[tokio::test] +async fn v1_responses_should_log_refresh_failure_after_401() { + let server = MockServer::start().await; + 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"} + }))) + .mount(&server) + .await; + let imported = build_imported_app_with_token_refresher( + server.uri(), + FailingTokenRefresher { + failure: RefreshFailure::InvalidGrant, + }, + ) + .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_failed") + .body(Body::from(r#"{"model":"gpt-5.5","input":[]}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_imported") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(account_status.0, "expired"); + let refresh_event: (String, String, String) = sqlx::query_as( + "select level, message, metadata_json from event_logs where request_id = ? and kind = 'account.refresh'", + ) + .bind("req_refresh_failed") + .fetch_one(&imported.pool) + .await + .unwrap(); + let metadata: Value = serde_json::from_str(&refresh_event.2).unwrap(); + assert_eq!(refresh_event.0, "warn"); + assert_eq!(refresh_event.1, "account refresh failed after upstream 401"); + assert_eq!(metadata["failure"], "invalidGrant"); + assert_eq!(metadata["accountStatus"], "expired"); + assert_eq!(metadata["trigger"], "upstream_401"); +} diff --git a/tests/v1_upstream_route_test.rs b/tests/v1_upstream_route_test.rs deleted file mode 100644 index 87a51d674..000000000 --- a/tests/v1_upstream_route_test.rs +++ /dev/null @@ -1,2115 +0,0 @@ -use std::collections::BTreeMap; - -use async_trait::async_trait; -use axum::{ - body::{to_bytes, Body}, - http::{Request, StatusCode}, -}; -use chrono::Utc; -use futures::{SinkExt, StreamExt}; -use secrecy::ExposeSecret; -use serde_json::{json, Value}; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, - sync::oneshot, -}; -use tokio_tungstenite::{ - accept_async, accept_hdr_async, - tungstenite::{ - handshake::server::{Request as WsRequest, Response as WsResponse}, - Message, - }, -}; -use tower::ServiceExt; -use wiremock::{ - matchers::{body_json, header, method, path}, - Mock, MockServer, ResponseTemplate, -}; - -use codex_proxy_rs::{ - accounts::repository::AccountRepository, - app::build_router, - 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, - UsageStatsConfig, - }, - cookies::repository::CookieRepository, - crypto::SecretBox, - state::AppState, - storage::db::connect_sqlite, -}; - -struct ImportedApp { - app: axum::Router, - pool: sqlx::SqlitePool, - secret_box: SecretBox, - client_api_key: String, - _tempdir: tempfile::TempDir, -} - -#[derive(Clone)] -struct StaticTokenRefresher { - access_token: String, - refresh_token: Option, -} - -#[derive(Clone)] -struct FailingTokenRefresher { - failure: RefreshFailure, -} - -struct ImportAccount { - id: &'static str, - account_id: &'static str, - token: &'static str, - refresh_token: &'static str, -} - -#[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(), - }) - } -} - -#[async_trait] -impl TokenRefresher for FailingTokenRefresher { - async fn refresh(&self, _refresh_token: &str) -> Result { - Err(self.failure) - } -} - -fn test_config(database_url: String, base_url: String) -> AppConfig { - AppConfig { - server: ServerConfig { - host: "127.0.0.1".to_string(), - port: 0, - }, - api: ApiConfig { base_url }, - 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, - }, - } -} - -async fn build_imported_app(base_url: String) -> ImportedApp { - build_imported_app_with_accounts( - base_url, - &[ImportAccount { - id: "acct_imported", - account_id: "chatgpt-account", - token: "access-secret", - refresh_token: "refresh-secret", - }], - ) - .await -} - -async fn build_imported_app_with_accounts( - base_url: String, - accounts: &[ImportAccount], -) -> ImportedApp { - let tempdir = tempfile::tempdir().unwrap(); - let db = tempdir.path().join("v1-upstream.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([21u8; 32]); - let hasher = ApiKeyHasher::new([22u8; 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_and_api_key_hasher( - test_config(url, base_url), - pool.clone(), - secret_box.clone(), - hasher, - )); - let accounts = accounts - .iter() - .map(|account| { - json!({ - "id": account.id, - "accountId": account.account_id, - "token": account.token, - "refreshToken": account.refresh_token, - "status": "active" - }) - }) - .collect::>(); - - 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": accounts }).to_string())) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(import_response.status(), StatusCode::OK); - - ImportedApp { - app, - pool, - secret_box, - client_api_key: generated.plaintext, - _tempdir: tempdir, - } -} - -async fn build_imported_app_with_refresher(base_url: String, access_token: &str) -> ImportedApp { - build_imported_app_with_token_refresher( - base_url, - StaticTokenRefresher { - access_token: access_token.to_string(), - refresh_token: None, - }, - ) - .await -} - -async fn build_imported_app_with_token_refresher( - base_url: String, - token_refresher: C, -) -> ImportedApp -where - C: TokenRefresher, -{ - build_imported_app_with_accounts_and_token_refresher( - base_url, - &[ImportAccount { - id: "acct_imported", - account_id: "chatgpt-account", - token: "access-secret", - refresh_token: "refresh-secret", - }], - token_refresher, - ) - .await -} - -async fn build_imported_app_with_accounts_and_token_refresher( - base_url: String, - accounts: &[ImportAccount], - token_refresher: C, -) -> ImportedApp -where - C: TokenRefresher, -{ - 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, - token_refresher, - ), - ); - let accounts = accounts - .iter() - .map(|account| { - json!({ - "id": account.id, - "accountId": account.account_id, - "token": account.token, - "refreshToken": account.refresh_token, - "status": "active" - }) - }) - .collect::>(); - - 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": accounts }).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; - let sse_body = concat!( - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"usage\":{\"input_tokens\":7,\"output_tokens\":4,\"input_tokens_details\":{\"cached_tokens\":2}}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-secret")) - .and(header("chatgpt-account-id", "chatgpt-account")) - .and(body_json(json!({ - "model": "gpt-5.5", - "instructions": "", - "input": [], - "stream": true, - "store": false - }))) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .insert_header( - "set-cookie", - "cf_clearance=new; Domain=.chatgpt.com; Path=/", - ) - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .clone() - .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_non_stream") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_1"); - assert_eq!(body["usage"]["input_tokens"], 7); - 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, 7, 4, 2)); - let cookie_header = CookieRepository::new(imported.pool.clone(), imported.secret_box) - .cookie_header("acct_imported", "chatgpt.com") - .await - .unwrap(); - assert_eq!(cookie_header.as_deref(), Some("cf_clearance=new")); - let event = fetch_v1_event_log(&imported.pool, "req_non_stream").await; - assert_eq!(event.0, "acct_imported"); - assert_eq!(event.1, "/v1/responses"); - assert_eq!(event.2, "gpt-5.5"); - assert_eq!(event.3, 200); - assert_eq!(event.4["stream"], false); - assert_eq!(event.4["usage"]["inputTokens"], 7); -} - -#[tokio::test] -async fn v1_responses_should_retry_next_account_after_429_retry_after() { - let server = MockServer::start().await; - let success_sse = concat!( - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_after_429\",\"object\":\"response\",\"usage\":{\"input_tokens\":5,\"output_tokens\":2}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with( - ResponseTemplate::new(429) - .insert_header("retry-after", "120") - .set_body_json(json!({ - "error": { - "message": "rate limited", - "resets_in_seconds": 120 - } - })), - ) - .up_to_n_times(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-b")) - .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_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_b", - account_id: "chatgpt-b", - token: "access-b", - refresh_token: "refresh-b", - }, - ], - ) - .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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_429"); - let usage_a: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", - ) - .bind("acct_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(usage_a, (1, 0, 0)); - let usage_b: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", - ) - .bind("acct_b") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(usage_b, (1, 5, 2)); -} - -#[tokio::test] -async fn v1_responses_stream_should_retry_next_account_after_429_retry_after() { - let server = MockServer::start().await; - let success_sse = concat!( - "event: response.output_text.delta\n", - "data: {\"delta\":\"ok\"}\n", - "\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_stream_after_429\",\"object\":\"response\",\"usage\":{\"input_tokens\":4,\"output_tokens\":1}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with( - ResponseTemplate::new(429) - .insert_header("retry-after", "90") - .set_body_json(json!({ - "error": {"message": "rate limited"} - })), - ) - .up_to_n_times(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-b")) - .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_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_b", - account_id: "chatgpt-b", - token: "access-b", - refresh_token: "refresh-b", - }, - ], - ) - .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") - .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("resp_stream_after_429")); -} - -#[tokio::test] -async fn v1_responses_should_mark_quota_exhausted_after_402_and_retry_next_account() { - let server = MockServer::start().await; - let success_sse = concat!( - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_after_402\",\"object\":\"response\",\"usage\":{\"input_tokens\":6,\"output_tokens\":3}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(402).set_body_json(json!({ - "error": {"message": "quota exhausted"} - }))) - .up_to_n_times(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-b")) - .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_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_b", - account_id: "chatgpt-b", - token: "access-b", - refresh_token: "refresh-b", - }, - ], - ) - .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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_402"); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "quota_exhausted"); -} - -#[tokio::test] -async fn v1_responses_should_mark_banned_after_403_and_retry_next_account() { - let server = MockServer::start().await; - let success_sse = concat!( - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_after_403\",\"object\":\"response\",\"usage\":{\"input_tokens\":6,\"output_tokens\":3}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(403).set_body_json(json!({ - "error": {"message": "account banned"} - }))) - .up_to_n_times(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-b")) - .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_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_b", - account_id: "chatgpt-b", - token: "access-b", - refresh_token: "refresh-b", - }, - ], - ) - .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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_403"); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "banned"); -} - -#[tokio::test] -async fn v1_responses_should_cool_down_cloudflare_403_and_retry_next_account() { - let server = MockServer::start().await; - let success_sse = concat!( - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_after_cf\",\"object\":\"response\",\"usage\":{\"input_tokens\":6,\"output_tokens\":3}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(403).set_body_string( - "Just a moment...cf_chl challenge", - )) - .up_to_n_times(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-b")) - .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_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_b", - account_id: "chatgpt-b", - token: "access-b", - refresh_token: "refresh-b", - }, - ], - ) - .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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_cf"); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "active"); -} - -#[tokio::test] -async fn v1_responses_should_forward_parity_fields_and_context_headers() { - let server = MockServer::start().await; - let sse_body = concat!( - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_fields\",\"object\":\"response\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-secret")) - .and(header("chatgpt-account-id", "chatgpt-account")) - .and(header("x-codex-turn-state", "turn-body")) - .and(header("x-codex-turn-metadata", "meta-header")) - .and(header("x-codex-beta-features", "beta-header")) - .and(header("x-responsesapi-include-timing-metrics", "true")) - .and(header("version", "2026-06-12")) - .and(header("x-codex-window-id", "window-1")) - .and(header("x-codex-parent-thread-id", "parent-1")) - .and(body_json(json!({ - "model": "gpt-5.5", - "instructions": "", - "input": [], - "stream": true, - "store": false, - "reasoning": { - "effort": "high", - "summary": "auto" - }, - "service_tier": "priority", - "tool_choice": { - "type": "function", - "function": {"name": "lookup"} - }, - "parallel_tool_calls": true, - "text": { - "format": { - "type": "json_schema", - "name": "Answer", - "schema": {"type": "object"}, - "strict": true - } - }, - "prompt_cache_key": "pcache", - "include": ["reasoning.encrypted_content"], - "client_metadata": { - "safe": "yes" - } - }))) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).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-codex-turn-state", "turn-header") - .header("x-codex-turn-metadata", "meta-header") - .header("x-codex-beta-features", "beta-header") - .header("x-responsesapi-include-timing-metrics", "true") - .header("version", "2026-06-12") - .header("x-codex-window-id", "window-1") - .header("x-codex-parent-thread-id", "parent-1") - .body(Body::from( - json!({ - "model": "gpt-5.5-fast", - "stream": false, - "input": [], - "reasoning": {"effort": "high"}, - "service_tier": "fast", - "tool_choice": { - "type": "function", - "function": {"name": "lookup"} - }, - "parallel_tool_calls": true, - "text": { - "format": { - "type": "json_schema", - "name": "Answer", - "schema": {"type": "object"}, - "strict": true - } - }, - "prompt_cache_key": "pcache", - "include": ["reasoning.encrypted_content"], - "client_metadata": { - "safe": "yes" - }, - "turnState": "turn-body" - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_fields"); -} - -#[tokio::test] -async fn v1_responses_should_include_encrypted_reasoning_by_default() { - let server = MockServer::start().await; - let sse_body = concat!( - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_reasoning_include\",\"object\":\"response\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(body_json(json!({ - "model": "gpt-5.5", - "instructions": "", - "input": [], - "stream": true, - "store": false, - "reasoning": { - "effort": "high", - "summary": "auto" - }, - "include": ["reasoning.encrypted_content"] - }))) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).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") - .body(Body::from( - r#"{"model":"gpt-5.5","stream":false,"input":[],"reasoning":{"effort":"high"}}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_reasoning_include"); -} - -#[tokio::test] -async fn v1_responses_should_reconstruct_non_stream_output_text_from_sse_deltas() { - let server = MockServer::start().await; - let sse_body = concat!( - "event: response.output_text.delta\n", - "data: {\"delta\":\"你好，我是\"}\n", - "\n", - "event: response.output_text.delta\n", - "data: {\"delta\":\"一个中文助手。\"}\n", - "\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_text\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":6,\"output_tokens\":4}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_text"); - assert_eq!( - body["output"][0]["content"][0]["text"], - "你好，我是一个中文助手。" - ); - assert_eq!(body["output_text"], "你好，我是一个中文助手。"); - assert_eq!(body["output"][0]["role"], "assistant"); -} - -#[tokio::test] -async fn v1_responses_should_use_done_output_items_when_completed_output_is_empty() { - let server = MockServer::start().await; - let sse_body = concat!( - "event: response.output_item.done\n", - "data: {\"item\":{\"type\":\"message\",\"status\":\"completed\",\"role\":\"assistant\",\"content\":[{\"type\":\"output_text\",\"text\":\"来自 done item\"}]}}\n", - "\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_item\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":5,\"output_tokens\":3}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_item"); - assert_eq!(body["output"][0]["content"][0]["text"], "来自 done item"); - assert_eq!(body["output_text"], "来自 done item"); -} - -#[tokio::test] -async fn v1_responses_non_stream_should_return_upstream_error_when_sse_error_event_arrives() { - let server = MockServer::start().await; - let sse_body = concat!( - "event: error\n", - "data: {\"error\":{\"message\":\"model overloaded\",\"type\":\"server_error\",\"code\":\"server_error\"}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_GATEWAY); - let body = response_json(response).await; - assert_eq!(body["error"]["code"], "upstream_error"); - assert_eq!(body["error"]["type"], "server_error"); - assert!(body["error"]["message"] - .as_str() - .unwrap() - .contains("model overloaded")); -} - -#[tokio::test] -async fn v1_responses_non_stream_should_return_upstream_error_when_response_failed_arrives() { - let server = MockServer::start().await; - let sse_body = concat!( - "event: response.failed\n", - "data: {\"response\":{\"id\":\"resp_failed\",\"object\":\"response\",\"status\":\"failed\",\"error\":{\"code\":\"rate_limit_exceeded\",\"message\":\"quota reached\"}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_GATEWAY); - let body = response_json(response).await; - assert_eq!(body["error"]["code"], "upstream_error"); - assert!(body["error"]["message"] - .as_str() - .unwrap() - .contains("quota reached")); -} - -#[tokio::test] -async fn v1_responses_should_passthrough_stream_and_record_usage_and_log() { - let server = MockServer::start().await; - let sse_body = concat!( - "event: response.output_text.delta\n", - "data: {\"delta\":\"pong\"}\n", - "\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_stream\",\"object\":\"response\",\"usage\":{\"input_tokens\":3,\"output_tokens\":5,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(body_json(json!({ - "model": "gpt-5.5", - "instructions": "", - "input": [], - "stream": true, - "store": false - }))) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).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_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.output_text.delta")); - 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, 3, 5, 1)); - let event = fetch_v1_event_log(&imported.pool, "req_stream").await; - assert_eq!(event.0, "acct_imported"); - assert_eq!(event.1, "/v1/responses"); - assert_eq!(event.2, "gpt-5.5"); - assert_eq!(event.3, 200); - assert_eq!(event.4["stream"], true); - assert_eq!(event.4["usage"]["outputTokens"], 5); -} - -#[tokio::test] -async fn v1_responses_stream_should_passthrough_error_event_and_log_failure() { - let server = MockServer::start().await; - let sse_body = concat!( - "event: error\n", - "data: {\"error\":{\"message\":\"upstream stream failed\",\"type\":\"server_error\",\"code\":\"upstream_error\"}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).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_stream_error_event") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":true}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_text(response).await; - assert!(body.contains("event: error")); - assert!(body.contains("upstream stream failed")); - let event = fetch_v1_event_log(&imported.pool, "req_stream_error_event").await; - assert_eq!(event.3, 502); - assert_eq!(event.4["failureEvent"], "error"); -} - -#[tokio::test] -async fn v1_responses_stream_should_passthrough_response_failed_and_log_failure() { - let server = MockServer::start().await; - let sse_body = concat!( - "event: response.failed\n", - "data: {\"response\":{\"id\":\"resp_stream_failed\",\"object\":\"response\",\"status\":\"failed\",\"error\":{\"code\":\"server_error\",\"message\":\"stream failed\"}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).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_stream_response_failed") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":true}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_text(response).await; - assert!(body.contains("event: response.failed")); - assert!(body.contains("resp_stream_failed")); - let event = fetch_v1_event_log(&imported.pool, "req_stream_response_failed").await; - assert_eq!(event.3, 502); - assert_eq!(event.4["failureEvent"], "response.failed"); -} - -#[tokio::test] -async fn v1_responses_should_default_to_streaming_when_stream_is_omitted() { - let server = MockServer::start().await; - let sse_body = concat!( - "event: response.output_text.delta\n", - "data: {\"delta\":\"pong\"}\n", - "\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_default_stream\",\"object\":\"response\",\"usage\":{\"input_tokens\":3,\"output_tokens\":5}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(body_json(json!({ - "model": "gpt-5.5", - "instructions": "", - "input": [], - "stream": true, - "store": false - }))) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).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") - .body(Body::from(r#"{"model":"gpt-5.5","input":[]}"#)) - .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.output_text.delta")); - assert!(body.contains("event: response.completed")); -} - -#[tokio::test] -async fn v1_responses_should_use_websocket_for_previous_response_id_streaming() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let (request_tx, request_rx) = oneshot::channel(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_async(stream).await.unwrap(); - let message = websocket.next().await.unwrap().unwrap(); - let request = serde_json::from_str::(&message.into_text().unwrap()).unwrap(); - request_tx.send(request).unwrap(); - websocket - .send(Message::Text( - json!({ - "type": "response.completed", - "response": { - "id": "resp_route_ws", - "object": "response", - "usage": { - "input_tokens": 8, - "output_tokens": 5 - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - let imported = build_imported_app(format!("http://{addr}")).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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"previous_response_id":"resp_prev"}"#, - )) - .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")); - assert!(body.contains("\"id\":\"resp_route_ws\"")); - let request = request_rx.await.unwrap(); - assert_eq!(request["type"], "response.create"); - assert_eq!(request["previous_response_id"], "resp_prev"); - server.await.unwrap(); -} - -#[tokio::test] -async fn v1_responses_previous_response_id_websocket_429_should_not_retry_different_account() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - reject_next_websocket_upgrade( - &listener, - "Bearer access-a", - 429, - "Too Many Requests", - Some(77), - r#"{"error":{"message":"history account rate limited"}}"#, - ) - .await; - }); - let imported = build_imported_app_with_accounts( - format!("http://{addr}"), - &[ - ImportAccount { - id: "acct_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_b", - account_id: "chatgpt-b", - token: "access-b", - refresh_token: "refresh-b", - }, - ], - ) - .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_ws_history_429") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"previous_response_id":"resp_prev"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); - server.await.unwrap(); - let account_b_usage = - sqlx::query_as::<_, (i64,)>("select count(*) from account_usage where account_id = ?") - .bind("acct_b") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_b_usage.0, 0); -} - -#[tokio::test] -async fn v1_responses_non_stream_previous_response_id_websocket_429_should_not_retry_different_account( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - reject_next_websocket_upgrade( - &listener, - "Bearer access-a", - 429, - "Too Many Requests", - Some(77), - r#"{"error":{"message":"history account rate limited"}}"#, - ) - .await; - }); - let imported = build_imported_app_with_accounts( - format!("http://{addr}"), - &[ - ImportAccount { - id: "acct_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_b", - account_id: "chatgpt-b", - token: "access-b", - refresh_token: "refresh-b", - }, - ], - ) - .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_ws_history_429_non_stream") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false,"previous_response_id":"resp_prev"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); - server.await.unwrap(); - let account_b_usage = - sqlx::query_as::<_, (i64,)>("select count(*) from account_usage where account_id = ?") - .bind("acct_b") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_b_usage.0, 0); -} - -#[tokio::test] -async fn v1_responses_websocket_without_history_should_fallback_and_refresh_fallback_account() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - reject_next_websocket_upgrade( - &listener, - "Bearer access-a", - 429, - "Too Many Requests", - Some(30), - r#"{"error":{"message":"rate limited"}}"#, - ) - .await; - reject_next_websocket_upgrade( - &listener, - "Bearer access-b", - 401, - "Unauthorized", - None, - r#"{"error":{"code":"token_revoked"}}"#, - ) - .await; - accept_successful_websocket_response(&listener, "Bearer refreshed-b", "resp_ws_fallback") - .await - }); - let imported = build_imported_app_with_accounts_and_token_refresher( - format!("http://{addr}"), - &[ - ImportAccount { - id: "acct_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_b", - account_id: "chatgpt-b", - token: "access-b", - refresh_token: "refresh-b", - }, - ], - StaticTokenRefresher { - access_token: "refreshed-b".to_string(), - refresh_token: None, - }, - ) - .await; - - let response = imported - .app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/responses") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":true,"use_websocket":true}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_text(response).await; - assert!(body.contains("\"id\":\"resp_ws_fallback\"")); - let websocket_request = server.await.unwrap(); - assert_eq!(websocket_request["type"], "response.create"); - assert!(websocket_request.get("previous_response_id").is_none()); - let repo = AccountRepository::new(imported.pool.clone(), imported.secret_box); - let account_b = repo.get("acct_b").await.unwrap().unwrap(); - assert_eq!(account_b.access_token.expose_secret(), "refreshed-b"); - assert_eq!( - account_b.refresh_token.unwrap().expose_secret(), - "refresh-b" - ); - let usage_a: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", - ) - .bind("acct_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(usage_a, (1, 0, 0)); - let usage_b: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", - ) - .bind("acct_b") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(usage_b, (1, 3, 1)); -} - -#[tokio::test] -async fn v1_responses_websocket_without_history_should_return_429_when_fallback_accounts_exhausted() -{ - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - reject_next_websocket_upgrade( - &listener, - "Bearer access-a", - 429, - "Too Many Requests", - Some(11), - r#"{"error":{"message":"first account limited"}}"#, - ) - .await; - reject_next_websocket_upgrade( - &listener, - "Bearer access-b", - 429, - "Too Many Requests", - Some(22), - r#"{"error":{"message":"second account limited"}}"#, - ) - .await; - }); - let imported = build_imported_app_with_accounts( - format!("http://{addr}"), - &[ - ImportAccount { - id: "acct_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_b", - account_id: "chatgpt-b", - token: "access-b", - refresh_token: "refresh-b", - }, - ], - ) - .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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":true,"use_websocket":true}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); - server.await.unwrap(); - let usage_a: (i64,) = - sqlx::query_as("select request_count from account_usage where account_id = ?") - .bind("acct_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(usage_a.0, 1); - let usage_b: (i64,) = - sqlx::query_as("select request_count from account_usage where account_id = ?") - .bind("acct_b") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(usage_b.0, 1); -} - -#[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":[],"stream":false}"#, - )) - .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)); -} - -#[tokio::test] -async fn v1_responses_should_log_refresh_failure_after_401() { - let server = MockServer::start().await; - 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"} - }))) - .mount(&server) - .await; - let imported = build_imported_app_with_token_refresher( - server.uri(), - FailingTokenRefresher { - failure: RefreshFailure::InvalidGrant, - }, - ) - .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_failed") - .body(Body::from(r#"{"model":"gpt-5.5","input":[]}"#)) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_imported") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "expired"); - let refresh_event: (String, String, String) = sqlx::query_as( - "select level, message, metadata_json from event_logs where request_id = ? and kind = 'account.refresh'", - ) - .bind("req_refresh_failed") - .fetch_one(&imported.pool) - .await - .unwrap(); - let metadata: Value = serde_json::from_str(&refresh_event.2).unwrap(); - assert_eq!(refresh_event.0, "warn"); - assert_eq!(refresh_event.1, "account refresh failed after upstream 401"); - assert_eq!(metadata["failure"], "invalidGrant"); - assert_eq!(metadata["accountStatus"], "expired"); - assert_eq!(metadata["trigger"], "upstream_401"); -} - -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 reject_next_websocket_upgrade( - listener: &TcpListener, - expected_authorization: &str, - status: u16, - reason: &str, - retry_after_seconds: Option, - body: &str, -) { - let (mut stream, _) = listener.accept().await.unwrap(); - let request = read_http_upgrade_request(&mut stream).await; - assert!(request.starts_with("GET /codex/responses HTTP/1.1")); - assert!( - request.contains(&format!("authorization: {expected_authorization}")), - "unexpected websocket authorization header in request:\n{request}" - ); - let retry_after = retry_after_seconds - .map(|seconds| format!("retry-after: {seconds}\r\n")) - .unwrap_or_default(); - let response = format!( - "HTTP/1.1 {status} {reason}\r\n{retry_after}content-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", - body.len() - ); - stream.write_all(response.as_bytes()).await.unwrap(); -} - -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn accept_successful_websocket_response( - listener: &TcpListener, - expected_authorization: &str, - response_id: &str, -) -> Value { - let (stream, _) = listener.accept().await.unwrap(); - let expected_authorization = expected_authorization.to_string(); - let mut websocket = - accept_hdr_async(stream, move |request: &WsRequest, response: WsResponse| { - assert_eq!( - request - .headers() - .get("authorization") - .and_then(|value| value.to_str().ok()), - Some(expected_authorization.as_str()) - ); - Ok(response) - }) - .await - .unwrap(); - let message = websocket.next().await.unwrap().unwrap(); - let request = serde_json::from_str::(&message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - json!({ - "type": "response.completed", - "response": { - "id": response_id, - "object": "response", - "usage": { - "input_tokens": 3, - "output_tokens": 1 - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - request -} - -async fn read_http_upgrade_request(stream: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0u8; 1024]; - loop { - let read = stream.read(&mut buffer).await.unwrap(); - if read == 0 { - break; - } - request.extend_from_slice(&buffer[..read]); - if request.windows(4).any(|window| window == b"\r\n\r\n") { - break; - } - } - String::from_utf8(request).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() -} - -async fn response_text(response: axum::response::Response) -> String { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - String::from_utf8(bytes.to_vec()).unwrap() -} - -async fn fetch_v1_event_log( - pool: &sqlx::SqlitePool, - request_id: &str, -) -> (String, String, String, i64, Value) { - let row: (String, String, String, i64, String) = sqlx::query_as( - "select account_id, route, model, status_code, metadata_json from event_logs where request_id = ? and kind = 'v1.response'", - ) - .bind(request_id) - .fetch_one(pool) - .await - .unwrap(); - ( - row.0, - row.1, - row.2, - row.3, - serde_json::from_str(&row.4).unwrap(), - ) -}