From 01b538c416da0de1efb888df775ddea3267a668b Mon Sep 17 00:00:00 2001 From: zyycn Date: Fri, 19 Jun 2026 20:53:26 +0800 Subject: [PATCH] docs: add architecture overview Co-authored-by: Codex --- .gitignore | 32 +- Cargo.lock | 215 +- Cargo.toml | 45 +- config.yaml | 8 +- crates/adapters/Cargo.toml | 38 + crates/adapters/src/codex/client.rs | 1497 +++ crates/adapters/src/codex/fingerprint.rs | 477 + crates/adapters/src/codex/mod.rs | 15 + crates/adapters/src/codex/models.rs | 3 + .../adapters/src/codex/websocket/connect.rs | 1078 +++ .../adapters/src/codex}/websocket/deflate.rs | 92 +- crates/adapters/src/codex/websocket/mod.rs | 6 + .../adapters/src/codex/websocket/opening.rs | 56 + .../adapters/src/codex}/websocket/pool.rs | 328 +- crates/adapters/src/lib.rs | 7 + crates/adapters/src/oauth/mod.rs | 4 + crates/adapters/src/oauth/openai.rs | 190 + crates/adapters/src/sqlite/account_tokens.rs | 10 + crates/adapters/src/sqlite/account_usage.rs | 21 + crates/adapters/src/sqlite/accounts.rs | 1671 ++++ .../adapters/src/sqlite/admin_sessions.rs | 74 +- crates/adapters/src/sqlite/client_keys.rs | 276 + .../adapters/src/sqlite/cookies.rs | 173 +- crates/adapters/src/sqlite/events.rs | 286 + crates/adapters/src/sqlite/mod.rs | 22 + crates/adapters/src/sqlite/models.rs | 102 + crates/adapters/src/sqlite/refresh_leases.rs | 80 + .../adapters/src/sqlite/session_affinity.rs | 170 + crates/adapters/tests/account_repository.rs | 322 + crates/adapters/tests/admin_sessions.rs | 82 + crates/adapters/tests/client_keys.rs | 89 + crates/adapters/tests/codex.rs | 3641 +++++++ crates/adapters/tests/cookies.rs | 93 + crates/adapters/tests/events.rs | 71 + crates/adapters/tests/models.rs | 93 + .../adapters/tests/oauth.rs | 24 +- crates/adapters/tests/refresh_leases.rs | 68 + crates/adapters/tests/session_affinity.rs | 102 + crates/assets/Cargo.toml | 19 + crates/assets/src/headers.rs | 60 + crates/assets/src/lib.rs | 4 + crates/assets/src/router.rs | 87 + crates/assets/tests/assets.rs | 94 + crates/core/Cargo.toml | 28 + crates/core/src/accounts/cloudflare.rs | 127 + crates/core/src/accounts/cookies.rs | 21 + .../codex => crates/core/src}/accounts/jwt.rs | 24 +- crates/core/src/accounts/lifecycle.rs | 20 + crates/core/src/accounts/mod.rs | 20 + .../core/src}/accounts/model.rs | 44 + .../core/src}/accounts/pool.rs | 104 +- crates/core/src/accounts/ports.rs | 100 + crates/core/src/accounts/service.rs | 44 + crates/core/src/accounts/usage.rs | 41 + crates/core/src/admin/auth.rs | 30 + crates/core/src/admin/client_keys.rs | 26 + crates/core/src/admin/mod.rs | 6 + crates/core/src/admin/ports.rs | 25 + crates/core/src/admin/settings.rs | 339 + crates/core/src/auth/mod.rs | 8 + crates/core/src/auth/oauth.rs | 417 + crates/core/src/auth/ports.rs | 30 + crates/core/src/auth/session.rs | 24 + crates/core/src/error.rs | 23 + crates/core/src/events/mod.rs | 8 + .../core/src/events/model.rs | 46 +- crates/core/src/events/ports.rs | 27 + crates/core/src/events/service.rs | 21 + .../core/src/gateway/conversation.rs | 6 +- .../core/src/gateway/fingerprint.rs | 16 + crates/core/src/gateway/installation.rs | 50 + crates/core/src/gateway/mod.rs | 10 + crates/core/src/gateway/ports.rs | 42 + crates/core/src/lib.rs | 14 + .../core/src}/models/catalog.rs | 137 +- crates/core/src/models/mod.rs | 10 + crates/core/src/models/model.rs | 183 + crates/core/src/models/ports.rs | 33 + crates/core/src/models/service.rs | 233 + crates/core/src/protocol/codex/chat.rs | 11 + .../core/src/protocol/codex/events.rs | 379 +- crates/core/src/protocol/codex/mod.rs | 14 + .../core/src/protocol/codex/responses.rs | 40 + .../core/src/protocol/codex/schema.rs | 4 + .../core/src/protocol/codex}/sse.rs | 54 +- .../core/src/protocol/codex/websocket.rs | 1369 +-- crates/core/src/protocol/mod.rs | 4 + crates/core/src/protocol/openai/chat.rs | 620 ++ crates/core/src/protocol/openai/errors.rs | 15 + crates/core/src/protocol/openai/mod.rs | 6 + crates/core/src/protocol/openai/models.rs | 25 + crates/core/src/protocol/openai/responses.rs | 745 ++ crates/core/src/serving/affinity.rs | 238 + crates/core/src/serving/chat.rs | 13 + crates/core/src/serving/errors.rs | 17 + crates/core/src/serving/fallback.rs | 16 + .../core/src/serving}/implicit_resume.rs | 69 +- crates/core/src/serving/mod.rs | 26 + .../core/src/serving}/quota.rs | 156 +- .../core/src/serving}/reasoning_replay.rs | 202 +- crates/core/src/serving/recovery.rs | 6 + crates/core/src/serving/responses.rs | 129 + crates/core/src/serving/routing.rs | 8 + crates/core/src/serving/stream.rs | 10 + crates/core/src/serving/usage.rs | 8 + crates/core/src/usage/mod.rs | 8 + crates/core/src/usage/model.rs | 28 + crates/core/src/usage/ports.rs | 27 + crates/core/src/usage/service.rs | 34 + .../core/tests/account_pool.rs | 72 +- crates/core/tests/admin_settings.rs | 85 + crates/core/tests/auth.rs | 59 + crates/core/tests/models.rs | 302 + crates/core/tests/protocol.rs | 1774 ++++ crates/core/tests/session_affinity.rs | 106 + crates/platform/Cargo.toml | 41 + crates/platform/src/config/loader.rs | 69 + crates/platform/src/config/mod.rs | 28 + {src => crates/platform/src}/config/types.rs | 63 + crates/platform/src/crypto/hash.rs | 25 + crates/platform/src/crypto/mod.rs | 7 + .../platform/src}/crypto/secret_box.rs | 18 + .../platform/src/identity/admin_password.rs | 4 +- .../platform/src}/identity/client_key.rs | 13 +- .../platform/src/identity/mod.rs | 14 + .../platform/src/json/mod.rs | 8 + .../platform/src/json}/pagination.rs | 6 + crates/platform/src/lib.rs | 10 + crates/platform/src/logging/mod.rs | 5 + .../platform/src}/logging/rotation.rs | 11 + crates/platform/src/storage/mod.rs | 11 + crates/platform/src/storage/paths.rs | 158 + .../platform/src}/storage/schema.sql | 0 .../platform/src/storage/sqlite.rs | 4 +- .../platform/tests/admin_password.rs | 7 +- .../platform/tests}/client_key_auth.rs | 2 +- {tests => crates/platform/tests}/config.rs | 37 +- .../platform/tests}/crypto.rs | 2 +- .../platform/tests}/log_rotation.rs | 2 +- .../platform/tests}/storage_schema.rs | 3 +- crates/runtime/Cargo.toml | 35 + crates/runtime/src/bootstrap.rs | 12 + crates/runtime/src/config.rs | 27 + crates/runtime/src/lib.rs | 9 + crates/runtime/src/repositories.rs | 54 + crates/runtime/src/services.rs | 6592 +++++++++++++ crates/runtime/src/state.rs | 316 + .../runtime/src}/tasks/cookie_cleanup.rs | 54 +- crates/runtime/src/tasks/coordinator.rs | 177 + .../runtime/src/tasks/fingerprint_update.rs | 35 + crates/runtime/src/tasks/mod.rs | 10 + .../runtime/src}/tasks/model_refresh.rs | 115 +- crates/runtime/src/tasks/quota_refresh.rs | 252 + .../src/tasks/session_affinity_cleanup.rs | 78 + crates/runtime/src/tasks/session_cleanup.rs | 75 + crates/runtime/src/tasks/token_refresh.rs | 834 ++ crates/runtime/src/upstream.rs | 50 + crates/runtime/tests/account_pool_restore.rs | 126 + crates/runtime/tests/fingerprint.rs | 198 + crates/runtime/tests/quota_refresh.rs | 218 + crates/runtime/tests/session_affinity.rs | 121 + crates/runtime/tests/tasks.rs | 604 ++ crates/runtime/tests/token_refresh.rs | 983 ++ crates/runtime/tests/upstream.rs | 132 + crates/server/Cargo.toml | 38 + .../server/src/admin_api}/accounts/cookies.rs | 103 +- .../server/src/admin_api/accounts/create.rs | 52 + .../server/src/admin_api/accounts/export.rs | 119 + .../server/src/admin_api}/accounts/health.rs | 195 +- .../server/src/admin_api/accounts/import.rs | 52 + .../src/admin_api/accounts/import_cli.rs | 75 + .../src/admin_api/accounts/lifecycle.rs | 336 + crates/server/src/admin_api/accounts/list.rs | 60 + crates/server/src/admin_api/accounts/mod.rs | 152 + crates/server/src/admin_api/accounts/oauth.rs | 378 + crates/server/src/admin_api/accounts/quota.rs | 201 + .../src/admin_api/client_keys/create.rs | 32 + .../src/admin_api/client_keys/export.rs | 56 + .../src/admin_api/client_keys/import.rs | 45 + .../src/admin_api}/client_keys/lifecycle.rs | 82 +- .../server/src/admin_api}/client_keys/list.rs | 30 +- .../server/src/admin_api}/client_keys/mod.rs | 227 +- .../server/src/admin_api}/diagnostics.rs | 11 +- .../server/src/admin_api}/logs/detail.rs | 15 +- crates/server/src/admin_api/logs/mod.rs | 90 + crates/server/src/admin_api/logs/query.rs | 82 + .../server/src/admin_api}/logs/state.rs | 75 +- crates/server/src/admin_api/mod.rs | 26 + crates/server/src/admin_api/models.rs | 94 + .../server/src/admin_api}/response.rs | 34 +- .../server/src/admin_api}/router.rs | 72 +- crates/server/src/admin_api/session.rs | 133 + .../server/src/admin_api}/settings.rs | 105 +- .../server/src/admin_api}/usage.rs | 222 +- crates/server/src/error/admin.rs | 6 + crates/server/src/error/mod.rs | 6 + crates/server/src/error/openai.rs | 6 + crates/server/src/lib.rs | 7 + crates/server/src/main.rs | 86 + crates/server/src/middleware/auth.rs | 5 + crates/server/src/middleware/cors.rs | 4 + crates/server/src/middleware/mod.rs | 10 + .../server/src/middleware}/request_id.rs | 9 +- crates/server/src/middleware/trace.rs | 81 + crates/server/src/openai_api/auth.rs | 29 + crates/server/src/openai_api/chat.rs | 197 + crates/server/src/openai_api/diagnostics.rs | 359 + crates/server/src/openai_api/error.rs | 6 + .../server/src/openai_api}/mod.rs | 7 +- .../server/src/openai_api}/models.rs | 41 +- crates/server/src/openai_api/responses.rs | 648 ++ .../server/src/openai_api}/router.rs | 9 +- crates/server/src/openai_api/sse.rs | 27 + crates/server/src/router.rs | 31 + crates/server/tests/admin_accounts_routes.rs | 3634 +++++++ .../server/tests/admin_api_contract.rs | 17 +- .../server/tests/admin_client_keys_routes.rs | 1169 +-- .../server/tests/admin_logs_routes.rs | 732 +- .../server/tests/admin_models_routes.rs | 454 +- .../server/tests/admin_session_routes.rs | 274 +- .../server/tests/admin_settings_routes.rs | 594 +- crates/server/tests/assets_routes.rs | 155 + .../server/tests/http_trace_middleware.rs | 7 +- crates/server/tests/openai_chat_upstream.rs | 8579 +++++++++++++++++ .../server/tests/openai_diagnostics_routes.rs | 323 + crates/server/tests/openai_models_auth.rs | 160 + .../server/tests/openai_responses_routes.rs | 244 + crates/xtask/Cargo.toml | 12 + crates/xtask/src/build_web.rs | 10 + crates/xtask/src/check_architecture.rs | 9 + crates/xtask/src/main.rs | 60 + crates/xtask/src/release.rs | 25 + docs/OPENAI_CODEX_PARITY_AUDIT.md | 17 +- docs/OPENAI_REQUEST_ALIGNMENT.md | 24 +- docs/api.md | 5 +- docs/architecture-audit.md | 7166 +++++++++++++- docs/architecture.md | 1096 ++- src/admin/api/accounts/create.rs | 42 - src/admin/api/accounts/delete.rs | 98 - src/admin/api/accounts/export.rs | 85 - src/admin/api/accounts/import.rs | 280 - src/admin/api/accounts/lifecycle.rs | 291 - src/admin/api/accounts/list.rs | 51 - src/admin/api/accounts/mod.rs | 319 - src/admin/api/accounts/oauth.rs | 338 - src/admin/api/accounts/quota.rs | 195 - src/admin/api/client_keys/create.rs | 28 - src/admin/api/client_keys/export.rs | 44 - src/admin/api/client_keys/import.rs | 40 - src/admin/api/logs/mod.rs | 60 - src/admin/api/logs/query.rs | 71 - src/admin/api/mod.rs | 20 - src/admin/api/models.rs | 114 - src/admin/api/session.rs | 338 - src/admin/client_keys/mod.rs | 1 - src/admin/client_keys/service.rs | 334 - src/admin/mod.rs | 5 - src/admin/session/mod.rs | 2 - src/admin/session/service.rs | 401 - src/admin/settings.rs | 290 - src/admin/tasks/mod.rs | 1 - src/admin/tasks/session_cleanup.rs | 59 - src/codex/accounts/cloudflare_challenge.rs | 143 - src/codex/accounts/cookies/jar.rs | 67 - src/codex/accounts/cookies/mod.rs | 2 - src/codex/accounts/lifecycle.rs | 9 - src/codex/accounts/mod.rs | 9 - src/codex/accounts/repository/accounts.rs | 754 -- src/codex/accounts/repository/leases.rs | 51 - src/codex/accounts/repository/mod.rs | 9 - src/codex/accounts/repository/quotas.rs | 146 - src/codex/accounts/repository/tokens.rs | 203 - src/codex/accounts/repository/usage.rs | 491 - src/codex/accounts/service/cookies.rs | 51 - src/codex/accounts/service/health.rs | 286 - src/codex/accounts/service/import.rs | 349 - src/codex/accounts/service/lifecycle.rs | 173 - src/codex/accounts/service/mod.rs | 363 - src/codex/accounts/service/pool_sync.rs | 78 - src/codex/accounts/service/refresh.rs | 294 - src/codex/accounts/usage_snapshots.rs | 290 - src/codex/events/mod.rs | 3 - src/codex/events/repository.rs | 213 - src/codex/events/service.rs | 203 - src/codex/gateway/fingerprint/mod.rs | 4 - src/codex/gateway/fingerprint/repository.rs | 122 - .../gateway/fingerprint/update_checker.rs | 246 - src/codex/gateway/fingerprint/updater.rs | 72 - src/codex/gateway/installation_id.rs | 112 - src/codex/gateway/mod.rs | 6 - src/codex/gateway/oauth/client.rs | 419 - src/codex/gateway/oauth/codex_cli.rs | 93 - src/codex/gateway/oauth/mod.rs | 9 - src/codex/gateway/oauth/refresh.rs | 142 - src/codex/gateway/oauth/token.rs | 5 - src/codex/gateway/protocol/codex_to_openai.rs | 567 -- src/codex/gateway/protocol/error.rs | 70 - src/codex/gateway/protocol/mod.rs | 5 - src/codex/gateway/protocol/openai_to_codex.rs | 311 - src/codex/gateway/protocol/schema.rs | 6 - src/codex/gateway/transport/custom_ca.rs | 361 - src/codex/gateway/transport/endpoints.rs | 61 - src/codex/gateway/transport/headers.rs | 92 - src/codex/gateway/transport/http_client.rs | 960 -- src/codex/gateway/transport/mod.rs | 10 - src/codex/gateway/transport/retry_after.rs | 143 - src/codex/gateway/transport/usage_events.rs | 122 - .../gateway/transport/websocket/audit.rs | 229 - src/codex/gateway/transport/websocket/mod.rs | 896 -- .../gateway/transport/websocket/opening.rs | 495 - src/codex/mod.rs | 7 - src/codex/models/mod.rs | 3 - src/codex/models/repository.rs | 62 - src/codex/models/service.rs | 191 - src/codex/serving/chat.rs | 313 - src/codex/serving/diagnostics.rs | 143 - src/codex/serving/dispatch/affinity.rs | 698 -- src/codex/serving/dispatch/fallback.rs | 632 -- src/codex/serving/dispatch/limits.rs | 91 - src/codex/serving/dispatch/mod.rs | 2063 ---- src/codex/serving/dispatch/routing.rs | 29 - src/codex/serving/dispatch/stream.rs | 638 -- src/codex/serving/dispatch/stream_audit.rs | 337 - src/codex/serving/dispatch/transition.rs | 92 - src/codex/serving/dispatch/usage.rs | 108 - src/codex/serving/http/auth.rs | 15 - src/codex/serving/http/chat.rs | 46 - src/codex/serving/http/diagnostics.rs | 296 - src/codex/serving/http/errors.rs | 119 - src/codex/serving/http/responses.rs | 70 - src/codex/serving/mod.rs | 5 - src/codex/serving/responses.rs | 1313 --- src/codex/tasks/mod.rs | 4 - src/codex/tasks/quota_refresh.rs | 117 - src/codex/tasks/token_refresh.rs | 882 -- src/codex/usage/mod.rs | 1 - src/codex/usage/service.rs | 47 - src/config/loader.rs | 19 - src/config/mod.rs | 6 - src/lib.rs | 8 - src/main.rs | 42 - src/platform/crypto/mod.rs | 3 - src/platform/http/auth.rs | 24 - src/platform/http/health.rs | 6 - src/platform/http/mod.rs | 3 - .../identity/client_key_repository.rs | 219 - src/platform/identity/mod.rs | 4 - src/platform/logging/mod.rs | 1 - src/platform/mod.rs | 5 - src/platform/storage/mod.rs | 2 - src/platform/storage/paths.rs | 14 - src/runtime/bootstrap.rs | 120 - src/runtime/mod.rs | 6 - src/runtime/router.rs | 120 - src/runtime/state.rs | 489 - src/runtime/tasks/coordinator.rs | 116 - src/runtime/tasks/mod.rs | 5 - src/runtime/tasks/types.rs | 38 - src/utils/mod.rs | 2 - src/web/assets.rs | 1 - src/web/mod.rs | 6 - src/web/router.rs | 1 - src/web/security.rs | 1 - src/web/shell.rs | 1 - tests/admin.rs | 21 - tests/admin/accounts/cookies_quota.rs | 662 -- tests/admin/accounts/import_export.rs | 863 -- tests/admin/accounts/lifecycle.rs | 330 - tests/admin/accounts/list.rs | 57 - tests/admin/accounts/oauth.rs | 678 -- tests/admin/session_repository.rs | 57 - tests/admin/usage_stats_route.rs | 246 - tests/architecture.rs | 12 +- tests/architecture/accounts_boundary.rs | 16 +- tests/architecture/admin_accounts_split.rs | 51 - tests/architecture/admin_boundary.rs | 9 +- tests/architecture/codex_activity_boundary.rs | 16 +- tests/architecture/dependency_direction.rs | 93 + tests/architecture/directory_shape.rs | 368 + tests/architecture/forbidden_imports.rs | 123 + tests/architecture/forbidden_legacy_paths.rs | 54 + tests/architecture/gateway_boundary.rs | 37 +- .../placeholder_implementations.rs | 128 + tests/architecture/platform_boundary.rs | 27 +- tests/architecture/runtime_tasks_boundary.rs | 6 +- tests/architecture/server_startup.rs | 64 + tests/architecture/serving_boundary.rs | 47 - tests/architecture/test_migration.rs | 449 + tests/architecture/workspace_scaffold.rs | 198 + tests/codex_accounts.rs | 10 - tests/codex_accounts/cookie_store.rs | 265 - tests/codex_accounts/refresh.rs | 160 - tests/codex_accounts/refresh_scheduler.rs | 17 - tests/codex_accounts/repository.rs | 609 -- tests/codex_accounts/service_refresh.rs | 171 - tests/codex_events.rs | 3 - tests/codex_events/logs_pagination.rs | 125 - tests/codex_gateway.rs | 9 - tests/codex_gateway/cli_auth_import.rs | 44 - tests/codex_gateway/fingerprint_update.rs | 72 - tests/codex_gateway/headers.rs | 176 - tests/codex_gateway/http_client.rs | 554 -- tests/codex_gateway/usage_events.rs | 245 - tests/codex_gateway/websocket.rs | 4983 ---------- tests/codex_gateway/websocket/pool.rs | 559 -- tests/codex_models.rs | 3 - tests/codex_models/catalog.rs | 145 - tests/codex_serving.rs | 12 - tests/codex_serving/chat_completions.rs | 910 -- tests/codex_serving/diagnostics_route.rs | 364 - tests/codex_serving/responses_http_sse.rs | 1766 ---- tests/codex_serving/responses_websocket.rs | 2863 ------ tests/codex_serving/routes_chat.rs | 35 - tests/codex_serving/routes_responses.rs | 372 - tests/codex_serving/upstream_errors.rs | 278 - tests/codex_serving/upstream_fallback.rs | 2068 ---- tests/platform.rs | 7 - tests/platform/http_auth.rs | 38 - tests/runtime.rs | 6 - tests/runtime/startup.rs | 374 - tests/support/admin_accounts.rs | 395 - tests/support/mod.rs | 98 - tests/support/upstream.rs | 399 - 423 files changed, 59073 insertions(+), 48538 deletions(-) create mode 100644 crates/adapters/Cargo.toml create mode 100644 crates/adapters/src/codex/client.rs create mode 100644 crates/adapters/src/codex/fingerprint.rs create mode 100644 crates/adapters/src/codex/mod.rs create mode 100644 crates/adapters/src/codex/models.rs create mode 100644 crates/adapters/src/codex/websocket/connect.rs rename {src/codex/gateway/transport => crates/adapters/src/codex}/websocket/deflate.rs (77%) create mode 100644 crates/adapters/src/codex/websocket/mod.rs create mode 100644 crates/adapters/src/codex/websocket/opening.rs rename {src/codex/gateway/transport => crates/adapters/src/codex}/websocket/pool.rs (70%) create mode 100644 crates/adapters/src/lib.rs create mode 100644 crates/adapters/src/oauth/mod.rs create mode 100644 crates/adapters/src/oauth/openai.rs create mode 100644 crates/adapters/src/sqlite/account_tokens.rs create mode 100644 crates/adapters/src/sqlite/account_usage.rs create mode 100644 crates/adapters/src/sqlite/accounts.rs rename src/admin/session/repository.rs => crates/adapters/src/sqlite/admin_sessions.rs (50%) create mode 100644 crates/adapters/src/sqlite/client_keys.rs rename src/codex/accounts/cookies/repository.rs => crates/adapters/src/sqlite/cookies.rs (71%) create mode 100644 crates/adapters/src/sqlite/events.rs create mode 100644 crates/adapters/src/sqlite/mod.rs create mode 100644 crates/adapters/src/sqlite/models.rs create mode 100644 crates/adapters/src/sqlite/refresh_leases.rs create mode 100644 crates/adapters/src/sqlite/session_affinity.rs create mode 100644 crates/adapters/tests/account_repository.rs create mode 100644 crates/adapters/tests/admin_sessions.rs create mode 100644 crates/adapters/tests/client_keys.rs create mode 100644 crates/adapters/tests/codex.rs create mode 100644 crates/adapters/tests/cookies.rs create mode 100644 crates/adapters/tests/events.rs create mode 100644 crates/adapters/tests/models.rs rename tests/codex_gateway/oauth_refresh.rs => crates/adapters/tests/oauth.rs (84%) create mode 100644 crates/adapters/tests/refresh_leases.rs create mode 100644 crates/adapters/tests/session_affinity.rs create mode 100644 crates/assets/Cargo.toml create mode 100644 crates/assets/src/headers.rs create mode 100644 crates/assets/src/lib.rs create mode 100644 crates/assets/src/router.rs create mode 100644 crates/assets/tests/assets.rs create mode 100644 crates/core/Cargo.toml create mode 100644 crates/core/src/accounts/cloudflare.rs create mode 100644 crates/core/src/accounts/cookies.rs rename {src/codex => crates/core/src}/accounts/jwt.rs (71%) create mode 100644 crates/core/src/accounts/lifecycle.rs create mode 100644 crates/core/src/accounts/mod.rs rename {src/codex => crates/core/src}/accounts/model.rs (65%) rename {src/codex => crates/core/src}/accounts/pool.rs (87%) create mode 100644 crates/core/src/accounts/ports.rs create mode 100644 crates/core/src/accounts/service.rs create mode 100644 crates/core/src/accounts/usage.rs create mode 100644 crates/core/src/admin/auth.rs create mode 100644 crates/core/src/admin/client_keys.rs create mode 100644 crates/core/src/admin/mod.rs create mode 100644 crates/core/src/admin/ports.rs create mode 100644 crates/core/src/admin/settings.rs create mode 100644 crates/core/src/auth/mod.rs create mode 100644 crates/core/src/auth/oauth.rs create mode 100644 crates/core/src/auth/ports.rs create mode 100644 crates/core/src/auth/session.rs create mode 100644 crates/core/src/error.rs create mode 100644 crates/core/src/events/mod.rs rename src/codex/events/event.rs => crates/core/src/events/model.rs (56%) create mode 100644 crates/core/src/events/ports.rs create mode 100644 crates/core/src/events/service.rs rename src/codex/gateway/conversation_identity.rs => crates/core/src/gateway/conversation.rs (97%) rename src/codex/gateway/fingerprint/model.rs => crates/core/src/gateway/fingerprint.rs (86%) create mode 100644 crates/core/src/gateway/installation.rs create mode 100644 crates/core/src/gateway/mod.rs create mode 100644 crates/core/src/gateway/ports.rs create mode 100644 crates/core/src/lib.rs rename {src/codex => crates/core/src}/models/catalog.rs (73%) create mode 100644 crates/core/src/models/mod.rs create mode 100644 crates/core/src/models/model.rs create mode 100644 crates/core/src/models/ports.rs create mode 100644 crates/core/src/models/service.rs create mode 100644 crates/core/src/protocol/codex/chat.rs rename src/codex/gateway/transport/rate_limits.rs => crates/core/src/protocol/codex/events.rs (55%) create mode 100644 crates/core/src/protocol/codex/mod.rs rename src/codex/gateway/transport/types.rs => crates/core/src/protocol/codex/responses.rs (74%) rename src/codex/gateway/protocol/tuple_schema.rs => crates/core/src/protocol/codex/schema.rs (99%) rename {src/codex/gateway/transport => crates/core/src/protocol/codex}/sse.rs (92%) rename src/codex/gateway/transport/websocket/codec.rs => crates/core/src/protocol/codex/websocket.rs (50%) create mode 100644 crates/core/src/protocol/mod.rs create mode 100644 crates/core/src/protocol/openai/chat.rs create mode 100644 crates/core/src/protocol/openai/errors.rs create mode 100644 crates/core/src/protocol/openai/mod.rs create mode 100644 crates/core/src/protocol/openai/models.rs create mode 100644 crates/core/src/protocol/openai/responses.rs create mode 100644 crates/core/src/serving/affinity.rs create mode 100644 crates/core/src/serving/chat.rs create mode 100644 crates/core/src/serving/errors.rs create mode 100644 crates/core/src/serving/fallback.rs rename {src/codex/serving/dispatch => crates/core/src/serving}/implicit_resume.rs (60%) create mode 100644 crates/core/src/serving/mod.rs rename {src/codex/accounts/service => crates/core/src/serving}/quota.rs (53%) rename {src/codex/serving/dispatch => crates/core/src/serving}/reasoning_replay.rs (62%) create mode 100644 crates/core/src/serving/recovery.rs create mode 100644 crates/core/src/serving/responses.rs create mode 100644 crates/core/src/serving/routing.rs create mode 100644 crates/core/src/serving/stream.rs create mode 100644 crates/core/src/serving/usage.rs create mode 100644 crates/core/src/usage/mod.rs create mode 100644 crates/core/src/usage/model.rs create mode 100644 crates/core/src/usage/ports.rs create mode 100644 crates/core/src/usage/service.rs rename tests/codex_accounts/pool_scheduling.rs => crates/core/tests/account_pool.rs (86%) create mode 100644 crates/core/tests/admin_settings.rs create mode 100644 crates/core/tests/auth.rs create mode 100644 crates/core/tests/models.rs create mode 100644 crates/core/tests/protocol.rs create mode 100644 crates/core/tests/session_affinity.rs create mode 100644 crates/platform/Cargo.toml create mode 100644 crates/platform/src/config/loader.rs create mode 100644 crates/platform/src/config/mod.rs rename {src => crates/platform/src}/config/types.rs (66%) create mode 100644 crates/platform/src/crypto/hash.rs create mode 100644 crates/platform/src/crypto/mod.rs rename {src/platform => crates/platform/src}/crypto/secret_box.rs (85%) rename src/platform/identity/admin_session.rs => crates/platform/src/identity/admin_password.rs (85%) rename {src/platform => crates/platform/src}/identity/client_key.rs (86%) rename src/platform/identity/error.rs => crates/platform/src/identity/mod.rs (60%) rename src/utils/json.rs => crates/platform/src/json/mod.rs (66%) rename {src/utils => crates/platform/src/json}/pagination.rs (82%) create mode 100644 crates/platform/src/lib.rs create mode 100644 crates/platform/src/logging/mod.rs rename {src/platform => crates/platform/src}/logging/rotation.rs (88%) create mode 100644 crates/platform/src/storage/mod.rs create mode 100644 crates/platform/src/storage/paths.rs rename {src/platform => crates/platform/src}/storage/schema.sql (100%) rename src/platform/storage/db.rs => crates/platform/src/storage/sqlite.rs (93%) rename tests/admin/session.rs => crates/platform/tests/admin_password.rs (60%) rename {tests/platform => crates/platform/tests}/client_key_auth.rs (88%) rename {tests => crates/platform/tests}/config.rs (77%) rename {tests/platform => crates/platform/tests}/crypto.rs (95%) rename {tests/platform => crates/platform/tests}/log_rotation.rs (90%) rename {tests/platform => crates/platform/tests}/storage_schema.rs (99%) create mode 100644 crates/runtime/Cargo.toml create mode 100644 crates/runtime/src/bootstrap.rs create mode 100644 crates/runtime/src/config.rs create mode 100644 crates/runtime/src/lib.rs create mode 100644 crates/runtime/src/repositories.rs create mode 100644 crates/runtime/src/services.rs create mode 100644 crates/runtime/src/state.rs rename {src/codex => crates/runtime/src}/tasks/cookie_cleanup.rs (51%) create mode 100644 crates/runtime/src/tasks/coordinator.rs create mode 100644 crates/runtime/src/tasks/fingerprint_update.rs create mode 100644 crates/runtime/src/tasks/mod.rs rename {src/codex => crates/runtime/src}/tasks/model_refresh.rs (52%) create mode 100644 crates/runtime/src/tasks/quota_refresh.rs create mode 100644 crates/runtime/src/tasks/session_affinity_cleanup.rs create mode 100644 crates/runtime/src/tasks/session_cleanup.rs create mode 100644 crates/runtime/src/tasks/token_refresh.rs create mode 100644 crates/runtime/src/upstream.rs create mode 100644 crates/runtime/tests/account_pool_restore.rs create mode 100644 crates/runtime/tests/fingerprint.rs create mode 100644 crates/runtime/tests/quota_refresh.rs create mode 100644 crates/runtime/tests/session_affinity.rs create mode 100644 crates/runtime/tests/tasks.rs create mode 100644 crates/runtime/tests/token_refresh.rs create mode 100644 crates/runtime/tests/upstream.rs create mode 100644 crates/server/Cargo.toml rename {src/admin/api => crates/server/src/admin_api}/accounts/cookies.rs (55%) create mode 100644 crates/server/src/admin_api/accounts/create.rs create mode 100644 crates/server/src/admin_api/accounts/export.rs rename {src/admin/api => crates/server/src/admin_api}/accounts/health.rs (50%) create mode 100644 crates/server/src/admin_api/accounts/import.rs create mode 100644 crates/server/src/admin_api/accounts/import_cli.rs create mode 100644 crates/server/src/admin_api/accounts/lifecycle.rs create mode 100644 crates/server/src/admin_api/accounts/list.rs create mode 100644 crates/server/src/admin_api/accounts/mod.rs create mode 100644 crates/server/src/admin_api/accounts/oauth.rs create mode 100644 crates/server/src/admin_api/accounts/quota.rs create mode 100644 crates/server/src/admin_api/client_keys/create.rs create mode 100644 crates/server/src/admin_api/client_keys/export.rs create mode 100644 crates/server/src/admin_api/client_keys/import.rs rename {src/admin/api => crates/server/src/admin_api}/client_keys/lifecycle.rs (58%) rename {src/admin/api => crates/server/src/admin_api}/client_keys/list.rs (52%) rename {src/admin/api => crates/server/src/admin_api}/client_keys/mod.rs (53%) rename {src/admin/api => crates/server/src/admin_api}/diagnostics.rs (67%) rename {src/admin/api => crates/server/src/admin_api}/logs/detail.rs (72%) create mode 100644 crates/server/src/admin_api/logs/mod.rs create mode 100644 crates/server/src/admin_api/logs/query.rs rename {src/admin/api => crates/server/src/admin_api}/logs/state.rs (60%) create mode 100644 crates/server/src/admin_api/mod.rs create mode 100644 crates/server/src/admin_api/models.rs rename {src/admin/api => crates/server/src/admin_api}/response.rs (77%) rename {src/admin/api => crates/server/src/admin_api}/router.rs (74%) create mode 100644 crates/server/src/admin_api/session.rs rename {src/admin/api => crates/server/src/admin_api}/settings.rs (64%) rename {src/admin/api => crates/server/src/admin_api}/usage.rs (64%) create mode 100644 crates/server/src/error/admin.rs create mode 100644 crates/server/src/error/mod.rs create mode 100644 crates/server/src/error/openai.rs create mode 100644 crates/server/src/lib.rs create mode 100644 crates/server/src/main.rs create mode 100644 crates/server/src/middleware/auth.rs create mode 100644 crates/server/src/middleware/cors.rs create mode 100644 crates/server/src/middleware/mod.rs rename {src/platform/http => crates/server/src/middleware}/request_id.rs (82%) create mode 100644 crates/server/src/middleware/trace.rs create mode 100644 crates/server/src/openai_api/auth.rs create mode 100644 crates/server/src/openai_api/chat.rs create mode 100644 crates/server/src/openai_api/diagnostics.rs create mode 100644 crates/server/src/openai_api/error.rs rename {src/codex/serving/http => crates/server/src/openai_api}/mod.rs (63%) rename {src/codex/serving/http => crates/server/src/openai_api}/models.rs (73%) create mode 100644 crates/server/src/openai_api/responses.rs rename {src/codex/serving/http => crates/server/src/openai_api}/router.rs (71%) create mode 100644 crates/server/src/openai_api/sse.rs create mode 100644 crates/server/src/router.rs create mode 100644 crates/server/tests/admin_accounts_routes.rs rename tests/admin/api_contract.rs => crates/server/tests/admin_api_contract.rs (81%) rename tests/admin/client_keys_route.rs => crates/server/tests/admin_client_keys_routes.rs (74%) rename tests/admin/logs_route.rs => crates/server/tests/admin_logs_routes.rs (50%) rename tests/admin/models_route.rs => crates/server/tests/admin_models_routes.rs (63%) rename tests/admin/session_login_route.rs => crates/server/tests/admin_session_routes.rs (78%) rename tests/admin/settings_route.rs => crates/server/tests/admin_settings_routes.rs (71%) create mode 100644 crates/server/tests/assets_routes.rs rename tests/runtime/http_trace.rs => crates/server/tests/http_trace_middleware.rs (91%) create mode 100644 crates/server/tests/openai_chat_upstream.rs create mode 100644 crates/server/tests/openai_diagnostics_routes.rs create mode 100644 crates/server/tests/openai_models_auth.rs create mode 100644 crates/server/tests/openai_responses_routes.rs create mode 100644 crates/xtask/Cargo.toml create mode 100644 crates/xtask/src/build_web.rs create mode 100644 crates/xtask/src/check_architecture.rs create mode 100644 crates/xtask/src/main.rs create mode 100644 crates/xtask/src/release.rs delete mode 100644 src/admin/api/accounts/create.rs delete mode 100644 src/admin/api/accounts/delete.rs delete mode 100644 src/admin/api/accounts/export.rs delete mode 100644 src/admin/api/accounts/import.rs delete mode 100644 src/admin/api/accounts/lifecycle.rs delete mode 100644 src/admin/api/accounts/list.rs delete mode 100644 src/admin/api/accounts/mod.rs delete mode 100644 src/admin/api/accounts/oauth.rs delete mode 100644 src/admin/api/accounts/quota.rs delete mode 100644 src/admin/api/client_keys/create.rs delete mode 100644 src/admin/api/client_keys/export.rs delete mode 100644 src/admin/api/client_keys/import.rs delete mode 100644 src/admin/api/logs/mod.rs delete mode 100644 src/admin/api/logs/query.rs delete mode 100644 src/admin/api/mod.rs delete mode 100644 src/admin/api/models.rs delete mode 100644 src/admin/api/session.rs delete mode 100644 src/admin/client_keys/mod.rs delete mode 100644 src/admin/client_keys/service.rs delete mode 100644 src/admin/mod.rs delete mode 100644 src/admin/session/mod.rs delete mode 100644 src/admin/session/service.rs delete mode 100644 src/admin/settings.rs delete mode 100644 src/admin/tasks/mod.rs delete mode 100644 src/admin/tasks/session_cleanup.rs delete mode 100644 src/codex/accounts/cloudflare_challenge.rs delete mode 100644 src/codex/accounts/cookies/jar.rs delete mode 100644 src/codex/accounts/cookies/mod.rs delete mode 100644 src/codex/accounts/lifecycle.rs delete mode 100644 src/codex/accounts/mod.rs delete mode 100644 src/codex/accounts/repository/accounts.rs delete mode 100644 src/codex/accounts/repository/leases.rs delete mode 100644 src/codex/accounts/repository/mod.rs delete mode 100644 src/codex/accounts/repository/quotas.rs delete mode 100644 src/codex/accounts/repository/tokens.rs delete mode 100644 src/codex/accounts/repository/usage.rs delete mode 100644 src/codex/accounts/service/cookies.rs delete mode 100644 src/codex/accounts/service/health.rs delete mode 100644 src/codex/accounts/service/import.rs delete mode 100644 src/codex/accounts/service/lifecycle.rs delete mode 100644 src/codex/accounts/service/mod.rs delete mode 100644 src/codex/accounts/service/pool_sync.rs delete mode 100644 src/codex/accounts/service/refresh.rs delete mode 100644 src/codex/accounts/usage_snapshots.rs delete mode 100644 src/codex/events/mod.rs delete mode 100644 src/codex/events/repository.rs delete mode 100644 src/codex/events/service.rs delete mode 100644 src/codex/gateway/fingerprint/mod.rs delete mode 100644 src/codex/gateway/fingerprint/repository.rs delete mode 100644 src/codex/gateway/fingerprint/update_checker.rs delete mode 100644 src/codex/gateway/fingerprint/updater.rs delete mode 100644 src/codex/gateway/installation_id.rs delete mode 100644 src/codex/gateway/mod.rs delete mode 100644 src/codex/gateway/oauth/client.rs delete mode 100644 src/codex/gateway/oauth/codex_cli.rs delete mode 100644 src/codex/gateway/oauth/mod.rs delete mode 100644 src/codex/gateway/oauth/refresh.rs delete mode 100644 src/codex/gateway/oauth/token.rs delete mode 100644 src/codex/gateway/protocol/codex_to_openai.rs delete mode 100644 src/codex/gateway/protocol/error.rs delete mode 100644 src/codex/gateway/protocol/mod.rs delete mode 100644 src/codex/gateway/protocol/openai_to_codex.rs delete mode 100644 src/codex/gateway/protocol/schema.rs delete mode 100644 src/codex/gateway/transport/custom_ca.rs delete mode 100644 src/codex/gateway/transport/endpoints.rs delete mode 100644 src/codex/gateway/transport/headers.rs delete mode 100644 src/codex/gateway/transport/http_client.rs delete mode 100644 src/codex/gateway/transport/mod.rs delete mode 100644 src/codex/gateway/transport/retry_after.rs delete mode 100644 src/codex/gateway/transport/usage_events.rs delete mode 100644 src/codex/gateway/transport/websocket/audit.rs delete mode 100644 src/codex/gateway/transport/websocket/mod.rs delete mode 100644 src/codex/gateway/transport/websocket/opening.rs delete mode 100644 src/codex/mod.rs delete mode 100644 src/codex/models/mod.rs delete mode 100644 src/codex/models/repository.rs delete mode 100644 src/codex/models/service.rs delete mode 100644 src/codex/serving/chat.rs delete mode 100644 src/codex/serving/diagnostics.rs delete mode 100644 src/codex/serving/dispatch/affinity.rs delete mode 100644 src/codex/serving/dispatch/fallback.rs delete mode 100644 src/codex/serving/dispatch/limits.rs delete mode 100644 src/codex/serving/dispatch/mod.rs delete mode 100644 src/codex/serving/dispatch/routing.rs delete mode 100644 src/codex/serving/dispatch/stream.rs delete mode 100644 src/codex/serving/dispatch/stream_audit.rs delete mode 100644 src/codex/serving/dispatch/transition.rs delete mode 100644 src/codex/serving/dispatch/usage.rs delete mode 100644 src/codex/serving/http/auth.rs delete mode 100644 src/codex/serving/http/chat.rs delete mode 100644 src/codex/serving/http/diagnostics.rs delete mode 100644 src/codex/serving/http/errors.rs delete mode 100644 src/codex/serving/http/responses.rs delete mode 100644 src/codex/serving/mod.rs delete mode 100644 src/codex/serving/responses.rs delete mode 100644 src/codex/tasks/mod.rs delete mode 100644 src/codex/tasks/quota_refresh.rs delete mode 100644 src/codex/tasks/token_refresh.rs delete mode 100644 src/codex/usage/mod.rs delete mode 100644 src/codex/usage/service.rs delete mode 100644 src/config/loader.rs delete mode 100644 src/config/mod.rs delete mode 100644 src/lib.rs delete mode 100644 src/main.rs delete mode 100644 src/platform/crypto/mod.rs delete mode 100644 src/platform/http/auth.rs delete mode 100644 src/platform/http/health.rs delete mode 100644 src/platform/http/mod.rs delete mode 100644 src/platform/identity/client_key_repository.rs delete mode 100644 src/platform/identity/mod.rs delete mode 100644 src/platform/logging/mod.rs delete mode 100644 src/platform/mod.rs delete mode 100644 src/platform/storage/mod.rs delete mode 100644 src/platform/storage/paths.rs delete mode 100644 src/runtime/bootstrap.rs delete mode 100644 src/runtime/mod.rs delete mode 100644 src/runtime/router.rs delete mode 100644 src/runtime/state.rs delete mode 100644 src/runtime/tasks/coordinator.rs delete mode 100644 src/runtime/tasks/mod.rs delete mode 100644 src/runtime/tasks/types.rs delete mode 100644 src/utils/mod.rs delete mode 100644 src/web/assets.rs delete mode 100644 src/web/mod.rs delete mode 100644 src/web/router.rs delete mode 100644 src/web/security.rs delete mode 100644 src/web/shell.rs delete mode 100644 tests/admin.rs delete mode 100644 tests/admin/accounts/cookies_quota.rs delete mode 100644 tests/admin/accounts/import_export.rs delete mode 100644 tests/admin/accounts/lifecycle.rs delete mode 100644 tests/admin/accounts/list.rs delete mode 100644 tests/admin/accounts/oauth.rs delete mode 100644 tests/admin/session_repository.rs delete mode 100644 tests/admin/usage_stats_route.rs delete mode 100644 tests/architecture/admin_accounts_split.rs create mode 100644 tests/architecture/dependency_direction.rs create mode 100644 tests/architecture/directory_shape.rs create mode 100644 tests/architecture/forbidden_imports.rs create mode 100644 tests/architecture/forbidden_legacy_paths.rs create mode 100644 tests/architecture/placeholder_implementations.rs create mode 100644 tests/architecture/server_startup.rs delete mode 100644 tests/architecture/serving_boundary.rs create mode 100644 tests/architecture/test_migration.rs create mode 100644 tests/architecture/workspace_scaffold.rs delete mode 100644 tests/codex_accounts.rs delete mode 100644 tests/codex_accounts/cookie_store.rs delete mode 100644 tests/codex_accounts/refresh.rs delete mode 100644 tests/codex_accounts/refresh_scheduler.rs delete mode 100644 tests/codex_accounts/repository.rs delete mode 100644 tests/codex_accounts/service_refresh.rs delete mode 100644 tests/codex_events.rs delete mode 100644 tests/codex_events/logs_pagination.rs delete mode 100644 tests/codex_gateway.rs delete mode 100644 tests/codex_gateway/cli_auth_import.rs delete mode 100644 tests/codex_gateway/fingerprint_update.rs delete mode 100644 tests/codex_gateway/headers.rs delete mode 100644 tests/codex_gateway/http_client.rs delete mode 100644 tests/codex_gateway/usage_events.rs delete mode 100644 tests/codex_gateway/websocket.rs delete mode 100644 tests/codex_gateway/websocket/pool.rs delete mode 100644 tests/codex_models.rs delete mode 100644 tests/codex_models/catalog.rs delete mode 100644 tests/codex_serving.rs delete mode 100644 tests/codex_serving/chat_completions.rs delete mode 100644 tests/codex_serving/diagnostics_route.rs delete mode 100644 tests/codex_serving/responses_http_sse.rs delete mode 100644 tests/codex_serving/responses_websocket.rs delete mode 100644 tests/codex_serving/routes_chat.rs delete mode 100644 tests/codex_serving/routes_responses.rs delete mode 100644 tests/codex_serving/upstream_errors.rs delete mode 100644 tests/codex_serving/upstream_fallback.rs delete mode 100644 tests/platform.rs delete mode 100644 tests/platform/http_auth.rs delete mode 100644 tests/runtime.rs delete mode 100644 tests/runtime/startup.rs delete mode 100644 tests/support/admin_accounts.rs delete mode 100644 tests/support/mod.rs delete mode 100644 tests/support/upstream.rs diff --git a/.gitignore b/.gitignore index 9ce0dd4f0..93b6f71b0 100644 --- a/.gitignore +++ b/.gitignore @@ -1,22 +1,32 @@ +# Rust build output. /target/ -/data/ -/logs/ -/local.yaml -/local.yml -/.worktrees/ -/worktrees/ -/.codex-ws-audit/ + +# Application runtime state. +/.runtime/ *.sqlite *.sqlite3 *.db *.db-shm *.db-wal *.log + +# Local configuration. +/local.yaml +/local.yml + +# Local workspaces and tooling state. +/.worktrees/ +/worktrees/ +/.codex-ws-audit/ +/.codegraph/ *.tmp *.swp *~ *.codegraph -*/node_modules -*/dist -.agents -.cargo +/.agents/ +/.cargo/ +/.reasonix/ + +# Frontend dependencies and build output. +node_modules/ +dist/ diff --git a/Cargo.lock b/Cargo.lock index 45b52f8da..ffe1fa8f1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -60,9 +60,9 @@ checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" [[package]] name = "alloc-stdlib" -version = "0.2.2" +version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" dependencies = [ "alloc-no-stdlib", ] @@ -262,9 +262,9 @@ dependencies = [ [[package]] name = "brotli" -version = "8.0.3" +version = "8.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8119e4516436f5708bbc474a9d395bf12f1b5395e93a92a56e647ac3388c8610" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" dependencies = [ "alloc-no-stdlib", "alloc-stdlib", @@ -273,9 +273,9 @@ dependencies = [ [[package]] name = "brotli-decompressor" -version = "5.0.1" +version = "5.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5962523e1b92ce1b5e793d9169b9943eece10d39f62550bc04bb605d75b94924" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" dependencies = [ "alloc-no-stdlib", "alloc-stdlib", @@ -365,27 +365,78 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" [[package]] -name = "codex-proxy-rs" +name = "codex-proxy-adapters" +version = "0.1.0" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "codex-proxy-core", + "codex-proxy-platform", + "flate2", + "futures", + "indexmap", + "reqwest", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "secrecy", + "serde", + "serde_json", + "sqlx", + "tempfile", + "thiserror 2.0.18", + "tokio", + "tokio-rustls", + "tokio-tungstenite", + "tracing", + "uuid", + "wiremock", +] + +[[package]] +name = "codex-proxy-assets" +version = "0.1.0" +dependencies = [ + "axum", + "tempfile", + "tokio", + "tower", +] + +[[package]] +name = "codex-proxy-core" +version = "0.1.0" +dependencies = [ + "async-trait", + "base64", + "chrono", + "hex", + "indexmap", + "rand 0.10.1", + "serde", + "serde_json", + "sha2 0.11.0", + "thiserror 2.0.18", + "tokio", + "tracing", + "uuid", +] + +[[package]] +name = "codex-proxy-platform" version = "0.1.0" dependencies = [ "aes-gcm", - "anyhow", "argon2", - "async-trait", - "axum", "base64", "chrono", "config", "dirs", "flate2", - "futures", "hex", "hmac", - "indexmap", - "insta", - "pretty_assertions", "rand 0.10.1", - "reqwest", "rustls", "rustls-native-certs", "rustls-pki-types", @@ -398,18 +449,77 @@ dependencies = [ "tempfile", "thiserror 2.0.18", "tokio", - "tokio-rustls", - "tokio-tungstenite", - "tower", - "tower-http", "tracing", "tracing-appender", "tracing-subscriber", "uuid", - "wiremock", "zeroize", ] +[[package]] +name = "codex-proxy-rs" +version = "0.1.0" +dependencies = [ + "codex-proxy-core", + "codex-proxy-runtime", + "codex-proxy-server", +] + +[[package]] +name = "codex-proxy-runtime" +version = "0.1.0" +dependencies = [ + "async-trait", + "base64", + "bytes", + "chrono", + "codex-proxy-adapters", + "codex-proxy-core", + "codex-proxy-platform", + "futures", + "reqwest", + "secrecy", + "serde_json", + "sqlx", + "tempfile", + "thiserror 2.0.18", + "tokio", + "tokio-tungstenite", + "tracing", + "uuid", + "wiremock", +] + +[[package]] +name = "codex-proxy-server" +version = "0.1.0" +dependencies = [ + "async-trait", + "axum", + "base64", + "chrono", + "codex-proxy-adapters", + "codex-proxy-assets", + "codex-proxy-core", + "codex-proxy-platform", + "codex-proxy-runtime", + "futures", + "secrecy", + "serde", + "serde_json", + "serde_yml", + "sqlx", + "tempfile", + "tokio", + "tokio-tungstenite", + "tower", + "tower-http", + "tracing", + "tracing-subscriber", + "uuid", + "wiremock", +] + [[package]] name = "compression-codecs" version = "0.4.38" @@ -459,17 +569,6 @@ dependencies = [ "yaml-rust2", ] -[[package]] -name = "console" -version = "0.16.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d64e8af5551369d19cf50138de61f1c42074ab970f74e99be916646777f8fc87" -dependencies = [ - "encode_unicode", - "libc", - "windows-sys 0.61.2", -] - [[package]] name = "const-oid" version = "0.10.2" @@ -693,12 +792,6 @@ dependencies = [ "powerfmt", ] -[[package]] -name = "diff" -version = "0.1.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "56254986775e3233ffa9c4d7d3faaf6d36a2c09d30b20687e9f88bc8bafc16c8" - [[package]] name = "digest" version = "0.10.7" @@ -787,12 +880,6 @@ dependencies = [ "serde", ] -[[package]] -name = "encode_unicode" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" - [[package]] name = "encoding_rs" version = "0.8.35" @@ -1072,9 +1159,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.14" +version = "0.4.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "171fefbc92fe4a4de27e0698d6a5b392d6a0e333506bc49133760b3bcf948733" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" dependencies = [ "atomic-waker", "bytes", @@ -1435,18 +1522,6 @@ dependencies = [ "generic-array", ] -[[package]] -name = "insta" -version = "1.47.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b4a6248eb93a4401ed2f37dfe8ea592d3cf05b7cf4f8efa867b6895af7e094e" -dependencies = [ - "console", - "once_cell", - "similar", - "tempfile", -] - [[package]] name = "ipnet" version = "2.12.0" @@ -1850,16 +1925,6 @@ dependencies = [ "zerocopy", ] -[[package]] -name = "pretty_assertions" -version = "1.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ae130e2f271fbc2ac3a40fb1d07180839cdbbe443c7a27e1e3c13c5cac0116d" -dependencies = [ - "diff", - "yansi", -] - [[package]] name = "prettyplease" version = "0.2.37" @@ -2460,12 +2525,6 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" -[[package]] -name = "similar" -version = "2.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbbb5d9659141646ae647b42fe094daf6c6192d1620870b449d9557f748b2daa" - [[package]] name = "slab" version = "0.4.12" @@ -3771,6 +3830,10 @@ version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +[[package]] +name = "xtask" +version = "0.1.0" + [[package]] name = "yaml-rust2" version = "0.11.0" @@ -3782,12 +3845,6 @@ dependencies = [ "hashlink", ] -[[package]] -name = "yansi" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049" - [[package]] name = "yoke" version = "0.8.3" diff --git a/Cargo.toml b/Cargo.toml index 23152d9b4..0c96fd4f4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,25 +1,46 @@ [package] name = "codex-proxy-rs" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +rust-version.workspace = true +description = "A lean Rust Codex proxy backed only by ChatGPT/Codex accounts." +repository = "https://example.invalid/codex-proxy-rs" +autobins = false + +[dev-dependencies] +codex-proxy-core = { path = "crates/core" } +codex-proxy-runtime = { path = "crates/runtime" } +codex-proxy-server = { path = "crates/server" } + +[lints] +workspace = true + +[workspace] +members = ["crates/*"] +resolver = "2" + +[workspace.package] version = "0.1.0" edition = "2021" license = "MIT" publish = false -description = "A lean Rust Codex proxy backed only by ChatGPT/Codex accounts." -repository = "https://example.invalid/codex-proxy-rs" rust-version = "1.95" -[dependencies] +[workspace.dependencies] aes-gcm = "0.10.3" anyhow = "1.0.102" argon2 = "0.5.3" async-trait = "0.1.89" axum = "0.8.9" base64 = "0.22.1" +bytes = "1.11.1" chrono = { version = "0.4.45", features = ["serde", "clock"] } config = "0.15.23" dirs = "5.0.1" -futures = "0.3.32" flate2 = "1.1.9" +futures = "0.3.32" hex = "0.4.3" hmac = "0.13.0" indexmap = { version = "2.0", features = ["serde"] } @@ -46,19 +67,19 @@ tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] } uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] } zeroize = "1.8.2" -[dev-dependencies] -insta = "1.47.2" -pretty_assertions = "1.4.1" -tempfile = "3.27.0" -wiremock = "0.6.5" - -[lints.rust] +[workspace.lints.rust] future_incompatible = { level = "warn", priority = -1 } nonstandard_style = { level = "deny", priority = -1 } unsafe_code = "forbid" -[lints.clippy] +[workspace.lints.clippy] all = { level = "deny", priority = 10 } redundant_clone = { level = "deny", priority = 9 } large_enum_variant = { level = "warn", priority = 8 } needless_collect = { level = "deny", priority = 7 } + +[profile.dev] +debug = true + +[profile.release] +debug = true diff --git a/config.yaml b/config.yaml index ebfcbe23b..9e094018b 100644 --- a/config.yaml +++ b/config.yaml @@ -32,10 +32,10 @@ quota: usage_stats: history_retention_days: null database: - url: sqlite://data/codex-proxy-rs.sqlite + url: sqlite://.runtime/data/codex-proxy-rs.sqlite security: - master_key_file: data/master.key - api_key_pepper_file: data/api-key-pepper.key + master_key_file: .runtime/data/master.key + api_key_pepper_file: .runtime/data/api-key-pepper.key tls: force_http11: false ws_pool: @@ -47,7 +47,7 @@ admin: default_username: admin default_password: admin logging: - directory: logs + directory: .runtime/logs retention_days: 14 enabled: true capacity: 2000 diff --git a/crates/adapters/Cargo.toml b/crates/adapters/Cargo.toml new file mode 100644 index 000000000..c8445de59 --- /dev/null +++ b/crates/adapters/Cargo.toml @@ -0,0 +1,38 @@ +[package] +name = "codex-proxy-adapters" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +rust-version.workspace = true + +[dependencies] +async-trait.workspace = true +bytes.workspace = true +codex-proxy-core = { path = "../core" } +codex-proxy-platform = { path = "../platform" } +chrono.workspace = true +flate2.workspace = true +futures.workspace = true +indexmap.workspace = true +reqwest.workspace = true +rustls.workspace = true +rustls-native-certs.workspace = true +rustls-pki-types.workspace = true +secrecy.workspace = true +serde.workspace = true +serde_json.workspace = true +sqlx.workspace = true +thiserror.workspace = true +tokio.workspace = true +tokio-rustls.workspace = true +tokio-tungstenite.workspace = true +tracing.workspace = true +uuid.workspace = true + +[dev-dependencies] +tempfile = "3.27.0" +wiremock = "0.6.5" + +[lints] +workspace = true diff --git a/crates/adapters/src/codex/client.rs b/crates/adapters/src/codex/client.rs new file mode 100644 index 000000000..b2f4839e8 --- /dev/null +++ b/crates/adapters/src/codex/client.rs @@ -0,0 +1,1497 @@ +//! Codex HTTP/SSE 上游客户端与请求头构造。 + +use std::{ + collections::HashMap, + env, fs, io, + path::PathBuf, + pin::Pin, + sync::{Arc, Mutex, OnceLock}, + time::Duration, +}; + +use bytes::Bytes; +use futures::{Stream, TryStreamExt}; +use indexmap::IndexMap; +use reqwest::{ + header::{ + HeaderMap, HeaderName, HeaderValue, ACCEPT, ACCEPT_ENCODING, CONTENT_TYPE, COOKIE, + RETRY_AFTER, SET_COOKIE, + }, + Client, Response as ReqwestResponse, StatusCode, +}; +use rustls::{ClientConfig, RootCertStore}; +use rustls_pki_types::{ + pem::{self, PemObject, SectionKind}, + CertificateDer, +}; +use serde_json::{Map, Value}; +use thiserror::Error; +use uuid::Uuid; + +use codex_proxy_core::{ + gateway::{ + fingerprint::Fingerprint, + ports::{CodexModelCatalogClient, CodexModelCatalogClientError, CodexModelCatalogRequest}, + }, + models::model::BackendModelEntry, + protocol::codex::{ + events::{extract_sse_usage, retry_after_seconds_from_body, TokenUsage}, + responses::{CodexCompactRequest, CodexResponsesRequest}, + sse::SseError, + websocket::{websocket_audit_artifact_from_attempt, websocket_payload_audit_snapshot}, + }, + serving::responses::{http_sse_fallback_allowed, transport_for_request, CodexTransport}, +}; + +use async_trait::async_trait; +use tokio_tungstenite::tungstenite::handshake::client::generate_key; + +use super::websocket::connect::{ + execute_response_create_request, execute_response_create_request_stream_with_pool, + execute_response_create_request_with_pool, CodexWebSocketConnection, + CodexWebSocketExchangeError, CodexWebSocketRateLimitHeaderUpdates, + CodexWebSocketTurnStateUpdate, +}; +use super::websocket::opening::write_websocket_audit_artifact_from_env; +use super::websocket::pool::{CodexWebSocketPool, CodexWebSocketPoolKey}; + +const MAX_UPSTREAM_ERROR_BODY_BYTES: usize = 1024 * 1024; +const CA_CERT_HINT: &str = "If you set CODEX_CA_CERTIFICATE or SSL_CERT_FILE, ensure it points to a PEM file containing one or more CERTIFICATE blocks, or unset it to use system roots."; + +type PemSection = (SectionKind, Vec); + +type ReqwestClientCacheKey = (bool, Option); +type ReqwestClientCache = Mutex>; + +/// Codex HTTP 客户端类型。 +pub type HttpClient = Client; + +/// `/codex/responses` +pub const CODEX_RESPONSES_PATH: &str = "/codex/responses"; +/// `/codex/responses/compact` +pub const CODEX_RESPONSES_COMPACT_PATH: &str = "/codex/responses/compact"; +/// `/codex/usage` +pub const CODEX_USAGE_PATH: &str = "/codex/usage"; +/// `/api/codex/usage` +pub const CODEX_USAGE_API_PATH: &str = "/api/codex/usage"; +/// `/wham/usage` +pub const WHAM_USAGE_PATH: &str = "/wham/usage"; +/// 自定义 CA 证书环境变量名。 +pub const CODEX_CA_CERT_ENV: &str = "CODEX_CA_CERTIFICATE"; +/// 系统 CA 文件环境变量名。 +pub const SSL_CERT_FILE_ENV: &str = "SSL_CERT_FILE"; + +/// Codex 上游 HTTP 客户端错误。 +#[derive(Debug, Error)] +pub enum CodexClientError { + /// Reqwest 传输失败。 + #[error("http transport error: {0}")] + Http(#[from] reqwest::Error), + /// 自定义 CA 构建失败。 + #[error("custom CA transport error: {0}")] + CustomCa(#[from] CustomCaError), + /// 请求头名字无效。 + #[error("invalid request header name: {0}")] + InvalidHeaderName(#[from] reqwest::header::InvalidHeaderName), + /// 请求头值无效。 + #[error("invalid request header value: {0}")] + InvalidHeaderValue(#[from] reqwest::header::InvalidHeaderValue), + /// SSE 响应解析失败。 + #[error("invalid upstream SSE response: {0}")] + InvalidSse(#[from] SseError), + /// WebSocket 请求编码失败。 + #[error("failed to encode websocket request: {0}")] + WebSocketEncode(#[source] serde_json::Error), + /// WebSocket 请求失败。 + #[error("websocket request failed: {0}")] + WebSocket(#[from] CodexWebSocketExchangeError), + /// 上游返回非成功响应。 + #[error("upstream returned status {status}: {body}")] + Upstream { + /// 上游状态码。 + status: StatusCode, + /// 上游错误体。 + body: String, + /// 推导出的重试秒数。 + retry_after_seconds: Option, + }, +} + +/// Codex 客户端结果类型。 +pub type CodexClientResult = Result; + +/// Codex SSE 字节流。 +pub type CodexBackendSseStream = + Pin> + Send + 'static>>; + +/// 自定义 CA 错误。 +#[derive(Debug, Error)] +pub enum CustomCaError { + /// 读取 CA 证书文件失败。 + #[error( + "Failed to read CA certificate file {} selected by {}: {source}. {hint}", + path.display(), + source_env, + hint = CA_CERT_HINT + )] + ReadCaFile { + /// 来源环境变量名。 + source_env: &'static str, + /// 证书路径。 + path: PathBuf, + /// 底层 IO 错误。 + source: io::Error, + }, + /// CA 证书文件格式无效。 + #[error( + "Failed to load CA certificates from {} selected by {}: {detail}. {hint}", + path.display(), + source_env, + hint = CA_CERT_HINT + )] + InvalidCaFile { + /// 来源环境变量名。 + source_env: &'static str, + /// 证书路径。 + path: PathBuf, + /// 详细错误。 + detail: String, + }, + /// 证书无法注册为 reqwest 根证书。 + #[error( + "Failed to parse certificate #{certificate_index} from {} selected by {}: {source}. {hint}", + path.display(), + source_env, + hint = CA_CERT_HINT + )] + RegisterCertificate { + /// 来源环境变量名。 + source_env: &'static str, + /// 证书路径。 + path: PathBuf, + /// 证书序号。 + certificate_index: usize, + /// 底层错误。 + source: reqwest::Error, + }, + /// 证书无法注册到 rustls root store。 + #[error( + "Failed to register certificate #{certificate_index} from {} selected by {} in rustls root store: {source}. {hint}", + path.display(), + source_env, + hint = CA_CERT_HINT + )] + RegisterRustlsCertificate { + /// 来源环境变量名。 + source_env: &'static str, + /// 证书路径。 + path: PathBuf, + /// 证书序号。 + certificate_index: usize, + /// 底层错误。 + source: rustls::Error, + }, + /// 使用自定义 CA 构建 reqwest client 失败。 + #[error("Failed to build HTTP client while using CA bundle from {} ({}): {source}", source_env, path.display())] + BuildClientWithCustomCa { + /// 来源环境变量名。 + source_env: &'static str, + /// 证书路径。 + path: PathBuf, + /// 底层错误。 + source: reqwest::Error, + }, + /// 使用系统根证书构建 reqwest client 失败。 + #[error("Failed to build HTTP client while using system root certificates: {0}")] + BuildClientWithSystemRoots(reqwest::Error), + /// 读取系统根证书失败。 + #[error("Failed to load native root certificates for custom CA transport: {0}")] + LoadNativeRoots(io::Error), +} + +/// 自定义 CA 结果类型。 +pub type CustomCaResult = Result; + +/// 拼接完整 endpoint URL。 +pub fn endpoint_url(base_url: &str, endpoint_path: &str) -> String { + format!( + "{}/{}", + base_url.trim_end_matches('/'), + endpoint_path.trim_start_matches('/') + ) +} + +/// 计算请求层路径。 +pub fn endpoint_request_path(base_url: &str, endpoint_path: &str) -> String { + let endpoint_path = endpoint_path.trim_start_matches('/'); + let base_path = reqwest::Url::parse(base_url) + .ok() + .map(|url| url.path().trim_end_matches('/').to_string()) + .filter(|path| !path.is_empty()) + .unwrap_or_default(); + + if base_path.is_empty() { + format!("/{endpoint_path}") + } else { + format!("{base_path}/{endpoint_path}") + } +} + +/// 返回 usage 相关 endpoint 的完整 URL 列表。 +pub fn usage_endpoint_urls(base_url: &str) -> Vec { + usage_endpoint_paths(base_url) + .into_iter() + .map(|path| endpoint_url(base_url, path)) + .collect() +} + +/// 返回 usage 主请求路径。 +pub fn primary_usage_request_path(base_url: &str) -> String { + let endpoint_path = usage_endpoint_paths(base_url) + .into_iter() + .next() + .unwrap_or(CODEX_USAGE_API_PATH); + endpoint_request_path(base_url, endpoint_path) +} + +/// 单次 Codex 上游请求的上下文。 +#[derive(Debug, Clone, Copy)] +pub struct CodexRequestContext<'a> { + /// Access token。 + pub access_token: &'a str, + /// ChatGPT account id。 + pub account_id: Option<&'a str>, + /// 请求 ID。 + pub request_id: &'a str, + /// x-codex-turn-state。 + pub turn_state: Option<&'a str>, + /// x-codex-turn-metadata。 + pub turn_metadata: Option<&'a str>, + /// x-codex-beta-features。 + pub beta_features: Option<&'a str>, + /// x-responsesapi-include-timing-metrics。 + pub include_timing_metrics: Option<&'a str>, + /// version。 + pub version: Option<&'a str>, + /// x-codex-window-id。 + pub codex_window_id: Option<&'a str>, + /// x-codex-parent-thread-id。 + pub parent_thread_id: Option<&'a str>, + /// cookie 头。 + pub cookie_header: Option<&'a str>, + /// x-codex-installation-id。 + pub installation_id: Option<&'a str>, + /// session_id。 + pub session_id: Option<&'a str>, +} + +/// Codex Responses 上游响应。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexBackendResponse { + /// 完整 SSE 文本。 + pub body: String, + /// 从 SSE 中提取出的最终 usage。 + pub usage: Option, + /// 响应头里的最新 turn state。 + pub turn_state: Option, + /// 上游透传的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游透传的限流头。 + pub rate_limit_headers: Vec<(String, String)>, +} + +/// Codex Responses 实际使用的上游传输。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CodexBackendTransport { + /// HTTP SSE transport. + HttpSse, + /// WebSocket transport. + WebSocket, +} + +/// Live stream rate-limit updates captured after the response headers. +pub type CodexRateLimitHeaderUpdates = CodexWebSocketRateLimitHeaderUpdates; + +/// Live stream turn-state updates captured after the response headers. +pub type CodexTurnStateUpdate = CodexWebSocketTurnStateUpdate; + +/// Codex Responses 上游 live SSE 响应。 +pub struct CodexBackendStreamingResponse { + /// 上游 SSE 字节流。 + pub body: CodexBackendSseStream, + /// 实际使用的上游传输。 + pub transport: CodexBackendTransport, + /// 响应头里的最新 turn state。 + pub turn_state: Option, + /// 上游透传的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游透传的限流头。 + pub rate_limit_headers: Vec<(String, String)>, + /// live stream 期间捕获的限流头更新。 + pub rate_limit_header_updates: Option, + /// live stream 期间捕获的 turn-state 更新。 + pub turn_state_update: Option, +} + +/// Codex compact 端点响应。 +#[derive(Debug, Clone, PartialEq)] +pub struct CodexCompactResponse { + /// 上游返回的 JSON。 + pub body: Value, + /// 上游透传的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游透传的限流头。 + pub rate_limit_headers: Vec<(String, String)>, +} + +/// 上游模型端点探测结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexConnectivityProbe { + /// 请求的完整端点。 + pub endpoint: String, + /// 返回状态码。 + pub status: StatusCode, +} + +/// Codex HTTP/SSE 上游客户端。 +#[derive(Clone)] +pub struct CodexBackendClient { + client: Client, + base_url: String, + fingerprint: Fingerprint, + websocket_pool: Option>, +} + +impl CodexBackendClient { + /// 构造客户端。 + pub fn new(client: Client, base_url: impl Into, fingerprint: Fingerprint) -> Self { + Self { + client, + base_url: base_url.into().trim_end_matches('/').to_string(), + fingerprint, + websocket_pool: None, + } + } + + /// 为 Responses WebSocket 请求启用连接池。 + pub fn with_websocket_pool(mut self, pool: Arc) -> Self { + self.websocket_pool = Some(pool); + self + } + + /// 驱逐指定账号的 Responses WebSocket 池连接。 + pub async fn evict_websocket_account(&self, account_id: &str) { + if let Some(pool) = &self.websocket_pool { + pool.evict_account(account_id).await; + } + } + + /// 发送 Responses SSE 请求并读取完整响应。 + pub async fn create_response( + &self, + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let upstream_request = response_upstream_request(request, context); + match transport_for_request(&upstream_request) { + CodexTransport::HttpSse => { + self.create_response_http_sse(&upstream_request, context) + .await + } + CodexTransport::WebSocketPreferred | CodexTransport::WebSocketRequired => { + match self + .create_response_websocket(&upstream_request, context) + .await + { + Ok(response) => Ok(response), + Err(error) + if http_sse_fallback_allowed(&upstream_request) + && websocket_error_allows_http_fallback(&error) => + { + tracing::warn!(error = %error, "websocket response failed; falling back to HTTP SSE"); + self.create_response_http_sse(&upstream_request, context) + .await + } + Err(error) => Err(error), + } + } + } + } + + /// 发送 Responses HTTP SSE 请求并返回 live 字节流。 + pub async fn create_response_stream( + &self, + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let upstream_request = response_upstream_request(request, context); + match transport_for_request(&upstream_request) { + CodexTransport::HttpSse => { + self.create_response_http_sse_stream(&upstream_request, context) + .await + } + CodexTransport::WebSocketPreferred | CodexTransport::WebSocketRequired => { + match self + .create_response_websocket_stream(&upstream_request, context) + .await + { + Ok(response) => Ok(response), + Err(error) + if http_sse_fallback_allowed(&upstream_request) + && websocket_error_allows_http_fallback(&error) => + { + tracing::warn!(error = %error, "websocket response stream failed; falling back to HTTP SSE"); + self.create_response_http_sse_stream(&upstream_request, context) + .await + } + Err(error) => Err(error), + } + } + } + } + + async fn create_response_http_sse( + &self, + upstream_request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let headers = self.request_headers_for_http_response(upstream_request, context)?; + let response = self + .client + .post(endpoint_url(&self.base_url, CODEX_RESPONSES_PATH)) + .headers(headers) + .json(&upstream_request) + .send() + .await?; + let status = response.status(); + let turn_state = turn_state(&response); + let set_cookie_headers = set_cookie_headers(&response); + let rate_limit_headers = rate_limit_headers(response.headers()); + let retry_after_seconds = retry_after_seconds(response.headers(), None); + + if !status.is_success() { + let body = read_capped_error_body(response).await?; + return Err(CodexClientError::Upstream { + status, + retry_after_seconds: retry_after_seconds + .or_else(|| retry_after_seconds_from_body(&body)), + body, + }); + } + + let body = response.text().await?; + let usage = extract_sse_usage(&body).map_err(CodexClientError::InvalidSse)?; + Ok(CodexBackendResponse { + body, + usage, + turn_state, + set_cookie_headers, + rate_limit_headers, + }) + } + + async fn create_response_http_sse_stream( + &self, + upstream_request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let headers = self.request_headers_for_http_response(upstream_request, context)?; + let response = self + .client + .post(endpoint_url(&self.base_url, CODEX_RESPONSES_PATH)) + .headers(headers) + .json(&upstream_request) + .send() + .await?; + let status = response.status(); + let turn_state = turn_state(&response); + let set_cookie_headers = set_cookie_headers(&response); + let rate_limit_headers = rate_limit_headers(response.headers()); + let retry_after_seconds = retry_after_seconds(response.headers(), None); + + if !status.is_success() { + let body = read_capped_error_body(response).await?; + return Err(CodexClientError::Upstream { + status, + retry_after_seconds: retry_after_seconds + .or_else(|| retry_after_seconds_from_body(&body)), + body, + }); + } + + Ok(CodexBackendStreamingResponse { + body: Box::pin(response.bytes_stream().map_err(CodexClientError::Http)), + transport: CodexBackendTransport::HttpSse, + turn_state, + set_cookie_headers, + rate_limit_headers, + rate_limit_header_updates: None, + turn_state_update: None, + }) + } + + async fn create_response_websocket( + &self, + upstream_request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let headers = self.request_headers_for_http_response(upstream_request, context)?; + let prepared = CodexWebSocketConnection::responses_create_request( + &self.base_url, + &generate_key(), + websocket_header_pairs(&headers, context), + upstream_request, + ) + .map_err(CodexClientError::WebSocketEncode)?; + let artifact = websocket_audit_artifact_from_attempt( + upstream_request, + prepared.connection().opening_audit_snapshot(), + websocket_payload_audit_snapshot(upstream_request), + ); + if let Err(error) = write_websocket_audit_artifact_from_env(&artifact).await { + tracing::warn!(error = %error, "failed to write Codex WebSocket audit artifact"); + } + let pool_key = self.websocket_pool_key(upstream_request, context); + let exchange = match (self.websocket_pool.as_deref(), pool_key) { + (Some(pool), Some(key)) => { + execute_response_create_request_with_pool(&prepared, Some((pool, key))).await + } + _ => execute_response_create_request(&prepared).await, + } + .map_err(websocket_exchange_error_to_client_error)?; + + Ok(CodexBackendResponse { + body: exchange.body, + usage: exchange.usage, + turn_state: exchange.turn_state, + set_cookie_headers: exchange.set_cookie_headers, + rate_limit_headers: exchange.rate_limit_headers, + }) + } + + async fn create_response_websocket_stream( + &self, + upstream_request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let headers = self.request_headers_for_http_response(upstream_request, context)?; + let prepared = CodexWebSocketConnection::responses_create_request( + &self.base_url, + &generate_key(), + websocket_header_pairs(&headers, context), + upstream_request, + ) + .map_err(CodexClientError::WebSocketEncode)?; + let artifact = websocket_audit_artifact_from_attempt( + upstream_request, + prepared.connection().opening_audit_snapshot(), + websocket_payload_audit_snapshot(upstream_request), + ); + if let Err(error) = write_websocket_audit_artifact_from_env(&artifact).await { + tracing::warn!(error = %error, "failed to write Codex WebSocket audit artifact"); + } + let pool_key = self.websocket_pool_key(upstream_request, context); + let exchange = match (self.websocket_pool.as_deref(), pool_key) { + (Some(pool), Some(key)) => { + execute_response_create_request_stream_with_pool(&prepared, Some((pool, key))).await + } + _ => execute_response_create_request_stream_with_pool(&prepared, None).await, + } + .map_err(websocket_exchange_error_to_client_error)?; + + Ok(CodexBackendStreamingResponse { + body: Box::pin( + exchange + .body + .map_err(websocket_exchange_error_to_client_error), + ), + transport: CodexBackendTransport::WebSocket, + turn_state: exchange.turn_state, + set_cookie_headers: exchange.set_cookie_headers, + rate_limit_headers: exchange.rate_limit_headers, + rate_limit_header_updates: Some(exchange.rate_limit_header_updates), + turn_state_update: Some(exchange.turn_state_update), + }) + } + + fn websocket_pool_key( + &self, + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> Option { + let account_id = context.account_id?; + let conversation_id = request + .prompt_cache_key + .as_deref() + .or(request.client_conversation_id.as_deref()) + .or(request.previous_response_id.as_deref())?; + Some(CodexWebSocketPoolKey::new( + &self.base_url, + account_id, + conversation_id, + )) + } + + /// 发送 compact JSON 请求并读取完整响应。 + pub async fn create_compact_response( + &self, + request: &CodexCompactRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let headers = self.compact_request_headers(context)?; + let response = self + .client + .post(endpoint_url(&self.base_url, CODEX_RESPONSES_COMPACT_PATH)) + .headers(headers) + .json(request) + .send() + .await?; + let status = response.status(); + let set_cookie_headers = set_cookie_headers(&response); + let rate_limit_headers = rate_limit_headers(response.headers()); + let retry_after_seconds = retry_after_seconds(response.headers(), None); + let body = response.text().await?; + + if !status.is_success() { + return Err(CodexClientError::Upstream { + status, + retry_after_seconds: retry_after_seconds + .or_else(|| retry_after_seconds_from_body(&body)), + body, + }); + } + + let parsed = + serde_json::from_str::(&body).map_err(|_| CodexClientError::Upstream { + status: StatusCode::BAD_GATEWAY, + retry_after_seconds: None, + body: format!( + "Compact response is not valid JSON: {}", + truncate_for_error(&body) + ), + })?; + + Ok(CodexCompactResponse { + body: parsed, + set_cookie_headers, + rate_limit_headers, + }) + } + + /// 获取 Codex usage JSON。 + pub async fn fetch_usage(&self, context: CodexRequestContext<'_>) -> CodexClientResult { + let mut last_invalid_body = None; + + for endpoint in usage_endpoint_urls(&self.base_url) { + let headers = self.auxiliary_request_headers(context)?; + let response = self.client.get(endpoint).headers(headers).send().await?; + let status = response.status(); + let retry_after_seconds = retry_after_seconds(response.headers(), None); + + if status == StatusCode::NOT_FOUND { + last_invalid_body = Some(read_capped_error_body(response).await?); + continue; + } + if !status.is_success() { + let body = read_capped_error_body(response).await?; + return Err(CodexClientError::Upstream { + status, + retry_after_seconds: retry_after_seconds + .or_else(|| retry_after_seconds_from_body(&body)), + body, + }); + } + + let body = response.text().await?; + match serde_json::from_str::(&body) { + Ok(parsed) if parsed.get("rate_limit").is_some() => return Ok(parsed), + _ => last_invalid_body = Some(body), + } + } + + Err(CodexClientError::Upstream { + status: StatusCode::BAD_GATEWAY, + retry_after_seconds: None, + body: last_invalid_body + .map(|body| format!("invalid usage response: {}", truncate_for_error(&body))) + .unwrap_or_else(|| "usage endpoint is unavailable".to_string()), + }) + } + + /// 获取后端模型目录条目。 + pub async fn fetch_models( + &self, + context: CodexRequestContext<'_>, + ) -> CodexClientResult> { + let endpoints = [ + format!( + "{}/codex/models?client_version={}", + self.base_url, self.fingerprint.app_version + ), + format!("{}/models", self.base_url), + format!("{}/sentinel/chat-requirements", self.base_url), + ]; + + for endpoint in endpoints { + let headers = self.auxiliary_request_headers(context)?; + let response = self.client.get(endpoint).headers(headers).send().await?; + if !response.status().is_success() { + continue; + } + let parsed = response.json::().await?; + let models = extract_backend_model_entries(&parsed); + if !models.is_empty() { + return Ok(models); + } + } + + Err(CodexClientError::Upstream { + status: StatusCode::BAD_GATEWAY, + retry_after_seconds: None, + body: "backend model catalog is unavailable".to_string(), + }) + } + + /// 探测主模型端点是否可达。 + pub async fn probe_models_endpoint( + &self, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let endpoint = format!( + "{}/codex/models?client_version={}", + self.base_url, self.fingerprint.app_version + ); + let headers = self.auxiliary_request_headers(context)?; + let response = self.client.get(&endpoint).headers(headers).send().await?; + + Ok(CodexConnectivityProbe { + endpoint, + status: response.status(), + }) + } + + fn request_headers_for_http_response( + &self, + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let mut headers = self.request_headers(context)?; + if let Some(subagent) = openai_subagent_from_metadata(request.client_metadata.as_ref()) { + headers.insert( + HeaderName::from_static("x-openai-subagent"), + HeaderValue::from_str(&subagent)?, + ); + } + Ok(headers) + } + + fn request_headers(&self, context: CodexRequestContext<'_>) -> CodexClientResult { + let request_id = context.session_id.unwrap_or(context.request_id); + let ordered_headers = build_ordered_codex_base_headers( + &self.fingerprint, + context.access_token, + context.account_id, + ); + + let mut headers = HeaderMap::new(); + insert_ordered_headers(&mut headers, &ordered_headers)?; + headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + insert_optional_header(&mut headers, "cookie", context.cookie_header)?; + headers.insert(ACCEPT, HeaderValue::from_static("text/event-stream")); + headers.insert( + HeaderName::from_static("openai-beta"), + HeaderValue::from_static("responses_websockets=2026-02-06"), + ); + headers.insert( + HeaderName::from_static("x-openai-internal-codex-residency"), + HeaderValue::from_static("us"), + ); + headers.insert( + HeaderName::from_static("x-client-request-id"), + HeaderValue::from_str(request_id)?, + ); + insert_optional_header( + &mut headers, + "x-codex-installation-id", + context.installation_id, + )?; + insert_optional_header(&mut headers, "session_id", context.session_id)?; + insert_optional_header(&mut headers, "x-codex-window-id", context.codex_window_id)?; + insert_optional_header(&mut headers, "x-codex-turn-state", context.turn_state)?; + insert_optional_header(&mut headers, "x-codex-turn-metadata", context.turn_metadata)?; + insert_optional_header(&mut headers, "x-codex-beta-features", context.beta_features)?; + insert_optional_header( + &mut headers, + "x-responsesapi-include-timing-metrics", + context.include_timing_metrics, + )?; + insert_optional_header(&mut headers, "version", context.version)?; + insert_optional_header( + &mut headers, + "x-codex-parent-thread-id", + context.parent_thread_id, + )?; + + Ok(headers) + } + + fn auxiliary_request_headers( + &self, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let ordered_headers = build_ordered_codex_base_headers( + &self.fingerprint, + context.access_token, + context.account_id, + ); + + let mut headers = HeaderMap::new(); + insert_ordered_headers(&mut headers, &ordered_headers)?; + if let Some(cookie_header) = context.cookie_header { + headers.insert(COOKIE, HeaderValue::from_str(cookie_header)?); + } + headers.insert(ACCEPT, HeaderValue::from_static("application/json")); + headers.insert(ACCEPT_ENCODING, HeaderValue::from_static("gzip, deflate")); + insert_optional_header( + &mut headers, + "x-codex-installation-id", + context.installation_id, + )?; + Ok(headers) + } + + fn compact_request_headers( + &self, + context: CodexRequestContext<'_>, + ) -> CodexClientResult { + let ordered_headers = build_ordered_codex_base_headers( + &self.fingerprint, + context.access_token, + context.account_id, + ); + + let mut headers = HeaderMap::new(); + insert_ordered_headers(&mut headers, &ordered_headers)?; + headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + insert_optional_header(&mut headers, "cookie", context.cookie_header)?; + headers.insert( + HeaderName::from_static("openai-beta"), + HeaderValue::from_static("responses_websockets=2026-02-06"), + ); + headers.insert( + HeaderName::from_static("x-openai-internal-codex-residency"), + HeaderValue::from_static("us"), + ); + headers.insert( + HeaderName::from_static("x-client-request-id"), + HeaderValue::from_str(&Uuid::new_v4().to_string())?, + ); + insert_optional_header( + &mut headers, + "x-codex-installation-id", + context.installation_id, + )?; + + Ok(headers) + } +} + +#[async_trait] +impl CodexModelCatalogClient for CodexBackendClient { + async fn fetch_models( + &self, + request: &CodexModelCatalogRequest<'_>, + ) -> Result, CodexModelCatalogClientError> { + Self::fetch_models( + self, + CodexRequestContext { + access_token: request.access_token, + account_id: request.account_id, + request_id: request.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, + installation_id: request.installation_id, + session_id: None, + }, + ) + .await + .map_err(|error| CodexModelCatalogClientError::RequestFailed { + message: error.to_string(), + }) + } +} + +/// 构建基础的 Codex 请求头集合。 +pub fn build_codex_base_headers( + fingerprint: &Fingerprint, + access_token: &str, + account_id: Option<&str>, +) -> IndexMap { + let mut headers = IndexMap::new(); + + headers.insert( + "authorization".to_string(), + format!("Bearer {access_token}"), + ); + if let Some(account_id) = account_id { + headers.insert("chatgpt-account-id".to_string(), account_id.to_string()); + } + headers.insert("originator".to_string(), fingerprint.originator.clone()); + headers.insert("user-agent".to_string(), fingerprint.user_agent()); + headers.insert("sec-ch-ua".to_string(), fingerprint.sec_ch_ua()); + + for (key, value) in &fingerprint.default_headers { + let key_lower = key.to_ascii_lowercase(); + if !headers.contains_key(&key_lower) { + headers.insert(key_lower, value.clone()); + } + } + + headers +} + +/// 构造 Responses HTTP/SSE 请求头集合。 +pub fn build_codex_headers( + fingerprint: &Fingerprint, + access_token: &str, + account_id: Option<&str>, + turn_state: Option<&str>, + request_id: &str, +) -> IndexMap { + let mut headers = build_codex_base_headers(fingerprint, access_token, account_id); + headers.insert( + "x-openai-internal-codex-residency".to_string(), + "us".to_string(), + ); + headers.insert("x-client-request-id".to_string(), request_id.to_string()); + + if let Some(turn_state) = turn_state { + headers.insert("x-codex-turn-state".to_string(), turn_state.to_string()); + } + + headers.insert("accept".to_string(), "text/event-stream".to_string()); + headers +} + +/// 按指纹声明顺序重排请求头。 +pub fn order_headers( + headers: IndexMap, + order: &[String], +) -> IndexMap { + let mut ordered = IndexMap::new(); + + for key in order { + if let Some(value) = headers.get(key) { + ordered.insert(key.clone(), value.clone()); + } + } + + for (key, value) in headers { + if !ordered.contains_key(&key) { + ordered.insert(key, value); + } + } + + ordered +} + +/// 构造并按指纹顺序重排 Responses 请求头。 +pub fn build_ordered_codex_headers( + fingerprint: &Fingerprint, + access_token: &str, + account_id: Option<&str>, + turn_state: Option<&str>, + request_id: &str, +) -> IndexMap { + let headers = build_codex_headers( + fingerprint, + access_token, + account_id, + turn_state, + request_id, + ); + order_headers(headers, &fingerprint.header_order) +} + +/// 构造并按指纹顺序重排基础请求头。 +pub fn build_ordered_codex_base_headers( + fingerprint: &Fingerprint, + access_token: &str, + account_id: Option<&str>, +) -> IndexMap { + let headers = build_codex_base_headers(fingerprint, access_token, account_id); + order_headers(headers, &fingerprint.header_order) +} + +/// 按 TLS 配置和自定义 CA 环境复用 reqwest client。 +pub fn build_reqwest_client(force_http11: bool) -> Result { + let cache_key = (force_http11, custom_ca_env_cache_key()); + static CLIENTS: OnceLock = OnceLock::new(); + let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new())); + let mut clients = cache + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + + if let Some(client) = clients.get(&cache_key) { + return Ok(client.clone()); + } + + let builder = Client::builder() + .use_rustls_tls() + .no_proxy() + .pool_max_idle_per_host(4) + .tcp_keepalive(Duration::from_secs(30)) + .gzip(true) + .brotli(true) + .zstd(true) + .deflate(true); + let builder = if force_http11 { + builder.http1_only() + } else { + builder + }; + let client = build_reqwest_client_with_custom_ca(builder)?; + clients.insert(cache_key, client.clone()); + Ok(client) +} + +fn response_upstream_request( + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, +) -> CodexResponsesRequest { + let mut upstream = request.clone(); + if let Some(session_id) = context.session_id { + upstream.prompt_cache_key = Some(session_id.to_string()); + } + upstream.client_metadata = response_client_metadata(request.client_metadata.as_ref(), context); + upstream +} + +fn response_client_metadata( + client_metadata: Option<&Value>, + context: CodexRequestContext<'_>, +) -> Option { + let mut metadata = Map::new(); + if let Some(Value::Object(input)) = client_metadata { + for (key, value) in input { + if let Some(value) = value.as_str() { + metadata.insert(key.clone(), Value::String(value.to_string())); + } + } + } + + insert_metadata_string( + &mut metadata, + "x-codex-installation-id", + context.installation_id, + ); + insert_metadata_string(&mut metadata, "x-codex-window-id", context.codex_window_id); + insert_metadata_string( + &mut metadata, + "x-codex-turn-metadata", + context.turn_metadata, + ); + insert_metadata_string( + &mut metadata, + "x-codex-parent-thread-id", + context.parent_thread_id, + ); + + if metadata.is_empty() { + None + } else { + Some(Value::Object(metadata)) + } +} + +fn insert_metadata_string(metadata: &mut Map, key: &str, value: Option<&str>) { + if let Some(value) = value.filter(|value| !value.trim().is_empty()) { + metadata.insert(key.to_string(), Value::String(value.to_string())); + } +} + +fn openai_subagent_from_metadata(client_metadata: Option<&Value>) -> Option { + let value = client_metadata? + .as_object()? + .get("x-openai-subagent")? + .as_str()? + .trim(); + if matches!( + value, + "review" | "compact" | "memory_consolidation" | "collab_spawn" + ) { + Some(value.to_string()) + } else { + None + } +} + +fn extract_backend_model_entries(value: &Value) -> Vec { + let Some(models) = value + .pointer("/chat_models/models") + .or_else(|| value.get("models")) + .or_else(|| value.get("data")) + .or_else(|| value.get("categories")) + .and_then(Value::as_array) + else { + return Vec::new(); + }; + + let mut entries = Vec::new(); + for model in models { + if let Some(nested) = model.get("models").and_then(Value::as_array) { + entries.extend(nested.iter().filter_map(parse_backend_model_entry)); + } else if let Some(entry) = parse_backend_model_entry(model) { + entries.push(entry); + } + } + entries +} + +fn parse_backend_model_entry(value: &Value) -> Option { + let entry = serde_json::from_value::(value.clone()).ok()?; + (entry.slug.is_some() + || entry.id.is_some() + || entry.name.is_some() + || entry.display_name.is_some() + || entry.title.is_some()) + .then_some(entry) +} + +fn insert_optional_header( + headers: &mut HeaderMap, + name: &'static str, + value: Option<&str>, +) -> CodexClientResult<()> { + let Some(value) = value else { + return Ok(()); + }; + headers.insert(HeaderName::from_static(name), HeaderValue::from_str(value)?); + Ok(()) +} + +fn insert_ordered_headers( + headers: &mut HeaderMap, + ordered_headers: &IndexMap, +) -> CodexClientResult<()> { + for (name, value) in ordered_headers { + headers.insert( + HeaderName::from_bytes(name.as_bytes())?, + HeaderValue::from_str(value)?, + ); + } + Ok(()) +} + +fn header_pairs(headers: &HeaderMap) -> Vec<(String, String)> { + headers + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.as_str().to_string(), value.to_string())) + }) + .collect() +} + +fn websocket_header_pairs( + headers: &HeaderMap, + context: CodexRequestContext<'_>, +) -> Vec<(String, String)> { + let mut pairs = header_pairs(headers) + .into_iter() + .filter(|(name, _)| { + !name.eq_ignore_ascii_case("content-type") + && !name.eq_ignore_ascii_case("accept") + && !name.eq_ignore_ascii_case("session_id") + }) + .collect::>(); + + if let Some(session_id) = context.session_id { + pairs.push(("session-id".to_string(), session_id.to_string())); + pairs.push(("thread-id".to_string(), session_id.to_string())); + } + + pairs +} + +fn websocket_exchange_error_to_client_error( + error: CodexWebSocketExchangeError, +) -> CodexClientError { + match error { + CodexWebSocketExchangeError::Upstream { + status_code, + retry_after_seconds, + body, + } => CodexClientError::Upstream { + status: StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_GATEWAY), + body, + retry_after_seconds, + }, + error => CodexClientError::WebSocket(error), + } +} + +fn websocket_error_allows_http_fallback(error: &CodexClientError) -> bool { + !matches!(error, CodexClientError::Upstream { .. }) +} + +async fn read_capped_error_body(response: ReqwestResponse) -> Result { + let body = response.bytes().await?; + let len = body.len().min(MAX_UPSTREAM_ERROR_BODY_BYTES); + Ok(String::from_utf8_lossy(&body[..len]).into_owned()) +} + +fn turn_state(response: &ReqwestResponse) -> Option { + response + .headers() + .get("x-codex-turn-state") + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string) +} + +fn set_cookie_headers(response: &ReqwestResponse) -> Vec { + response + .headers() + .get_all(SET_COOKIE) + .iter() + .filter_map(|value| value.to_str().ok().map(ToString::to_string)) + .collect() +} + +fn rate_limit_headers(headers: &HeaderMap) -> Vec<(String, String)> { + headers + .iter() + .filter(|(name, _)| is_rate_limit_header(name.as_str())) + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.as_str().to_string(), value.to_string())) + }) + .collect() +} + +fn is_rate_limit_header(name: &str) -> bool { + let name = name.to_ascii_lowercase(); + name == "retry-after" + || name.contains("ratelimit") + || name.contains("rate-limit") + || name.starts_with("x-codex-primary-") + || name.starts_with("x-codex-secondary-") + || name.starts_with("x-codex-code-review-") + || name.starts_with("x-codex-review-") + || name.starts_with("x-code-review-") +} + +fn retry_after_seconds(headers: &HeaderMap, body: Option<&str>) -> Option { + headers + .get(RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.trim().parse::().ok()) + .filter(|seconds| *seconds > 0) + .or_else(|| body.and_then(retry_after_seconds_from_body)) +} + +fn truncate_for_error(body: &str) -> String { + body.chars().take(200).collect() +} + +fn usage_endpoint_paths(base_url: &str) -> Vec<&'static str> { + if has_backend_api_base_path(base_url) { + vec![WHAM_USAGE_PATH, CODEX_USAGE_PATH] + } else { + vec![CODEX_USAGE_API_PATH, CODEX_USAGE_PATH] + } +} + +fn has_backend_api_base_path(base_url: &str) -> bool { + reqwest::Url::parse(base_url).ok().is_some_and(|url| { + url.path() + .split('/') + .any(|segment| segment == "backend-api") + }) +} + +/// 在 reqwest builder 上应用自定义 CA。 +pub fn build_reqwest_client_with_custom_ca( + builder: reqwest::ClientBuilder, +) -> CustomCaResult { + build_reqwest_client_with_env(&ProcessEnv, builder) +} + +/// 返回当前自定义 CA 的缓存键。 +pub fn custom_ca_env_cache_key() -> Option { + ProcessEnv + .configured_ca_bundle() + .map(|bundle| format!("{}={}", bundle.source_env, bundle.path.display())) +} + +/// 构建 rustls client config，若未配置自定义 CA 则返回 `None`。 +pub fn maybe_build_rustls_client_config_with_custom_ca( +) -> CustomCaResult>> { + maybe_build_rustls_client_config_with_env(&ProcessEnv) +} + +fn build_reqwest_client_with_env( + env_source: &dyn EnvSource, + mut builder: reqwest::ClientBuilder, +) -> CustomCaResult { + let Some(bundle) = env_source.configured_ca_bundle() else { + return builder + .build() + .map_err(CustomCaError::BuildClientWithSystemRoots); + }; + + builder = builder.use_rustls_tls(); + for (idx, cert) in bundle.load_certificates()?.iter().enumerate() { + let certificate = reqwest::Certificate::from_der(cert.as_ref()).map_err(|source| { + CustomCaError::RegisterCertificate { + source_env: bundle.source_env, + path: bundle.path.clone(), + certificate_index: idx + 1, + source, + } + })?; + builder = builder.add_root_certificate(certificate); + } + + builder + .build() + .map_err(|source| CustomCaError::BuildClientWithCustomCa { + source_env: bundle.source_env, + path: bundle.path, + source, + }) +} + +fn maybe_build_rustls_client_config_with_env( + env_source: &dyn EnvSource, +) -> CustomCaResult>> { + let Some(bundle) = env_source.configured_ca_bundle() else { + return Ok(None); + }; + + let mut root_store = native_root_store().map_err(CustomCaError::LoadNativeRoots)?; + for (idx, cert) in bundle.load_certificates()?.into_iter().enumerate() { + root_store + .add(cert) + .map_err(|source| CustomCaError::RegisterRustlsCertificate { + source_env: bundle.source_env, + path: bundle.path.clone(), + certificate_index: idx + 1, + source, + })?; + } + + Ok(Some(std::sync::Arc::new( + ClientConfig::builder() + .with_root_certificates(root_store) + .with_no_client_auth(), + ))) +} + +pub(crate) fn native_root_store() -> Result { + let mut root_store = RootCertStore::empty(); + let rustls_native_certs::CertificateResult { certs, errors, .. } = + rustls_native_certs::load_native_certs(); + if !errors.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("failed to load native root certificates: {errors:?}"), + )); + } + + let (added, _) = root_store.add_parsable_certificates(certs); + if added == 0 { + return Err(io::Error::new( + io::ErrorKind::NotFound, + "no native root certificates found", + )); + } + + Ok(root_store) +} + +trait EnvSource { + fn var(&self, key: &str) -> Option; + + fn non_empty_path(&self, key: &str) -> Option { + self.var(key) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + } + + fn configured_ca_bundle(&self) -> Option { + self.non_empty_path(CODEX_CA_CERT_ENV) + .map(|path| ConfiguredCaBundle { + source_env: CODEX_CA_CERT_ENV, + path, + }) + .or_else(|| { + self.non_empty_path(SSL_CERT_FILE_ENV) + .map(|path| ConfiguredCaBundle { + source_env: SSL_CERT_FILE_ENV, + path, + }) + }) + } +} + +struct ProcessEnv; + +impl EnvSource for ProcessEnv { + fn var(&self, key: &str) -> Option { + env::var(key).ok() + } +} + +struct ConfiguredCaBundle { + source_env: &'static str, + path: PathBuf, +} + +impl ConfiguredCaBundle { + fn load_certificates(&self) -> CustomCaResult>> { + let pem_data = fs::read(&self.path).map_err(|source| CustomCaError::ReadCaFile { + source_env: self.source_env, + path: self.path.clone(), + source, + })?; + let normalized = normalize_trusted_certificate_labels(&pem_data); + let mut certificates = Vec::new(); + for section in PemSection::pem_slice_iter(normalized.as_bytes()) { + let (kind, der) = section.map_err(|error| self.pem_parse_error(&error))?; + if kind == SectionKind::Certificate { + certificates.push(CertificateDer::from(der)); + } + } + if certificates.is_empty() { + return Err(self.pem_parse_error(&pem::Error::NoItemsFound)); + } + Ok(certificates) + } + + fn pem_parse_error(&self, error: &pem::Error) -> CustomCaError { + let detail = match error { + pem::Error::NoItemsFound => "no certificates found in PEM file".to_string(), + _ => format!("failed to parse PEM file: {error}"), + }; + CustomCaError::InvalidCaFile { + source_env: self.source_env, + path: self.path.clone(), + detail, + } + } +} + +fn normalize_trusted_certificate_labels(input: &[u8]) -> String { + String::from_utf8_lossy(input).replace("TRUSTED CERTIFICATE", "CERTIFICATE") +} diff --git a/crates/adapters/src/codex/fingerprint.rs b/crates/adapters/src/codex/fingerprint.rs new file mode 100644 index 000000000..449f334bb --- /dev/null +++ b/crates/adapters/src/codex/fingerprint.rs @@ -0,0 +1,477 @@ +//! 指纹更新与持久化适配器。 + +use std::{path::PathBuf, sync::Arc, time::Duration}; + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use sqlx::{Row, SqlitePool}; +use thiserror::Error; +use tokio::sync::Mutex; +use tracing::{info, warn}; +use uuid::Uuid; + +use codex_proxy_core::gateway::fingerprint::Fingerprint; + +/// 指纹历史记录的来源标识。 +pub const CODEX_DESKTOP_UPDATE_SOURCE: &str = "codex_desktop_update_source"; + +const AUTO_UPDATED_FINGERPRINT_ID: &str = "auto_updated"; +const AUTO_UPDATE_SOURCE: &str = "auto_update"; +const AUTO_UPDATE_PLATFORM: &str = "darwin"; +const AUTO_UPDATE_ARCH: &str = "arm64"; +const AUTO_UPDATE_USER_AGENT_TEMPLATE: &str = "Codex Desktop/{version} ({platform}; {arch})"; +const DEFAULT_AUTO_UPDATE_CHROMIUM_VERSION: &str = "146"; +const APPCAST_TIMEOUT: Duration = Duration::from_secs(30); +const POLL_INTERVAL: Duration = Duration::from_secs(3 * 24 * 60 * 60); + +type AppcastFields = (Option, Option, Option); + +/// 更新清单解析结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FingerprintUpdate { + /// 应用版本。 + pub app_version: String, + /// 构建号。 + pub build_number: String, +} + +/// 指纹更新错误。 +#[derive(Debug, Error)] +pub enum FingerprintError { + /// 更新清单 JSON 无效。 + #[error("invalid update manifest: {0}")] + InvalidManifest(#[from] serde_json::Error), + /// HTTP 请求失败。 + #[error("failed to fetch update manifest: {0}")] + Http(#[from] reqwest::Error), + /// 数据库存储失败。 + #[error("failed to persist fingerprint update: {0}")] + Database(#[from] sqlx::Error), +} + +/// Appcast 检查错误。 +#[derive(Debug, Error)] +pub enum UpdateError { + /// HTTP 请求失败。 + #[error("HTTP 请求失败: {0}")] + Http(#[from] reqwest::Error), + /// 获取 appcast 失败。 + #[error("获取 appcast 失败，状态码: {0}")] + AppcastFetch(u16), + /// 解析 appcast 失败。 + #[error("解析 appcast 失败")] + AppcastParse, + /// JSON 序列化失败。 + #[error("JSON 序列化失败: {0}")] + Json(#[from] serde_json::Error), + /// 文件操作失败。 + #[error("文件操作失败: {0}")] + Io(#[from] std::io::Error), + /// 数据库存储失败。 + #[error("数据库操作失败: {0}")] + Database(#[from] sqlx::Error), +} + +#[derive(Deserialize)] +struct Manifest { + version: String, + build_number: String, +} + +/// 从更新清单 JSON 中提取版本信息。 +pub fn parse_update_manifest(input: &str) -> Result { + // 自动更新只同步桌面端指纹字段，不把远端配置当作运行时业务配置执行。 + let manifest: Manifest = serde_json::from_str(input)?; + Ok(FingerprintUpdate { + app_version: manifest.version, + build_number: manifest.build_number, + }) +} + +/// 指纹历史记录。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StoredFingerprint { + /// 应用版本。 + pub app_version: String, + /// 构建号。 + pub build_number: String, + /// 来源。 + pub source: String, +} + +/// 指纹自动更新状态。 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct UpdateState { + /// 最近检查时间。 + pub last_check: DateTime, + /// 最新版本。 + pub latest_version: Option, + /// 最新构建号。 + pub latest_build: Option, + /// 下载地址。 + pub download_url: Option, + /// 是否有可用更新。 + pub update_available: bool, + /// 当前版本。 + pub current_version: String, + /// 当前构建号。 + pub current_build: String, +} + +/// SQLite 指纹仓储。 +#[derive(Clone)] +pub struct FingerprintRepository { + pool: SqlitePool, +} + +impl FingerprintRepository { + /// 使用给定连接池构造仓储。 + pub fn new(pool: SqlitePool) -> Self { + Self { pool } + } + + /// 插入一条指纹更新历史。 + pub async fn insert_update(&self, update: &FingerprintUpdate) -> Result<(), sqlx::Error> { + let mut fingerprint = Fingerprint::default_codex_desktop(); + fingerprint.app_version.clone_from(&update.app_version); + fingerprint.build_number.clone_from(&update.build_number); + + sqlx::query( + "insert into fingerprints (id, app_version, build_number, platform, arch, chromium_version, user_agent_template, source, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(Uuid::new_v4().to_string()) + .bind(fingerprint.app_version) + .bind(fingerprint.build_number) + .bind(fingerprint.platform) + .bind(fingerprint.arch) + .bind(fingerprint.chromium_version) + .bind(fingerprint.user_agent_template) + .bind(CODEX_DESKTOP_UPDATE_SOURCE) + .bind(Utc::now().to_rfc3339()) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// 更新自动应用的最新版本快照。 + pub async fn upsert_auto_update( + &self, + app_version: &str, + build_number: &str, + chromium_version: Option<&str>, + ) -> Result<(), sqlx::Error> { + sqlx::query( + "insert into fingerprints (id, app_version, build_number, platform, arch, chromium_version, user_agent_template, source, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?) \ + on conflict(id) do update set app_version = excluded.app_version, build_number = excluded.build_number, platform = excluded.platform, arch = excluded.arch, chromium_version = excluded.chromium_version, user_agent_template = excluded.user_agent_template, source = excluded.source, created_at = excluded.created_at", + ) + .bind(AUTO_UPDATED_FINGERPRINT_ID) + .bind(app_version) + .bind(build_number) + .bind(AUTO_UPDATE_PLATFORM) + .bind(AUTO_UPDATE_ARCH) + .bind(chromium_version.unwrap_or(DEFAULT_AUTO_UPDATE_CHROMIUM_VERSION)) + .bind(AUTO_UPDATE_USER_AGENT_TEMPLATE) + .bind(AUTO_UPDATE_SOURCE) + .bind(Utc::now().to_rfc3339()) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// 读取最新一条指纹历史记录。 + pub async fn latest(&self) -> Result, sqlx::Error> { + let row = sqlx::query( + "select app_version, build_number, source from fingerprints order by created_at desc, id desc limit 1", + ) + .fetch_optional(&self.pool) + .await?; + + Ok(row.map(|row| StoredFingerprint { + app_version: row.get("app_version"), + build_number: row.get("build_number"), + source: row.get("source"), + })) + } + + /// 读取自动更新应用后的最新指纹。 + pub async fn load_latest_auto_updated(&self) -> Result, sqlx::Error> { + let row = sqlx::query( + r#" + select app_version, build_number, platform, arch, chromium_version, user_agent_template + from fingerprints + where source = ? + order by created_at desc + limit 1 + "#, + ) + .bind(AUTO_UPDATE_SOURCE) + .fetch_optional(&self.pool) + .await?; + + let Some(row) = row else { + return Ok(None); + }; + + Ok(Some(Fingerprint { + originator: "Codex Desktop".to_string(), + app_version: row.get("app_version"), + build_number: row.get("build_number"), + platform: row.get("platform"), + arch: row.get("arch"), + chromium_version: row.get("chromium_version"), + user_agent_template: row.get("user_agent_template"), + default_headers: Fingerprint::default_headers(), + header_order: Fingerprint::default_header_order(), + })) + } +} + +/// 通过 HTTP 拉取更新清单并写入历史记录。 +#[derive(Clone)] +pub struct FingerprintUpdater { + client: reqwest::Client, + repository: FingerprintRepository, + update_url: String, +} + +impl FingerprintUpdater { + /// 构造更新器。 + pub fn new( + client: reqwest::Client, + repository: FingerprintRepository, + update_url: impl Into, + ) -> Self { + Self { + client, + repository, + update_url: update_url.into(), + } + } + + /// 执行一次更新清单拉取并持久化历史记录。 + pub async fn poll_once(&self) -> Result { + let manifest = self + .client + .get(&self.update_url) + .timeout(APPCAST_TIMEOUT) + .send() + .await? + .error_for_status()? + .text() + .await?; + let update = parse_update_manifest(&manifest)?; + self.repository.insert_update(&update).await?; + Ok(update) + } +} + +/// Appcast 指纹更新检查器。 +#[derive(Clone)] +pub struct UpdateChecker { + repository: Option, + client: reqwest::Client, + appcast_url: String, + extracted_fingerprint_path: PathBuf, + state: Arc>, +} + +struct InternalState { + current_version: String, + current_build: String, +} + +impl UpdateChecker { + /// 使用自定义 client 和路径构造检查器。 + pub fn with_client( + repository: Option, + client: reqwest::Client, + appcast_url: impl Into, + extracted_fingerprint_path: PathBuf, + current_version: impl Into, + current_build: impl Into, + ) -> Self { + Self { + repository, + client, + appcast_url: appcast_url.into(), + extracted_fingerprint_path, + state: Arc::new(Mutex::new(InternalState { + current_version: current_version.into(), + current_build: current_build.into(), + })), + } + } + + /// 检查 appcast 是否存在更新。 + pub async fn check_for_update(&self) -> Result { + let state = self.state.lock().await; + let current_version = state.current_version.clone(); + let current_build = state.current_build.clone(); + drop(state); + + let xml = self.fetch_appcast().await?; + let (version, build, download_url) = parse_appcast(&xml)?; + + let update_available = version + .as_ref() + .is_some_and(|value| value != ¤t_version) + || build.as_ref().is_some_and(|value| value != ¤t_build); + + Ok(UpdateState { + last_check: Utc::now(), + latest_version: version, + latest_build: build, + download_url, + update_available, + current_version, + current_build, + }) + } + + /// 检查并在需要时应用更新。 + pub async fn check_and_apply_update(&self) -> Result { + let update_state = self.check_for_update().await?; + + if !update_state.update_available { + return Ok(false); + } + + let (Some(version), Some(build)) = ( + update_state.latest_version.as_deref(), + update_state.latest_build.as_deref(), + ) else { + return Ok(false); + }; + + let state = self.state.lock().await; + let current_version = state.current_version.clone(); + let current_build = state.current_build.clone(); + drop(state); + + info!( + version = %version, + build = %build, + current_version = %current_version, + current_build = %current_build, + "发现新的 fingerprint 版本" + ); + + self.apply_version_update(version, build).await?; + + let mut state = self.state.lock().await; + state.current_version = version.to_string(); + state.current_build = build.to_string(); + + Ok(true) + } + + /// 启动后台轮询任务。 + pub fn start_background_checker(self) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + let mut ticker = tokio::time::interval(POLL_INTERVAL); + + info!( + interval_secs = POLL_INTERVAL.as_secs(), + "fingerprint 后台版本检查器已启动" + ); + + if let Err(error) = self.check_and_apply_update().await { + warn!(error = %error, "fingerprint 首次版本检查失败"); + } + + loop { + ticker.tick().await; + if let Err(error) = self.check_and_apply_update().await { + warn!(error = %error, "fingerprint 定期版本检查失败"); + } + } + }) + } + + async fn fetch_appcast(&self) -> Result { + let response = self + .client + .get(&self.appcast_url) + .timeout(APPCAST_TIMEOUT) + .send() + .await?; + + if !response.status().is_success() { + return Err(UpdateError::AppcastFetch(response.status().as_u16())); + } + + response.text().await.map_err(UpdateError::from) + } + + async fn apply_version_update(&self, version: &str, build: &str) -> Result<(), UpdateError> { + let chromium_version = + load_matching_chromium_version(&self.extracted_fingerprint_path, version, build); + + if let Some(repository) = &self.repository { + repository + .upsert_auto_update(version, build, chromium_version.as_deref()) + .await?; + } + + Ok(()) + } +} + +fn parse_appcast(xml: &str) -> Result { + let item_start = xml.find("").ok_or(UpdateError::AppcastParse)?; + let item_end = xml.find("").ok_or(UpdateError::AppcastParse)?; + let item = &xml[item_start..item_end]; + + let version = extract_sparkle_field(item, "shortVersionString"); + let build = extract_sparkle_field(item, "version"); + let download_url = item.find("url=\"").and_then(|position| { + let start = position + 5; + item[start..] + .find('"') + .map(|end| item[start..start + end].to_string()) + }); + + Ok((version, build, download_url)) +} + +fn extract_sparkle_field(item: &str, field: &str) -> Option { + let attr_pattern = format!("sparkle:{}=\"", field); + if let Some(position) = item.find(&attr_pattern) { + let start = position + attr_pattern.len(); + if let Some(end) = item[start..].find('"') { + return Some(item[start..start + end].to_string()); + } + } + + let elem_start = format!("", field); + let elem_end = format!("", field); + if let Some(start_position) = item.find(&elem_start) { + let content_start = start_position + elem_start.len(); + if let Some(end_position) = item[content_start..].find(&elem_end) { + return Some(item[content_start..content_start + end_position].to_string()); + } + } + + None +} + +fn load_matching_chromium_version( + extracted_fingerprint_path: &PathBuf, + version: &str, + build: &str, +) -> Option { + let content = std::fs::read_to_string(extracted_fingerprint_path).ok()?; + let parsed: serde_json::Value = serde_json::from_str(&content).ok()?; + + let extracted_version = parsed.get("app_version")?.as_str()?; + let extracted_build = parsed.get("build_number")?.as_str()?; + + if extracted_version == version && extracted_build == build { + parsed + .get("chromium_version")? + .as_str() + .map(ToString::to_string) + } else { + None + } +} diff --git a/crates/adapters/src/codex/mod.rs b/crates/adapters/src/codex/mod.rs new file mode 100644 index 000000000..270494205 --- /dev/null +++ b/crates/adapters/src/codex/mod.rs @@ -0,0 +1,15 @@ +//! Codex 上游适配器。 + +pub mod client; +pub mod fingerprint; +pub mod models; +pub mod websocket; + +pub use client::{ + endpoint_request_path, endpoint_url, primary_usage_request_path, usage_endpoint_urls, +}; + +/// 构造默认 Codex HTTP 客户端。 +pub fn default_http_client() -> client::HttpClient { + client::HttpClient::new() +} diff --git a/crates/adapters/src/codex/models.rs b/crates/adapters/src/codex/models.rs new file mode 100644 index 000000000..a243e6296 --- /dev/null +++ b/crates/adapters/src/codex/models.rs @@ -0,0 +1,3 @@ +//! Codex 模型目录适配器。 + +pub use crate::codex::client::CodexBackendClient; diff --git a/crates/adapters/src/codex/websocket/connect.rs b/crates/adapters/src/codex/websocket/connect.rs new file mode 100644 index 000000000..58afdb4bb --- /dev/null +++ b/crates/adapters/src/codex/websocket/connect.rs @@ -0,0 +1,1078 @@ +//! Codex WebSocket 连接建立。 + +use std::{io, pin::Pin, sync::Arc, time::Duration, time::Instant}; + +use bytes::Bytes; +use codex_proxy_core::protocol::codex::{ + events::{ + extract_sse_usage, parse_rate_limits_event_raw, rate_limits_to_header_pairs, + retry_after_seconds_from_body, TokenUsage, + }, + responses::CodexResponsesRequest, + sse::SseError, + websocket::{ + classify_websocket_error_frame, is_terminal_websocket_event, + retry_after_seconds_from_wrapped_error_headers, websocket_event_to_sse_frame, + websocket_incomplete_response_reason, websocket_metadata_turn_state, + websocket_response_completed_parse_error, websocket_response_create_payload_text, + OpeningAuditHeader, OpeningAuditSnapshot, WebSocketErrorClassificationProfile, + }, +}; +use futures::{channel::mpsc, SinkExt, Stream, StreamExt}; +use rustls_pki_types::ServerName; +use serde_json::Value; +use thiserror::Error; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpStream, + sync::Mutex, + time::timeout, +}; +use tokio_rustls::TlsConnector; +use tokio_tungstenite::{ + tungstenite::{ + self, + error::{CapacityError, ProtocolError, TlsError, UrlError}, + handshake::derive_accept_key, + http::{HeaderName, HeaderValue, Response as WsResponse, StatusCode as WsStatusCode, Uri}, + protocol::Role, + Message, + }, + MaybeTlsStream, WebSocketStream, +}; + +use crate::codex::client::{maybe_build_rustls_client_config_with_custom_ca, native_root_store}; + +use super::deflate::{permessage_deflate_enabled, PerMessageDeflateStream}; +use super::pool::{ + CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey, CodexWsStream, + PooledWebSocketConnection, WebSocketPoolAcquire, +}; + +const REDACTED_HEADER_VALUE: &str = ""; +const CODEX_RESPONSES_PATH: &str = "/codex/responses"; +const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits"; +const MAX_OPENING_RESPONSE_BYTES: usize = 64 * 1024; +const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20); + +/// WebSocket 连接适配器。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexWebSocketConnection { + endpoint: String, + headers: Vec<(String, String)>, +} + +/// Prepared Responses WebSocket request descriptor. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexWebSocketRequest { + connection: CodexWebSocketConnection, + payload_text: String, +} + +/// Responses WebSocket exchange result. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexWebSocketExchange { + /// 完整 SSE 文本。 + pub body: String, + /// 从 SSE 中提取的 usage。 + pub usage: Option, + /// 上游 metadata 帧中的最新 turn state。 + pub turn_state: Option, + /// 上游握手响应里的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游握手响应里的限流头。 + pub rate_limit_headers: Vec<(String, String)>, + /// 打开握手响应状态码。 + pub handshake_status: u16, +} + +/// Responses WebSocket live SSE exchange result. +pub struct CodexWebSocketStreamingExchange { + /// Live SSE bytes converted from WebSocket events. + pub body: CodexWebSocketSseStream, + /// 上游 metadata 帧中的最新 turn state。 + pub turn_state: Option, + /// 上游握手响应里的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游握手响应里的限流头。 + pub rate_limit_headers: Vec<(String, String)>, + /// 上游内部 rate-limit 事件里的动态更新。 + pub rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates, + /// 上游内部 metadata 事件里的动态 turn state。 + pub turn_state_update: CodexWebSocketTurnStateUpdate, + /// 打开握手响应状态码。 + pub handshake_status: u16, +} + +/// Responses WebSocket live SSE byte stream. +pub type CodexWebSocketSseStream = + Pin> + Send + 'static>>; +/// WebSocket live stream rate-limit header updates. +pub type CodexWebSocketRateLimitHeaderUpdates = Arc>>; +/// WebSocket live stream turn-state update. +pub type CodexWebSocketTurnStateUpdate = Arc>>; + +/// Responses WebSocket exchange error. +#[derive(Debug, Error)] +pub enum CodexWebSocketExchangeError { + /// opening request 无法构造。 + #[error("invalid websocket request: {0}")] + InvalidRequest(#[from] tungstenite::http::Error), + /// WebSocket 传输失败。 + #[error("websocket transport error: {0}")] + Transport(#[from] tungstenite::Error), + /// SSE 聚合结果无法解析。 + #[error("invalid websocket SSE response: {0}")] + InvalidSse(#[from] SseError), + /// 上游 WebSocket 错误帧。 + #[error("websocket upstream error {status_code}: {body}")] + Upstream { + /// HTTP-style upstream status code. + status_code: u16, + /// 推导出的重试秒数。 + retry_after_seconds: Option, + /// 原始错误帧。 + body: String, + }, + /// 上游返回 `response.incomplete`。 + #[error("Incomplete response returned, reason: {reason}")] + IncompleteResponse { + /// incomplete_details.reason。 + reason: String, + }, + /// 上游返回无法按官方形状解析的 `response.completed`。 + #[error("{message}")] + InvalidCompletedResponse { + /// 解析失败说明。 + message: String, + }, + /// 上游在 terminal 事件前关闭。 + #[error("websocket closed before terminal event")] + ClosedBeforeTerminal, + /// 上游在指定时间内没有发送任何事件。 + #[error("websocket receive idle timeout after {timeout:?}")] + ReceiveIdleTimeout { + /// 超时时长。 + timeout: Duration, + }, + /// 上游返回非文本事件帧。 + #[error("unexpected binary websocket event")] + UnexpectedBinaryEvent, +} + +impl CodexWebSocketRequest { + /// 返回连接描述。 + pub fn connection(&self) -> &CodexWebSocketConnection { + &self.connection + } + + /// 返回将要发送的首个文本帧。 + pub fn payload_text(&self) -> &str { + &self.payload_text + } +} + +/// 执行一次 prepared Responses WebSocket 请求并聚合为 SSE。 +pub async fn execute_response_create_request( + request: &CodexWebSocketRequest, +) -> Result { + execute_response_create_request_with_pool(request, None).await +} + +pub(crate) async fn execute_response_create_request_with_pool( + request: &CodexWebSocketRequest, + pool: Option<(&CodexWebSocketPool, CodexWebSocketPoolKey)>, +) -> Result { + let Some((pool, key)) = pool else { + return execute_fresh_response_create_request(request).await; + }; + + match pool.acquire(&key).await { + WebSocketPoolAcquire::Reused(connection) => { + execute_pooled_response_create_request(request, pool, key, *connection).await + } + WebSocketPoolAcquire::FreshReserved => { + let result = + execute_fresh_pooled_response_create_request(request, pool, key.clone()).await; + if result.is_err() { + pool.discard(&key).await; + } + result + } + WebSocketPoolAcquire::Bypass => execute_fresh_response_create_request(request).await, + } +} + +pub(crate) async fn execute_response_create_request_stream_with_pool( + request: &CodexWebSocketRequest, + pool: Option<(&CodexWebSocketPool, CodexWebSocketPoolKey)>, +) -> Result { + let Some((pool, key)) = pool else { + return execute_fresh_response_create_request_stream(request).await; + }; + + match pool.acquire(&key).await { + WebSocketPoolAcquire::Reused(connection) => { + execute_pooled_response_create_request_stream(request, pool.clone(), key, *connection) + .await + } + WebSocketPoolAcquire::FreshReserved => { + let result = execute_fresh_pooled_response_create_request_stream( + request, + pool.clone(), + key.clone(), + ) + .await; + if result.is_err() { + pool.discard(&key).await; + } + result + } + WebSocketPoolAcquire::Bypass => execute_fresh_response_create_request_stream(request).await, + } +} + +async fn execute_fresh_response_create_request( + request: &CodexWebSocketRequest, +) -> Result { + let (mut websocket, response) = connect_websocket(request.connection()).await?; + websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await?; + + let metadata = websocket_connection_metadata(&response); + let (exchange, _websocket, _metadata) = collect_websocket_response(websocket, metadata).await?; + Ok(exchange) +} + +async fn execute_fresh_response_create_request_stream( + request: &CodexWebSocketRequest, +) -> Result { + let (mut websocket, response) = connect_websocket(request.connection()).await?; + websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await?; + + let metadata = websocket_connection_metadata(&response); + Ok(stream_websocket_response(websocket, metadata, None)) +} + +async fn execute_fresh_pooled_response_create_request( + request: &CodexWebSocketRequest, + pool: &CodexWebSocketPool, + key: CodexWebSocketPoolKey, +) -> Result { + let (mut websocket, response) = connect_websocket(request.connection()).await?; + websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await?; + + let now = Instant::now(); + let metadata = websocket_connection_metadata(&response); + let (exchange, websocket, metadata) = collect_websocket_response(websocket, metadata).await?; + let last_activity_at = Instant::now(); + pool.put( + key, + PooledWebSocketConnection { + websocket, + metadata, + created_at: now, + last_activity_at, + last_ping_at: None, + }, + ) + .await; + Ok(exchange) +} + +async fn execute_fresh_pooled_response_create_request_stream( + request: &CodexWebSocketRequest, + pool: CodexWebSocketPool, + key: CodexWebSocketPoolKey, +) -> Result { + let (mut websocket, response) = connect_websocket(request.connection()).await?; + websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await?; + + let now = Instant::now(); + let metadata = websocket_connection_metadata(&response); + Ok(stream_websocket_response( + websocket, + metadata, + Some(WebSocketStreamPoolReturn { + pool, + key, + created_at: now, + last_ping_at: None, + }), + )) +} + +async fn execute_pooled_response_create_request( + request: &CodexWebSocketRequest, + pool: &CodexWebSocketPool, + key: CodexWebSocketPoolKey, + mut connection: PooledWebSocketConnection, +) -> Result { + if let Err(error) = connection + .websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await + { + pool.discard(&key).await; + return Err(error.into()); + } + + let created_at = connection.created_at; + let last_ping_at = connection.last_ping_at; + match collect_websocket_response(connection.websocket, connection.metadata).await { + Ok((exchange, websocket, metadata)) => { + let last_activity_at = Instant::now(); + pool.put( + key, + PooledWebSocketConnection { + websocket, + metadata, + created_at, + last_activity_at, + last_ping_at, + }, + ) + .await; + Ok(exchange) + } + Err(error) => { + pool.discard(&key).await; + Err(error) + } + } +} + +async fn execute_pooled_response_create_request_stream( + request: &CodexWebSocketRequest, + pool: CodexWebSocketPool, + key: CodexWebSocketPoolKey, + mut connection: PooledWebSocketConnection, +) -> Result { + if let Err(error) = connection + .websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await + { + pool.discard(&key).await; + return Err(error.into()); + } + + Ok(stream_websocket_response( + connection.websocket, + connection.metadata, + Some(WebSocketStreamPoolReturn { + pool, + key, + created_at: connection.created_at, + last_ping_at: connection.last_ping_at, + }), + )) +} + +impl CodexWebSocketConnection { + /// 构造待打开的 WebSocket 连接描述。 + pub fn new(endpoint: impl Into, headers: Vec<(String, String)>) -> Self { + Self { + endpoint: endpoint.into(), + headers, + } + } + + /// 构造 Responses WebSocket 连接描述。 + pub fn responses( + base_url: &str, + websocket_key: &str, + business_headers: Vec<(String, String)>, + ) -> Self { + let endpoint = responses_websocket_endpoint(base_url); + let mut headers = Vec::new(); + if let Some(host) = websocket_host_header(&endpoint) { + headers.push(("Host".to_string(), host)); + } + headers.extend([ + ("Connection".to_string(), "Upgrade".to_string()), + ("Upgrade".to_string(), "websocket".to_string()), + ("Sec-WebSocket-Version".to_string(), "13".to_string()), + ("Sec-WebSocket-Key".to_string(), websocket_key.to_string()), + ]); + headers.extend(business_headers); + headers.push(( + "sec-websocket-extensions".to_string(), + WEBSOCKET_EXTENSIONS.to_string(), + )); + Self { endpoint, headers } + } + + /// 构造 Responses WebSocket opening 与首个 `response.create` 文本帧。 + pub fn responses_create_request( + base_url: &str, + websocket_key: &str, + business_headers: Vec<(String, String)>, + request: &CodexResponsesRequest, + ) -> Result { + Ok(CodexWebSocketRequest { + connection: Self::responses(base_url, websocket_key, business_headers), + payload_text: websocket_response_create_payload_text(request)?, + }) + } + + /// 返回 WebSocket endpoint。 + pub fn endpoint(&self) -> &str { + &self.endpoint + } + + /// 返回按发送顺序保存的请求头。 + pub fn headers(&self) -> &[(String, String)] { + &self.headers + } + + /// 生成真实 WebSocket opening 请求文本，用于 capture/audit parity。 + pub fn opening_request_text(&self) -> String { + String::from_utf8_lossy(&opening_request_bytes(self)).into_owned() + } + + /// 生成打开握手审计快照。 + pub fn opening_audit_snapshot(&self) -> OpeningAuditSnapshot { + OpeningAuditSnapshot { + request_line: request_line_for_endpoint(&self.endpoint), + header_order: self.headers.iter().map(|(name, _)| name.clone()).collect(), + headers: self + .headers + .iter() + .map(|(name, value)| OpeningAuditHeader { + name: name.clone(), + value: audit_header_value(name, value), + }) + .collect(), + } + } +} + +/// 将 Codex backend base URL 转换为 Responses WebSocket endpoint。 +pub fn responses_websocket_endpoint(base_url: &str) -> String { + let endpoint = format!("{}{}", base_url.trim_end_matches('/'), CODEX_RESPONSES_PATH); + if let Some(rest) = endpoint.strip_prefix("https://") { + format!("wss://{rest}") + } else if let Some(rest) = endpoint.strip_prefix("http://") { + format!("ws://{rest}") + } else { + endpoint + } +} + +fn websocket_host_header(endpoint: &str) -> Option { + let url = reqwest::Url::parse(endpoint).ok()?; + let host = url.host_str()?; + Some(match url.port() { + Some(port) => format!("{host}:{port}"), + None => host.to_string(), + }) +} + +async fn connect_websocket( + connection: &CodexWebSocketConnection, +) -> Result<(CodexWsStream, WsResponse>>), CodexWebSocketExchangeError> { + let uri = connection + .endpoint() + .parse::() + .map_err(tungstenite::Error::from)?; + let mut stream = connect_stream(&uri).await?; + let request_bytes = opening_request_bytes(connection); + stream + .write_all(&request_bytes) + .await + .map_err(tungstenite::Error::Io)?; + stream.flush().await.map_err(tungstenite::Error::Io)?; + + let (response, buffered) = read_opening_response(&mut stream).await?; + if response.status() != WsStatusCode::SWITCHING_PROTOCOLS { + return Err(websocket_opening_error(response)); + } + verify_accept_key(connection, &response)?; + let permessage_deflate = response_accepts_permessage_deflate(&response); + let stream = PerMessageDeflateStream::new(stream, permessage_deflate, buffered); + let websocket = + WebSocketStream::from_partially_read(stream, Vec::new(), Role::Client, None).await; + Ok((websocket, response)) +} + +fn websocket_opening_error(response: WsResponse>>) -> CodexWebSocketExchangeError { + let status_code = response.status().as_u16(); + let body = response + .body() + .as_ref() + .map(|body| String::from_utf8_lossy(body).into_owned()) + .unwrap_or_default(); + let retry_after_seconds = response + .headers() + .get("retry-after") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .filter(|seconds| *seconds > 0) + .or_else(|| retry_after_seconds_from_body(&body)); + CodexWebSocketExchangeError::Upstream { + status_code, + retry_after_seconds, + body, + } +} + +async fn connect_stream(uri: &Uri) -> Result, tungstenite::Error> { + let host = uri.host().ok_or(UrlError::NoHostName)?; + if host.is_empty() { + return Err(UrlError::EmptyHostName.into()); + } + let port = uri + .port_u16() + .or_else(|| match uri.scheme_str() { + Some("wss") => Some(443), + Some("ws") => Some(80), + _ => None, + }) + .ok_or(UrlError::UnsupportedUrlScheme)?; + + let socket = TcpStream::connect((host, port)).await?; + match uri.scheme_str() { + Some("ws") => Ok(MaybeTlsStream::Plain(socket)), + Some("wss") => connect_tls_stream(socket, host).await, + _ => Err(UrlError::UnsupportedUrlScheme.into()), + } +} + +async fn connect_tls_stream( + socket: TcpStream, + host: &str, +) -> Result, tungstenite::Error> { + let server_name = + ServerName::try_from(host.to_string()).map_err(|_| TlsError::InvalidDnsName)?; + let config = if let Some(config) = maybe_build_rustls_client_config_with_custom_ca() + .map_err(|error| tungstenite::Error::Io(io::Error::other(error)))? + { + config + } else { + Arc::new( + rustls::ClientConfig::builder() + .with_root_certificates(native_root_store().map_err(tungstenite::Error::Io)?) + .with_no_client_auth(), + ) + }; + let stream = TlsConnector::from(config) + .connect(server_name, socket) + .await + .map_err(tungstenite::Error::Io)?; + Ok(MaybeTlsStream::Rustls(stream)) +} + +fn opening_request_bytes(connection: &CodexWebSocketConnection) -> Vec { + let mut bytes = Vec::new(); + bytes.extend_from_slice(request_line_for_endpoint(connection.endpoint()).as_bytes()); + bytes.extend_from_slice(b"\r\n"); + for (name, value) in connection.headers() { + bytes.extend_from_slice(name.as_bytes()); + bytes.extend_from_slice(b": "); + bytes.extend_from_slice(value.as_bytes()); + bytes.extend_from_slice(b"\r\n"); + } + bytes.extend_from_slice(b"\r\n"); + bytes +} + +async fn read_opening_response( + stream: &mut MaybeTlsStream, +) -> Result<(WsResponse>>, Vec), tungstenite::Error> { + let mut bytes = Vec::new(); + let mut chunk = [0_u8; 1024]; + let header_end = loop { + let read = stream.read(&mut chunk).await?; + if read == 0 { + return Err(ProtocolError::HandshakeIncomplete.into()); + } + bytes.extend_from_slice(&chunk[..read]); + if bytes.len() > MAX_OPENING_RESPONSE_BYTES { + return Err(CapacityError::MessageTooLong { + size: bytes.len(), + max_size: MAX_OPENING_RESPONSE_BYTES, + } + .into()); + } + if let Some(index) = find_header_end(&bytes) { + break index; + } + }; + + let body_start = header_end + 4; + let mut body = bytes[body_start..].to_vec(); + let mut response = parse_response_head(&bytes[..header_end])?; + if response.status() != WsStatusCode::SWITCHING_PROTOCOLS { + read_response_body(stream, response.headers(), &mut body).await?; + *response.body_mut() = Some(body); + return Ok((response, Vec::new())); + } + + Ok((response, body)) +} + +fn find_header_end(bytes: &[u8]) -> Option { + bytes.windows(4).position(|window| window == b"\r\n\r\n") +} + +fn parse_response_head(head: &[u8]) -> Result>>, tungstenite::Error> { + let head = std::str::from_utf8(head)?; + let mut lines = head.split("\r\n"); + let status_line = lines + .next() + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "empty response"))?; + let status = status_line + .split_whitespace() + .nth(1) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "missing response status"))? + .parse::() + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?; + + let mut response = WsResponse::new(None); + *response.status_mut() = WsStatusCode::from_u16(status)?; + for line in lines { + let Some((name, value)) = line.split_once(':') else { + continue; + }; + response.headers_mut().append( + HeaderName::from_bytes(name.trim().as_bytes())?, + HeaderValue::from_bytes(value.trim_start().as_bytes())?, + ); + } + Ok(response) +} + +async fn read_response_body( + stream: &mut MaybeTlsStream, + headers: &tokio_tungstenite::tungstenite::http::HeaderMap, + body: &mut Vec, +) -> Result<(), tungstenite::Error> { + let Some(content_length) = headers + .get("content-length") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + else { + return Ok(()); + }; + + while body.len() < content_length { + let remaining = content_length - body.len(); + let mut chunk = vec![0_u8; remaining.min(1024)]; + let read = stream.read(&mut chunk).await?; + if read == 0 { + break; + } + body.extend_from_slice(&chunk[..read]); + } + body.truncate(content_length); + Ok(()) +} + +fn verify_accept_key( + connection: &CodexWebSocketConnection, + response: &WsResponse>>, +) -> Result<(), tungstenite::Error> { + let key = connection + .headers() + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("sec-websocket-key")) + .map(|(_, value)| value.as_str()) + .ok_or(ProtocolError::MissingSecWebSocketKey)?; + let expected = derive_accept_key(key.as_bytes()); + let actual = response + .headers() + .get("sec-websocket-accept") + .ok_or(ProtocolError::SecWebSocketAcceptKeyMismatch)?; + + if actual == expected.as_str() { + Ok(()) + } else { + Err(ProtocolError::SecWebSocketAcceptKeyMismatch.into()) + } +} + +fn response_accepts_permessage_deflate(response: &WsResponse>>) -> bool { + response + .headers() + .get_all("sec-websocket-extensions") + .iter() + .filter_map(|value| value.to_str().ok()) + .any(permessage_deflate_enabled) +} + +async fn collect_websocket_response( + mut websocket: CodexWsStream, + mut metadata: CodexWebSocketConnectionMetadata, +) -> Result< + ( + CodexWebSocketExchange, + CodexWsStream, + CodexWebSocketConnectionMetadata, + ), + CodexWebSocketExchangeError, +> { + let mut body = String::new(); + + loop { + let message = next_websocket_message(&mut websocket).await?; + let Some(message) = message else { + break; + }; + let text = match message { + Message::Text(text) => text, + Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent), + Message::Close(_) => break, + _ => continue, + }; + let raw = text.to_string(); + if let Some(classified) = + classify_websocket_error_frame(&raw, WebSocketErrorClassificationProfile::OneShot) + { + let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) + .or_else(|| retry_after_seconds_from_body(&raw)); + return Err(CodexWebSocketExchangeError::Upstream { + status_code: classified.status_code, + retry_after_seconds, + body: raw, + }); + } + if let Some(reason) = websocket_incomplete_response_reason(&raw) { + return Err(CodexWebSocketExchangeError::IncompleteResponse { reason }); + } + if let Some(message) = websocket_response_completed_parse_error(&raw) { + return Err(CodexWebSocketExchangeError::InvalidCompletedResponse { message }); + } + if let Some(headers) = websocket_rate_limit_event_headers(&raw) { + metadata.rate_limit_headers.extend(headers); + continue; + } + if let Some(metadata_turn_state) = websocket_metadata_turn_state(&raw) { + metadata.turn_state = Some(metadata_turn_state); + continue; + } + let event = websocket_event_type(&raw); + if event.as_deref() == Some("error") { + continue; + } + let forwarded = if let Some(frame) = websocket_event_to_sse_frame(&raw) { + body.push_str(&frame); + true + } else { + false + }; + if forwarded && event.as_deref().is_some_and(is_terminal_websocket_event) { + let usage = extract_sse_usage(&body)?; + let exchange = CodexWebSocketExchange { + body, + usage, + turn_state: metadata.turn_state.clone(), + set_cookie_headers: metadata.set_cookie_headers.clone(), + rate_limit_headers: metadata.rate_limit_headers.clone(), + handshake_status: metadata.handshake_status, + }; + return Ok((exchange, websocket, metadata)); + } + } + + Err(CodexWebSocketExchangeError::ClosedBeforeTerminal) +} + +fn websocket_connection_metadata( + response: &WsResponse>>, +) -> CodexWebSocketConnectionMetadata { + CodexWebSocketConnectionMetadata { + turn_state: websocket_response_turn_state(response.headers()), + set_cookie_headers: websocket_set_cookie_headers(response.headers()), + rate_limit_headers: websocket_rate_limit_headers(response.headers()), + handshake_status: response.status().as_u16(), + } +} + +fn websocket_rate_limit_event_headers(raw: &str) -> Option> { + parse_rate_limits_event_raw(raw).map(|parsed| rate_limits_to_header_pairs(&parsed)) +} + +struct WebSocketStreamPoolReturn { + pool: CodexWebSocketPool, + key: CodexWebSocketPoolKey, + created_at: Instant, + last_ping_at: Option, +} + +fn stream_websocket_response( + websocket: CodexWsStream, + metadata: CodexWebSocketConnectionMetadata, + pool_return: Option, +) -> CodexWebSocketStreamingExchange { + let response_metadata = metadata.clone(); + let rate_limit_header_updates = Arc::new(Mutex::new(Vec::new())); + let rate_limit_header_updates_for_task = Arc::clone(&rate_limit_header_updates); + let turn_state_update = Arc::new(Mutex::new(metadata.turn_state.clone())); + let turn_state_update_for_task = Arc::clone(&turn_state_update); + let (tx, rx) = mpsc::unbounded(); + tokio::spawn(async move { + forward_websocket_response_stream( + websocket, + metadata, + pool_return, + rate_limit_header_updates_for_task, + turn_state_update_for_task, + tx, + ) + .await; + }); + + CodexWebSocketStreamingExchange { + body: Box::pin(rx), + turn_state: response_metadata.turn_state, + set_cookie_headers: response_metadata.set_cookie_headers, + rate_limit_headers: response_metadata.rate_limit_headers, + rate_limit_header_updates, + turn_state_update, + handshake_status: response_metadata.handshake_status, + } +} + +async fn forward_websocket_response_stream( + mut websocket: CodexWsStream, + mut metadata: CodexWebSocketConnectionMetadata, + pool_return: Option, + rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates, + turn_state_update: CodexWebSocketTurnStateUpdate, + tx: mpsc::UnboundedSender>, +) { + let mut pool_return = pool_return; + loop { + let message = match next_websocket_message(&mut websocket).await { + Ok(message) => message, + Err(error) => { + discard_stream_websocket(websocket, pool_return).await; + let _ = tx.unbounded_send(Err(error)); + return; + } + }; + let Some(message) = message else { + break; + }; + let text = match message { + Message::Text(text) => text, + Message::Binary(_) => { + discard_stream_websocket(websocket, pool_return).await; + let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent)); + return; + } + Message::Close(_) => { + discard_stream_websocket(websocket, pool_return).await; + let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)); + return; + } + _ => continue, + }; + let raw = text.to_string(); + if let Some(classified) = + classify_websocket_error_frame(&raw, WebSocketErrorClassificationProfile::OneShot) + { + discard_stream_websocket(websocket, pool_return).await; + let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) + .or_else(|| retry_after_seconds_from_body(&raw)); + let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::Upstream { + status_code: classified.status_code, + retry_after_seconds, + body: raw, + })); + return; + } + if let Some(reason) = websocket_incomplete_response_reason(&raw) { + discard_stream_websocket(websocket, pool_return).await; + let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::IncompleteResponse { + reason, + })); + return; + } + if let Some(message) = websocket_response_completed_parse_error(&raw) { + discard_stream_websocket(websocket, pool_return).await; + let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::InvalidCompletedResponse { + message, + })); + return; + } + if let Some(headers) = websocket_rate_limit_event_headers(&raw) { + metadata.rate_limit_headers.extend(headers.iter().cloned()); + rate_limit_header_updates.lock().await.extend(headers); + continue; + } + if let Some(metadata_turn_state) = websocket_metadata_turn_state(&raw) { + metadata.turn_state = Some(metadata_turn_state); + *turn_state_update.lock().await = metadata.turn_state.clone(); + continue; + } + let event = websocket_event_type(&raw); + if event.as_deref() == Some("error") { + continue; + } + let Some(frame) = websocket_event_to_sse_frame(&raw) else { + continue; + }; + let terminal = event.as_deref().is_some_and(is_terminal_websocket_event); + if tx.unbounded_send(Ok(Bytes::from(frame))).is_err() { + discard_stream_websocket(websocket, pool_return).await; + return; + } + if terminal { + finish_stream_websocket(websocket, metadata, pool_return.take()).await; + return; + } + } + + discard_stream_websocket(websocket, pool_return).await; + let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)); +} + +async fn next_websocket_message( + websocket: &mut CodexWsStream, +) -> Result, CodexWebSocketExchangeError> { + match timeout(WEBSOCKET_RECEIVE_IDLE_TIMEOUT, websocket.next()).await { + Ok(message) => message.transpose().map_err(Into::into), + Err(_) => Err(CodexWebSocketExchangeError::ReceiveIdleTimeout { + timeout: WEBSOCKET_RECEIVE_IDLE_TIMEOUT, + }), + } +} + +async fn finish_stream_websocket( + mut websocket: CodexWsStream, + metadata: CodexWebSocketConnectionMetadata, + pool_return: Option, +) { + let Some(pool_return) = pool_return else { + let _ = websocket.close(None).await; + return; + }; + pool_return + .pool + .put( + pool_return.key, + PooledWebSocketConnection { + websocket, + metadata, + created_at: pool_return.created_at, + last_activity_at: Instant::now(), + last_ping_at: pool_return.last_ping_at, + }, + ) + .await; +} + +async fn discard_stream_websocket( + mut websocket: CodexWsStream, + pool_return: Option, +) { + if let Some(pool_return) = pool_return { + pool_return.pool.discard(&pool_return.key).await; + } + let _ = websocket.close(None).await; +} + +fn websocket_event_type(raw: &str) -> Option { + serde_json::from_str::(raw).ok().and_then(|value| { + value + .get("type") + .and_then(Value::as_str) + .map(ToString::to_string) + }) +} + +fn websocket_response_turn_state( + headers: &tokio_tungstenite::tungstenite::http::HeaderMap, +) -> Option { + headers + .get("x-codex-turn-state") + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string) +} + +fn websocket_set_cookie_headers( + headers: &tokio_tungstenite::tungstenite::http::HeaderMap, +) -> Vec { + headers + .get_all("set-cookie") + .iter() + .filter_map(|value| value.to_str().ok().map(ToString::to_string)) + .collect() +} + +fn websocket_rate_limit_headers( + headers: &tokio_tungstenite::tungstenite::http::HeaderMap, +) -> Vec<(String, String)> { + headers + .iter() + .filter(|(name, _)| is_rate_limit_header(name.as_str())) + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.as_str().to_string(), value.to_string())) + }) + .collect() +} + +fn is_rate_limit_header(name: &str) -> bool { + let name = name.to_ascii_lowercase(); + name == "retry-after" + || name.contains("ratelimit") + || name.contains("rate-limit") + || name.starts_with("x-codex-primary-") + || name.starts_with("x-codex-secondary-") + || name.starts_with("x-codex-code-review-") + || name.starts_with("x-codex-review-") + || name.starts_with("x-code-review-") +} + +fn request_line_for_endpoint(endpoint: &str) -> String { + let path = reqwest::Url::parse(endpoint) + .ok() + .map(|url| { + let mut path = url.path().to_string(); + if let Some(query) = url.query() { + path.push('?'); + path.push_str(query); + } + path + }) + .filter(|path| !path.is_empty()) + .unwrap_or_else(|| endpoint.to_string()); + format!("GET {path} HTTP/1.1") +} + +fn audit_header_value(name: &str, value: &str) -> String { + if is_sensitive_opening_header(name) { + REDACTED_HEADER_VALUE.to_string() + } else { + value.to_string() + } +} + +fn is_sensitive_opening_header(name: &str) -> bool { + matches!( + name.to_ascii_lowercase().as_str(), + "authorization" + | "chatgpt-account-id" + | "cookie" + | "session_id" + | "session-id" + | "thread-id" + | "x-client-request-id" + | "x-codex-window-id" + | "x-codex-turn-metadata" + | "x-codex-turn-state" + | "x-codex-parent-thread-id" + ) +} diff --git a/src/codex/gateway/transport/websocket/deflate.rs b/crates/adapters/src/codex/websocket/deflate.rs similarity index 77% rename from src/codex/gateway/transport/websocket/deflate.rs rename to crates/adapters/src/codex/websocket/deflate.rs index 47ab62187..ab6810f42 100644 --- a/src/codex/gateway/transport/websocket/deflate.rs +++ b/crates/adapters/src/codex/websocket/deflate.rs @@ -1,19 +1,24 @@ +//! WebSocket permessage-deflate 处理。 + use std::{ collections::VecDeque, - io, + io::{self, Read}, pin::Pin, task::{Context, Poll}, }; -use flate2::{Decompress, FlushDecompress, Status}; +use flate2::read::DeflateDecoder; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; -const SERVER_FRAME_CHUNK_BYTES: usize = 8192; const PERMESSAGE_DEFLATE_TAIL: [u8; 4] = [0x00, 0x00, 0xff, 0xff]; const OPCODE_CONTINUATION: u8 = 0x0; const OPCODE_TEXT: u8 = 0x1; const OPCODE_BINARY: u8 = 0x2; -pub(super) struct PerMessageDeflateStream { +const SERVER_FRAME_CHUNK_BYTES: usize = 8192; + +/// Async stream wrapper that inflates server permessage-deflate data frames +/// before tungstenite sees them. +pub(crate) struct PerMessageDeflateStream { inner: S, enabled: bool, raw_input: Vec, @@ -26,6 +31,7 @@ struct CompressedMessage { payload: Vec, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] struct FrameParts { fin: bool, rsv1: bool, @@ -37,7 +43,7 @@ struct FrameParts { } impl PerMessageDeflateStream { - pub(super) fn new(inner: S, enabled: bool, preloaded: Vec) -> Self { + pub(crate) fn new(inner: S, enabled: bool, preloaded: Vec) -> Self { Self { inner, enabled, @@ -67,8 +73,10 @@ where return Poll::Ready(Ok(())); } - if this.rewrite_next_frame()? { - continue; + match this.rewrite_next_frame() { + Ok(true) => continue, + Ok(false) => {} + Err(error) => return Poll::Ready(Err(error)), } let mut chunk = [0_u8; SERVER_FRAME_CHUNK_BYTES]; @@ -122,9 +130,9 @@ impl PerMessageDeflateStream { return Ok(true); } - if frame.rsv1 && is_data_frame(frame.opcode) { + if frame.rsv1 && is_initial_data_frame(frame.opcode) { if frame.fin { - let inflated = inflate_message(payload)?; + let inflated = inflate_permessage_deflate_message(payload)?; self.decoded_input .extend(encode_server_frame(frame.opcode, false, &inflated)); } else { @@ -143,7 +151,7 @@ impl PerMessageDeflateStream { let message = self.compressed_message.take().ok_or_else(|| { io::Error::new(io::ErrorKind::InvalidData, "missing compressed message") })?; - let inflated = inflate_message(&message.payload)?; + let inflated = inflate_permessage_deflate_message(&message.payload)?; self.decoded_input.extend(encode_server_frame( message.opcode, false, @@ -170,6 +178,48 @@ fn copy_decoded_input(decoded_input: &mut VecDeque, buffer: &mut ReadBuf<'_> true } +/// 判断扩展响应是否启用了 permessage-deflate。 +pub fn permessage_deflate_enabled(extension_header: &str) -> bool { + extension_header + .split(',') + .any(|value| value.trim().starts_with("permessage-deflate")) +} + +/// 解压单条 permessage-deflate 消息 payload。 +pub fn inflate_permessage_deflate_message(payload: &[u8]) -> io::Result> { + let mut compressed = Vec::with_capacity(payload.len() + PERMESSAGE_DEFLATE_TAIL.len()); + compressed.extend_from_slice(payload); + compressed.extend_from_slice(&PERMESSAGE_DEFLATE_TAIL); + + let mut decoder = DeflateDecoder::new(compressed.as_slice()); + let mut inflated = Vec::new(); + decoder.read_to_end(&mut inflated)?; + if inflated.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "permessage-deflate frame produced no output", + )); + } + Ok(inflated) +} + +/// 将启用 RSV1 的压缩服务端 data frame 重写为普通未压缩 frame。 +pub fn rewrite_permessage_deflate_server_frame(frame: &[u8]) -> io::Result>> { + let Some(parts) = parse_frame_parts(frame)? else { + return Ok(None); + }; + if parts.frame_len != frame.len() || parts.masked || !parts.rsv1 || !parts.fin { + return Ok(None); + } + if !is_data_frame(parts.opcode) { + return Ok(None); + } + + let payload = &frame[parts.payload_offset..parts.payload_offset + parts.payload_len]; + let inflated = inflate_permessage_deflate_message(payload)?; + Ok(Some(encode_server_frame(parts.opcode, false, &inflated))) +} + fn parse_frame_parts(bytes: &[u8]) -> io::Result> { if bytes.len() < 2 { return Ok(None); @@ -231,24 +281,6 @@ fn parse_frame_parts(bytes: &[u8]) -> io::Result> { })) } -fn inflate_message(payload: &[u8]) -> io::Result> { - let mut compressed = Vec::with_capacity(payload.len() + PERMESSAGE_DEFLATE_TAIL.len()); - compressed.extend_from_slice(payload); - compressed.extend_from_slice(&PERMESSAGE_DEFLATE_TAIL); - - let mut decompressor = Decompress::new(false); - let mut inflated = Vec::with_capacity(payload.len().saturating_mul(4).max(1024)); - match decompressor.decompress_vec(&compressed, &mut inflated, FlushDecompress::Sync) { - Ok(Status::Ok | Status::StreamEnd) => Ok(inflated), - Ok(Status::BufError) if !inflated.is_empty() => Ok(inflated), - Ok(Status::BufError) => Err(io::Error::new( - io::ErrorKind::InvalidData, - "permessage-deflate frame produced no output", - )), - Err(error) => Err(io::Error::new(io::ErrorKind::InvalidData, error)), - } -} - fn encode_server_frame(opcode: u8, rsv1: bool, payload: &[u8]) -> Vec { let mut frame = Vec::new(); let rsv1_bit = if rsv1 { 0x40 } else { 0 }; @@ -269,5 +301,9 @@ fn encode_server_frame(opcode: u8, rsv1: bool, payload: &[u8]) -> Vec { } fn is_data_frame(opcode: u8) -> bool { + matches!(opcode, OPCODE_TEXT | OPCODE_BINARY | OPCODE_CONTINUATION) +} + +fn is_initial_data_frame(opcode: u8) -> bool { matches!(opcode, OPCODE_TEXT | OPCODE_BINARY) } diff --git a/crates/adapters/src/codex/websocket/mod.rs b/crates/adapters/src/codex/websocket/mod.rs new file mode 100644 index 000000000..a1bc44fb3 --- /dev/null +++ b/crates/adapters/src/codex/websocket/mod.rs @@ -0,0 +1,6 @@ +//! Codex WebSocket 适配器。 + +pub mod connect; +pub mod deflate; +pub mod opening; +pub mod pool; diff --git a/crates/adapters/src/codex/websocket/opening.rs b/crates/adapters/src/codex/websocket/opening.rs new file mode 100644 index 000000000..611cb7bbd --- /dev/null +++ b/crates/adapters/src/codex/websocket/opening.rs @@ -0,0 +1,56 @@ +//! WebSocket 打开握手。 + +use std::{ + io, + path::{Path, PathBuf}, +}; + +use chrono::Utc; +use codex_proxy_core::protocol::codex::websocket::{OpeningAuditSnapshot, WebSocketAuditArtifact}; +use uuid::Uuid; + +/// WebSocket audit artifact 输出目录环境变量。 +pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR"; + +/// 构造握手审计快照。 +pub fn opening_audit_snapshot(header_order: Vec) -> OpeningAuditSnapshot { + OpeningAuditSnapshot { + header_order, + ..OpeningAuditSnapshot::default() + } +} + +/// 显式写入 WebSocket audit artifact。 +pub async fn write_websocket_audit_artifact_for_dir( + dir: Option<&Path>, + artifact: &WebSocketAuditArtifact, +) -> io::Result> { + let Some(dir) = dir.filter(|dir| !dir.as_os_str().is_empty()) else { + return Ok(None); + }; + + tokio::fs::create_dir_all(dir).await?; + let path = dir.join(websocket_audit_file_name()); + let body = serde_json::to_vec_pretty(artifact).map_err(io::Error::other)?; + tokio::fs::write(&path, body).await?; + Ok(Some(path)) +} + +/// 按环境变量配置写入 WebSocket audit artifact。 +pub async fn write_websocket_audit_artifact_from_env( + artifact: &WebSocketAuditArtifact, +) -> io::Result> { + let Some(dir) = std::env::var_os(WS_AUDIT_DIR_ENV) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + else { + return Ok(None); + }; + + write_websocket_audit_artifact_for_dir(Some(&dir), artifact).await +} + +fn websocket_audit_file_name() -> String { + let timestamp = Utc::now().format("%Y%m%dT%H%M%S%.3fZ"); + format!("codex-ws-audit-{timestamp}-{}.json", Uuid::new_v4()) +} diff --git a/src/codex/gateway/transport/websocket/pool.rs b/crates/adapters/src/codex/websocket/pool.rs similarity index 70% rename from src/codex/gateway/transport/websocket/pool.rs rename to crates/adapters/src/codex/websocket/pool.rs index f42020c50..5547c51f3 100644 --- a/src/codex/gateway/transport/websocket/pool.rs +++ b/crates/adapters/src/codex/websocket/pool.rs @@ -1,3 +1,5 @@ +//! Codex WebSocket 连接池。 + use std::{ collections::HashMap, sync::Arc, @@ -6,15 +8,18 @@ use std::{ use futures::{SinkExt, StreamExt}; use tokio::{net::TcpStream, sync::Mutex, time::timeout}; -use tokio_tungstenite::{ - tungstenite::{http::HeaderMap as WsHeaderMap, Message}, - MaybeTlsStream, WebSocketStream, -}; +use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream}; -use super::codec::{rate_limit_headers, set_cookie_headers, turn_state}; use super::deflate::PerMessageDeflateStream; -use super::OpeningAuditSnapshot; +const DEFAULT_MAX_PER_ACCOUNT: usize = 8; +const DEFAULT_MAX_AGE: Duration = Duration::from_secs(55 * 60); +const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25); +const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25); +const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5); +const DEFAULT_LIVENESS_TIMEOUT: Duration = Duration::from_millis(62_500); + +/// WebSocket 连接池 key。 #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct CodexWebSocketPoolKey { base_url: String, @@ -23,6 +28,7 @@ pub struct CodexWebSocketPoolKey { } impl CodexWebSocketPoolKey { + /// 构造连接池 key。 pub fn new( base_url: impl Into, account_id: impl Into, @@ -35,44 +41,63 @@ impl CodexWebSocketPoolKey { } } - fn account_id(&self) -> &str { + pub(crate) fn account_id(&self) -> &str { &self.account_id } } +/// WebSocket 连接池。 #[derive(Clone)] pub struct CodexWebSocketPool { inner: Arc>, config: CodexWebSocketPoolConfig, } +impl Default for CodexWebSocketPool { + fn default() -> Self { + Self::with_config(CodexWebSocketPoolConfig::default()) + } +} + +/// WebSocket 连接池配置。 #[derive(Debug, Clone, Copy)] pub struct CodexWebSocketPoolConfig { + /// 是否启用连接池。 pub enabled: bool, + /// 单个 socket 的最大生命周期。 pub max_age: Duration, + /// 单个账号允许占用的最大池 slot 数。 pub max_per_account: usize, + /// 后台维护间隔；`None` 表示不启动后台任务。 pub maintenance_interval: Option, + /// idle socket 探活 ping 间隔；`None` 表示维护时只做过期清理。 pub ping_interval: Option, + /// 发送 ping 后等待上游响应的超时时间。 pub ping_timeout: Duration, + /// idle socket 无活动多久后视为失活。 pub liveness_timeout: Option, } -impl CodexWebSocketPool { - const DEFAULT_MAX_AGE: Duration = Duration::from_secs(55 * 60); - const DEFAULT_MAX_PER_ACCOUNT: usize = 8; - const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25); - const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25); - const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5); - const DEFAULT_LIVENESS_TIMEOUT: Duration = Duration::from_millis(62_500); - - pub fn with_default_max_age() -> Self { - Self::with_config(CodexWebSocketPoolConfig::default()) +impl Default for CodexWebSocketPoolConfig { + fn default() -> Self { + Self { + enabled: true, + max_age: DEFAULT_MAX_AGE, + max_per_account: DEFAULT_MAX_PER_ACCOUNT, + maintenance_interval: Some(DEFAULT_MAINTENANCE_INTERVAL), + ping_interval: Some(DEFAULT_PING_INTERVAL), + ping_timeout: DEFAULT_PING_TIMEOUT, + liveness_timeout: Some(DEFAULT_LIVENESS_TIMEOUT), + } } +} - pub fn with_limits(max_age: Duration, max_per_account: usize) -> Self { +impl CodexWebSocketPool { + /// 构造连接池策略和状态。 + pub fn new(max_per_account: usize, max_age: Duration) -> Self { Self::with_config(CodexWebSocketPoolConfig { + max_per_account: max_per_account.max(1), max_age, - max_per_account, maintenance_interval: None, ping_interval: None, liveness_timeout: None, @@ -80,6 +105,7 @@ impl CodexWebSocketPool { }) } + /// 使用完整配置构造连接池。 pub fn with_config(config: CodexWebSocketPoolConfig) -> Self { let pool = Self { inner: Arc::new(Mutex::new(WebSocketPoolState::default())), @@ -89,46 +115,47 @@ impl CodexWebSocketPool { pool } - pub async fn evict_account(&self, account_id: &str) { - let mut idle_connections = Vec::new(); - { - let mut state = self.inner.lock().await; - let keys = state - .slots - .keys() - .filter(|key| key.account_id() == account_id) - .cloned() - .collect::>(); - for key in keys { - if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) { - idle_connections.push(*connection); - } - } + /// 使用默认配置构造连接池。 + pub fn with_default_max_age() -> Self { + Self::with_config(CodexWebSocketPoolConfig::default()) + } + + /// 使用最大存活时间与账号容量限制构造连接池。 + pub fn with_limits(max_age: Duration, max_per_account: usize) -> Self { + Self { + inner: Arc::new(Mutex::new(WebSocketPoolState::default())), + config: CodexWebSocketPoolConfig { + max_age, + max_per_account: max_per_account.max(1), + maintenance_interval: None, + ping_interval: None, + liveness_timeout: None, + ..CodexWebSocketPoolConfig::default() + }, } - close_idle_connections(idle_connections).await; } - pub async fn shutdown(&self) { - let idle_connections = { - let mut state = self.inner.lock().await; - state.shutting_down = true; - state - .slots - .drain() - .filter_map(|(_, slot)| match slot { - WebSocketPoolSlot::Idle(connection) => Some(*connection), - WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => None, - }) - .collect::>() - }; - close_idle_connections(idle_connections).await; + /// 返回单账号最大连接数。 + pub fn max_per_account(&self) -> usize { + self.config.max_per_account } - pub async fn gc_sweep(&self) { - self.maintain_idle_connections().await; + /// 返回连接最大存活时长。 + pub fn max_age(&self) -> Duration { + self.config.max_age } - pub(super) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire { + /// 判断当前账号是否还能打开新连接。 + pub fn permits_new_connection(&self, current_connections: usize) -> bool { + current_connections < self.config.max_per_account + } + + /// 判断连接是否已达到回收年龄。 + pub fn should_recycle(&self, age: Duration) -> bool { + age >= self.config.max_age + } + + pub(crate) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire { let mut expired_connection = None; let acquire = { let mut state = self.inner.lock().await; @@ -143,7 +170,7 @@ impl CodexWebSocketPool { let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(key) else { return WebSocketPoolAcquire::Bypass; }; - if connection.created_at.elapsed() <= self.config.max_age { + if connection.created_at.elapsed() < self.config.max_age { state.slots.insert(key.clone(), WebSocketPoolSlot::Busy); return WebSocketPoolAcquire::Reused(connection); } @@ -163,13 +190,13 @@ impl CodexWebSocketPool { }; if let Some(connection) = expired_connection { - close_idle_connections(vec![connection]).await; + close_pooled_connection(connection).await; } acquire } - pub(super) async fn put( + pub(crate) async fn put( &self, key: CodexWebSocketPoolKey, connection: PooledWebSocketConnection, @@ -179,7 +206,7 @@ impl CodexWebSocketPool { let mut state = self.inner.lock().await; let expired = connection .as_ref() - .is_some_and(|connection| connection.created_at.elapsed() > self.config.max_age); + .is_some_and(|connection| connection.created_at.elapsed() >= self.config.max_age); if expired || state.shutting_down || !self.config.enabled { state.slots.remove(&key); } else if matches!(state.slots.get(&key), Some(WebSocketPoolSlot::Busy)) { @@ -195,7 +222,7 @@ impl CodexWebSocketPool { } } - pub(super) async fn discard(&self, key: &CodexWebSocketPoolKey) { + pub(crate) async fn discard(&self, key: &CodexWebSocketPoolKey) { let idle_connection = { let mut state = self.inner.lock().await; match state.slots.remove(key) { @@ -204,7 +231,59 @@ impl CodexWebSocketPool { } }; if let Some(connection) = idle_connection { - close_idle_connections(vec![connection]).await; + close_pooled_connection(connection).await; + } + } + + /// 驱逐指定账号的 idle 连接，并阻止已占用 slot 回收到池中。 + pub async fn evict_account(&self, account_id: &str) { + let mut idle_connections = Vec::new(); + { + let mut state = self.inner.lock().await; + let keys = state + .slots + .keys() + .filter(|key| key.account_id() == account_id) + .cloned() + .collect::>(); + for key in keys { + if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) { + idle_connections.push(*connection); + } + } + } + close_pooled_connections(idle_connections).await; + } + + /// 关闭连接池，关闭所有 idle 连接，并让后续 acquire 直接绕过池。 + pub async fn shutdown(&self) { + let idle_connections = { + let mut state = self.inner.lock().await; + state.shutting_down = true; + state + .slots + .drain() + .filter_map(|(_, slot)| match slot { + WebSocketPoolSlot::Idle(connection) => Some(*connection), + WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => None, + }) + .collect::>() + }; + close_pooled_connections(idle_connections).await; + } + + /// 维护 idle 连接：关闭过期/失活连接，并按需 ping 探活。 + pub async fn gc_sweep(&self) { + self.maintain_idle_connections().await; + } + + /// 维护 idle 连接：关闭过期/失活连接，并按需 ping 探活。 + pub async fn maintain_idle_connections(&self) { + let maintenance = self.take_idle_connections_for_maintenance().await; + close_pooled_connections(maintenance.close).await; + for (key, connection) in maintenance.probe { + let connection = probe_idle_connection(connection, self.config.ping_timeout).await; + self.return_maintained_connection(key, connection).await; } } @@ -212,11 +291,11 @@ impl CodexWebSocketPool { let Some(interval_duration) = self.config.maintenance_interval else { return; }; - let inner = Arc::downgrade(&self.inner); - let config = self.config; let Ok(_handle) = tokio::runtime::Handle::try_current() else { return; }; + let inner = Arc::downgrade(&self.inner); + let config = self.config; tokio::spawn(async move { let mut interval = tokio::time::interval(interval_duration); loop { @@ -237,15 +316,6 @@ impl CodexWebSocketPool { self.inner.lock().await.shutting_down } - pub async fn maintain_idle_connections(&self) { - let maintenance = self.take_idle_connections_for_maintenance().await; - close_idle_connections(maintenance.close).await; - for (key, connection) in maintenance.probe { - let connection = probe_idle_connection(connection, self.config.ping_timeout).await; - self.return_maintained_connection(key, connection).await; - } - } - async fn take_idle_connections_for_maintenance(&self) -> WebSocketPoolMaintenance { let mut close = Vec::new(); let mut probe = Vec::new(); @@ -255,8 +325,26 @@ impl CodexWebSocketPool { if state.shutting_down || !self.config.enabled { return WebSocketPoolMaintenance { close, probe }; } - let expired_keys = expired_idle_keys(&state.slots, self.config.max_age); - for key in expired_keys { + let keys = state + .slots + .iter() + .filter_map(|(key, slot)| match slot { + WebSocketPoolSlot::Idle(connection) + if should_close_idle_connection( + connection, + now, + self.config.max_age, + self.config.liveness_timeout, + ) => + { + Some(key.clone()) + } + WebSocketPoolSlot::Idle(_) + | WebSocketPoolSlot::Busy + | WebSocketPoolSlot::Checking => None, + }) + .collect::>(); + for key in keys { if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) { close.push(*connection); } @@ -332,59 +420,28 @@ impl CodexWebSocketPool { } } -impl Default for CodexWebSocketPool { - fn default() -> Self { - Self::with_default_max_age() - } -} - -impl Default for CodexWebSocketPoolConfig { - fn default() -> Self { - Self { - enabled: true, - max_age: CodexWebSocketPool::DEFAULT_MAX_AGE, - max_per_account: CodexWebSocketPool::DEFAULT_MAX_PER_ACCOUNT, - maintenance_interval: Some(CodexWebSocketPool::DEFAULT_MAINTENANCE_INTERVAL), - ping_interval: Some(CodexWebSocketPool::DEFAULT_PING_INTERVAL), - ping_timeout: CodexWebSocketPool::DEFAULT_PING_TIMEOUT, - liveness_timeout: Some(CodexWebSocketPool::DEFAULT_LIVENESS_TIMEOUT), - } - } -} - #[derive(Default)] struct WebSocketPoolState { slots: HashMap, shutting_down: bool, } -pub(super) type CodexWsStream = WebSocketStream>>; +pub(crate) type CodexWsStream = WebSocketStream>>; -#[derive(Debug, Clone, PartialEq, Eq)] -pub(super) struct CodexWebSocketConnectionMetadata { - pub(super) turn_state: Option, - pub(super) set_cookie_headers: Vec, - pub(super) rate_limit_headers: Vec<(String, String)>, - pub(super) opening_audit: OpeningAuditSnapshot, +#[derive(Clone)] +pub(crate) struct CodexWebSocketConnectionMetadata { + pub(crate) turn_state: Option, + pub(crate) set_cookie_headers: Vec, + pub(crate) rate_limit_headers: Vec<(String, String)>, + pub(crate) handshake_status: u16, } -impl CodexWebSocketConnectionMetadata { - pub(super) fn from_headers(headers: &WsHeaderMap, opening_audit: OpeningAuditSnapshot) -> Self { - Self { - turn_state: turn_state(headers), - set_cookie_headers: set_cookie_headers(headers), - rate_limit_headers: rate_limit_headers(headers), - opening_audit, - } - } -} - -pub(super) struct PooledWebSocketConnection { - pub(super) websocket: CodexWsStream, - pub(super) metadata: CodexWebSocketConnectionMetadata, - pub(super) created_at: Instant, - pub(super) last_activity_at: Instant, - pub(super) last_ping_at: Option, +pub(crate) struct PooledWebSocketConnection { + pub(crate) websocket: CodexWsStream, + pub(crate) metadata: CodexWebSocketConnectionMetadata, + pub(crate) created_at: Instant, + pub(crate) last_activity_at: Instant, + pub(crate) last_ping_at: Option, } enum WebSocketPoolSlot { @@ -393,64 +450,46 @@ enum WebSocketPoolSlot { Checking, } -pub(super) enum WebSocketPoolAcquire { +pub(crate) enum WebSocketPoolAcquire { Reused(Box), FreshReserved, Bypass, } fn account_slot_count( - inner: &HashMap, + slots: &HashMap, account_id: &str, ) -> usize { - inner + slots .keys() .filter(|key| key.account_id() == account_id) .count() } -async fn close_idle_connections(connections: Vec) { +async fn close_pooled_connection(mut connection: PooledWebSocketConnection) { + let _ = connection.websocket.send(Message::Close(None)).await; +} + +async fn close_pooled_connections(connections: Vec) { for connection in connections { close_pooled_connection(connection).await; } } -async fn close_pooled_connection(connection: PooledWebSocketConnection) { - let mut websocket = connection.websocket; - let _ = websocket.close(None).await; -} - struct WebSocketPoolMaintenance { close: Vec, probe: Vec<(CodexWebSocketPoolKey, PooledWebSocketConnection)>, } -fn expired_idle_keys( - slots: &HashMap, - max_age: Duration, -) -> Vec { - slots - .iter() - .filter_map(|(key, slot)| match slot { - WebSocketPoolSlot::Idle(connection) if connection.created_at.elapsed() > max_age => { - Some(key.clone()) - } - WebSocketPoolSlot::Idle(_) | WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => { - None - } - }) - .collect() -} - fn should_close_idle_connection( connection: &PooledWebSocketConnection, now: Instant, max_age: Duration, liveness_timeout: Option, ) -> bool { - connection.created_at.elapsed() > max_age + now.duration_since(connection.created_at) >= max_age || liveness_timeout.is_some_and(|timeout| { - !timeout.is_zero() && now.duration_since(connection.last_activity_at) > timeout + !timeout.is_zero() && now.duration_since(connection.last_activity_at) >= timeout }) } @@ -469,8 +508,7 @@ async fn probe_idle_connection( mut connection: PooledWebSocketConnection, ping_timeout: Duration, ) -> Option { - let now = Instant::now(); - connection.last_ping_at = Some(now); + connection.last_ping_at = Some(Instant::now()); if connection .websocket .send(Message::Ping(Vec::new().into())) diff --git a/crates/adapters/src/lib.rs b/crates/adapters/src/lib.rs new file mode 100644 index 000000000..bcf4885d9 --- /dev/null +++ b/crates/adapters/src/lib.rs @@ -0,0 +1,7 @@ +#![deny(missing_docs)] + +//! 适配器层，承载具体存储、HTTP、OAuth 和 WebSocket 实现。 + +pub mod codex; +pub mod oauth; +pub mod sqlite; diff --git a/crates/adapters/src/oauth/mod.rs b/crates/adapters/src/oauth/mod.rs new file mode 100644 index 000000000..3c9d93f33 --- /dev/null +++ b/crates/adapters/src/oauth/mod.rs @@ -0,0 +1,4 @@ +//! OAuth 上游适配器。 + +/// OpenAI OAuth 的 Reqwest 实现。 +pub mod openai; diff --git a/crates/adapters/src/oauth/openai.rs b/crates/adapters/src/oauth/openai.rs new file mode 100644 index 000000000..4f7d57e4e --- /dev/null +++ b/crates/adapters/src/oauth/openai.rs @@ -0,0 +1,190 @@ +//! OpenAI OAuth Reqwest 适配器。 + +use async_trait::async_trait; +use reqwest::{Client, StatusCode}; +use serde::Deserialize; + +use codex_proxy_core::auth::{ + oauth::{DeviceCode, OAuthConfig, OAuthError, RefreshFailure, TokenPair}, + ports::{OAuthClient, TokenRefresher}, +}; + +/// OpenAI OAuth 上游客户端。 +#[derive(Clone)] +pub struct OpenAiOAuthClient { + client: Client, + config: OAuthConfig, +} + +impl OpenAiOAuthClient { + /// 使用 Reqwest 客户端和静态配置构造 OAuth 客户端。 + pub fn new(client: Client, config: OAuthConfig) -> Self { + Self { client, config } + } +} + +/// 使用默认 Reqwest 客户端构造 OpenAI OAuth 客户端。 +pub fn default_openai_oauth_client(config: OAuthConfig) -> OpenAiOAuthClient { + OpenAiOAuthClient::new(Client::new(), config) +} + +#[derive(Deserialize)] +struct RefreshTokenResponse { + access_token: String, + refresh_token: Option, +} + +#[derive(Deserialize)] +struct OAuthErrorResponse { + error: Option, + error_description: Option, +} + +#[async_trait] +impl OAuthClient for OpenAiOAuthClient { + async fn exchange_code( + &self, + code: &str, + code_verifier: &str, + redirect_uri: &str, + ) -> Result { + let response = self + .client + .post(&self.config.token_endpoint) + .form(&[ + ("grant_type", "authorization_code"), + ("client_id", self.config.client_id.as_str()), + ("code", code), + ("redirect_uri", redirect_uri), + ("code_verifier", code_verifier), + ]) + .send() + .await + .map_err(|_| OAuthError::Transport)?; + let status = response.status(); + let body = response.text().await.map_err(|_| OAuthError::Transport)?; + if !status.is_success() { + return Err(OAuthError::Rejected(format!( + "Token exchange failed ({status}): {body}" + ))); + } + parse_token_pair(&body).map_err(|_| OAuthError::Transport) + } + + async fn request_device_code(&self) -> Result { + let response = self + .client + .post(&self.config.device_code_endpoint) + .form(&[ + ("client_id", self.config.client_id.as_str()), + ("scope", "openid profile email offline_access"), + ]) + .send() + .await + .map_err(|_| OAuthError::Transport)?; + let status = response.status(); + let body = response.text().await.map_err(|_| OAuthError::Transport)?; + if !status.is_success() { + return Err(OAuthError::Rejected(format!( + "Device code request failed ({status}): {body}" + ))); + } + serde_json::from_str::(&body).map_err(|_| OAuthError::Transport) + } + + async fn poll_device_token(&self, device_code: &str) -> Result { + let response = self + .client + .post(&self.config.token_endpoint) + .form(&[ + ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"), + ("device_code", device_code), + ("client_id", self.config.client_id.as_str()), + ]) + .send() + .await + .map_err(|_| OAuthError::Transport)?; + let status = response.status(); + let body = response.text().await.map_err(|_| OAuthError::Transport)?; + if !status.is_success() { + return Err(classify_device_poll_failure(status, &body)); + } + parse_token_pair(&body).map_err(|_| OAuthError::Transport) + } +} + +#[async_trait] +impl TokenRefresher for OpenAiOAuthClient { + async fn refresh(&self, refresh_token: &str) -> Result { + let response = self + .client + .post(&self.config.token_endpoint) + .form(&[ + ("grant_type", "refresh_token"), + ("client_id", self.config.client_id.as_str()), + ("refresh_token", refresh_token), + ]) + .send() + .await + .map_err(|_| RefreshFailure::Transport)?; + let status = response.status(); + let body = response + .text() + .await + .map_err(|_| RefreshFailure::Transport)?; + if !status.is_success() { + return Err(classify_refresh_failure(status, &body)); + } + parse_token_pair(&body).map_err(|_| RefreshFailure::Transport) + } +} + +fn parse_token_pair(body: &str) -> Result { + let tokens = serde_json::from_str::(body).map_err(|_| ())?; + if tokens.access_token.trim().is_empty() { + return Err(()); + } + Ok(TokenPair { + access_token: tokens.access_token, + refresh_token: tokens.refresh_token, + }) +} + +fn classify_device_poll_failure(status: StatusCode, body: &str) -> OAuthError { + let parsed = serde_json::from_str::(body).ok(); + match parsed.as_ref().and_then(|data| data.error.as_deref()) { + Some("authorization_pending") => OAuthError::AuthorizationPending, + Some("slow_down") => OAuthError::SlowDown, + Some(code) => OAuthError::Rejected( + parsed + .as_ref() + .and_then(|data| data.error_description.clone()) + .unwrap_or_else(|| code.to_string()), + ), + None => OAuthError::Rejected(format!("Device token poll failed ({status}): {body}")), + } +} + +fn classify_refresh_failure(status: StatusCode, body: &str) -> RefreshFailure { + let lower = body.to_ascii_lowercase(); + if lower.contains("quota") { + return RefreshFailure::QuotaExhausted; + } + if lower.contains("account has been deactivated") + || lower.contains("refresh_token_reused") + || lower.contains("banned") + { + return RefreshFailure::Banned; + } + if status == StatusCode::BAD_REQUEST + || status == StatusCode::UNAUTHORIZED + || lower.contains("invalid_grant") + || lower.contains("invalid_token") + || lower.contains("access_denied") + || lower.contains("refresh_token_expired") + || lower.contains("token_revoked") + { + return RefreshFailure::InvalidGrant; + } + RefreshFailure::Transport +} diff --git a/crates/adapters/src/sqlite/account_tokens.rs b/crates/adapters/src/sqlite/account_tokens.rs new file mode 100644 index 000000000..740f734d0 --- /dev/null +++ b/crates/adapters/src/sqlite/account_tokens.rs @@ -0,0 +1,10 @@ +//! SQLite 账号 token 加解密辅助。 + +use secrecy::SecretString; + +use codex_proxy_platform::crypto::{CryptoResult, SecretBox}; + +/// 加密账号 token。 +pub fn encrypt_token(secret_box: &SecretBox, token: &SecretString) -> CryptoResult { + secret_box.encrypt(token) +} diff --git a/crates/adapters/src/sqlite/account_usage.rs b/crates/adapters/src/sqlite/account_usage.rs new file mode 100644 index 000000000..524d685c3 --- /dev/null +++ b/crates/adapters/src/sqlite/account_usage.rs @@ -0,0 +1,21 @@ +//! SQLite 账号用量存储。 + +use sqlx::SqlitePool; + +/// SQLite 用量存储。 +#[derive(Clone)] +pub struct SqliteUsageStore { + pool: SqlitePool, +} + +impl SqliteUsageStore { + /// 构造存储。 + pub fn new(pool: SqlitePool) -> Self { + Self { pool } + } + + /// 返回底层连接池。 + pub fn pool(&self) -> &SqlitePool { + &self.pool + } +} diff --git a/crates/adapters/src/sqlite/accounts.rs b/crates/adapters/src/sqlite/accounts.rs new file mode 100644 index 000000000..8308cf20d --- /dev/null +++ b/crates/adapters/src/sqlite/accounts.rs @@ -0,0 +1,1671 @@ +//! SQLite 账号仓储适配器。 + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use secrecy::{ExposeSecret, SecretString}; +use sqlx::{sqlite::SqliteRow, Row, SqlitePool}; +use thiserror::Error; + +use codex_proxy_core::accounts::{ + model::{Account, AccountStatus}, + ports::{AccountStore, AccountStoreError, AccountStoreResult}, + usage::AccountUsageDelta, +}; +use codex_proxy_platform::{ + crypto::{CryptoError, SecretBox}, + json::{decode_cursor, encode_cursor, Page}, +}; + +const INSERT_ACCOUNT_SQL: &str = r" +insert into accounts ( + id, + email, + account_id, + user_id, + label, + plan_type, + access_token_cipher, + refresh_token_cipher, + access_token_expires_at, + status, + added_at, + updated_at +) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"; + +const SELECT_STORED_ACCOUNT_BY_ID_SQL: &str = r" +select + id, + email, + account_id, + user_id, + label, + plan_type, + access_token_cipher, + refresh_token_cipher, + access_token_expires_at, + status, + added_at, + updated_at +from accounts +where id = ?"; + +const SELECT_STORED_ACCOUNT_BY_CHATGPT_IDENTITY_SQL: &str = r" +select + id, + email, + account_id, + user_id, + label, + plan_type, + access_token_cipher, + refresh_token_cipher, + access_token_expires_at, + status, + added_at, + updated_at +from accounts +where account_id = ? + and ((user_id is null and ? is null) or user_id = ?) +order by added_at asc +limit 1"; + +const LIST_STORED_ACCOUNTS_AFTER_CURSOR_SQL: &str = r" +select + id, + email, + account_id, + user_id, + label, + plan_type, + access_token_cipher, + refresh_token_cipher, + access_token_expires_at, + status, + added_at, + updated_at +from accounts +where added_at < ? + or (added_at = ? and id < ?) +order by added_at desc, id desc +limit ?"; + +const LIST_STORED_ACCOUNTS_SQL: &str = r" +select + id, + email, + account_id, + user_id, + label, + plan_type, + access_token_cipher, + refresh_token_cipher, + access_token_expires_at, + status, + added_at, + updated_at +from accounts +order by added_at desc, id desc +limit ?"; + +const LIST_ACCOUNT_METADATA_AFTER_CURSOR_SQL: &str = r" +select + id, + email, + account_id, + user_id, + label, + plan_type, + access_token_expires_at, + status, + added_at, + updated_at +from accounts +where added_at < ? + or (added_at = ? and id < ?) +order by added_at desc, id desc +limit ?"; + +const LIST_ACCOUNT_METADATA_SQL: &str = r" +select + id, + email, + account_id, + user_id, + label, + plan_type, + access_token_expires_at, + status, + added_at, + updated_at +from accounts +order by added_at desc, id desc +limit ?"; + +const LIST_POOL_ACCOUNTS_SQL: &str = r" +select + a.id, + a.email, + a.account_id, + a.user_id, + a.label, + a.plan_type, + a.access_token_cipher, + a.refresh_token_cipher, + a.access_token_expires_at, + a.status, + a.quota_limit_reached, + a.quota_verify_required, + a.quota_cooldown_until, + a.cloudflare_cooldown_until, + a.added_at, + coalesce(au.request_count, 0) as usage_request_count, + coalesce(au.empty_response_count, 0) as usage_empty_response_count, + coalesce(au.image_input_tokens, 0) as usage_image_input_tokens, + coalesce(au.image_output_tokens, 0) as usage_image_output_tokens, + coalesce(au.image_request_count, 0) as usage_image_request_count, + coalesce(au.image_request_failed_count, 0) as usage_image_request_failed_count, + coalesce(au.window_request_count, 0) as usage_window_request_count, + coalesce(au.window_input_tokens, 0) as usage_window_input_tokens, + coalesce(au.window_output_tokens, 0) as usage_window_output_tokens, + coalesce(au.window_cached_tokens, 0) as usage_window_cached_tokens, + coalesce(au.window_image_input_tokens, 0) as usage_window_image_input_tokens, + coalesce(au.window_image_output_tokens, 0) as usage_window_image_output_tokens, + coalesce(au.window_image_request_count, 0) as usage_window_image_request_count, + coalesce(au.window_image_request_failed_count, 0) as usage_window_image_request_failed_count, + au.window_started_at as usage_window_started_at, + au.window_reset_at as usage_window_reset_at, + au.limit_window_seconds as usage_limit_window_seconds, + au.last_used_at as usage_last_used_at +from accounts a +left join account_usage au on au.account_id = a.id +order by a.added_at desc, a.id desc"; + +const GET_POOL_ACCOUNT_SQL: &str = r" +select + a.id, + a.email, + a.account_id, + a.user_id, + a.label, + a.plan_type, + a.access_token_cipher, + a.refresh_token_cipher, + a.access_token_expires_at, + a.status, + a.quota_limit_reached, + a.quota_verify_required, + a.quota_cooldown_until, + a.cloudflare_cooldown_until, + a.added_at, + coalesce(au.request_count, 0) as usage_request_count, + coalesce(au.empty_response_count, 0) as usage_empty_response_count, + coalesce(au.image_input_tokens, 0) as usage_image_input_tokens, + coalesce(au.image_output_tokens, 0) as usage_image_output_tokens, + coalesce(au.image_request_count, 0) as usage_image_request_count, + coalesce(au.image_request_failed_count, 0) as usage_image_request_failed_count, + coalesce(au.window_request_count, 0) as usage_window_request_count, + coalesce(au.window_input_tokens, 0) as usage_window_input_tokens, + coalesce(au.window_output_tokens, 0) as usage_window_output_tokens, + coalesce(au.window_cached_tokens, 0) as usage_window_cached_tokens, + coalesce(au.window_image_input_tokens, 0) as usage_window_image_input_tokens, + coalesce(au.window_image_output_tokens, 0) as usage_window_image_output_tokens, + coalesce(au.window_image_request_count, 0) as usage_window_image_request_count, + coalesce(au.window_image_request_failed_count, 0) as usage_window_image_request_failed_count, + au.window_started_at as usage_window_started_at, + au.window_reset_at as usage_window_reset_at, + au.limit_window_seconds as usage_limit_window_seconds, + au.last_used_at as usage_last_used_at +from accounts a +left join account_usage au on au.account_id = a.id +where a.id = ?"; + +const RECORD_USAGE_SQL: &str = r" +insert into account_usage ( + account_id, + request_count, + empty_response_count, + input_tokens, + output_tokens, + cached_tokens, + image_input_tokens, + image_output_tokens, + image_request_count, + image_request_failed_count, + window_request_count, + window_input_tokens, + window_output_tokens, + window_cached_tokens, + window_image_input_tokens, + window_image_output_tokens, + window_image_request_count, + window_image_request_failed_count, + last_used_at +) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +on conflict(account_id) do update set + request_count = request_count + excluded.request_count, + empty_response_count = empty_response_count + excluded.empty_response_count, + input_tokens = input_tokens + excluded.input_tokens, + output_tokens = output_tokens + excluded.output_tokens, + cached_tokens = cached_tokens + excluded.cached_tokens, + image_input_tokens = image_input_tokens + excluded.image_input_tokens, + image_output_tokens = image_output_tokens + excluded.image_output_tokens, + image_request_count = image_request_count + excluded.image_request_count, + image_request_failed_count = image_request_failed_count + excluded.image_request_failed_count, + window_request_count = window_request_count + excluded.window_request_count, + window_input_tokens = window_input_tokens + excluded.window_input_tokens, + window_output_tokens = window_output_tokens + excluded.window_output_tokens, + window_cached_tokens = window_cached_tokens + excluded.window_cached_tokens, + window_image_input_tokens = window_image_input_tokens + excluded.window_image_input_tokens, + window_image_output_tokens = window_image_output_tokens + excluded.window_image_output_tokens, + window_image_request_count = window_image_request_count + excluded.window_image_request_count, + window_image_request_failed_count = window_image_request_failed_count + excluded.window_image_request_failed_count, + window_started_at = case + when account_usage.window_started_at is null + and (account_usage.window_reset_at is not null or account_usage.limit_window_seconds is not null) + then excluded.last_used_at + else account_usage.window_started_at + end, + last_used_at = excluded.last_used_at"; + +const GET_USAGE_SQL: &str = r" +select + account_id, + request_count, + empty_response_count, + input_tokens, + output_tokens, + cached_tokens, + image_input_tokens, + image_output_tokens, + image_request_count, + image_request_failed_count, + window_request_count, + window_input_tokens, + window_output_tokens, + window_cached_tokens, + window_image_input_tokens, + window_image_output_tokens, + window_image_request_count, + window_image_request_failed_count, + window_started_at, + window_reset_at, + limit_window_seconds, + last_used_at +from account_usage +where account_id = ?"; + +const LIST_USAGE_AFTER_CURSOR_SQL: &str = r" +select + au.account_id, + a.email, + a.label, + a.plan_type, + au.request_count, + au.empty_response_count, + au.input_tokens, + au.output_tokens, + au.cached_tokens, + au.image_input_tokens, + au.image_output_tokens, + au.image_request_count, + au.image_request_failed_count, + au.last_used_at +from account_usage au +left join accounts a on a.id = au.account_id +where au.last_used_at < ? + or (au.last_used_at = ? and au.account_id < ?) +order by au.last_used_at desc, au.account_id desc +limit ?"; + +const LIST_USAGE_SQL: &str = r" +select + au.account_id, + a.email, + a.label, + a.plan_type, + au.request_count, + au.empty_response_count, + au.input_tokens, + au.output_tokens, + au.cached_tokens, + au.image_input_tokens, + au.image_output_tokens, + au.image_request_count, + au.image_request_failed_count, + au.last_used_at +from account_usage au +left join accounts a on a.id = au.account_id +order by au.last_used_at desc, au.account_id desc +limit ?"; + +const USAGE_SUMMARY_SQL: &str = r" +select + count(*) as account_count, + coalesce(sum(request_count), 0) as request_count, + coalesce(sum(empty_response_count), 0) as empty_response_count, + coalesce(sum(input_tokens), 0) as input_tokens, + coalesce(sum(output_tokens), 0) as output_tokens, + coalesce(sum(cached_tokens), 0) as cached_tokens, + coalesce(sum(image_input_tokens), 0) as image_input_tokens, + coalesce(sum(image_output_tokens), 0) as image_output_tokens, + coalesce(sum(image_request_count), 0) as image_request_count, + coalesce(sum(image_request_failed_count), 0) as image_request_failed_count +from account_usage"; + +const RESET_USAGE_SQL: &str = r" +update account_usage +set + request_count = 0, + empty_response_count = 0, + input_tokens = 0, + output_tokens = 0, + cached_tokens = 0, + image_input_tokens = 0, + image_output_tokens = 0, + image_request_count = 0, + image_request_failed_count = 0, + window_request_count = 0, + window_input_tokens = 0, + window_output_tokens = 0, + window_cached_tokens = 0, + window_image_input_tokens = 0, + window_image_output_tokens = 0, + window_image_request_count = 0, + window_image_request_failed_count = 0, + window_started_at = null, + last_used_at = null +where account_id = ?"; + +const LIST_QUOTA_SNAPSHOTS_SQL: &str = r" +select + id, + email, + quota_json, + quota_fetched_at +from accounts +where quota_json is not null + and trim(quota_json) <> '' +order by coalesce(quota_fetched_at, '') desc, id desc"; + +const UPDATE_QUOTA_JSON_SQL: &str = r" +update accounts +set + quota_json = ?, + quota_fetched_at = ?, + updated_at = ? +where id = ?"; + +const SELECT_RATE_LIMIT_WINDOW_SQL: &str = r" +select + window_reset_at, + limit_window_seconds +from account_usage +where account_id = ?"; + +const SYNC_RATE_LIMIT_WINDOW_RESET_SQL: &str = r" +insert into account_usage ( + account_id, + window_request_count, + window_input_tokens, + window_output_tokens, + window_cached_tokens, + window_image_input_tokens, + window_image_output_tokens, + window_image_request_count, + window_image_request_failed_count, + window_started_at, + window_reset_at, + limit_window_seconds +) values (?, 0, 0, 0, 0, 0, 0, 0, 0, ?, ?, ?) +on conflict(account_id) do update set + window_request_count = 0, + window_input_tokens = 0, + window_output_tokens = 0, + window_cached_tokens = 0, + window_image_input_tokens = 0, + window_image_output_tokens = 0, + window_image_request_count = 0, + window_image_request_failed_count = 0, + window_started_at = excluded.window_started_at, + window_reset_at = excluded.window_reset_at, + limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)"; + +const SYNC_RATE_LIMIT_WINDOW_SQL: &str = r" +insert into account_usage ( + account_id, + window_reset_at, + limit_window_seconds +) values (?, ?, ?) +on conflict(account_id) do update set + window_reset_at = excluded.window_reset_at, + limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)"; + +const MARK_QUOTA_LIMITED_UNTIL_SQL: &str = r" +update accounts +set + quota_limit_reached = 1, + quota_verify_required = 0, + quota_cooldown_until = ?, + updated_at = ? +where id = ?"; + +const SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL: &str = r" +update accounts +set + cloudflare_cooldown_until = ?, + updated_at = ? +where id = ?"; + +const UPDATE_ACCOUNT_CLAIMS_WITH_REFRESH_SQL: &str = r" +update accounts +set + email = ?, + account_id = ?, + user_id = ?, + plan_type = ?, + access_token_cipher = ?, + refresh_token_cipher = ?, + access_token_expires_at = ?, + status = ?, + updated_at = ? +where id = ?"; + +const UPDATE_ACCOUNT_CLAIMS_PRESERVING_REFRESH_SQL: &str = r" +update accounts +set + email = ?, + account_id = ?, + user_id = ?, + plan_type = ?, + access_token_cipher = ?, + access_token_expires_at = ?, + status = ?, + updated_at = ? +where id = ?"; + +/// SQLite 账号仓储错误。 +#[derive(Debug, Error)] +pub enum SqliteAccountStoreError { + /// 数据库错误。 + #[error("sqlite account store database error: {0}")] + Database(#[from] sqlx::Error), + /// 加解密错误。 + #[error("sqlite account store crypto error: {0}")] + Crypto(#[from] CryptoError), + /// 时间格式错误。 + #[error("sqlite account store timestamp error: {0}")] + Timestamp(#[from] chrono::ParseError), + /// 账号状态非法。 + #[error("sqlite account store status error: {0}")] + InvalidStatus(String), + /// 分页游标非法。 + #[error("invalid account pagination cursor")] + InvalidCursor, +} + +/// SQLite 账号仓储结果。 +pub type SqliteAccountStoreResult = Result; + +/// 新建账号数据。 +#[derive(Debug)] +pub struct NewAccount { + /// 账号 ID。 + pub id: String, + /// 邮箱。 + pub email: Option, + /// ChatGPT 账号 ID。 + pub account_id: Option, + /// ChatGPT 用户 ID。 + pub user_id: Option, + /// 管理标签。 + pub label: Option, + /// 计划类型。 + pub plan_type: Option, + /// access token 明文。 + pub access_token: SecretString, + /// refresh token 明文。 + pub refresh_token: Option, + /// access token 过期时间。 + pub access_token_expires_at: Option>, + /// 初始状态。 + pub status: AccountStatus, +} + +/// 通过 JWT claims 更新已有账号。 +#[derive(Debug)] +pub struct AccountClaimsUpdate { + /// 展示邮箱。 + pub email: Option, + /// ChatGPT 账号 ID。 + pub account_id: Option, + /// ChatGPT 用户 ID。 + pub user_id: Option, + /// 计划类型。 + pub plan_type: Option, + /// 新 access token。 + pub access_token: SecretString, + /// 可选的新 refresh token。 + pub refresh_token: Option, + /// access token 过期时间。 + pub access_token_expires_at: Option>, + /// 新状态。 + pub status: AccountStatus, +} + +/// 已存储账号数据。 +#[derive(Debug, Clone)] +pub struct StoredAccount { + /// 账号 ID。 + pub id: String, + /// 邮箱。 + pub email: Option, + /// ChatGPT 账号 ID。 + pub account_id: Option, + /// ChatGPT 用户 ID。 + pub user_id: Option, + /// 管理标签。 + pub label: Option, + /// 计划类型。 + pub plan_type: Option, + /// access token 明文。 + pub access_token: SecretString, + /// refresh token 明文。 + pub refresh_token: Option, + /// access token 过期时间。 + pub access_token_expires_at: Option>, + /// 账号状态。 + pub status: AccountStatus, + /// 创建时间。 + pub added_at: DateTime, + /// 更新时间。 + pub updated_at: DateTime, +} + +/// 已存储账号元数据，不包含任何密文或明文 token。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StoredAccountMetadata { + /// 账号 ID。 + pub id: String, + /// 邮箱。 + pub email: Option, + /// ChatGPT 账号 ID。 + pub account_id: Option, + /// ChatGPT 用户 ID。 + pub user_id: Option, + /// 管理标签。 + pub label: Option, + /// 计划类型。 + pub plan_type: Option, + /// access token 过期时间。 + pub access_token_expires_at: Option>, + /// 账号状态。 + pub status: AccountStatus, + /// 创建时间。 + pub added_at: DateTime, + /// 更新时间。 + pub updated_at: DateTime, +} + +/// 账号用量增量。 +#[derive(Debug, Clone, Copy)] +pub struct UsageDelta { + /// 请求数增量。 + pub request_count: i64, + /// 输入 token 增量。 + pub input_tokens: i64, + /// 输出 token 增量。 + pub output_tokens: i64, + /// 缓存 token 增量。 + pub cached_tokens: i64, + /// 图片输入 token 增量。 + pub image_input_tokens: i64, + /// 图片输出 token 增量。 + pub image_output_tokens: i64, + /// 图片请求数增量。 + pub image_request_count: i64, + /// 图片请求失败数增量。 + pub image_request_failed_count: i64, + /// 空响应数增量。 + pub empty_response_count: i64, +} + +impl Default for UsageDelta { + fn default() -> Self { + Self { + request_count: 1, + input_tokens: 0, + output_tokens: 0, + cached_tokens: 0, + image_input_tokens: 0, + image_output_tokens: 0, + image_request_count: 0, + image_request_failed_count: 0, + empty_response_count: 0, + } + } +} + +/// 账号用量记录。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountUsageRecord { + /// 账号 ID。 + pub account_id: String, + /// 请求数。 + pub request_count: i64, + /// 空响应数。 + pub empty_response_count: i64, + /// 输入 token 数。 + pub input_tokens: i64, + /// 输出 token 数。 + pub output_tokens: i64, + /// 缓存 token 数。 + pub cached_tokens: i64, + /// 图片输入 token 数。 + pub image_input_tokens: i64, + /// 图片输出 token 数。 + pub image_output_tokens: i64, + /// 图片请求数。 + pub image_request_count: i64, + /// 图片请求失败数。 + pub image_request_failed_count: i64, + /// 当前窗口请求数。 + pub window_request_count: i64, + /// 当前窗口输入 token 数。 + pub window_input_tokens: i64, + /// 当前窗口输出 token 数。 + pub window_output_tokens: i64, + /// 当前窗口缓存 token 数。 + pub window_cached_tokens: i64, + /// 当前窗口图片输入 token 数。 + pub window_image_input_tokens: i64, + /// 当前窗口图片输出 token 数。 + pub window_image_output_tokens: i64, + /// 当前窗口图片请求数。 + pub window_image_request_count: i64, + /// 当前窗口图片请求失败数。 + pub window_image_request_failed_count: i64, + /// 当前窗口起始时间。 + pub window_started_at: Option>, + /// 当前窗口重置时间。 + pub window_reset_at: Option>, + /// 限流窗口秒数。 + pub limit_window_seconds: Option, + /// 最近使用时间。 + pub last_used_at: Option>, +} + +/// 管理端账号用量列表记录。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountUsageListRecord { + /// 账号 ID。 + pub account_id: String, + /// 邮箱。 + pub email: Option, + /// 管理标签。 + pub label: Option, + /// 计划类型。 + pub plan_type: Option, + /// 请求数。 + pub request_count: i64, + /// 空响应数。 + pub empty_response_count: i64, + /// 输入 token 数。 + pub input_tokens: i64, + /// 输出 token 数。 + pub output_tokens: i64, + /// 缓存 token 数。 + pub cached_tokens: i64, + /// 图片输入 token 数。 + pub image_input_tokens: i64, + /// 图片输出 token 数。 + pub image_output_tokens: i64, + /// 图片请求数。 + pub image_request_count: i64, + /// 图片请求失败数。 + pub image_request_failed_count: i64, + /// 最近使用时间。 + pub last_used_at: Option>, +} + +/// 管理端账号用量汇总。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AccountUsageSummary { + /// 有用量记录的账号数。 + pub account_count: i64, + /// 请求总数。 + pub request_count: i64, + /// 空响应总数。 + pub empty_response_count: i64, + /// 输入 token 总数。 + pub input_tokens: i64, + /// 输出 token 总数。 + pub output_tokens: i64, + /// 缓存 token 总数。 + pub cached_tokens: i64, + /// 图片输入 token 总数。 + pub image_input_tokens: i64, + /// 图片输出 token 总数。 + pub image_output_tokens: i64, + /// 图片请求总数。 + pub image_request_count: i64, + /// 图片请求失败总数。 + pub image_request_failed_count: i64, +} + +/// 账号配额快照。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountQuotaSnapshot { + /// 账号 ID。 + pub account_id: String, + /// 邮箱。 + pub email: Option, + /// 配额 JSON。 + pub quota_json: String, + /// 配额拉取时间。 + pub quota_fetched_at: Option>, +} + +/// SQLite 账号仓储。 +#[derive(Clone)] +pub struct SqliteAccountStore { + pool: SqlitePool, + secret_box: SecretBox, +} + +#[derive(Debug)] +struct TokenWrite { + access_token_cipher: String, + refresh_token_cipher: Option, + access_token_expires_at: Option, + updated_at: String, +} + +impl SqliteAccountStore { + /// 构造适配器。 + pub fn new(pool: SqlitePool, secret_box: SecretBox) -> Self { + Self { pool, secret_box } + } + + /// 插入账号。 + pub async fn insert(&self, account: NewAccount) -> SqliteAccountStoreResult<()> { + let now = Utc::now().to_rfc3339(); + let access_token_cipher = self.secret_box.encrypt(&account.access_token)?; + let refresh_token_cipher = account + .refresh_token + .as_ref() + .map(|token| self.secret_box.encrypt(token)) + .transpose()?; + sqlx::query(INSERT_ACCOUNT_SQL) + .bind(account.id) + .bind(account.email) + .bind(account.account_id) + .bind(account.user_id) + .bind(account.label) + .bind(account.plan_type) + .bind(access_token_cipher) + .bind(refresh_token_cipher) + .bind( + account + .access_token_expires_at + .map(|value| value.to_rfc3339()), + ) + .bind(status_to_db(account.status)) + .bind(&now) + .bind(now) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// 按 ID 读取账号。 + pub async fn get(&self, id: &str) -> SqliteAccountStoreResult> { + let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_ID_SQL) + .bind(id) + .fetch_optional(&self.pool) + .await?; + row.map(|row| self.account_from_row(&row)).transpose() + } + + /// 按 ChatGPT 账号身份读取账号。 + pub async fn find_by_chatgpt_identity( + &self, + account_id: &str, + user_id: Option<&str>, + ) -> SqliteAccountStoreResult> { + let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_CHATGPT_IDENTITY_SQL) + .bind(account_id) + .bind(user_id) + .bind(user_id) + .fetch_optional(&self.pool) + .await?; + row.map(|row| self.account_from_row(&row)).transpose() + } + + /// 用 JWT claims 更新账号，并在缺少新 refresh token 时保留旧值。 + pub async fn update_from_claims( + &self, + id: &str, + update: AccountClaimsUpdate, + ) -> SqliteAccountStoreResult { + let TokenWrite { + access_token_cipher, + refresh_token_cipher, + access_token_expires_at, + updated_at, + } = self.prepare_token_write( + &update.access_token, + update.refresh_token.as_ref(), + update.access_token_expires_at, + )?; + let status = status_to_db(update.status); + + let result = if let Some(refresh_token_cipher) = refresh_token_cipher { + sqlx::query(UPDATE_ACCOUNT_CLAIMS_WITH_REFRESH_SQL) + .bind(update.email) + .bind(update.account_id) + .bind(update.user_id) + .bind(update.plan_type) + .bind(access_token_cipher) + .bind(refresh_token_cipher) + .bind(access_token_expires_at) + .bind(status) + .bind(updated_at) + .bind(id) + .execute(&self.pool) + .await? + } else { + sqlx::query(UPDATE_ACCOUNT_CLAIMS_PRESERVING_REFRESH_SQL) + .bind(update.email) + .bind(update.account_id) + .bind(update.user_id) + .bind(update.plan_type) + .bind(access_token_cipher) + .bind(access_token_expires_at) + .bind(status) + .bind(updated_at) + .bind(id) + .execute(&self.pool) + .await? + }; + Ok(result.rows_affected() > 0) + } + + /// 按创建时间倒序分页列出账号。 + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> SqliteAccountStoreResult> { + let fetch_limit = i64::from(limit) + 1; + let rows = if let Some(cursor) = cursor { + let (created_at, id) = + decode_cursor(&cursor).ok_or(SqliteAccountStoreError::InvalidCursor)?; + sqlx::query(LIST_STORED_ACCOUNTS_AFTER_CURSOR_SQL) + .bind(&created_at) + .bind(created_at) + .bind(id) + .bind(fetch_limit) + .fetch_all(&self.pool) + .await? + } else { + sqlx::query(LIST_STORED_ACCOUNTS_SQL) + .bind(fetch_limit) + .fetch_all(&self.pool) + .await? + }; + + let has_next = rows.len() > limit as usize; + let take_count = rows.len().min(limit as usize); + let mut items = Vec::with_capacity(take_count); + for row in rows.into_iter().take(take_count) { + items.push(self.account_from_row(&row)?); + } + let next_cursor = if has_next { + items + .last() + .map(|account| encode_cursor(&account.added_at.to_rfc3339(), &account.id)) + } else { + None + }; + Ok(Page { items, next_cursor }) + } + + /// 按创建时间倒序分页列出账号元数据，不解密 token。 + pub async fn list_metadata( + &self, + cursor: Option, + limit: u32, + ) -> SqliteAccountStoreResult> { + let fetch_limit = i64::from(limit) + 1; + let rows = if let Some(cursor) = cursor { + let (created_at, id) = + decode_cursor(&cursor).ok_or(SqliteAccountStoreError::InvalidCursor)?; + sqlx::query(LIST_ACCOUNT_METADATA_AFTER_CURSOR_SQL) + .bind(&created_at) + .bind(created_at) + .bind(id) + .bind(fetch_limit) + .fetch_all(&self.pool) + .await? + } else { + sqlx::query(LIST_ACCOUNT_METADATA_SQL) + .bind(fetch_limit) + .fetch_all(&self.pool) + .await? + }; + + let has_next = rows.len() > limit as usize; + let take_count = rows.len().min(limit as usize); + let items = rows + .into_iter() + .take(take_count) + .map(|row| account_metadata_from_row(&row)) + .collect::>>()?; + let next_cursor = if has_next { + items + .last() + .map(|account| encode_cursor(&account.added_at.to_rfc3339(), &account.id)) + } else { + None + }; + Ok(Page { items, next_cursor }) + } + + /// 更新账号状态。 + pub async fn set_status( + &self, + id: &str, + status: AccountStatus, + ) -> SqliteAccountStoreResult { + let result = sqlx::query("update accounts set status = ?, updated_at = ? where id = ?") + .bind(status_to_db(status)) + .bind(Utc::now().to_rfc3339()) + .bind(id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 更新账号标签。 + pub async fn set_label( + &self, + id: &str, + label: Option, + ) -> SqliteAccountStoreResult { + let result = sqlx::query("update accounts set label = ?, updated_at = ? where id = ?") + .bind(label) + .bind(Utc::now().to_rfc3339()) + .bind(id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 删除账号。 + pub async fn delete(&self, id: &str) -> SqliteAccountStoreResult { + let result = sqlx::query("delete from accounts where id = ?") + .bind(id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 删除全部账号及账号关联状态。 + pub async fn delete_all(&self) -> SqliteAccountStoreResult { + sqlx::query("delete from account_cookies") + .execute(&self.pool) + .await?; + sqlx::query("delete from account_usage") + .execute(&self.pool) + .await?; + sqlx::query("delete from account_refresh_leases") + .execute(&self.pool) + .await?; + let result = sqlx::query("delete from accounts") + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } + + /// 记录账号用量。 + pub async fn record_usage( + &self, + account_id: &str, + usage: UsageDelta, + ) -> SqliteAccountStoreResult<()> { + let now = Utc::now().to_rfc3339(); + sqlx::query(RECORD_USAGE_SQL) + .bind(account_id) + .bind(usage.request_count) + .bind(usage.empty_response_count) + .bind(usage.input_tokens) + .bind(usage.output_tokens) + .bind(usage.cached_tokens) + .bind(usage.image_input_tokens) + .bind(usage.image_output_tokens) + .bind(usage.image_request_count) + .bind(usage.image_request_failed_count) + .bind(usage.request_count) + .bind(usage.input_tokens) + .bind(usage.output_tokens) + .bind(usage.cached_tokens) + .bind(usage.image_input_tokens) + .bind(usage.image_output_tokens) + .bind(usage.image_request_count) + .bind(usage.image_request_failed_count) + .bind(now) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// 读取账号用量。 + pub async fn get_usage( + &self, + account_id: &str, + ) -> SqliteAccountStoreResult> { + let row = sqlx::query(GET_USAGE_SQL) + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + row.map(|row| usage_from_row(&row)).transpose() + } + + /// 分页列出账号用量。 + pub async fn list_usage( + &self, + cursor: Option, + limit: u32, + ) -> SqliteAccountStoreResult> { + let fetch_limit = i64::from(limit) + 1; + let rows = if let Some(cursor) = cursor { + let (last_used_at, account_id) = + decode_cursor(&cursor).ok_or(SqliteAccountStoreError::InvalidCursor)?; + sqlx::query(LIST_USAGE_AFTER_CURSOR_SQL) + .bind(&last_used_at) + .bind(last_used_at) + .bind(account_id) + .bind(fetch_limit) + .fetch_all(&self.pool) + .await? + } else { + sqlx::query(LIST_USAGE_SQL) + .bind(fetch_limit) + .fetch_all(&self.pool) + .await? + }; + + let has_next = rows.len() > limit as usize; + let take_count = rows.len().min(limit as usize); + let items = rows + .into_iter() + .take(take_count) + .map(|row| usage_list_from_row(&row)) + .collect::>>()?; + let next_cursor = if has_next { + items.last().map(|usage| { + encode_cursor( + usage + .last_used_at + .as_ref() + .map(DateTime::to_rfc3339) + .as_deref() + .unwrap_or(""), + &usage.account_id, + ) + }) + } else { + None + }; + Ok(Page { items, next_cursor }) + } + + /// 汇总账号用量。 + pub async fn usage_summary(&self) -> SqliteAccountStoreResult { + let row = sqlx::query(USAGE_SUMMARY_SQL).fetch_one(&self.pool).await?; + Ok(AccountUsageSummary { + account_count: row.get("account_count"), + request_count: row.get("request_count"), + empty_response_count: row.get("empty_response_count"), + input_tokens: row.get("input_tokens"), + output_tokens: row.get("output_tokens"), + cached_tokens: row.get("cached_tokens"), + image_input_tokens: row.get("image_input_tokens"), + image_output_tokens: row.get("image_output_tokens"), + image_request_count: row.get("image_request_count"), + image_request_failed_count: row.get("image_request_failed_count"), + }) + } + + /// 清零账号累计和当前窗口用量，保留窗口重置时间。 + pub async fn reset_usage(&self, account_id: &str) -> SqliteAccountStoreResult { + let result = sqlx::query(RESET_USAGE_SQL) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 列出已缓存的账号配额快照。 + pub async fn list_quota_snapshots( + &self, + ) -> SqliteAccountStoreResult> { + let rows = sqlx::query(LIST_QUOTA_SNAPSHOTS_SQL) + .fetch_all(&self.pool) + .await?; + rows.iter().map(quota_snapshot_from_row).collect() + } + + /// 写入账号配额 JSON 快照。 + pub async fn update_quota_json( + &self, + account_id: &str, + quota_json: &str, + ) -> SqliteAccountStoreResult { + let now = Utc::now().to_rfc3339(); + let result = sqlx::query(UPDATE_QUOTA_JSON_SQL) + .bind(quota_json) + .bind(&now) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 标记账号进入配额冷却期。 + pub async fn mark_quota_limited_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> SqliteAccountStoreResult { + let now = Utc::now().to_rfc3339(); + let result = sqlx::query(MARK_QUOTA_LIMITED_UNTIL_SQL) + .bind(cooldown_until.to_rfc3339()) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 标记账号进入 Cloudflare 冷却期。 + pub async fn set_cloudflare_cooldown_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> SqliteAccountStoreResult { + let now = Utc::now().to_rfc3339(); + let result = sqlx::query(SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL) + .bind(cooldown_until.to_rfc3339()) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 读取账号配额 JSON 快照。 + pub async fn get_quota_json( + &self, + account_id: &str, + ) -> SqliteAccountStoreResult> { + let row = sqlx::query("select quota_json from accounts where id = ?") + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + Ok(row.and_then(|row| row.get("quota_json"))) + } + + /// 同步账号当前 rate-limit 统计窗口。 + pub async fn sync_rate_limit_window( + &self, + account_id: &str, + reset_at: DateTime, + limit_window_seconds: Option, + ) -> SqliteAccountStoreResult<()> { + let row = sqlx::query(SELECT_RATE_LIMIT_WINDOW_SQL) + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + let existing_reset_at = row + .as_ref() + .map(|row| parse_optional_rfc3339(row.get::, _>("window_reset_at"))) + .transpose()? + .flatten(); + let existing_limit_window_seconds = row + .as_ref() + .and_then(|row| optional_positive_i64_to_u64(row.get("limit_window_seconds"))); + let reset_at_db = reset_at.to_rfc3339(); + let limit_window_seconds_db = limit_window_seconds.map(u64_to_i64_saturating); + + if should_reset_usage_window( + existing_reset_at, + existing_limit_window_seconds, + reset_at, + limit_window_seconds, + ) { + sqlx::query(SYNC_RATE_LIMIT_WINDOW_RESET_SQL) + .bind(account_id) + .bind(Utc::now().to_rfc3339()) + .bind(reset_at_db) + .bind(limit_window_seconds_db) + .execute(&self.pool) + .await?; + } else { + sqlx::query(SYNC_RATE_LIMIT_WINDOW_SQL) + .bind(account_id) + .bind(reset_at_db) + .bind(limit_window_seconds_db) + .execute(&self.pool) + .await?; + } + Ok(()) + } + + fn account_from_row( + &self, + row: &sqlx::sqlite::SqliteRow, + ) -> SqliteAccountStoreResult { + let access_token_cipher = row.get::("access_token_cipher"); + let access_token = self.secret_box.decrypt(&access_token_cipher)?; + let refresh_token = match row.get::, _>("refresh_token_cipher") { + Some(cipher) => Some(self.secret_box.decrypt(&cipher)?), + None => None, + }; + Ok(StoredAccount { + id: row.get("id"), + email: row.get("email"), + account_id: row.get("account_id"), + user_id: row.get("user_id"), + label: row.get("label"), + plan_type: row.get("plan_type"), + access_token, + refresh_token, + access_token_expires_at: parse_optional_rfc3339( + row.get::, _>("access_token_expires_at"), + )?, + status: status_from_db(&row.get::("status"))?, + added_at: parse_rfc3339(&row.get::("added_at"))?, + updated_at: parse_rfc3339(&row.get::("updated_at"))?, + }) + } + + fn prepare_token_write( + &self, + access_token: &SecretString, + refresh_token: Option<&SecretString>, + access_token_expires_at: Option>, + ) -> SqliteAccountStoreResult { + Ok(TokenWrite { + access_token_cipher: self.secret_box.encrypt(access_token)?, + refresh_token_cipher: refresh_token + .map(|token| self.secret_box.encrypt(token)) + .transpose()?, + access_token_expires_at: access_token_expires_at.map(|value| value.to_rfc3339()), + updated_at: Utc::now().to_rfc3339(), + }) + } +} + +fn account_metadata_from_row( + row: &sqlx::sqlite::SqliteRow, +) -> SqliteAccountStoreResult { + Ok(StoredAccountMetadata { + id: row.get("id"), + email: row.get("email"), + account_id: row.get("account_id"), + user_id: row.get("user_id"), + label: row.get("label"), + plan_type: row.get("plan_type"), + access_token_expires_at: parse_optional_rfc3339( + row.get::, _>("access_token_expires_at"), + )?, + status: status_from_db(&row.get::("status"))?, + added_at: parse_rfc3339(&row.get::("added_at"))?, + updated_at: parse_rfc3339(&row.get::("updated_at"))?, + }) +} + +#[async_trait] +impl AccountStore for SqliteAccountStore { + async fn list_pool_accounts(&self) -> AccountStoreResult> { + list_pool_accounts(self) + .await + .map_err(map_account_store_error) + } + + async fn get_pool_account(&self, account_id: &str) -> AccountStoreResult> { + get_pool_account(self, account_id) + .await + .map_err(map_account_store_error) + } + + async fn mark_quota_limited_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> AccountStoreResult { + SqliteAccountStore::mark_quota_limited_until(self, account_id, cooldown_until) + .await + .map_err(map_account_store_error) + } + + async fn set_cloudflare_cooldown_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> AccountStoreResult { + SqliteAccountStore::set_cloudflare_cooldown_until(self, account_id, cooldown_until) + .await + .map_err(map_account_store_error) + } + + async fn set_status( + &self, + account_id: &str, + status: AccountStatus, + ) -> AccountStoreResult { + SqliteAccountStore::set_status(self, account_id, status) + .await + .map_err(map_account_store_error) + } + + async fn record_usage_delta( + &self, + account_id: &str, + usage: AccountUsageDelta, + ) -> AccountStoreResult<()> { + self.record_usage(account_id, sqlite_usage_delta(usage)) + .await + .map_err(map_account_store_error) + } + + async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult> { + SqliteAccountStore::get_quota_json(self, account_id) + .await + .map_err(map_account_store_error) + } + + async fn update_quota_json( + &self, + account_id: &str, + quota_json: &str, + ) -> AccountStoreResult { + SqliteAccountStore::update_quota_json(self, account_id, quota_json) + .await + .map_err(map_account_store_error) + } + + async fn sync_rate_limit_window( + &self, + account_id: &str, + reset_at: DateTime, + limit_window_seconds: Option, + ) -> AccountStoreResult<()> { + SqliteAccountStore::sync_rate_limit_window(self, account_id, reset_at, limit_window_seconds) + .await + .map_err(map_account_store_error) + } +} + +async fn list_pool_accounts(store: &SqliteAccountStore) -> SqliteAccountStoreResult> { + let rows = sqlx::query(LIST_POOL_ACCOUNTS_SQL) + .fetch_all(&store.pool) + .await?; + let mut accounts = Vec::with_capacity(rows.len()); + + for row in rows { + accounts.push(pool_account_from_row(store, &row)?); + } + + Ok(accounts) +} + +async fn get_pool_account( + store: &SqliteAccountStore, + account_id: &str, +) -> SqliteAccountStoreResult> { + let row = sqlx::query(GET_POOL_ACCOUNT_SQL) + .bind(account_id) + .fetch_optional(&store.pool) + .await?; + row.map(|row| pool_account_from_row(store, &row)) + .transpose() +} + +fn pool_account_from_row( + store: &SqliteAccountStore, + row: &SqliteRow, +) -> SqliteAccountStoreResult { + let access_token_cipher = row.get::("access_token_cipher"); + let access_token = store.secret_box.decrypt(&access_token_cipher)?; + let refresh_token = row + .get::, _>("refresh_token_cipher") + .map(|cipher| store.secret_box.decrypt(&cipher)) + .transpose()? + .map(|token| token.expose_secret().to_string()); + + Ok(Account { + id: row.get("id"), + email: row.get("email"), + account_id: row.get("account_id"), + user_id: row.get("user_id"), + label: row.get("label"), + plan_type: row.get("plan_type"), + access_token: access_token.expose_secret().to_string(), + refresh_token, + access_token_expires_at: parse_optional_rfc3339( + row.get::, _>("access_token_expires_at"), + )?, + status: status_from_db(&row.get::("status"))?, + quota_limit_reached: row.get::("quota_limit_reached") != 0, + quota_verify_required: row.get::("quota_verify_required") != 0, + quota_cooldown_until: parse_optional_rfc3339( + row.get::, _>("quota_cooldown_until"), + )?, + cloudflare_cooldown_until: parse_optional_rfc3339( + row.get::, _>("cloudflare_cooldown_until"), + )?, + request_count: nonnegative_i64_to_u64(row.get::("usage_request_count")), + empty_response_count: nonnegative_i64_to_u64( + row.get::("usage_empty_response_count"), + ), + image_input_tokens: nonnegative_i64_to_u64(row.get::("usage_image_input_tokens")), + image_output_tokens: nonnegative_i64_to_u64(row.get::("usage_image_output_tokens")), + image_request_count: nonnegative_i64_to_u64(row.get::("usage_image_request_count")), + image_request_failed_count: nonnegative_i64_to_u64( + row.get::("usage_image_request_failed_count"), + ), + window_request_count: nonnegative_i64_to_u64( + row.get::("usage_window_request_count"), + ), + window_input_tokens: nonnegative_i64_to_u64(row.get::("usage_window_input_tokens")), + window_output_tokens: nonnegative_i64_to_u64( + row.get::("usage_window_output_tokens"), + ), + window_cached_tokens: nonnegative_i64_to_u64( + row.get::("usage_window_cached_tokens"), + ), + window_image_input_tokens: nonnegative_i64_to_u64( + row.get::("usage_window_image_input_tokens"), + ), + window_image_output_tokens: nonnegative_i64_to_u64( + row.get::("usage_window_image_output_tokens"), + ), + window_image_request_count: nonnegative_i64_to_u64( + row.get::("usage_window_image_request_count"), + ), + window_image_request_failed_count: nonnegative_i64_to_u64( + row.get::("usage_window_image_request_failed_count"), + ), + window_started_at: parse_optional_rfc3339( + row.get::, _>("usage_window_started_at"), + )?, + window_reset_at: parse_optional_rfc3339( + row.get::, _>("usage_window_reset_at"), + )?, + limit_window_seconds: optional_positive_i64_to_u64( + row.get::, _>("usage_limit_window_seconds"), + ), + added_at: row.get("added_at"), + last_used_at: row.get("usage_last_used_at"), + }) +} + +fn map_account_store_error(error: impl ToString) -> AccountStoreError { + AccountStoreError::OperationFailed { + message: error.to_string(), + } +} + +fn sqlite_usage_delta(usage: AccountUsageDelta) -> UsageDelta { + UsageDelta { + request_count: u64_to_i64_saturating(usage.requests), + 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), + image_input_tokens: u64_to_i64_saturating(usage.image_input_tokens), + image_output_tokens: u64_to_i64_saturating(usage.image_output_tokens), + image_request_count: u64_to_i64_saturating(usage.image_requests), + image_request_failed_count: u64_to_i64_saturating(usage.image_request_failures), + empty_response_count: u64_to_i64_saturating(usage.empty_responses), + } +} + +fn should_reset_usage_window( + existing_reset_at: Option>, + existing_limit_window_seconds: Option, + new_reset_at: DateTime, + new_limit_window_seconds: Option, +) -> bool { + let Some(existing_reset_at) = existing_reset_at else { + return false; + }; + if existing_reset_at == new_reset_at { + return false; + } + let drift = existing_reset_at + .signed_duration_since(new_reset_at) + .num_seconds() + .unsigned_abs(); + let window_seconds = new_limit_window_seconds + .or(existing_limit_window_seconds) + .unwrap_or(0); + let threshold = if window_seconds > 0 { + window_seconds / 2 + } else { + 3_600 + }; + drift >= threshold +} + +fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> SqliteAccountStoreResult { + Ok(AccountUsageRecord { + account_id: row.get("account_id"), + request_count: row.get("request_count"), + empty_response_count: row.get("empty_response_count"), + input_tokens: row.get("input_tokens"), + output_tokens: row.get("output_tokens"), + cached_tokens: row.get("cached_tokens"), + image_input_tokens: row.get("image_input_tokens"), + image_output_tokens: row.get("image_output_tokens"), + image_request_count: row.get("image_request_count"), + image_request_failed_count: row.get("image_request_failed_count"), + window_request_count: row.get("window_request_count"), + window_input_tokens: row.get("window_input_tokens"), + window_output_tokens: row.get("window_output_tokens"), + window_cached_tokens: row.get("window_cached_tokens"), + window_image_input_tokens: row.get("window_image_input_tokens"), + window_image_output_tokens: row.get("window_image_output_tokens"), + window_image_request_count: row.get("window_image_request_count"), + window_image_request_failed_count: row.get("window_image_request_failed_count"), + window_started_at: parse_optional_rfc3339( + row.get::, _>("window_started_at"), + )?, + window_reset_at: parse_optional_rfc3339(row.get::, _>("window_reset_at"))?, + limit_window_seconds: optional_positive_i64_to_u64( + row.get::, _>("limit_window_seconds"), + ), + last_used_at: parse_optional_rfc3339(row.get::, _>("last_used_at"))?, + }) +} + +fn usage_list_from_row( + row: &sqlx::sqlite::SqliteRow, +) -> SqliteAccountStoreResult { + Ok(AccountUsageListRecord { + account_id: row.get("account_id"), + email: row.get("email"), + label: row.get("label"), + plan_type: row.get("plan_type"), + request_count: row.get("request_count"), + empty_response_count: row.get("empty_response_count"), + input_tokens: row.get("input_tokens"), + output_tokens: row.get("output_tokens"), + cached_tokens: row.get("cached_tokens"), + image_input_tokens: row.get("image_input_tokens"), + image_output_tokens: row.get("image_output_tokens"), + image_request_count: row.get("image_request_count"), + image_request_failed_count: row.get("image_request_failed_count"), + last_used_at: parse_optional_rfc3339(row.get::, _>("last_used_at"))?, + }) +} + +fn quota_snapshot_from_row( + row: &sqlx::sqlite::SqliteRow, +) -> SqliteAccountStoreResult { + Ok(AccountQuotaSnapshot { + account_id: row.get("id"), + email: row.get("email"), + quota_json: row.get("quota_json"), + quota_fetched_at: parse_optional_rfc3339(row.get::, _>("quota_fetched_at"))?, + }) +} + +fn status_to_db(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", + } +} + +fn status_from_db(value: &str) -> SqliteAccountStoreResult { + match value { + "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(SqliteAccountStoreError::InvalidStatus(other.to_string())), + } +} + +fn parse_optional_rfc3339( + value: Option, +) -> SqliteAccountStoreResult>> { + value.as_deref().map(parse_rfc3339).transpose() +} + +fn parse_rfc3339(value: &str) -> SqliteAccountStoreResult> { + Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc)) +} + +fn nonnegative_i64_to_u64(value: i64) -> u64 { + value.max(0) as u64 +} + +fn u64_to_i64_saturating(value: u64) -> i64 { + value.min(i64::MAX as u64) as i64 +} + +fn optional_positive_i64_to_u64(value: Option) -> Option { + value + .and_then(|value| u64::try_from(value).ok()) + .filter(|value| *value > 0) +} diff --git a/src/admin/session/repository.rs b/crates/adapters/src/sqlite/admin_sessions.rs similarity index 50% rename from src/admin/session/repository.rs rename to crates/adapters/src/sqlite/admin_sessions.rs index 2de2cb827..14827e6bf 100644 --- a/src/admin/session/repository.rs +++ b/crates/adapters/src/sqlite/admin_sessions.rs @@ -1,66 +1,81 @@ -use chrono::{DateTime, Utc}; -use sqlx::{Row, SqlitePool}; -use uuid::Uuid; +//! SQLite 管理员会话存储。 +use chrono::{DateTime, Utc}; +use sqlx::SqlitePool; + +/// 已持久化的管理员用户。 #[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminUserRecord { +pub struct StoredAdminUser { + /// 管理员用户 ID。 pub id: String, + /// 管理员密码哈希。 pub password_hash: String, } +/// SQLite 管理员会话存储。 #[derive(Clone)] -pub struct AdminAuthRepository { +pub struct SqliteAdminSessionStore { pool: SqlitePool, } -impl AdminAuthRepository { +/// SQLite 管理员会话存储结果。 +pub type SqliteAdminSessionStoreResult = Result; + +impl SqliteAdminSessionStore { + /// 构造存储。 pub fn new(pool: SqlitePool) -> Self { Self { pool } } - pub async fn ensure_default_admin(&self, password_hash: &str) -> Result { - let mut tx = self.pool.begin().await?; - let count: (i64,) = sqlx::query_as("select count(*) from admin_users") - .fetch_one(&mut *tx) + /// 返回底层连接池。 + pub fn pool(&self) -> &SqlitePool { + &self.pool + } + + /// 如果还没有管理员用户，则创建默认管理员。 + pub async fn ensure_default_admin( + &self, + password_hash: &str, + ) -> SqliteAdminSessionStoreResult { + let existing: (i64,) = sqlx::query_as("select count(*) from admin_users") + .fetch_one(&self.pool) .await?; - if count.0 > 0 { - tx.commit().await?; + if existing.0 > 0 { return Ok(false); } - let admin_id = format!("admin_{}", Uuid::new_v4().simple()); let now = Utc::now().to_rfc3339(); sqlx::query( "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", ) - .bind(admin_id) + .bind("admin_1") .bind(password_hash) .bind(&now) .bind(&now) - .execute(&mut *tx) + .execute(&self.pool) .await?; - tx.commit().await?; + Ok(true) } - pub async fn load_first_admin(&self) -> Result, sqlx::Error> { - let row = sqlx::query( - "select id, password_hash from admin_users order by created_at asc limit 1", + /// 读取第一个管理员用户。 + pub async fn load_first_admin(&self) -> SqliteAdminSessionStoreResult> { + let row = sqlx::query_as::<_, (String, String)>( + "select id, password_hash from admin_users order by created_at asc, id asc limit 1", ) .fetch_optional(&self.pool) .await?; - Ok(row.map(|row| AdminUserRecord { - id: row.get("id"), - password_hash: row.get("password_hash"), - })) + + Ok(row.map(|(id, password_hash)| StoredAdminUser { id, password_hash })) } + /// 创建管理员会话。 pub async fn create_session( &self, session_id: &str, user_id: &str, expires_at: DateTime, - ) -> Result<(), sqlx::Error> { + ) -> SqliteAdminSessionStoreResult<()> { let now = Utc::now().to_rfc3339(); sqlx::query( "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", @@ -74,7 +89,8 @@ impl AdminAuthRepository { Ok(()) } - pub async fn validate_session(&self, session_id: &str) -> Result { + /// 校验管理员会话是否存在且未过期。 + pub async fn validate_session(&self, session_id: &str) -> SqliteAdminSessionStoreResult { let now = Utc::now().to_rfc3339(); let row = sqlx::query("select 1 from admin_sessions where id = ? and expires_at > ?") .bind(session_id) @@ -84,8 +100,12 @@ impl AdminAuthRepository { Ok(row.is_some()) } - pub async fn cleanup_expired_sessions(&self, now: DateTime) -> Result { - let result = sqlx::query("delete from admin_sessions where expires_at < ?") + /// 删除已经过期的管理员会话。 + pub async fn cleanup_expired_sessions( + &self, + now: DateTime, + ) -> SqliteAdminSessionStoreResult { + let result = sqlx::query("delete from admin_sessions where expires_at <= ?") .bind(now.to_rfc3339()) .execute(&self.pool) .await?; diff --git a/crates/adapters/src/sqlite/client_keys.rs b/crates/adapters/src/sqlite/client_keys.rs new file mode 100644 index 000000000..93f860c71 --- /dev/null +++ b/crates/adapters/src/sqlite/client_keys.rs @@ -0,0 +1,276 @@ +//! SQLite 客户端 API Key 存储适配器。 + +use async_trait::async_trait; +use chrono::Utc; +use sqlx::{Row, SqlitePool}; +use thiserror::Error; +use uuid::Uuid; + +use codex_proxy_core::admin::ports::{ClientKeyStore, ClientKeyStoreError, ClientKeyStoreResult}; +use codex_proxy_platform::{ + identity::{ApiKeyHasher, AuthError}, + json::{decode_cursor, encode_cursor, Page}, +}; + +/// SQLite 客户端 API Key 存储错误。 +#[derive(Debug, Error)] +pub enum SqliteClientKeyStoreError { + /// 数据库错误。 + #[error("sqlite client key database error: {0}")] + Database(#[from] sqlx::Error), + /// API Key 验证错误。 + #[error("sqlite client key auth error: {0}")] + Auth(#[from] AuthError), + /// 分页游标无效。 + #[error("invalid client key pagination cursor")] + InvalidCursor, +} + +/// 已持久化的客户端 API Key 元数据。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StoredClientApiKey { + /// API Key 记录 ID。 + pub id: String, + /// API Key 名称。 + pub name: String, + /// 管理员可见标签。 + pub label: Option, + /// 明文 API Key 的短前缀。 + pub prefix: String, + /// 是否允许用于 `/v1` 认证。 + pub enabled: bool, + /// 创建时间。 + pub created_at: String, + /// 最近一次成功使用时间。 + pub last_used_at: Option, +} + +/// 新建客户端 API Key 后的一次性结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CreatedClientApiKey { + /// API Key 记录 ID。 + pub id: String, + /// API Key 名称。 + pub name: String, + /// 管理员可见标签。 + pub label: Option, + /// 明文 API Key 的短前缀。 + pub prefix: String, + /// 是否允许用于 `/v1` 认证。 + pub enabled: bool, + /// 创建时间。 + pub created_at: String, + /// 最近一次成功使用时间。 + pub last_used_at: Option, + /// 仅返回一次的明文 API Key。 + pub plaintext: String, +} + +/// SQLite 客户端 API Key 存储。 +#[derive(Clone)] +pub struct SqliteClientKeyStore { + pool: SqlitePool, + hasher: ApiKeyHasher, +} + +impl SqliteClientKeyStore { + /// 构造存储适配器。 + pub fn new(pool: SqlitePool, hasher: ApiKeyHasher) -> Self { + Self { pool, hasher } + } + + /// 暴露底层连接池，供集成测试和运行时组合层复用。 + pub fn pool(&self) -> &SqlitePool { + &self.pool + } + + /// 创建新的本地客户端 API Key。 + pub async fn create( + &self, + name: &str, + ) -> Result { + let generated = self.hasher.generate_client_api_key(name); + let id = format!("key_{}", Uuid::new_v4().simple()); + let now = Utc::now().to_rfc3339(); + sqlx::query( + "insert into client_api_keys (id, name, label, prefix, key_hash, enabled, created_at, last_used_at) values (?, ?, null, ?, ?, 1, ?, null)", + ) + .bind(&id) + .bind(name) + .bind(&generated.prefix) + .bind(&generated.key_hash) + .bind(&now) + .execute(&self.pool) + .await?; + + Ok(CreatedClientApiKey { + id, + name: name.to_string(), + label: None, + prefix: generated.prefix, + enabled: true, + created_at: now, + last_used_at: None, + plaintext: generated.plaintext, + }) + } + + /// 按创建时间倒序分页列出客户端 API Key。 + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> Result, SqliteClientKeyStoreError> { + let fetch_limit = i64::from(limit) + 1; + let rows = if let Some(cursor) = cursor { + let (created_at, id) = + decode_cursor(&cursor).ok_or(SqliteClientKeyStoreError::InvalidCursor)?; + sqlx::query( + "select id, name, label, prefix, enabled, created_at, last_used_at from client_api_keys where created_at < ? or (created_at = ? and id < ?) order by created_at desc, id desc limit ?", + ) + .bind(&created_at) + .bind(created_at) + .bind(id) + .bind(fetch_limit) + .fetch_all(&self.pool) + .await? + } else { + sqlx::query( + "select id, name, label, prefix, enabled, created_at, last_used_at from client_api_keys order by created_at desc, id desc limit ?", + ) + .bind(fetch_limit) + .fetch_all(&self.pool) + .await? + }; + + let has_next = rows.len() > limit as usize; + let take_count = rows.len().min(limit as usize); + let items = rows + .into_iter() + .take(take_count) + .map(|row| key_from_row(&row)) + .collect::>(); + let next_cursor = if has_next { + items + .last() + .map(|key| encode_cursor(&key.created_at, &key.id)) + } else { + None + }; + + Ok(Page { items, next_cursor }) + } + + /// 按 ID 读取客户端 API Key 元数据。 + pub async fn get( + &self, + id: &str, + ) -> Result, SqliteClientKeyStoreError> { + let row = sqlx::query( + "select id, name, label, prefix, enabled, created_at, last_used_at from client_api_keys where id = ?", + ) + .bind(id) + .fetch_optional(&self.pool) + .await?; + + Ok(row.map(|row| key_from_row(&row))) + } + + /// 更新客户端 API Key 启用状态。 + pub async fn set_enabled( + &self, + id: &str, + enabled: bool, + ) -> Result { + let result = sqlx::query("update client_api_keys set enabled = ? where id = ?") + .bind(if enabled { 1_i64 } else { 0_i64 }) + .bind(id) + .execute(&self.pool) + .await?; + + Ok(result.rows_affected() > 0) + } + + /// 更新客户端 API Key 标签并返回最新元数据。 + pub async fn set_label( + &self, + id: &str, + label: Option, + ) -> Result, SqliteClientKeyStoreError> { + let result = sqlx::query("update client_api_keys set label = ? where id = ?") + .bind(label) + .bind(id) + .execute(&self.pool) + .await?; + if result.rows_affected() == 0 { + return Ok(None); + } + + self.get(id).await + } + + /// 删除客户端 API Key。 + pub async fn delete(&self, id: &str) -> Result { + let result = sqlx::query("delete from client_api_keys where id = ?") + .bind(id) + .execute(&self.pool) + .await?; + + Ok(result.rows_affected() > 0) + } +} + +#[async_trait] +impl ClientKeyStore for SqliteClientKeyStore { + async fn verify_and_touch(&self, plaintext: &str) -> ClientKeyStoreResult { + verify_and_touch(&self.pool, &self.hasher, plaintext) + .await + .map_err(map_client_key_store_error) + } +} + +async fn verify_and_touch( + pool: &SqlitePool, + hasher: &ApiKeyHasher, + plaintext: &str, +) -> Result { + let prefix = plaintext.chars().take(12).collect::(); + let rows = + sqlx::query("select id, key_hash from client_api_keys where prefix = ? and enabled = 1") + .bind(prefix) + .fetch_all(pool) + .await?; + + for row in rows { + let key_hash = row.get::("key_hash"); + if hasher.verify_client_api_key(plaintext, &key_hash)? { + let id = row.get::("id"); + sqlx::query("update client_api_keys set last_used_at = ? where id = ?") + .bind(Utc::now().to_rfc3339()) + .bind(id) + .execute(pool) + .await?; + return Ok(true); + } + } + + Ok(false) +} + +fn key_from_row(row: &sqlx::sqlite::SqliteRow) -> StoredClientApiKey { + StoredClientApiKey { + id: row.get("id"), + name: row.get("name"), + label: row.get("label"), + prefix: row.get("prefix"), + enabled: row.get::("enabled") != 0, + created_at: row.get("created_at"), + last_used_at: row.get("last_used_at"), + } +} + +fn map_client_key_store_error(error: SqliteClientKeyStoreError) -> ClientKeyStoreError { + ClientKeyStoreError::OperationFailed { + message: error.to_string(), + } +} diff --git a/src/codex/accounts/cookies/repository.rs b/crates/adapters/src/sqlite/cookies.rs similarity index 71% rename from src/codex/accounts/cookies/repository.rs rename to crates/adapters/src/sqlite/cookies.rs index c17c24f8f..b455c2c30 100644 --- a/src/codex/accounts/cookies/repository.rs +++ b/crates/adapters/src/sqlite/cookies.rs @@ -1,35 +1,63 @@ +//! SQLite Cookie 存储。 + use chrono::{DateTime, Duration, Utc}; use secrecy::{ExposeSecret, SecretString}; use sqlx::{Row, SqlitePool}; use thiserror::Error; use uuid::Uuid; -use crate::platform::crypto::{CryptoError, SecretBox}; +use codex_proxy_platform::crypto::{CryptoError, SecretBox}; +const DEFAULT_COOKIE_DOMAIN: &str = "chatgpt.com"; const CAPTURABLE_COOKIES: &[&str] = &["cf_clearance"]; +/// SQLite Cookie 存储错误。 #[derive(Debug, Error)] -pub enum CookieError { - #[error("database error: {0}")] +pub enum SqliteCookieStoreError { + /// 数据库错误。 + #[error("sqlite cookie store database error: {0}")] Database(#[from] sqlx::Error), - #[error("cookie encryption error: {0}")] + /// 加解密错误。 + #[error("sqlite cookie store crypto error: {0}")] Crypto(#[from] CryptoError), } -pub type CookieResult = Result; +/// SQLite Cookie 存储结果。 +pub type SqliteCookieStoreResult = Result; +/// SQLite Cookie 存储。 #[derive(Clone)] -pub struct CookieRepository { +pub struct SqliteCookieStore { pool: SqlitePool, secret_box: SecretBox, } -impl CookieRepository { +impl SqliteCookieStore { + /// 构造存储。 pub fn new(pool: SqlitePool, secret_box: SecretBox) -> Self { Self { pool, secret_box } } - pub async fn capture_set_cookie(&self, account_id: &str, raw: &str) -> CookieResult<()> { + /// 返回底层连接池。 + pub fn pool(&self) -> &SqlitePool { + &self.pool + } + + /// 检查账号是否存在。 + pub async fn account_exists(&self, account_id: &str) -> SqliteCookieStoreResult { + let row = sqlx::query("select 1 from accounts where id = ?") + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + Ok(row.is_some()) + } + + /// 捕获上游 `Set-Cookie` 响应头中允许持久化的 Cookie。 + pub async fn capture_set_cookie( + &self, + account_id: &str, + raw: &str, + ) -> SqliteCookieStoreResult<()> { let Some(parsed) = parse_set_cookie(raw) else { return Ok(()); }; @@ -39,7 +67,12 @@ impl CookieRepository { self.upsert_cookie(account_id, parsed).await } - pub async fn set_cookie_header(&self, account_id: &str, raw: &str) -> CookieResult { + /// 将 Cookie 请求头写入账号 Cookie 存储。 + pub async fn set_cookie_header( + &self, + account_id: &str, + raw: &str, + ) -> SqliteCookieStoreResult { let parsed = parse_cookie_header(raw); let count = parsed.len(); for cookie in parsed { @@ -48,42 +81,23 @@ impl CookieRepository { Ok(count) } - async fn upsert_cookie(&self, account_id: &str, parsed: ParsedCookie) -> CookieResult<()> { - let now = Utc::now().to_rfc3339(); - let value_cipher = self - .secret_box - .encrypt(&SecretString::new(parsed.value.into()))?; - sqlx::query( - "insert into account_cookies (id, account_id, domain, name, value_cipher, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?) on conflict(account_id, domain, name, path) do update set value_cipher = excluded.value_cipher, expires_at = excluded.expires_at, updated_at = excluded.updated_at", - ) - .bind(Uuid::new_v4().to_string()) - .bind(account_id) - .bind(parsed.domain) - .bind(parsed.name) - .bind(value_cipher) - .bind(parsed.path) - .bind(parsed.expires_at) - .bind(now) - .execute(&self.pool) - .await?; - Ok(()) - } - + /// 为请求域名读取账号 Cookie 请求头。 pub async fn cookie_header( &self, account_id: &str, request_domain: &str, - ) -> CookieResult> { + ) -> SqliteCookieStoreResult> { self.cookie_header_for_request(account_id, request_domain, "/") .await } + /// 为请求域名和路径读取账号 Cookie 请求头。 pub async fn cookie_header_for_request( &self, account_id: &str, request_domain: &str, request_path: &str, - ) -> CookieResult> { + ) -> SqliteCookieStoreResult> { let rows = sqlx::query( "select domain, name, value_cipher, path, expires_at from account_cookies where account_id = ?", ) @@ -101,7 +115,8 @@ impl CookieRepository { if !path_matches(request_path, &path) { continue; } - if cookie_is_expired(row.get::, _>("expires_at").as_deref(), now) { + let expires_at = row.get::, _>("expires_at"); + if cookie_is_expired(expires_at.as_deref(), now) { continue; } let name = row.get::("name"); @@ -132,7 +147,8 @@ impl CookieRepository { } } - pub async fn delete_account_cookies(&self, account_id: &str) -> CookieResult { + /// 删除账号全部 Cookie。 + pub async fn delete_account_cookies(&self, account_id: &str) -> SqliteCookieStoreResult { let result = sqlx::query("delete from account_cookies where account_id = ?") .bind(account_id) .execute(&self.pool) @@ -140,27 +156,44 @@ impl CookieRepository { Ok(result.rows_affected()) } - pub async fn cleanup_expired(&self, now: DateTime) -> CookieResult { - let rows = - sqlx::query("select id, expires_at from account_cookies where expires_at is not null") - .fetch_all(&self.pool) - .await?; - let mut deleted = 0; - for row in rows { - let expires_at = row.get::("expires_at"); - if !cookie_is_expired(Some(&expires_at), now) { - continue; - } - let result = sqlx::query("delete from account_cookies where id = ?") - .bind(row.get::("id")) - .execute(&self.pool) - .await?; - deleted += result.rows_affected(); - } - Ok(deleted) + /// 删除已经过期的 Cookie。 + pub async fn cleanup_expired(&self, now: DateTime) -> SqliteCookieStoreResult { + let result = sqlx::query( + "delete from account_cookies where expires_at is not null and expires_at <= ?", + ) + .bind(now.to_rfc3339()) + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } + + async fn upsert_cookie( + &self, + account_id: &str, + parsed: ParsedCookie, + ) -> SqliteCookieStoreResult<()> { + let now = Utc::now().to_rfc3339(); + let value_cipher = self + .secret_box + .encrypt(&SecretString::new(parsed.value.into()))?; + sqlx::query( + "insert into account_cookies (id, account_id, domain, name, value_cipher, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?) on conflict(account_id, domain, name, path) do update set value_cipher = excluded.value_cipher, expires_at = excluded.expires_at, updated_at = excluded.updated_at", + ) + .bind(Uuid::new_v4().to_string()) + .bind(account_id) + .bind(parsed.domain) + .bind(parsed.name) + .bind(value_cipher) + .bind(parsed.path) + .bind(parsed.expires_at) + .bind(now) + .execute(&self.pool) + .await?; + Ok(()) } } +#[derive(Debug, Clone, PartialEq, Eq)] struct ParsedCookie { domain: String, name: String, @@ -169,6 +202,7 @@ struct ParsedCookie { expires_at: Option, } +#[derive(Debug, Clone, PartialEq, Eq)] struct CookieHeaderPair { path_len: usize, name: String, @@ -178,7 +212,13 @@ struct CookieHeaderPair { fn parse_set_cookie(raw: &str) -> Option { let mut parts = raw.split(';').map(str::trim); let (name, value) = parts.next()?.split_once('=')?; - let mut domain = "chatgpt.com".to_string(); + let name = name.trim(); + let value = value.trim(); + if name.is_empty() || value.is_empty() { + return None; + } + + let mut domain = DEFAULT_COOKIE_DOMAIN.to_string(); let mut path = "/".to_string(); let mut expires_at = None; for part in parts { @@ -186,18 +226,19 @@ fn parse_set_cookie(raw: &str) -> Option { continue; }; match attribute.trim().to_ascii_lowercase().as_str() { - "domain" => domain = value.trim_start_matches('.').to_string(), + "domain" => domain = value.trim().trim_start_matches('.').to_string(), "path" => path = normalize_cookie_path(value), "max-age" => { - if let Ok(seconds) = value.parse::() { + if let Ok(seconds) = value.trim().parse::() { expires_at = Some(max_age_expires_at(seconds)); } break; } - "expires" => expires_at = Some(value.to_string()), + "expires" => expires_at = Some(value.trim().to_string()), _ => {} } } + Some(ParsedCookie { domain, name: name.to_string(), @@ -226,7 +267,7 @@ fn parse_cookie_header(raw: &str) -> Vec { return None; } Some(ParsedCookie { - domain: "chatgpt.com".to_string(), + domain: DEFAULT_COOKIE_DOMAIN.to_string(), name: name.to_string(), value: value.to_string(), path: "/".to_string(), @@ -243,15 +284,6 @@ fn domain_matches(request_domain: &str, cookie_domain: &str) -> bool { .is_some_and(|prefix| prefix.ends_with('.')) } -fn normalize_cookie_path(path: &str) -> String { - let path = path.trim(); - if path.starts_with('/') { - path.to_string() - } else { - "/".to_string() - } -} - fn path_matches(request_path: &str, cookie_path: &str) -> bool { let request_path = normalize_request_path(request_path); let cookie_path = normalize_cookie_path(cookie_path); @@ -273,6 +305,15 @@ fn normalize_request_path(path: &str) -> String { } } +fn normalize_cookie_path(path: &str) -> String { + let path = path.trim(); + if path.starts_with('/') { + path.to_string() + } else { + "/".to_string() + } +} + fn cookie_is_expired(expires_at: Option<&str>, now: DateTime) -> bool { expires_at .and_then(parse_cookie_expires_at) diff --git a/crates/adapters/src/sqlite/events.rs b/crates/adapters/src/sqlite/events.rs new file mode 100644 index 000000000..23b58bc12 --- /dev/null +++ b/crates/adapters/src/sqlite/events.rs @@ -0,0 +1,286 @@ +//! SQLite 事件日志存储。 + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use sqlx::{QueryBuilder, Row, Sqlite, SqlitePool}; +use thiserror::Error; + +use codex_proxy_core::events::{ + model::{EventLevel, EventLog}, + ports::{EventLogStore, EventLogStoreError, EventLogStoreResult}, +}; +use codex_proxy_platform::json::{decode_cursor, encode_cursor, Page}; + +/// SQLite 事件日志错误。 +#[derive(Debug, Error)] +pub enum SqliteEventLogStoreError { + /// 数据库错误。 + #[error("sqlite event log database error: {0}")] + Database(#[from] sqlx::Error), + /// JSON 错误。 + #[error("sqlite event log json error: {0}")] + Json(#[from] serde_json::Error), + /// 时间格式错误。 + #[error("sqlite event log timestamp error: {0}")] + Timestamp(#[from] chrono::ParseError), + /// 事件等级非法。 + #[error("invalid event level: {0}")] + InvalidLevel(String), + /// 分页游标非法。 + #[error("invalid event log pagination cursor")] + InvalidCursor, +} + +/// SQLite 事件日志结果。 +pub type SqliteEventLogStoreResult = Result; + +/// 事件日志查询过滤器。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct EventLogFilter { + /// 事件类别。 + pub kind: Option, + /// 事件等级。 + pub level: Option, + /// 请求 ID。 + pub request_id: Option, + /// 账号 ID。 + pub account_id: Option, + /// 路由。 + pub route: Option, + /// 模型。 + pub model: Option, + /// HTTP 状态码。 + pub status_code: Option, + /// 搜索关键词。 + pub search: Option, +} + +/// SQLite 事件日志存储。 +#[derive(Clone)] +pub struct SqliteEventLogStore { + pool: SqlitePool, +} + +impl SqliteEventLogStore { + /// 构造存储。 + pub fn new(pool: SqlitePool) -> Self { + Self { pool } + } + + /// 写入事件日志。 + pub async fn append(&self, event: &EventLog) -> SqliteEventLogStoreResult<()> { + append_event(&self.pool, event).await + } + + /// 保留最新的指定数量事件日志。 + pub async fn trim_to_capacity(&self, capacity: u32) -> SqliteEventLogStoreResult { + trim_to_capacity(&self.pool, capacity).await + } + + /// 分页查询事件日志。 + pub async fn list( + &self, + filter: EventLogFilter, + cursor: Option, + limit: u32, + ) -> SqliteEventLogStoreResult> { + let fetch_limit = i64::from(limit) + 1; + let mut builder = QueryBuilder::::new( + "select id, request_id, kind, level, account_id, route, model, status_code, latency_ms, message, metadata_json, created_at from event_logs", + ); + push_filter(&mut builder, &filter, cursor.as_deref())?; + builder.push(" order by created_at desc, id desc limit "); + builder.push_bind(fetch_limit); + + let rows = builder.build().fetch_all(&self.pool).await?; + let has_next = rows.len() > limit as usize; + let take_count = rows.len().min(limit as usize); + let items = rows + .into_iter() + .take(take_count) + .map(|row| event_from_row(&row)) + .collect::>>()?; + let next_cursor = if has_next { + items + .last() + .map(|event| encode_cursor(&event.created_at.to_rfc3339(), &event.id)) + } else { + None + }; + + Ok(Page { items, next_cursor }) + } + + /// 按 ID 读取事件日志。 + pub async fn get(&self, id: &str) -> SqliteEventLogStoreResult> { + let row = sqlx::query( + "select id, request_id, kind, level, account_id, route, model, status_code, latency_ms, message, metadata_json, created_at from event_logs where id = ?", + ) + .bind(id) + .fetch_optional(&self.pool) + .await?; + row.map(|row| event_from_row(&row)).transpose() + } + + /// 统计事件日志数量。 + pub async fn count(&self) -> SqliteEventLogStoreResult { + let count: (i64,) = sqlx::query_as("select count(*) from event_logs") + .fetch_one(&self.pool) + .await?; + Ok(count.0.max(0) as u64) + } + + /// 清空事件日志。 + pub async fn clear(&self) -> SqliteEventLogStoreResult { + let result = sqlx::query("delete from event_logs") + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } +} + +#[async_trait] +impl EventLogStore for SqliteEventLogStore { + async fn append(&self, event: &EventLog) -> EventLogStoreResult<()> { + append_event(&self.pool, event).await.map_err(|error| { + EventLogStoreError::OperationFailed { + message: error.to_string(), + } + })?; + Ok(()) + } +} + +async fn append_event(pool: &SqlitePool, event: &EventLog) -> SqliteEventLogStoreResult<()> { + sqlx::query( + "insert into event_logs (id, request_id, kind, level, account_id, route, model, status_code, latency_ms, message, metadata_json, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(&event.id) + .bind(&event.request_id) + .bind(&event.kind) + .bind(level_to_db(event.level)) + .bind(&event.account_id) + .bind(&event.route) + .bind(&event.model) + .bind(event.status_code) + .bind(event.latency_ms) + .bind(&event.message) + .bind(event.metadata.to_string()) + .bind(event.created_at.to_rfc3339()) + .execute(pool) + .await?; + Ok(()) +} + +async fn trim_to_capacity(pool: &SqlitePool, capacity: u32) -> SqliteEventLogStoreResult { + let result = sqlx::query( + "delete from event_logs where id in ( + select id from event_logs + order by created_at desc, id desc + limit -1 offset ? + )", + ) + .bind(i64::from(capacity)) + .execute(pool) + .await?; + Ok(result.rows_affected()) +} + +fn push_filter( + builder: &mut QueryBuilder, + filter: &EventLogFilter, + cursor: Option<&str>, +) -> SqliteEventLogStoreResult<()> { + let mut separated = builder.separated(" and "); + separated.push(" where 1 = 1"); + + if let Some(kind) = filter.kind.as_deref() { + separated.push("kind = "); + separated.push_bind_unseparated(kind); + } + if let Some(level) = filter.level { + separated.push("level = "); + separated.push_bind_unseparated(level_to_db(level)); + } + if let Some(request_id) = filter.request_id.as_deref() { + separated.push("request_id = "); + separated.push_bind_unseparated(request_id); + } + if let Some(account_id) = filter.account_id.as_deref() { + separated.push("account_id = "); + separated.push_bind_unseparated(account_id); + } + if let Some(route) = filter.route.as_deref() { + separated.push("route = "); + separated.push_bind_unseparated(route); + } + if let Some(model) = filter.model.as_deref() { + separated.push("model = "); + separated.push_bind_unseparated(model); + } + if let Some(status_code) = filter.status_code { + separated.push("status_code = "); + separated.push_bind_unseparated(status_code); + } + if let Some(search) = filter.search.as_deref() { + let pattern = format!("%{search}%"); + separated.push("(message like "); + separated.push_bind_unseparated(pattern.clone()); + separated.push_unseparated(" or metadata_json like "); + separated.push_bind_unseparated(pattern); + separated.push_unseparated(")"); + } + if let Some(cursor) = cursor { + let (created_at, id) = + decode_cursor(cursor).ok_or(SqliteEventLogStoreError::InvalidCursor)?; + separated.push("(created_at < "); + separated.push_bind_unseparated(created_at.clone()); + separated.push_unseparated(" or (created_at = "); + separated.push_bind_unseparated(created_at); + separated.push_unseparated(" and id < "); + separated.push_bind_unseparated(id); + separated.push_unseparated("))"); + } + + Ok(()) +} + +fn event_from_row(row: &sqlx::sqlite::SqliteRow) -> SqliteEventLogStoreResult { + Ok(EventLog { + id: row.get("id"), + request_id: row.get("request_id"), + kind: row.get("kind"), + level: level_from_db(&row.get::("level"))?, + account_id: row.get("account_id"), + route: row.get("route"), + model: row.get("model"), + status_code: row.get("status_code"), + latency_ms: row.get("latency_ms"), + message: row.get("message"), + metadata: serde_json::from_str(&row.get::("metadata_json"))?, + created_at: parse_rfc3339(&row.get::("created_at"))?, + }) +} + +fn level_to_db(level: EventLevel) -> &'static str { + match level { + EventLevel::Debug => "debug", + EventLevel::Info => "info", + EventLevel::Warn => "warn", + EventLevel::Error => "error", + } +} + +fn level_from_db(value: &str) -> SqliteEventLogStoreResult { + match value { + "debug" => Ok(EventLevel::Debug), + "info" => Ok(EventLevel::Info), + "warn" => Ok(EventLevel::Warn), + "error" => Ok(EventLevel::Error), + other => Err(SqliteEventLogStoreError::InvalidLevel(other.to_string())), + } +} + +fn parse_rfc3339(value: &str) -> SqliteEventLogStoreResult> { + Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc)) +} diff --git a/crates/adapters/src/sqlite/mod.rs b/crates/adapters/src/sqlite/mod.rs new file mode 100644 index 000000000..345d67d03 --- /dev/null +++ b/crates/adapters/src/sqlite/mod.rs @@ -0,0 +1,22 @@ +//! SQLite 适配器。 + +/// 账号 token 存储辅助。 +pub mod account_tokens; +/// 账号用量存储。 +pub mod account_usage; +/// 账号读取适配器。 +pub mod accounts; +/// 管理员会话存储。 +pub mod admin_sessions; +/// 客户端 API Key 存储。 +pub mod client_keys; +/// Cookie 存储。 +pub mod cookies; +/// 事件日志存储。 +pub mod events; +/// 模型快照存储。 +pub mod models; +/// 刷新租约存储。 +pub mod refresh_leases; +/// 会话亲和性存储。 +pub mod session_affinity; diff --git a/crates/adapters/src/sqlite/models.rs b/crates/adapters/src/sqlite/models.rs new file mode 100644 index 000000000..b3dd11c1a --- /dev/null +++ b/crates/adapters/src/sqlite/models.rs @@ -0,0 +1,102 @@ +//! 模型快照 SQLite 仓储。 + +use async_trait::async_trait; +use chrono::Utc; +use sqlx::{Row, SqlitePool}; +use thiserror::Error; + +use codex_proxy_core::models::{ + model::ModelPlanSnapshot, + ports::{ModelSnapshotStore, ModelSnapshotStoreError, ModelSnapshotStoreResult}, +}; + +/// 模型快照仓储实现。 +#[derive(Debug, Clone)] +pub struct ModelSnapshotRepository { + pool: SqlitePool, +} + +/// 模型快照仓储错误。 +#[derive(Debug, Error)] +pub enum ModelSnapshotRepositoryError { + /// 存储层错误。 + #[error("model snapshot store error: {0}")] + Store(#[from] ModelSnapshotStoreError), + /// 数据库错误。 + #[error("model snapshot database error: {0}")] + Database(#[from] sqlx::Error), + /// JSON 错误。 + #[error("model snapshot json error: {0}")] + Json(#[from] serde_json::Error), +} + +impl ModelSnapshotRepository { + /// 构造仓储。 + pub fn new(pool: SqlitePool) -> Self { + Self { pool } + } + + /// 直接写入单个计划快照。 + pub async fn replace_plan_snapshot( + &self, + snapshot: &ModelPlanSnapshot, + ) -> ModelSnapshotStoreResult<()> { + let models_json = serde_json::to_string(&snapshot.models).map_err(|error| { + ModelSnapshotStoreError::OperationFailed { + message: error.to_string(), + } + })?; + let fetched_at = Utc::now().to_rfc3339(); + sqlx::query( + "insert into model_plan_snapshots (plan_type, models_json, fetched_at) values (?, ?, ?) \ + on conflict(plan_type) do update set models_json = excluded.models_json, fetched_at = excluded.fetched_at", + ) + .bind(&snapshot.plan_type) + .bind(models_json) + .bind(fetched_at) + .execute(&self.pool) + .await + .map_err(|error| ModelSnapshotStoreError::OperationFailed { + message: error.to_string(), + })?; + Ok(()) + } + + /// 列出所有快照。 + pub async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult> { + let rows = sqlx::query( + "select plan_type, models_json from model_plan_snapshots order by plan_type asc", + ) + .fetch_all(&self.pool) + .await + .map_err(|error| ModelSnapshotStoreError::OperationFailed { + message: error.to_string(), + })?; + rows.into_iter() + .map(|row| { + let plan_type: String = row.get("plan_type"); + let models_json: String = row.get("models_json"); + let models = serde_json::from_str(&models_json).map_err(|error| { + ModelSnapshotStoreError::OperationFailed { + message: error.to_string(), + } + })?; + Ok(ModelPlanSnapshot { plan_type, models }) + }) + .collect() + } +} + +#[async_trait] +impl ModelSnapshotStore for ModelSnapshotRepository { + async fn replace_plan_snapshot( + &self, + snapshot: &ModelPlanSnapshot, + ) -> ModelSnapshotStoreResult<()> { + Self::replace_plan_snapshot(self, snapshot).await + } + + async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult> { + Self::list_plan_snapshots(self).await + } +} diff --git a/crates/adapters/src/sqlite/refresh_leases.rs b/crates/adapters/src/sqlite/refresh_leases.rs new file mode 100644 index 000000000..647649f2b --- /dev/null +++ b/crates/adapters/src/sqlite/refresh_leases.rs @@ -0,0 +1,80 @@ +//! SQLite 刷新租约存储。 + +use chrono::{DateTime, Utc}; +use sqlx::SqlitePool; +use thiserror::Error; + +/// SQLite 刷新租约存储结果。 +pub type SqliteRefreshLeaseStoreResult = Result; + +/// SQLite 刷新租约存储错误。 +#[derive(Debug, Error)] +pub enum SqliteRefreshLeaseStoreError { + /// SQLite 操作失败。 + #[error("sqlite refresh lease query failed: {0}")] + Sqlx(#[from] sqlx::Error), +} + +/// SQLite 刷新租约存储。 +#[derive(Clone)] +pub struct SqliteRefreshLeaseStore { + pool: SqlitePool, +} + +impl SqliteRefreshLeaseStore { + /// 构造存储。 + pub fn new(pool: SqlitePool) -> Self { + Self { pool } + } + + /// 返回底层连接池。 + pub fn pool(&self) -> &SqlitePool { + &self.pool + } + + /// 尝试获取账号刷新租约。 + pub async fn try_acquire( + &self, + account_id: &str, + owner: &str, + expires_at: DateTime, + now: DateTime, + ) -> SqliteRefreshLeaseStoreResult { + let result = sqlx::query( + r#" +insert into account_refresh_leases (account_id, owner, expires_at, updated_at) +values (?, ?, ?, ?) +on conflict(account_id) do update set + owner = excluded.owner, + expires_at = excluded.expires_at, + updated_at = excluded.updated_at +where account_refresh_leases.expires_at <= ? + or account_refresh_leases.owner = ? +"#, + ) + .bind(account_id) + .bind(owner) + .bind(expires_at.to_rfc3339()) + .bind(now.to_rfc3339()) + .bind(now.to_rfc3339()) + .bind(owner) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 释放账号刷新租约。 + pub async fn release( + &self, + account_id: &str, + owner: &str, + ) -> SqliteRefreshLeaseStoreResult { + let result = + sqlx::query("delete from account_refresh_leases where account_id = ? and owner = ?") + .bind(account_id) + .bind(owner) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } +} diff --git a/crates/adapters/src/sqlite/session_affinity.rs b/crates/adapters/src/sqlite/session_affinity.rs new file mode 100644 index 000000000..b6480e0f5 --- /dev/null +++ b/crates/adapters/src/sqlite/session_affinity.rs @@ -0,0 +1,170 @@ +//! SQLite 会话亲和性存储。 + +use chrono::{DateTime, Duration, Utc}; +use codex_proxy_core::serving::affinity::{SessionAffinityEntry, StoredSessionAffinity}; +use sqlx::SqlitePool; +use thiserror::Error; + +/// SQLite 会话亲和性存储。 +#[derive(Clone)] +pub struct SqliteSessionAffinityStore { + pool: SqlitePool, +} + +/// SQLite 会话亲和性存储错误。 +#[derive(Debug, Error)] +pub enum SqliteSessionAffinityStoreError { + /// 数据库错误。 + #[error("database error: {0}")] + Database(#[from] sqlx::Error), + /// 时间戳格式无效。 + #[error("invalid session affinity timestamp: {0}")] + InvalidTimestamp(#[from] chrono::ParseError), + /// 函数调用 ID JSON 无效。 + #[error("invalid session affinity function call ids: {0}")] + InvalidFunctionCallIds(#[from] serde_json::Error), +} + +/// SQLite 会话亲和性存储结果。 +pub type SqliteSessionAffinityStoreResult = Result; + +impl SqliteSessionAffinityStore { + /// 构造存储。 + pub fn new(pool: SqlitePool) -> Self { + Self { pool } + } + + /// 返回底层连接池。 + pub fn pool(&self) -> &SqlitePool { + &self.pool + } + + /// 插入或更新响应 ID 的会话亲和性条目。 + pub async fn upsert( + &self, + response_id: &str, + entry: &SessionAffinityEntry, + ttl: Duration, + ) -> SqliteSessionAffinityStoreResult<()> { + let function_call_ids_json = serde_json::to_string(&entry.function_call_ids)?; + let expires_at = entry + .created_at + .checked_add_signed(ttl) + .unwrap_or(DateTime::::MAX_UTC); + sqlx::query( + r" +insert into session_affinities ( + response_id, + account_id, + conversation_id, + turn_state, + instructions_hash, + input_tokens, + function_call_ids_json, + variant_hash, + expires_at, + created_at +) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +on conflict(response_id) do update set + account_id = excluded.account_id, + conversation_id = excluded.conversation_id, + turn_state = excluded.turn_state, + instructions_hash = excluded.instructions_hash, + input_tokens = excluded.input_tokens, + function_call_ids_json = excluded.function_call_ids_json, + variant_hash = excluded.variant_hash, + expires_at = excluded.expires_at, + created_at = excluded.created_at", + ) + .bind(response_id) + .bind(&entry.account_id) + .bind(&entry.conversation_id) + .bind(&entry.turn_state) + .bind(&entry.instructions_hash) + .bind(entry.input_tokens.map(u64_to_i64_saturating)) + .bind(function_call_ids_json) + .bind(&entry.variant_hash) + .bind(expires_at.to_rfc3339()) + .bind(entry.created_at.to_rfc3339()) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// 列出当前未过期的会话亲和性条目。 + pub async fn list_active( + &self, + now: DateTime, + ) -> SqliteSessionAffinityStoreResult> { + let rows = sqlx::query( + r" +select + response_id, + account_id, + conversation_id, + turn_state, + instructions_hash, + input_tokens, + function_call_ids_json, + variant_hash, + expires_at, + created_at +from session_affinities +where expires_at > ? +order by created_at asc, response_id asc", + ) + .bind(now.to_rfc3339()) + .fetch_all(&self.pool) + .await?; + + rows.into_iter() + .map(|row| stored_session_affinity_from_row(&row)) + .collect() + } + + /// 删除已过期的会话亲和性条目。 + pub async fn delete_expired( + &self, + now: DateTime, + ) -> SqliteSessionAffinityStoreResult { + let result = sqlx::query("delete from session_affinities where expires_at <= ?") + .bind(now.to_rfc3339()) + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } +} + +fn stored_session_affinity_from_row( + row: &sqlx::sqlite::SqliteRow, +) -> SqliteSessionAffinityStoreResult { + use sqlx::Row as _; + + let function_call_ids_json = row.get::("function_call_ids_json"); + Ok(StoredSessionAffinity { + response_id: row.get("response_id"), + entry: SessionAffinityEntry { + account_id: row.get("account_id"), + conversation_id: row.get("conversation_id"), + turn_state: row.get("turn_state"), + instructions_hash: row.get("instructions_hash"), + input_tokens: optional_nonnegative_i64_to_u64(row.get("input_tokens")), + function_call_ids: serde_json::from_str(&function_call_ids_json)?, + variant_hash: row.get("variant_hash"), + created_at: parse_rfc3339(&row.get::("created_at"))?, + }, + expires_at: parse_rfc3339(&row.get::("expires_at"))?, + }) +} + +fn parse_rfc3339(value: &str) -> SqliteSessionAffinityStoreResult> { + Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc)) +} + +fn optional_nonnegative_i64_to_u64(value: Option) -> Option { + value.and_then(|value| u64::try_from(value).ok()) +} + +fn u64_to_i64_saturating(value: u64) -> i64 { + value.min(i64::MAX as u64) as i64 +} diff --git a/crates/adapters/tests/account_repository.rs b/crates/adapters/tests/account_repository.rs new file mode 100644 index 000000000..a6687155b --- /dev/null +++ b/crates/adapters/tests/account_repository.rs @@ -0,0 +1,322 @@ +use chrono::{Duration, Utc}; +use codex_proxy_adapters::sqlite::accounts::{NewAccount, SqliteAccountStore, UsageDelta}; +use codex_proxy_core::accounts::{model::AccountStatus, ports::AccountStore}; +use codex_proxy_platform::{crypto::SecretBox, storage::connect_sqlite}; +use secrecy::{ExposeSecret, SecretString}; + +#[tokio::test] +async fn account_repository_should_encrypt_tokens_and_load_decrypted_account() { + let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 4).await; + let repo = SqliteAccountStore::new(pool.clone(), secret_box); + let expires_at = Utc::now() + Duration::hours(1); + + repo.insert(NewAccount { + id: "acct_a".to_string(), + email: Some("user@example.com".to_string()), + account_id: Some("chatgpt-account".to_string()), + user_id: Some("chatgpt-user".to_string()), + label: Some("primary".to_string()), + plan_type: Some("plus".to_string()), + access_token: SecretString::new("access-secret".to_string().into()), + refresh_token: Some(SecretString::new("refresh-secret".to_string().into())), + access_token_expires_at: Some(expires_at), + status: AccountStatus::Active, + }) + .await + .unwrap(); + + let stored: (String, String) = sqlx::query_as( + "select access_token_cipher, refresh_token_cipher from accounts where id = ?", + ) + .bind("acct_a") + .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 loaded = repo.get("acct_a").await.unwrap().unwrap(); + assert_eq!(loaded.access_token.expose_secret(), "access-secret"); + assert_eq!( + loaded.refresh_token.unwrap().expose_secret(), + "refresh-secret" + ); + assert_eq!(loaded.email.as_deref(), Some("user@example.com")); + assert_eq!(loaded.status, AccountStatus::Active); +} + +#[tokio::test] +async fn account_repository_should_cursor_page_accounts_by_added_at_desc() { + let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 5).await; + let repo = SqliteAccountStore::new(pool.clone(), secret_box.clone()); + + seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await; + seed_repo_account(&pool, &secret_box, "acct_b", "2026-06-11T00:01:00Z").await; + seed_repo_account(&pool, &secret_box, "acct_c", "2026-06-11T00:02:00Z").await; + + let first_page = repo.list(None, 2).await.unwrap(); + assert_eq!( + first_page + .items + .iter() + .map(|account| account.id.as_str()) + .collect::>(), + ["acct_c", "acct_b"] + ); + + let second_page = repo.list(first_page.next_cursor, 2).await.unwrap(); + assert_eq!( + second_page + .items + .iter() + .map(|account| account.id.as_str()) + .collect::>(), + ["acct_a"] + ); +} + +#[tokio::test] +async fn account_repository_should_list_metadata_without_decrypting_tokens() { + let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 35).await; + let repo = SqliteAccountStore::new(pool.clone(), secret_box); + 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("2026-06-11T00:00:00Z") + .bind("2026-06-11T00:00:00Z") + .execute(&pool) + .await + .unwrap(); + + let page = repo.list_metadata(None, 10).await.unwrap(); + + assert_eq!(page.items[0].id, "acct_corrupt"); + assert_eq!(page.items[0].email.as_deref(), Some("user@example.com")); + assert_eq!(page.items[0].status, AccountStatus::Active); +} + +#[tokio::test] +async fn account_repository_should_update_status_and_label_without_rewriting_tokens() { + let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 6).await; + let repo = SqliteAccountStore::new(pool.clone(), secret_box); + + repo.insert(NewAccount { + id: "acct_a".to_string(), + email: None, + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-secret".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }) + .await + .unwrap(); + let before: (String,) = sqlx::query_as("select access_token_cipher from accounts where id = ?") + .bind("acct_a") + .fetch_one(&pool) + .await + .unwrap(); + + assert!(repo + .set_status("acct_a", AccountStatus::Disabled) + .await + .unwrap()); + assert!(repo + .set_label("acct_a", Some("work".to_string())) + .await + .unwrap()); + + let after: (String,) = sqlx::query_as("select access_token_cipher from accounts where id = ?") + .bind("acct_a") + .fetch_one(&pool) + .await + .unwrap(); + let loaded = repo.get("acct_a").await.unwrap().unwrap(); + assert_eq!(after.0, before.0); + assert_eq!(loaded.status, AccountStatus::Disabled); + assert_eq!(loaded.label.as_deref(), Some("work")); +} + +#[tokio::test] +async fn account_repository_should_accumulate_usage_counters() { + let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 7).await; + let repo = SqliteAccountStore::new(pool, secret_box); + seed_new_account(&repo, "acct_a").await; + + repo.record_usage( + "acct_a", + UsageDelta { + input_tokens: 10, + output_tokens: 4, + cached_tokens: 3, + empty_response_count: 0, + ..UsageDelta::default() + }, + ) + .await + .unwrap(); + repo.record_usage( + "acct_a", + UsageDelta { + input_tokens: 8, + output_tokens: 2, + cached_tokens: 1, + empty_response_count: 1, + ..UsageDelta::default() + }, + ) + .await + .unwrap(); + + let usage = repo.get_usage("acct_a").await.unwrap().unwrap(); + assert_eq!(usage.request_count, 2); + assert_eq!(usage.input_tokens, 18); + assert_eq!(usage.output_tokens, 6); + assert_eq!(usage.cached_tokens, 4); + assert_eq!(usage.empty_response_count, 1); + assert!(usage.last_used_at.is_some()); +} + +#[tokio::test] +async fn account_repository_should_restore_window_usage_into_runtime_pool_accounts() { + let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 33).await; + let repo = SqliteAccountStore::new(pool.clone(), secret_box.clone()); + let window_started_at = Utc::now() - Duration::minutes(2); + let window_reset_at = Utc::now() + Duration::minutes(3); + + seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await; + sqlx::query( + "insert into account_usage (account_id, request_count, image_input_tokens, image_output_tokens, image_request_count, image_request_failed_count, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_image_input_tokens, window_image_output_tokens, window_image_request_count, window_image_request_failed_count, window_started_at, window_reset_at, limit_window_seconds) values (?, 9, 31, 9, 2, 1, 7, 11, 13, 17, 19, 5, 1, 1, ?, ?, 300)", + ) + .bind("acct_a") + .bind(window_started_at.to_rfc3339()) + .bind(window_reset_at.to_rfc3339()) + .execute(&pool) + .await + .unwrap(); + + let account = repo.list_pool_accounts().await.unwrap().remove(0); + + assert_eq!(account.window_request_count, 7); + assert_eq!(account.image_input_tokens, 31); + assert_eq!(account.image_output_tokens, 9); + assert_eq!(account.image_request_count, 2); + assert_eq!(account.image_request_failed_count, 1); + assert_eq!(account.window_input_tokens, 11); + assert_eq!(account.window_output_tokens, 13); + assert_eq!(account.window_cached_tokens, 17); + assert_eq!(account.window_image_input_tokens, 19); + assert_eq!(account.window_image_output_tokens, 5); + assert_eq!(account.window_image_request_count, 1); + assert_eq!(account.window_image_request_failed_count, 1); + assert_eq!(account.window_started_at, Some(window_started_at)); + assert_eq!(account.window_reset_at, Some(window_reset_at)); + assert_eq!(account.limit_window_seconds, Some(300)); +} + +#[tokio::test] +async fn account_repository_should_restore_quota_verify_required_into_runtime_pool_accounts() { + let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 34).await; + let repo = SqliteAccountStore::new(pool.clone(), secret_box.clone()); + + seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await; + sqlx::query("update accounts set quota_verify_required = 1 where id = ?") + .bind("acct_a") + .execute(&pool) + .await + .unwrap(); + + let account = repo.list_pool_accounts().await.unwrap().remove(0); + + assert!(account.quota_verify_required); +} + +#[tokio::test] +async fn account_repository_should_update_quota_json_and_fetched_at() { + let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 36).await; + let repo = SqliteAccountStore::new(pool.clone(), secret_box.clone()); + seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await; + + let updated = repo + .update_quota_json("acct_a", r#"{"rate_limit":{"limit_reached":false}}"#) + .await + .unwrap(); + let missing_updated = repo + .update_quota_json("missing", r#"{"rate_limit":null}"#) + .await + .unwrap(); + let row: (String, Option) = + sqlx::query_as("select quota_json, quota_fetched_at from accounts where id = ?") + .bind("acct_a") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!( + (updated, missing_updated, row.0, row.1.is_some()), + ( + true, + false, + r#"{"rate_limit":{"limit_reached":false}}"#.to_string(), + true + ) + ); +} + +async fn sqlite_account_store_parts( + db_name: &str, + key_byte: u8, +) -> (sqlx::SqlitePool, SecretBox, 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(); + (pool, SecretBox::new([key_byte; 32]), dir) +} + +async fn seed_new_account(repo: &SqliteAccountStore, id: &str) { + repo.insert(NewAccount { + id: id.to_string(), + email: None, + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-secret".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }) + .await + .unwrap(); +} + +async fn seed_repo_account( + pool: &sqlx::SqlitePool, + secret_box: &SecretBox, + id: &str, + added_at: &str, +) { + let token = SecretString::new(format!("access-{id}").into()); + let cipher = secret_box.encrypt(&token).unwrap(); + sqlx::query( + "insert into accounts (id, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?)", + ) + .bind(id) + .bind(cipher) + .bind("active") + .bind(added_at) + .bind(added_at) + .execute(pool) + .await + .unwrap(); +} diff --git a/crates/adapters/tests/admin_sessions.rs b/crates/adapters/tests/admin_sessions.rs new file mode 100644 index 000000000..a8d3701d0 --- /dev/null +++ b/crates/adapters/tests/admin_sessions.rs @@ -0,0 +1,82 @@ +use chrono::{Duration, Utc}; +use codex_proxy_adapters::sqlite::admin_sessions::SqliteAdminSessionStore; +use codex_proxy_platform::{identity::hash_admin_password, storage::connect_sqlite}; + +#[tokio::test] +async fn admin_session_store_should_create_and_load_default_admin_once() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("admin-auth.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAdminSessionStore::new(pool); + let password_hash = hash_admin_password("admin-password").expect("password hash"); + + let created = store + .ensure_default_admin(&password_hash) + .await + .expect("default admin should be created"); + let skipped = store + .ensure_default_admin("different-hash") + .await + .expect("second default admin should be skipped"); + let admin = store + .load_first_admin() + .await + .expect("default admin should load") + .expect("default admin should exist"); + + assert!(created); + assert!(!skipped); + assert_eq!(admin.password_hash, password_hash); +} + +#[tokio::test] +async fn admin_session_store_should_create_validate_and_cleanup_sessions() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("admin-sessions.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAdminSessionStore::new(pool); + let password_hash = hash_admin_password("admin-password").expect("password hash"); + store + .ensure_default_admin(&password_hash) + .await + .expect("default admin should be created"); + let admin = store + .load_first_admin() + .await + .expect("admin should load") + .expect("admin should exist"); + let now = Utc::now(); + + store + .create_session("sess_future", &admin.id, now + Duration::minutes(5)) + .await + .expect("future session should be inserted"); + store + .create_session("sess_expired", &admin.id, now - Duration::minutes(5)) + .await + .expect("expired session should be inserted"); + + assert!(store + .validate_session("sess_future") + .await + .expect("future session should validate")); + assert!(!store + .validate_session("sess_expired") + .await + .expect("expired session should not validate")); + + let deleted = store + .cleanup_expired_sessions(now) + .await + .expect("expired session should be deleted"); + + assert_eq!(deleted, 1); + assert!(store + .validate_session("sess_future") + .await + .expect("future session should remain")); +} diff --git a/crates/adapters/tests/client_keys.rs b/crates/adapters/tests/client_keys.rs new file mode 100644 index 000000000..f802a1f16 --- /dev/null +++ b/crates/adapters/tests/client_keys.rs @@ -0,0 +1,89 @@ +use codex_proxy_adapters::sqlite::client_keys::SqliteClientKeyStore; +use codex_proxy_core::admin::ports::ClientKeyStore; +use codex_proxy_platform::{identity::ApiKeyHasher, storage::connect_sqlite}; + +#[tokio::test] +async fn client_key_store_should_create_list_disable_enable_and_delete_keys() { + let (store, _dir) = client_key_store("client-keys.sqlite", 10).await; + + let created = store.create("cursor").await.unwrap(); + assert_eq!(created.name, "cursor"); + assert!(created.plaintext.starts_with("cpr_")); + assert!(store.verify_and_touch(&created.plaintext).await.unwrap()); + + let first_page = store.list(None, 10).await.unwrap(); + assert_eq!(first_page.items.len(), 1); + assert_eq!(first_page.items[0].name, "cursor"); + assert!(first_page.items[0].last_used_at.is_some()); + + assert!(store.set_enabled(&created.id, false).await.unwrap()); + assert!(!store.verify_and_touch(&created.plaintext).await.unwrap()); + + assert!(store.set_enabled(&created.id, true).await.unwrap()); + assert!(store.verify_and_touch(&created.plaintext).await.unwrap()); + + assert!(store.delete(&created.id).await.unwrap()); + assert!(!store.verify_and_touch(&created.plaintext).await.unwrap()); +} + +#[tokio::test] +async fn client_key_store_should_page_keys_by_created_at_desc() { + let (store, _dir) = client_key_store("client-keys-page.sqlite", 11).await; + + let key_a = store.create("a").await.unwrap(); + let key_b = store.create("b").await.unwrap(); + let key_c = store.create("c").await.unwrap(); + sqlx::query("update client_api_keys set created_at = ? where id = ?") + .bind("2026-06-11T00:00:00Z") + .bind(&key_a.id) + .execute(store.pool()) + .await + .unwrap(); + sqlx::query("update client_api_keys set created_at = ? where id = ?") + .bind("2026-06-11T00:01:00Z") + .bind(&key_b.id) + .execute(store.pool()) + .await + .unwrap(); + sqlx::query("update client_api_keys set created_at = ? where id = ?") + .bind("2026-06-11T00:02:00Z") + .bind(&key_c.id) + .execute(store.pool()) + .await + .unwrap(); + + let first_page = store.list(None, 2).await.unwrap(); + assert_eq!( + first_page + .items + .iter() + .map(|key| key.name.as_str()) + .collect::>(), + ["c", "b"] + ); + + let second_page = store.list(first_page.next_cursor, 2).await.unwrap(); + assert_eq!( + second_page + .items + .iter() + .map(|key| key.name.as_str()) + .collect::>(), + ["a"] + ); +} + +async fn client_key_store( + db_name: &str, + key_byte: u8, +) -> (SqliteClientKeyStore, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join(db_name); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .unwrap(); + ( + SqliteClientKeyStore::new(pool, ApiKeyHasher::new([key_byte; 32])), + dir, + ) +} diff --git a/crates/adapters/tests/codex.rs b/crates/adapters/tests/codex.rs new file mode 100644 index 000000000..bc83ad53e --- /dev/null +++ b/crates/adapters/tests/codex.rs @@ -0,0 +1,3641 @@ +use std::{ + io::Write, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, + time::Duration, +}; + +use codex_proxy_adapters::codex::{ + client, + websocket::pool::{CodexWebSocketPool, CodexWebSocketPoolConfig}, +}; +use flate2::{write::DeflateEncoder, Compression}; +use futures::{SinkExt, StreamExt}; +use serde_json::json; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + time::timeout, +}; +use tokio_tungstenite::{ + accept_async, accept_hdr_async, + tungstenite::{ + handshake::{ + derive_accept_key, + server::{Request as WsRequest, Response as WsResponse}, + }, + Message, + }, +}; + +#[test] +fn endpoints_should_join_backend_paths() { + assert_eq!( + client::endpoint_url("https://api.example.com/", "/codex/responses"), + "https://api.example.com/codex/responses" + ); + assert_eq!( + client::endpoint_request_path("https://api.example.com/backend-api", "/codex/usage"), + "/backend-api/codex/usage" + ); +} + +#[test] +fn ordered_codex_headers_should_preserve_fingerprint_priority_and_request_fields() { + let fingerprint = codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(); + + let headers = client::build_ordered_codex_headers( + &fingerprint, + "access-token", + Some("acct-1"), + Some("turn-1"), + "req-1", + ); + let keys = headers.keys().cloned().collect::>(); + + assert_eq!(headers["authorization"], "Bearer access-token"); + assert_eq!(headers["chatgpt-account-id"], "acct-1"); + assert_eq!(headers["x-client-request-id"], "req-1"); + assert_eq!(headers["x-codex-turn-state"], "turn-1"); + assert_eq!(headers["accept"], "text/event-stream"); + assert_eq!(keys.first().map(String::as_str), Some("authorization")); +} + +#[test] +fn custom_ca_should_report_environment_cache_key_consistently() { + let cache_key = client::custom_ca_env_cache_key(); + + assert_eq!( + cache_key.is_some(), + std::env::var("CODEX_CA_CERTIFICATE").is_ok() || std::env::var("SSL_CERT_FILE").is_ok() + ); +} + +#[test] +fn reqwest_http2_feature_should_be_enabled_for_fallback_parity() { + let builder = reqwest::Client::builder().http2_prior_knowledge(); + + drop(builder); +} + +#[test] +fn websocket_connection_should_preserve_endpoint_and_header_order() { + let connection = codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::new( + "wss://chatgpt.com/backend-api/codex", + vec![ + ("authorization".to_string(), "Bearer token".to_string()), + ("user-agent".to_string(), "Codex Desktop/test".to_string()), + ], + ); + + assert_eq!( + ( + connection.endpoint(), + connection.opening_audit_snapshot().header_order, + ), + ( + "wss://chatgpt.com/backend-api/codex", + vec!["authorization".to_string(), "user-agent".to_string()], + ) + ); +} + +#[test] +fn websocket_connection_opening_audit_should_redact_sensitive_headers() { + let connection = codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::new( + "wss://chatgpt.com/backend-api/codex/responses?source=audit", + vec![ + ( + "authorization".to_string(), + "Bearer access-secret".to_string(), + ), + ("chatgpt-account-id".to_string(), "acct-secret".to_string()), + ("user-agent".to_string(), "Codex Desktop/test".to_string()), + ("x-client-request-id".to_string(), "req-secret".to_string()), + ( + "x-codex-turn-metadata".to_string(), + "{\"secret\":true}".to_string(), + ), + ], + ); + + let snapshot = connection.opening_audit_snapshot(); + + assert_eq!( + snapshot.request_line, + "GET /backend-api/codex/responses?source=audit HTTP/1.1" + ); + assert_eq!( + snapshot + .headers + .iter() + .map(|header| (header.name.as_str(), header.value.as_str())) + .collect::>(), + vec![ + ("authorization", ""), + ("chatgpt-account-id", ""), + ("user-agent", "Codex Desktop/test"), + ("x-client-request-id", ""), + ("x-codex-turn-metadata", ""), + ] + ); +} + +#[tokio::test] +async fn websocket_audit_artifact_should_require_explicit_directory() { + let connection = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses( + "https://chatgpt.com", + "test-websocket-key", + vec![ + ( + "authorization".to_string(), + "Bearer access-token".to_string(), + ), + ("chatgpt-account-id".to_string(), "acct-secret".to_string()), + ("cookie".to_string(), "session=secret".to_string()), + ], + ); + let mut payload_request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "private instructions", + vec![json!({"role": "user", "content": "private prompt"})], + ); + payload_request.prompt_cache_key = Some("cache-secret".to_string()); + payload_request.client_metadata = Some(json!({"thread_id": "thread-secret"})); + let mut artifact = + codex_proxy_core::protocol::codex::websocket::websocket_audit_artifact_from_attempt( + &payload_request, + connection.opening_audit_snapshot(), + codex_proxy_core::protocol::codex::websocket::websocket_payload_audit_snapshot( + &payload_request, + ), + ); + artifact.error = Some( + codex_proxy_core::protocol::codex::websocket::WebSocketAuditErrorSnapshot { + classification: "opening_failed".to_string(), + message: "connection refused".to_string(), + }, + ); + let dir = tempfile::tempdir().expect("temp dir"); + + let disabled = + codex_proxy_adapters::codex::websocket::opening::write_websocket_audit_artifact_for_dir( + None, &artifact, + ) + .await + .expect("disabled audit should be ok"); + + assert!(disabled.is_none()); + assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 0); + + let written = + codex_proxy_adapters::codex::websocket::opening::write_websocket_audit_artifact_for_dir( + Some(dir.path()), + &artifact, + ) + .await + .expect("enabled audit should write") + .expect("enabled audit path"); + let body = std::fs::read_to_string(&written).expect("audit file"); + let json = serde_json::from_str::(&body).expect("audit json"); + + assert_eq!(json["transport_mode"], "websocket_preferred"); + assert_eq!( + json["opening"]["request_line"], + "GET /codex/responses HTTP/1.1" + ); + assert_eq!(json["payload"]["top_level_keys"][1], "model"); + assert_eq!(json["error"]["classification"], "opening_failed"); + assert!(!body.contains("access-token")); + assert!(!body.contains("acct-secret")); + assert!(!body.contains("session=secret")); + assert!(!body.contains("private prompt")); + assert!(!body.contains("private instructions")); + assert!(!body.contains("cache-secret")); + assert!(!body.contains("thread-secret")); +} + +#[test] +fn websocket_responses_endpoint_should_convert_http_base_url_to_ws_endpoint() { + assert_eq!( + codex_proxy_adapters::codex::websocket::connect::responses_websocket_endpoint( + "https://chatgpt.com/backend-api" + ), + "wss://chatgpt.com/backend-api/codex/responses" + ); + assert_eq!( + codex_proxy_adapters::codex::websocket::connect::responses_websocket_endpoint( + "http://127.0.0.1:8080" + ), + "ws://127.0.0.1:8080/codex/responses" + ); +} + +#[test] +fn websocket_deflate_should_rewrite_compressed_server_text_frame() { + let payload = json!({ + "type": "response.completed", + "response": { + "id": "resp_deflate", + "object": "response" + } + }) + .to_string(); + let compressed = compressed_permessage_deflate_payload(payload.as_bytes()); + let frame = server_websocket_frame(0x1, true, &compressed); + + let rewritten = + codex_proxy_adapters::codex::websocket::deflate::rewrite_permessage_deflate_server_frame( + &frame, + ) + .expect("compressed frame should rewrite") + .expect("compressed frame should produce rewritten frame"); + + assert_eq!(rewritten[0] & 0x40, 0); + assert_eq!(rewritten[0] & 0x0f, 0x1); + assert_eq!(server_websocket_payload(&rewritten), payload.as_bytes()); +} + +#[test] +fn websocket_deflate_should_leave_uncompressed_server_frame_unchanged() { + let frame = server_websocket_frame(0x1, false, b"plain"); + + let rewritten = + codex_proxy_adapters::codex::websocket::deflate::rewrite_permessage_deflate_server_frame( + &frame, + ) + .expect("plain frame should parse"); + + assert_eq!(rewritten, None); +} + +#[tokio::test] +async fn codex_backend_client_should_decode_live_permessage_deflate_websocket_frame() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let request = read_http_request(&mut stream).await; + let websocket_key = websocket_opening_header(&request, "sec-websocket-key") + .expect("opening request should include sec-websocket-key"); + let accept_key = derive_accept_key(websocket_key.as_bytes()); + stream + .write_all( + format!( + "HTTP/1.1 101 Switching Protocols\r\n\ + Upgrade: websocket\r\n\ + Connection: Upgrade\r\n\ + Sec-WebSocket-Accept: {accept_key}\r\n\ + Sec-WebSocket-Extensions: permessage-deflate\r\n\ + \r\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + + let completed = json!({ + "type": "response.completed", + "response": { + "id": "resp_live_deflate", + "object": "response", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string(); + let compressed = compressed_permessage_deflate_payload(completed.as_bytes()); + stream + .write_all(&server_websocket_frame(0x1, true, &compressed)) + .await + .unwrap(); + stream.flush().await.unwrap(); + }); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + + let response = backend + .create_response( + &request, + request_context("req_live_deflate", Some("chatgpt-account")), + ) + .await + .expect("deflated websocket response should decode"); + server.await.unwrap(); + + assert!(response.body.contains("resp_live_deflate")); +} + +fn compressed_permessage_deflate_payload(payload: &[u8]) -> Vec { + let mut encoder = DeflateEncoder::new(Vec::new(), Compression::fast()); + encoder.write_all(payload).expect("payload should compress"); + encoder.finish().expect("compressed payload") +} + +fn server_websocket_frame(opcode: u8, rsv1: bool, payload: &[u8]) -> Vec { + let mut frame = Vec::new(); + frame.push(0x80 | if rsv1 { 0x40 } else { 0 } | opcode); + match payload.len() { + len @ 0..=125 => frame.push(len as u8), + len @ 126..=65_535 => { + frame.push(126); + frame.extend_from_slice(&(len as u16).to_be_bytes()); + } + len => { + frame.push(127); + frame.extend_from_slice(&(len as u64).to_be_bytes()); + } + } + frame.extend_from_slice(payload); + frame +} + +fn server_websocket_payload(frame: &[u8]) -> &[u8] { + let len_marker = frame[1] & 0x7f; + let (offset, len) = match len_marker { + len @ 0..=125 => (2, usize::from(len)), + 126 => (4, usize::from(u16::from_be_bytes([frame[2], frame[3]]))), + 127 => { + let len = u64::from_be_bytes([ + frame[2], frame[3], frame[4], frame[5], frame[6], frame[7], frame[8], frame[9], + ]) as usize; + (10, len) + } + _ => unreachable!("websocket length marker is masked to 7 bits"), + }; + &frame[offset..offset + len] +} + +#[test] +fn websocket_connection_should_build_standard_opening_headers_around_business_headers() { + let connection = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses( + "https://chatgpt.com/backend-api", + "test-websocket-key", + vec![ + ( + "authorization".to_string(), + "Bearer access-token".to_string(), + ), + ("user-agent".to_string(), "Codex Desktop/test".to_string()), + ( + "openai-beta".to_string(), + "responses_websockets=2026-02-06".to_string(), + ), + ], + ); + + let snapshot = connection.opening_audit_snapshot(); + + assert_eq!( + connection.endpoint(), + "wss://chatgpt.com/backend-api/codex/responses" + ); + assert_eq!( + snapshot.header_order, + vec![ + "Host", + "Connection", + "Upgrade", + "Sec-WebSocket-Version", + "Sec-WebSocket-Key", + "authorization", + "user-agent", + "openai-beta", + "sec-websocket-extensions", + ] + ); + assert_eq!( + snapshot + .headers + .iter() + .map(|header| (header.name.as_str(), header.value.as_str())) + .collect::>(), + vec![ + ("Host", "chatgpt.com"), + ("Connection", "Upgrade"), + ("Upgrade", "websocket"), + ("Sec-WebSocket-Version", "13"), + ("Sec-WebSocket-Key", "test-websocket-key"), + ("authorization", ""), + ("user-agent", "Codex Desktop/test"), + ("openai-beta", "responses_websockets=2026-02-06"), + ( + "sec-websocket-extensions", + "permessage-deflate; client_max_window_bits" + ), + ] + ); +} + +#[test] +fn websocket_connection_should_render_raw_opening_bytes_for_capture_parity() { + let connection = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses( + "https://chatgpt.com/backend-api", + "test-websocket-key", + vec![ + ( + "chatgpt-account-id".to_string(), + "chatgpt-account".to_string(), + ), + ( + "authorization".to_string(), + "Bearer access-token".to_string(), + ), + ( + "user-agent".to_string(), + "Codex Desktop/26.519.81530 (darwin; arm64)".to_string(), + ), + ("originator".to_string(), "Codex Desktop".to_string()), + ( + "openai-beta".to_string(), + "responses_websockets=2026-02-06".to_string(), + ), + ( + "x-codex-beta-features".to_string(), + "terminal_resize_reflow,memories,network_proxy,prevent_idle_sleep,remote_compaction_v2".to_string(), + ), + ("x-client-request-id".to_string(), "session-1".to_string()), + ("session-id".to_string(), "session-1".to_string()), + ("thread-id".to_string(), "session-1".to_string()), + ("x-codex-window-id".to_string(), "session-1:0".to_string()), + ( + "x-codex-turn-metadata".to_string(), + r#"{"installation_id":"install-1","session_id":"session-1","thread_id":"session-1","turn_id":"","window_id":"session-1:0","request_kind":"prewarm","sandbox":"seccomp"}"#.to_string(), + ), + ], + ); + + let opening = connection.opening_request_text(); + + assert!(opening.starts_with("GET /backend-api/codex/responses HTTP/1.1\r\n")); + assert_substrings_appear_in_order( + &opening, + &[ + "Host: chatgpt.com\r\n", + "Connection: Upgrade\r\n", + "Upgrade: websocket\r\n", + "Sec-WebSocket-Version: 13\r\n", + "Sec-WebSocket-Key: test-websocket-key\r\n", + "chatgpt-account-id: chatgpt-account\r\n", + "authorization: Bearer access-token\r\n", + "user-agent: Codex Desktop/26.519.81530 (darwin; arm64)\r\n", + "originator: Codex Desktop\r\n", + "openai-beta: responses_websockets=2026-02-06\r\n", + "x-codex-beta-features: terminal_resize_reflow,memories,network_proxy,prevent_idle_sleep,remote_compaction_v2\r\n", + "x-client-request-id: session-1\r\n", + "session-id: session-1\r\n", + "thread-id: session-1\r\n", + "x-codex-window-id: session-1:0\r\n", + r#"x-codex-turn-metadata: {"installation_id":"install-1","session_id":"session-1","thread_id":"session-1","turn_id":"","window_id":"session-1:0","request_kind":"prewarm","sandbox":"seccomp"}"#, + "sec-websocket-extensions: permessage-deflate; client_max_window_bits\r\n", + ], + ); + assert!(opening.ends_with("\r\n\r\n")); +} + +#[test] +fn websocket_connection_should_prepare_response_create_payload_text() { + let request = codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + vec![json!({ + "role": "user", + "content": "hello", + })], + ); + + let prepared = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + "https://chatgpt.com/backend-api", + "test-websocket-key", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize"); + let payload: serde_json::Value = + serde_json::from_str(prepared.payload_text()).expect("payload should be json"); + + assert_eq!( + prepared.connection().endpoint(), + "wss://chatgpt.com/backend-api/codex/responses" + ); + assert_eq!(payload["type"], "response.create"); + assert_eq!(payload["model"], "gpt-5.5"); + assert_eq!(payload["instructions"], "be brief"); + assert_eq!(payload["input"][0]["content"], "hello"); + assert_eq!(payload["stream"], true); +} + +#[test] +fn websocket_connection_should_prepare_capture_payload_with_old_field_order() { + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "private capture instructions", + vec![json!({ + "role": "user", + "content": "private capture prompt", + })], + ); + request.prompt_cache_key = Some("session-1".to_string()); + request.client_metadata = Some(json!({ + "thread_id": "capture-thread-secret", + "safe": "capture", + })); + request.generate = Some(false); + + let prepared = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + "https://chatgpt.com/backend-api", + "test-websocket-key", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize"); + + assert_substrings_appear_in_order( + prepared.payload_text(), + &[ + "\"type\":\"response.create\"", + "\"model\":\"gpt-5.5\"", + "\"instructions\":\"private capture instructions\"", + "\"input\":", + "\"tools\":[]", + "\"tool_choice\":\"auto\"", + "\"parallel_tool_calls\":true", + "\"reasoning\":null", + "\"store\":false", + "\"stream\":true", + "\"include\":[]", + "\"prompt_cache_key\":\"session-1\"", + "\"generate\":false", + "\"client_metadata\":", + ], + ); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_collect_completed_sse() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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 payload = serde_json::from_str::(&message.into_text().unwrap()) + .expect("client payload should be json"); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_ws_live", + "object": "response", + "usage": { + "input_tokens": 5, + "output_tokens": 2, + "total_tokens": 7 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + payload + }); + let request = codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + let prepared = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + &format!("http://{addr}"), + "dGhlIHNhbXBsZSBub25jZQ==", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize"); + + let response = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect("websocket exchange should succeed"); + let payload = server.await.unwrap(); + + assert_eq!(payload["type"], "response.create"); + assert!(response.body.contains("event: response.completed")); + assert!(response.body.contains("\"id\":\"resp_ws_live\"")); + assert_eq!(response.usage.expect("usage").input_tokens, 5); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_surface_response_failed_as_upstream_error( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Text( + json!({ + "type": "response.failed", + "response": { + "error": { + "code": "rate_limit_exceeded", + "message": "Rate limit reached. Please try again in 11.054s." + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let request = codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + let prepared = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + &format!("http://{addr}"), + "dGhlIHNhbXBsZSBub25jZQ==", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize"); + + let error = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect_err("response.failed should be surfaced as upstream error"); + server.await.unwrap(); + + let codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::Upstream { + status_code, + retry_after_seconds, + body, + } = error + else { + panic!("expected upstream websocket error"); + }; + assert_eq!(status_code, 429); + assert_eq!(retry_after_seconds, Some(12)); + assert!(body.contains("rate_limit_exceeded")); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_preserve_opening_error_status_body_and_retry_after( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let request = read_http_request(&mut stream).await; + assert!(request.starts_with("GET /codex/responses HTTP/1.1")); + assert!(request.contains("authorization: Bearer access-token")); + let body = r#"{"error":{"message":"rate limited"}}"#; + stream + .write_all( + format!( + "HTTP/1.1 429 Too Many Requests\r\nretry-after: 33\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .await + .unwrap(); + }); + let prepared = prepared_websocket_request(&format!("http://{addr}")); + + let error = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect_err("failed opening should surface upstream status"); + server.await.unwrap(); + + let codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::Upstream { + status_code, + retry_after_seconds, + body, + } = error + else { + panic!("expected upstream opening error"); + }; + assert_eq!(status_code, 429); + assert_eq!(retry_after_seconds, Some(33)); + assert!(body.contains("rate limited")); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_reject_binary_event() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Binary(b"unexpected-binary".to_vec().into())) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let prepared = prepared_websocket_request(&format!("http://{addr}")); + + let error = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect_err("binary websocket events should be rejected"); + server.await.unwrap(); + + assert!(matches!( + error, + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::UnexpectedBinaryEvent + )); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_surface_wrapped_error_status_and_retry_after( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Text( + json!({ + "type": "error", + "status": 409, + "headers": { + "retry-after": ["17"] + }, + "error": { + "code": "conflict", + "message": "wrapped conflict" + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let prepared = prepared_websocket_request(&format!("http://{addr}")); + + let error = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect_err("wrapped error should surface as upstream error"); + server.await.unwrap(); + + let codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::Upstream { + status_code, + retry_after_seconds, + body, + } = error + else { + panic!("expected wrapped upstream error"); + }; + assert_eq!(status_code, 409); + assert_eq!(retry_after_seconds, Some(17)); + assert!(body.contains("wrapped conflict")); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_surface_connection_limit_as_503() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Text( + json!({ + "type": "response.failed", + "response": { + "id": "resp_connection_limit", + "error": { + "code": "websocket_connection_limit_reached", + "message": "connection limit reached" + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let prepared = prepared_websocket_request(&format!("http://{addr}")); + + let error = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect_err("connection limit should surface as upstream error"); + server.await.unwrap(); + + let codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::Upstream { + status_code, + body, + .. + } = error + else { + panic!("expected connection limit upstream error"); + }; + assert_eq!(status_code, 503); + assert!(body.contains("websocket_connection_limit_reached")); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_skip_invalid_stream_events() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_after_invalid", + "object": "response" + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let request = codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + let prepared = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + &format!("http://{addr}"), + "dGhlIHNhbXBsZSBub25jZQ==", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize"); + + let response = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect("websocket exchange should succeed"); + server.await.unwrap(); + + assert!(!response.body.contains("response.output_text.delta")); + assert!(response.body.contains("event: response.completed")); + assert!(response.body.contains("\"id\":\"resp_after_invalid\"")); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_capture_internal_metadata_and_rate_limit_events( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Text( + json!({ + "type": "codex.rate_limits", + "rate_limits": { + "primary": { + "used_percent": 100, + "window_minutes": 5, + "reset_at": 1893456300 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.metadata", + "headers": { + "x-codex-turn-state": ["turn-from-metadata"] + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_internal_events", + "object": "response", + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let prepared = prepared_websocket_request(&format!("http://{addr}")); + + let response = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect("internal events should update metadata without forwarding"); + server.await.unwrap(); + + assert_eq!(response.turn_state.as_deref(), Some("turn-from-metadata")); + assert!(!response.body.contains("codex.rate_limits")); + assert!(!response.body.contains("response.metadata")); + assert!(response.body.contains("event: response.completed")); + assert!(response + .rate_limit_headers + .iter() + .any(|(name, value)| name == "x-codex-primary-used-percent" && value == "100")); + assert!(response + .rate_limit_headers + .iter() + .any(|(name, value)| name == "x-codex-primary-reset-at" && value == "1893456300")); +} + +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn websocket_execute_response_create_request_should_capture_handshake_headers() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = + accept_hdr_async(stream, |_request: &WsRequest, mut response: WsResponse| { + response + .headers_mut() + .insert("x-codex-turn-state", "turn-from-handshake".parse().unwrap()); + response.headers_mut().insert( + "set-cookie", + "cf_clearance=ws; Domain=.chatgpt.com; Path=/" + .parse() + .unwrap(), + ); + response + .headers_mut() + .insert("x-ratelimit-remaining-requests", "41".parse().unwrap()); + Ok(response) + }) + .await + .unwrap(); + let _message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_ws_headers", + "object": "response", + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let request = codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + let prepared = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + &format!("http://{addr}"), + "dGhlIHNhbXBsZSBub25jZQ==", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize"); + + let response = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect("websocket exchange should succeed"); + server.await.unwrap(); + + assert_eq!(response.turn_state.as_deref(), Some("turn-from-handshake")); + assert_eq!( + response.set_cookie_headers, + vec!["cf_clearance=ws; Domain=.chatgpt.com; Path=/".to_string()] + ); + assert!(response + .rate_limit_headers + .iter() + .any(|(name, value)| name == "x-ratelimit-remaining-requests" && value == "41")); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_surface_incomplete_response() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Text( + json!({ + "type": "response.incomplete", + "response": { + "id": "resp_incomplete", + "incomplete_details": { + "reason": "max_output_tokens" + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let request = codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + let prepared = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + &format!("http://{addr}"), + "dGhlIHNhbXBsZSBub25jZQ==", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize"); + + let error = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect_err("response.incomplete should be surfaced as websocket error"); + server.await.unwrap(); + + let codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::IncompleteResponse { + reason, + } = error + else { + panic!("expected incomplete websocket response"); + }; + assert_eq!(reason, "max_output_tokens"); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_reject_invalid_completed_response() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_invalid_completed", + "object": "response", + "usage": { + "input_tokens": "bad", + "output_tokens": 1, + "total_tokens": 1 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let request = codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + let prepared = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + &format!("http://{addr}"), + "dGhlIHNhbXBsZSBub25jZQ==", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize"); + + let error = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect_err("invalid response.completed should be rejected"); + server.await.unwrap(); + + let codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::InvalidCompletedResponse { + message, + } = error + else { + panic!("expected invalid completed websocket response"); + }; + assert!(message.contains("failed to parse ResponseCompleted")); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_ignore_completed_without_response_until_close( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let request = codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + let prepared = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + &format!("http://{addr}"), + "dGhlIHNhbXBsZSBub25jZQ==", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize"); + + let error = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect_err("response.completed without response should not finish the stream"); + server.await.unwrap(); + + assert!(matches!( + error, + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::ClosedBeforeTerminal + )); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_ignore_success_status_error_until_close() +{ + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Text( + json!({ + "type": "error", + "status": 200, + "error": { + "message": "non-terminal informational frame" + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let request = codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + let prepared = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + &format!("http://{addr}"), + "dGhlIHNhbXBsZSBub25jZQ==", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize"); + + let error = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect_err("unclassified success-status error frame should not finish the stream"); + server.await.unwrap(); + + assert!(matches!( + error, + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::ClosedBeforeTerminal + )); +} + +#[tokio::test(start_paused = true)] +async fn websocket_execute_response_create_request_should_timeout_when_upstream_is_silent() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (request_seen_tx, request_seen_rx) = tokio::sync::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(); + request_seen_tx.send(()).unwrap(); + futures::future::pending::<()>().await; + }); + let prepared = prepared_websocket_request(&format!("http://{addr}")); + + let response_task = tokio::spawn(async move { + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + }); + request_seen_rx.await.unwrap(); + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_secs(20)).await; + tokio::task::yield_now().await; + + let error = response_task + .await + .expect("websocket task should finish") + .expect_err("silent upstream should time out"); + server.abort(); + + let codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::ReceiveIdleTimeout { + timeout, + } = error + else { + panic!("expected receive idle timeout"); + }; + assert_eq!(timeout, Duration::from_secs(20)); +} + +#[tokio::test] +async fn websocket_execute_response_create_request_should_reply_to_server_ping_before_terminal() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Ping("codex-ping".as_bytes().to_vec().into())) + .await + .unwrap(); + let pong = timeout(Duration::from_secs(1), websocket.next()) + .await + .expect("client should reply to server ping before terminal response") + .expect("client should send a websocket frame") + .expect("client frame should be valid"); + assert_eq!(pong, Message::Pong("codex-ping".as_bytes().to_vec().into())); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_after_ping", + "object": "response", + "usage": { + "input_tokens": 2, + "output_tokens": 1, + "total_tokens": 3 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let request = codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + let prepared = + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + &format!("http://{addr}"), + "dGhlIHNhbXBsZSBub25jZQ==", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize"); + + let response = + codex_proxy_adapters::codex::websocket::connect::execute_response_create_request(&prepared) + .await + .expect("websocket exchange should succeed after ping/pong"); + server.await.unwrap(); + + assert!(response.body.contains("event: response.completed")); + assert!(response.body.contains("\"id\":\"resp_after_ping\"")); + assert_eq!(response.usage.expect("usage").input_tokens, 2); +} + +#[tokio::test] +async fn codex_backend_client_stream_should_reject_binary_websocket_event() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Binary(b"unexpected-binary".to_vec().into())) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + + let mut response = backend + .create_response_stream( + &request, + request_context("req_stream_binary", Some("chatgpt-account")), + ) + .await + .expect("websocket stream should open"); + let error = response + .body + .next() + .await + .expect("stream should yield binary error") + .expect_err("binary event should be rejected"); + server.await.unwrap(); + + assert!(matches!( + error, + client::CodexClientError::WebSocket( + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::UnexpectedBinaryEvent + ) + )); +} + +#[tokio::test] +async fn codex_backend_client_stream_should_error_when_websocket_closes_before_terminal() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + 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(); + websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "partial" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + + let mut response = backend + .create_response_stream( + &request, + request_context("req_stream_mid_close", Some("chatgpt-account")), + ) + .await + .expect("websocket stream should open"); + let first = response + .body + .next() + .await + .expect("stream should yield partial frame") + .expect("partial frame should be valid"); + assert!(std::str::from_utf8(&first).unwrap().contains("partial")); + let error = response + .body + .next() + .await + .expect("stream should yield close-before-terminal error") + .expect_err("close before terminal should be an error"); + server.await.unwrap(); + + assert!(matches!( + error, + client::CodexClientError::WebSocket( + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketExchangeError::ClosedBeforeTerminal + ) + )); +} + +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn codex_backend_client_should_use_websocket_when_previous_response_id_is_present() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = + accept_hdr_async(stream, |_request: &WsRequest, mut response: WsResponse| { + response + .headers_mut() + .insert("x-codex-turn-state", "turn-ws-client".parse().unwrap()); + response.headers_mut().insert( + "set-cookie", + "cf_clearance=client-ws; Domain=.chatgpt.com; Path=/" + .parse() + .unwrap(), + ); + response + .headers_mut() + .insert("x-ratelimit-remaining-requests", "17".parse().unwrap()); + Ok(response) + }) + .await + .unwrap(); + let message = websocket.next().await.unwrap().unwrap(); + let payload = serde_json::from_str::(&message.into_text().unwrap()) + .expect("client payload should be json"); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_ws_client", + "object": "response", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + payload + }); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + + let response = backend + .create_response( + &request, + client::CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_ws_client", + 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, + installation_id: None, + session_id: None, + }, + ) + .await + .expect("websocket response should succeed"); + let payload = server.await.unwrap(); + + assert_eq!(payload["type"], "response.create"); + assert_eq!(payload["previous_response_id"], "resp_previous"); + assert!(response.body.contains("event: response.completed")); + assert!(response.body.contains("\"id\":\"resp_ws_client\"")); + assert_eq!(response.usage.expect("usage").input_tokens, 3); + assert_eq!(response.turn_state.as_deref(), Some("turn-ws-client")); + assert_eq!( + response.set_cookie_headers, + vec!["cf_clearance=client-ws; Domain=.chatgpt.com; Path=/".to_string()] + ); + assert!(response + .rate_limit_headers + .iter() + .any(|(name, value)| name == "x-ratelimit-remaining-requests" && value == "17")); +} + +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn codex_backend_client_websocket_should_forward_security_chain_headers_and_payload_fields() { + let received_headers = Arc::new(Mutex::new(Vec::<(String, String)>::new())); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let headers_for_server = Arc::clone(&received_headers); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = + accept_hdr_async(stream, move |request: &WsRequest, response: WsResponse| { + let headers = request + .headers() + .iter() + .map(|(name, value)| { + ( + name.as_str().to_string(), + value.to_str().unwrap_or_default().to_string(), + ) + }) + .collect::>(); + *headers_for_server.lock().unwrap() = headers; + Ok(response) + }) + .await + .unwrap(); + let message = websocket.next().await.unwrap().unwrap(); + let payload = serde_json::from_str::(&message.into_text().unwrap()) + .expect("client payload should be json"); + websocket + .send(Message::Text( + completed_websocket_response("resp_ws_security", 1, 1).into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + payload + }); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.use_websocket = true; + request.prompt_cache_key = Some("client-thread".to_string()); + request.client_metadata = Some(json!({ + "safe": "yes", + "x-openai-subagent": "review", + "ignored_non_string": 42 + })); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + + let response = backend + .create_response( + &request, + client::CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_ws_security", + turn_state: Some("turn-state"), + turn_metadata: Some("{\"thread_source\":\"subagent\"}"), + beta_features: Some("feature-a"), + include_timing_metrics: Some("true"), + version: Some("26.318.11754"), + codex_window_id: Some("cw_derived"), + parent_thread_id: Some("parent-456"), + cookie_header: None, + installation_id: Some("install-123"), + session_id: Some("cp_derived"), + }, + ) + .await + .expect("websocket response should succeed"); + let payload = server.await.unwrap(); + + assert!(response.body.contains("resp_ws_security")); + assert_eq!(payload["prompt_cache_key"], "cp_derived"); + assert_eq!( + payload["client_metadata"], + json!({ + "safe": "yes", + "x-openai-subagent": "review", + "x-codex-installation-id": "install-123", + "x-codex-window-id": "cw_derived", + "x-codex-turn-metadata": "{\"thread_source\":\"subagent\"}", + "x-codex-parent-thread-id": "parent-456" + }) + ); + + let headers = received_headers.lock().unwrap().clone(); + assert!(headers + .iter() + .any(|(name, value)| name == "x-client-request-id" && value == "cp_derived")); + assert!(headers + .iter() + .any(|(name, value)| name == "x-openai-subagent" && value == "review")); + assert!(headers + .iter() + .any(|(name, value)| name == "session-id" && value == "cp_derived")); + assert!(headers + .iter() + .any(|(name, value)| name == "thread-id" && value == "cp_derived")); + assert!(headers.iter().all(|(name, _)| name != "content-type")); + assert!(headers.iter().all(|(name, _)| name != "accept")); + assert!(headers.iter().all(|(name, _)| name != "session_id")); +} + +#[test] +fn websocket_pool_should_apply_capacity_and_age_policy() { + let pool = codex_proxy_adapters::codex::websocket::pool::CodexWebSocketPool::new( + 2, + Duration::from_secs(60), + ); + + assert_eq!( + ( + pool.permits_new_connection(1), + pool.permits_new_connection(2), + pool.should_recycle(Duration::from_secs(59)), + pool.should_recycle(Duration::from_secs(60)), + ), + (true, false, false, true) + ); +} + +#[tokio::test] +async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and_conversation() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut websocket = accept_async(stream).await.unwrap(); + for response_id in ["resp_pool_first", "resp_pool_second"] { + let _message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "object": "response", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + } + websocket.close(None).await.unwrap(); + }); + let pool = Arc::new( + codex_proxy_adapters::codex::websocket::pool::CodexWebSocketPool::new( + 8, + Duration::from_secs(60), + ), + ); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool)); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + request.prompt_cache_key = Some("conversation-pool".to_string()); + + let first = backend + .create_response( + &request, + request_context("req_pool_first", Some("chatgpt-account")), + ) + .await + .expect("first pooled websocket response should succeed"); + let second = backend + .create_response( + &request, + request_context("req_pool_second", Some("chatgpt-account")), + ) + .await + .expect("second pooled websocket response should succeed"); + server.await.unwrap(); + + assert!(first.body.contains("resp_pool_first")); + assert!(second.body.contains("resp_pool_second")); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (release_first_tx, release_first_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + let mut first_websocket = accept_async(first_stream).await.unwrap(); + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "first connection is still busy" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + + let (second_stream, _) = listener.accept().await.unwrap(); + let mut second_websocket = accept_async(second_stream).await.unwrap(); + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + completed_websocket_response("resp_busy_second", 2, 1).into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + + let (third_stream, _) = listener.accept().await.unwrap(); + let mut third_websocket = accept_async(third_stream).await.unwrap(); + let _third_message = third_websocket.next().await.unwrap().unwrap(); + third_websocket + .send(Message::Text( + completed_websocket_response("resp_busy_third", 2, 1).into(), + )) + .await + .unwrap(); + third_websocket.close(None).await.unwrap(); + + release_first_rx.await.unwrap(); + first_websocket + .send(Message::Text( + completed_websocket_response("resp_busy_first", 2, 1).into(), + )) + .await + .unwrap(); + first_websocket.close(None).await.unwrap(); + }); + let pool = Arc::new(CodexWebSocketPool::with_default_max_age()); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(pool); + let request = pooled_websocket_request("conversation-busy"); + + let mut first = backend + .create_response_stream( + &request, + request_context("req_busy_first", Some("chatgpt-account")), + ) + .await + .expect("first pooled websocket stream should start") + .body; + let first_chunk = first + .next() + .await + .expect("first stream should yield an initial chunk") + .expect("first stream chunk should be valid"); + let first_chunk = std::str::from_utf8(&first_chunk).unwrap(); + assert!(first_chunk.contains("first connection is still busy")); + + let second = backend + .create_response( + &request, + request_context("req_busy_second", Some("chatgpt-account")), + ) + .await + .expect("busy key should bypass with a one-shot second connection"); + let third = backend + .create_response( + &request, + request_context("req_busy_third", Some("chatgpt-account")), + ) + .await + .expect("busy key should bypass with a one-shot third connection"); + + release_first_tx.send(()).unwrap(); + while first.next().await.transpose().unwrap().is_some() {} + server.await.unwrap(); + + assert!(second.body.contains("resp_busy_second")); + assert!(third.body.contains("resp_busy_third")); +} + +#[tokio::test] +async fn websocket_pool_should_bypass_new_keys_after_account_cap() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let server = tokio::spawn(async move { + for response_id in ["resp_cap_first", "resp_cap_second", "resp_cap_third"] { + let (stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut websocket = accept_async(stream).await.unwrap(); + let _message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + completed_websocket_response(response_id, 2, 1).into(), + )) + .await + .unwrap(); + if response_id == "resp_cap_third" { + websocket.close(None).await.unwrap(); + } + } + }); + let pool = Arc::new(CodexWebSocketPool::with_limits(Duration::from_secs(60), 1)); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(pool); + let first_request = pooled_websocket_request("conversation-cap-one"); + let second_request = pooled_websocket_request("conversation-cap-two"); + + let first = backend + .create_response( + &first_request, + request_context("req_cap_first", Some("chatgpt-account")), + ) + .await + .expect("first capped websocket response should succeed"); + let second = backend + .create_response( + &second_request, + request_context("req_cap_second", Some("chatgpt-account")), + ) + .await + .expect("new key over account cap should use one-shot connection"); + let third = backend + .create_response( + &second_request, + request_context("req_cap_third", Some("chatgpt-account")), + ) + .await + .expect("capped key should keep bypassing instead of entering the pool"); + server.await.unwrap(); + + assert!(first.body.contains("resp_cap_first")); + assert!(second.body.contains("resp_cap_second")); + assert!(third.body.contains("resp_cap_third")); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 3); +} + +#[tokio::test] +async fn codex_backend_client_should_ping_idle_pooled_websocket_during_maintenance() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let ping_count = Arc::new(AtomicUsize::new(0)); + let ping_count_for_server = Arc::clone(&ping_count); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let _first_message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_pool_keepalive_first", + "object": "response", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + + let ping = timeout(Duration::from_secs(1), websocket.next()) + .await + .expect("pool maintenance should probe the idle websocket") + .expect("pool maintenance should send a websocket frame") + .expect("pool maintenance frame should be valid"); + let Message::Ping(payload) = ping else { + panic!("expected pool maintenance ping frame, got {ping:?}"); + }; + ping_count_for_server.fetch_add(1, Ordering::SeqCst); + websocket.send(Message::Pong(payload)).await.unwrap(); + + let _second_message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_pool_keepalive_second", + "object": "response", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let pool = Arc::new(CodexWebSocketPool::with_config( + websocket_pool_config_for_tests(None, Some(Duration::from_millis(1)), None), + )); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool)); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + request.prompt_cache_key = Some("conversation-pool-keepalive".to_string()); + + let first = backend + .create_response( + &request, + request_context("req_pool_keepalive_first", Some("chatgpt-account")), + ) + .await + .expect("first pooled websocket response should succeed"); + pool.maintain_idle_connections().await; + let second = backend + .create_response( + &request, + request_context("req_pool_keepalive_second", Some("chatgpt-account")), + ) + .await + .expect("second pooled websocket response should reuse the probed socket"); + server.await.unwrap(); + + assert!(first.body.contains("resp_pool_keepalive_first")); + assert!(second.body.contains("resp_pool_keepalive_second")); + assert_eq!(ping_count.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn websocket_pool_should_evict_idle_connection_when_ping_times_out() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut first_websocket = accept_async(first_stream).await.unwrap(); + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_websocket + .send(Message::Text( + completed_websocket_response("resp_no_pong_first", 2, 1).into(), + )) + .await + .unwrap(); + let (second_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut second_websocket = accept_async(second_stream).await.unwrap(); + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + completed_websocket_response("resp_no_pong_second", 2, 1).into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + drop(first_websocket); + }); + let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { + ping_interval: Some(Duration::from_millis(1)), + ping_timeout: Duration::from_millis(20), + maintenance_interval: None, + ..websocket_pool_config_for_tests(None, None, None) + })); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool)); + let request = pooled_websocket_request("conversation-no-pong"); + + let first = backend + .create_response( + &request, + request_context("req_no_pong_first", Some("chatgpt-account")), + ) + .await + .expect("first websocket response should succeed"); + pool.maintain_idle_connections().await; + let second = backend + .create_response( + &request, + request_context("req_no_pong_second", Some("chatgpt-account")), + ) + .await + .expect("second websocket response should use a fresh connection"); + server.await.unwrap(); + + assert!(first.body.contains("resp_no_pong_first")); + assert!(second.body.contains("resp_no_pong_second")); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn websocket_pool_should_gc_expired_idle_connections() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut first_websocket = accept_async(first_stream).await.unwrap(); + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_websocket + .send(Message::Text( + completed_websocket_response("resp_gc_first", 2, 1).into(), + )) + .await + .unwrap(); + let close = timeout(Duration::from_secs(1), first_websocket.next()) + .await + .expect("gc sweep should close the expired idle websocket") + .expect("gc sweep should send a close frame") + .expect("close frame should be valid"); + assert!(matches!(close, Message::Close(_))); + + let (second_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut second_websocket = accept_async(second_stream).await.unwrap(); + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + completed_websocket_response("resp_gc_second", 2, 1).into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + }); + let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { + max_age: Duration::from_millis(5), + maintenance_interval: None, + ping_interval: None, + liveness_timeout: None, + ..CodexWebSocketPoolConfig::default() + })); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool)); + let request = pooled_websocket_request("conversation-gc"); + + let first = backend + .create_response( + &request, + request_context("req_gc_first", Some("chatgpt-account")), + ) + .await + .expect("first websocket response should succeed"); + tokio::time::sleep(Duration::from_millis(15)).await; + pool.gc_sweep().await; + let second = backend + .create_response( + &request, + request_context("req_gc_second", Some("chatgpt-account")), + ) + .await + .expect("second websocket response should use a fresh connection after gc"); + server.await.unwrap(); + + assert!(first.body.contains("resp_gc_first")); + assert!(second.body.contains("resp_gc_second")); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn codex_backend_client_should_ping_idle_pooled_websocket_from_background_maintenance() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let ping_count = Arc::new(AtomicUsize::new(0)); + let ping_count_for_server = Arc::clone(&ping_count); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let _first_message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_pool_background_first", + "object": "response", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + + let ping = timeout(Duration::from_secs(1), websocket.next()) + .await + .expect("background maintenance should probe the idle websocket") + .expect("background maintenance should send a websocket frame") + .expect("background maintenance frame should be valid"); + let Message::Ping(payload) = ping else { + panic!("expected background maintenance ping frame, got {ping:?}"); + }; + ping_count_for_server.fetch_add(1, Ordering::SeqCst); + websocket.send(Message::Pong(payload)).await.unwrap(); + + let _second_message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_pool_background_second", + "object": "response", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + }); + let pool = Arc::new(CodexWebSocketPool::with_config( + websocket_pool_config_for_tests( + Some(Duration::from_millis(20)), + Some(Duration::from_secs(60)), + None, + ), + )); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool)); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + request.prompt_cache_key = Some("conversation-pool-background".to_string()); + + let first = backend + .create_response( + &request, + request_context("req_pool_background_first", Some("chatgpt-account")), + ) + .await + .expect("first pooled websocket response should succeed"); + tokio::time::sleep(Duration::from_millis(80)).await; + let second = backend + .create_response( + &request, + request_context("req_pool_background_second", Some("chatgpt-account")), + ) + .await + .expect("second pooled websocket response should reuse the background-probed socket"); + server.await.unwrap(); + pool.shutdown().await; + + assert!(first.body.contains("resp_pool_background_first")); + assert!(second.body.contains("resp_pool_background_second")); + assert_eq!(ping_count.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn codex_backend_client_should_close_idle_pooled_websocket_when_account_is_evicted() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut first_websocket = accept_async(first_stream).await.unwrap(); + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_pool_evict_first", + "object": "response", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + let close = timeout(Duration::from_secs(1), first_websocket.next()) + .await + .expect("evict_account should close the idle websocket") + .expect("evict_account should send a close frame") + .expect("close frame should be valid"); + assert!(matches!(close, Message::Close(_))); + + let (second_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut second_websocket = accept_async(second_stream).await.unwrap(); + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_pool_evict_second", + "object": "response", + "usage": { + "input_tokens": 4, + "output_tokens": 1, + "total_tokens": 5 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + }); + let pool = Arc::new(CodexWebSocketPool::with_config( + websocket_pool_config_for_tests(None, None, None), + )); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool)); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + request.prompt_cache_key = Some("conversation-pool-evict".to_string()); + + let first = backend + .create_response( + &request, + request_context("req_pool_evict_first", Some("chatgpt-account")), + ) + .await + .expect("first pooled websocket response should succeed"); + pool.evict_account("chatgpt-account").await; + let second = backend + .create_response( + &request, + request_context("req_pool_evict_second", Some("chatgpt-account")), + ) + .await + .expect("second websocket response should open a fresh socket after eviction"); + server.await.unwrap(); + + assert!(first.body.contains("resp_pool_evict_first")); + assert!(second.body.contains("resp_pool_evict_second")); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn codex_backend_client_should_stop_reusing_pooled_websockets_after_shutdown() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut first_websocket = accept_async(first_stream).await.unwrap(); + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_pool_shutdown_first", + "object": "response", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + let close = timeout(Duration::from_secs(1), first_websocket.next()) + .await + .expect("shutdown should close the idle websocket") + .expect("shutdown should send a close frame") + .expect("close frame should be valid"); + assert!(matches!(close, Message::Close(_))); + + let (second_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut second_websocket = accept_async(second_stream).await.unwrap(); + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_pool_shutdown_second", + "object": "response", + "usage": { + "input_tokens": 4, + "output_tokens": 1, + "total_tokens": 5 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + }); + let pool = Arc::new(CodexWebSocketPool::with_config( + websocket_pool_config_for_tests(None, None, None), + )); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool)); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + request.prompt_cache_key = Some("conversation-pool-shutdown".to_string()); + + let first = backend + .create_response( + &request, + request_context("req_pool_shutdown_first", Some("chatgpt-account")), + ) + .await + .expect("first pooled websocket response should succeed"); + pool.shutdown().await; + let second = backend + .create_response( + &request, + request_context("req_pool_shutdown_second", Some("chatgpt-account")), + ) + .await + .expect("second websocket response should bypass the shut down pool"); + server.await.unwrap(); + + assert!(first.body.contains("resp_pool_shutdown_first")); + assert!(second.body.contains("resp_pool_shutdown_second")); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn codex_backend_client_should_close_idle_pooled_websocket_after_liveness_timeout() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut first_websocket = accept_async(first_stream).await.unwrap(); + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_pool_liveness_first", + "object": "response", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + let close = timeout(Duration::from_secs(1), first_websocket.next()) + .await + .expect("liveness timeout should close the idle websocket") + .expect("liveness timeout should send a close frame") + .expect("close frame should be valid"); + assert!(matches!(close, Message::Close(_))); + + let (second_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut second_websocket = accept_async(second_stream).await.unwrap(); + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_pool_liveness_second", + "object": "response", + "usage": { + "input_tokens": 4, + "output_tokens": 1, + "total_tokens": 5 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + }); + let pool = Arc::new(CodexWebSocketPool::with_config( + websocket_pool_config_for_tests(None, None, Some(Duration::from_millis(1))), + )); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(Arc::clone(&pool)); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + request.prompt_cache_key = Some("conversation-pool-liveness".to_string()); + + let first = backend + .create_response( + &request, + request_context("req_pool_liveness_first", Some("chatgpt-account")), + ) + .await + .expect("first pooled websocket response should succeed"); + tokio::time::sleep(Duration::from_millis(10)).await; + pool.maintain_idle_connections().await; + let second = backend + .create_response( + &request, + request_context("req_pool_liveness_second", Some("chatgpt-account")), + ) + .await + .expect("second websocket response should open a fresh socket after liveness close"); + server.await.unwrap(); + + assert!(first.body.contains("resp_pool_liveness_first")); + assert!(second.body.contains("resp_pool_liveness_second")); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn codex_backend_client_should_discard_pooled_websocket_after_upstream_error() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut first_websocket = accept_async(first_stream).await.unwrap(); + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_websocket + .send(Message::Text( + json!({ + "type": "response.failed", + "response": { + "id": "resp_pool_rate_limit", + "error": { + "code": "rate_limit_exceeded", + "message": "Rate limit reached. Please try again in 1s." + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + first_websocket.close(None).await.unwrap(); + + let (second_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut second_websocket = accept_async(second_stream).await.unwrap(); + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_pool_after_error", + "object": "response", + "usage": { + "input_tokens": 5, + "output_tokens": 2, + "total_tokens": 7 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + }); + let pool = Arc::new( + codex_proxy_adapters::codex::websocket::pool::CodexWebSocketPool::new( + 8, + Duration::from_secs(60), + ), + ); + let backend = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ) + .with_websocket_pool(pool); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + request.prompt_cache_key = Some("conversation-pool".to_string()); + + let first_error = backend + .create_response( + &request, + request_context("req_pool_error", Some("chatgpt-account")), + ) + .await + .expect_err("first pooled websocket response should surface upstream error"); + let second = backend + .create_response( + &request, + request_context("req_pool_after_error", Some("chatgpt-account")), + ) + .await + .expect("second pooled websocket response should use a fresh connection"); + server.await.unwrap(); + + let client::CodexClientError::Upstream { status, body, .. } = first_error else { + panic!("expected upstream error from first pooled websocket response"); + }; + assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS); + assert!(body.contains("rate_limit_exceeded")); + assert!(second.body.contains("resp_pool_after_error")); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 2); +} + +#[test] +fn update_manifest_should_extract_version_and_build_number() { + let update = codex_proxy_adapters::codex::fingerprint::parse_update_manifest( + r#"{"version":"26.700.111","build_number":"5002"}"#, + ) + .expect("manifest should parse"); + + assert_eq!(update.app_version, "26.700.111"); + assert_eq!(update.build_number, "5002"); +} + +#[tokio::test] +async fn fingerprint_repository_should_upsert_auto_update_record() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("fingerprints.sqlite"); + let pool = codex_proxy_platform::storage::connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let repo = codex_proxy_adapters::codex::fingerprint::FingerprintRepository::new(pool.clone()); + + repo.upsert_auto_update("26.800.1", "6001", Some("147")) + .await + .expect("first upsert"); + repo.upsert_auto_update("26.800.2", "6002", None) + .await + .expect("second upsert"); + + let stored = repo + .load_latest_auto_updated() + .await + .expect("load latest") + .expect("stored fingerprint"); + let count: (i64,) = + sqlx::query_as("select count(*) from fingerprints where id = 'auto_updated'") + .fetch_one(&pool) + .await + .expect("count row"); + + assert_eq!(count.0, 1); + assert_eq!(stored.app_version, "26.800.2"); + assert_eq!(stored.build_number, "6002"); + assert_eq!(stored.chromium_version, "146"); +} + +#[tokio::test] +async fn fingerprint_updater_should_fetch_manifest_and_persist_history() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/desktop/update.json")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_raw( + r#"{"version":"26.700.111","build_number":"5002"}"#, + "application/json", + )) + .mount(&server) + .await; + + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("fingerprints.sqlite"); + let pool = codex_proxy_platform::storage::connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let repo = codex_proxy_adapters::codex::fingerprint::FingerprintRepository::new(pool); + let updater = codex_proxy_adapters::codex::fingerprint::FingerprintUpdater::new( + reqwest::Client::new(), + repo.clone(), + format!("{}/desktop/update.json", server.uri()), + ); + + updater.poll_once().await.expect("poll once"); + + let latest = repo + .latest() + .await + .expect("latest row") + .expect("stored history"); + assert_eq!(latest.app_version, "26.700.111"); + assert_eq!(latest.build_number, "5002"); + assert_eq!( + latest.source, + codex_proxy_adapters::codex::fingerprint::CODEX_DESKTOP_UPDATE_SOURCE + ); +} + +#[tokio::test] +async fn update_checker_should_report_available_update_from_appcast() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/appcast.xml")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_raw( + r#" + + + + + + + + "#, + "application/xml", + )) + .mount(&server) + .await; + + let checker = codex_proxy_adapters::codex::fingerprint::UpdateChecker::with_client( + None, + reqwest::Client::new(), + format!("{}/appcast.xml", server.uri()), + tempfile::tempdir() + .expect("temp dir") + .path() + .join("extracted-fingerprint.json"), + "26.800.1", + "6001", + ); + + let state = checker.check_for_update().await.expect("update state"); + + assert!(state.update_available); + assert_eq!(state.latest_version.as_deref(), Some("26.900.1")); + assert_eq!(state.latest_build.as_deref(), Some("7001")); +} + +#[tokio::test] +async fn update_checker_should_apply_available_update_to_repository() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/appcast.xml")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_raw( + r#" + + + + + + + + "#, + "application/xml", + )) + .mount(&server) + .await; + + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("fingerprints.sqlite"); + let pool = codex_proxy_platform::storage::connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let repo = codex_proxy_adapters::codex::fingerprint::FingerprintRepository::new(pool); + let extracted_path = dir.path().join("extracted-fingerprint.json"); + std::fs::write( + &extracted_path, + r#"{"app_version":"26.900.1","build_number":"7001","chromium_version":"147"}"#, + ) + .expect("write extracted fingerprint"); + + let checker = codex_proxy_adapters::codex::fingerprint::UpdateChecker::with_client( + Some(repo.clone()), + reqwest::Client::new(), + format!("{}/appcast.xml", server.uri()), + extracted_path, + "26.800.1", + "6001", + ); + + let applied = checker + .check_and_apply_update() + .await + .expect("apply update"); + let stored = repo + .load_latest_auto_updated() + .await + .expect("load latest") + .expect("stored fingerprint"); + + assert!(applied); + assert_eq!(stored.app_version, "26.900.1"); + assert_eq!(stored.build_number, "7001"); + assert_eq!(stored.chromium_version, "147"); +} + +#[tokio::test] +async fn codex_backend_client_should_send_desktop_headers_and_capture_response_metadata() { + let server = wiremock::MockServer::start().await; + let sse_body = concat!( + "event: response.completed\n", + "data: {\"response\":{\"id\":\"resp_1\",\"usage\":{\"input_tokens\":2,\"output_tokens\":3,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n", + "\n", + ); + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/codex/responses")) + .and(wiremock::matchers::header( + "authorization", + "Bearer access-token", + )) + .and(wiremock::matchers::header( + "chatgpt-account-id", + "chatgpt-account", + )) + .and(wiremock::matchers::header("originator", "Codex Desktop")) + .and(wiremock::matchers::header("x-client-request-id", "req_1")) + .and(wiremock::matchers::header("x-codex-turn-state", "turn_1")) + .and(wiremock::matchers::header("cookie", "cf_clearance=old")) + .and(wiremock::matchers::body_json(json!({ + "model": "gpt-5.5", + "instructions": "", + "input": [], + "stream": true, + "store": false + }))) + .respond_with( + wiremock::ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .insert_header( + "set-cookie", + "cf_clearance=new; Domain=.chatgpt.com; Path=/", + ) + .insert_header("x-codex-turn-state", "turn_2") + .set_body_string(sse_body), + ) + .mount(&server) + .await; + let client = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + server.uri(), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "", + Vec::new(), + ); + request.force_http_sse = true; + + let response = client + .create_response( + &request, + client::CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_1", + turn_state: Some("turn_1"), + turn_metadata: None, + beta_features: None, + include_timing_metrics: None, + version: None, + codex_window_id: None, + parent_thread_id: None, + cookie_header: Some("cf_clearance=old"), + installation_id: None, + session_id: None, + }, + ) + .await + .unwrap(); + + assert_eq!( + response.usage, + Some(codex_proxy_core::protocol::codex::events::TokenUsage { + input_tokens: 2, + output_tokens: 3, + cached_tokens: 1, + image_input_tokens: 0, + image_output_tokens: 0, + total_tokens: 5, + }) + ); + assert_eq!(response.turn_state.as_deref(), Some("turn_2")); + assert_eq!( + response.set_cookie_headers, + vec!["cf_clearance=new; Domain=.chatgpt.com; Path=/".to_string()] + ); +} + +#[tokio::test] +async fn codex_backend_client_usage_should_use_original_auxiliary_headers() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/api/codex/usage")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_json(json!({ + "rate_limit": { + "limit_reached": false + } + }))) + .mount(&server) + .await; + let client = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + server.uri(), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + + let usage = client + .fetch_usage(client::CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_aux", + turn_state: Some("turn-state"), + turn_metadata: Some("turn-meta"), + beta_features: Some("feature-a"), + include_timing_metrics: Some("true"), + version: Some("26.318.11754"), + codex_window_id: Some("cw_1"), + parent_thread_id: Some("parent-1"), + cookie_header: Some("cf_clearance=old"), + installation_id: Some("install-1"), + session_id: Some("session-1"), + }) + .await + .unwrap(); + + assert_eq!(usage["rate_limit"]["limit_reached"], false); + let requests = server.received_requests().await.unwrap(); + let headers = &requests[0].headers; + assert_eq!( + headers.get("accept").and_then(|value| value.to_str().ok()), + Some("application/json") + ); + assert_eq!( + headers + .get("accept-encoding") + .and_then(|value| value.to_str().ok()), + Some("gzip, deflate") + ); + assert!(headers.get("content-type").is_none()); + assert!(headers.get("openai-beta").is_none()); + assert!(headers.get("x-openai-internal-codex-residency").is_none()); + assert!(headers.get("x-client-request-id").is_none()); + assert_eq!( + headers + .get("x-codex-installation-id") + .and_then(|value| value.to_str().ok()), + Some("install-1") + ); + assert!(headers.get("session_id").is_none()); +} + +#[tokio::test] +async fn codex_backend_client_models_should_use_original_auxiliary_headers() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/codex/models")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_json(json!({ + "models": [ + {"slug": "gpt-5.5", "title": "GPT 5.5"} + ] + }))) + .mount(&server) + .await; + let client = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + server.uri(), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + + let models = client + .fetch_models(client::CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_models", + turn_state: Some("turn-state"), + turn_metadata: None, + beta_features: None, + include_timing_metrics: None, + version: None, + codex_window_id: None, + parent_thread_id: None, + cookie_header: None, + installation_id: Some("install-1"), + session_id: Some("session-1"), + }) + .await + .unwrap(); + + assert_eq!(models.len(), 1); + let requests = server.received_requests().await.unwrap(); + let models_request = requests + .iter() + .find(|request| request.url.path() == "/codex/models") + .unwrap(); + let headers = &models_request.headers; + assert_eq!( + headers.get("accept").and_then(|value| value.to_str().ok()), + Some("application/json") + ); + assert!(headers.get("content-type").is_none()); + assert!(headers.get("openai-beta").is_none()); + assert!(headers.get("x-openai-internal-codex-residency").is_none()); + assert!(headers.get("x-client-request-id").is_none()); + assert_eq!( + headers + .get("x-codex-installation-id") + .and_then(|value| value.to_str().ok()), + Some("install-1") + ); + assert!(headers.get("session_id").is_none()); +} + +#[tokio::test] +async fn codex_backend_client_probe_should_return_models_endpoint_status() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/codex/models")) + .respond_with(wiremock::ResponseTemplate::new(204)) + .mount(&server) + .await; + let client = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + server.uri(), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + + let probe = client + .probe_models_endpoint(client::CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_probe", + 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, + installation_id: None, + session_id: None, + }) + .await + .unwrap(); + + assert_eq!(probe.status, reqwest::StatusCode::NO_CONTENT); + assert!(probe + .endpoint + .ends_with("/codex/models?client_version=26.519.81530")); +} + +#[tokio::test] +async fn codex_backend_client_should_cap_non_success_error_body_at_one_mib() { + let server = wiremock::MockServer::start().await; + let large_error_body = "x".repeat(1024 * 1024 + 17); + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/codex/responses")) + .respond_with(wiremock::ResponseTemplate::new(500).set_body_string(large_error_body)) + .mount(&server) + .await; + let client = client::CodexBackendClient::new( + client::build_reqwest_client(false).unwrap(), + server.uri(), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "", + Vec::new(), + ); + request.force_http_sse = true; + + let result = client + .create_response( + &request, + client::CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_large_error", + 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, + installation_id: None, + session_id: None, + }, + ) + .await; + + let Err(client::CodexClientError::Upstream { status, body, .. }) = result else { + panic!("expected upstream error"); + }; + assert_eq!(status, reqwest::StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(body.len(), 1024 * 1024); +} + +#[tokio::test] +async fn codex_backend_client_should_send_http_sse_headers_in_fingerprint_order() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let request = read_http_request(&mut stream).await; + write_completed_sse_response(&mut stream).await; + request + }); + + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "", + Vec::new(), + ); + request.force_http_sse = true; + request.turn_metadata = Some("turn-meta".to_string()); + request.beta_features = Some("beta-a".to_string()); + request.include_timing_metrics = Some("true".to_string()); + request.version = Some("26.519.81530".to_string()); + request.codex_window_id = Some("cw_1".to_string()); + request.parent_thread_id = Some("parent-1".to_string()); + let client = client::CodexBackendClient::new( + client::build_reqwest_client(true).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + + client + .create_response( + &request, + client::CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_order", + turn_state: Some("turn-state"), + 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: Some("cf_clearance=old"), + installation_id: Some("install-1"), + session_id: Some("session-1"), + }, + ) + .await + .unwrap(); + + let raw_request = server.await.unwrap(); + let header_names = read_header_names(&raw_request); + assert_header_subsequence( + &header_names, + &[ + "authorization", + "chatgpt-account-id", + "originator", + "user-agent", + "sec-ch-ua", + "sec-ch-ua-mobile", + "sec-ch-ua-platform", + "accept-encoding", + "accept-language", + "sec-fetch-site", + "sec-fetch-mode", + "sec-fetch-dest", + "content-type", + "cookie", + "accept", + "openai-beta", + "x-openai-internal-codex-residency", + "x-client-request-id", + "x-codex-installation-id", + "session_id", + "x-codex-window-id", + "x-codex-turn-state", + "x-codex-turn-metadata", + "x-codex-beta-features", + "x-responsesapi-include-timing-metrics", + "version", + "x-codex-parent-thread-id", + ], + ); +} + +#[tokio::test] +async fn codex_backend_client_should_send_compact_headers_in_fingerprint_order() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let request = read_http_request(&mut stream).await; + write_compact_json_response(&mut stream).await; + request + }); + let client = client::CodexBackendClient::new( + client::build_reqwest_client(true).unwrap(), + format!("http://{addr}"), + codex_proxy_core::gateway::fingerprint::Fingerprint::default_for_tests(), + ); + + client + .create_compact_response( + &codex_proxy_core::protocol::codex::responses::CodexCompactRequest { + model: "gpt-5.5".to_string(), + input: Vec::new(), + instructions: String::new(), + tools: None, + parallel_tool_calls: None, + reasoning: None, + text: None, + }, + client::CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_compact", + 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: Some("cf_clearance=old"), + installation_id: Some("install-1"), + session_id: None, + }, + ) + .await + .unwrap(); + + let raw_request = server.await.unwrap(); + let header_names = read_header_names(&raw_request); + assert_header_subsequence( + &header_names, + &[ + "authorization", + "chatgpt-account-id", + "originator", + "user-agent", + "sec-ch-ua", + "sec-ch-ua-mobile", + "sec-ch-ua-platform", + "accept-encoding", + "accept-language", + "sec-fetch-site", + "sec-fetch-mode", + "sec-fetch-dest", + "content-type", + "cookie", + "openai-beta", + "x-openai-internal-codex-residency", + "x-client-request-id", + "x-codex-installation-id", + ], + ); +} + +#[tokio::test] +async fn build_reqwest_client_should_reuse_cached_connection_pool() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut first_stream, _) = listener.accept().await.unwrap(); + read_http_request(&mut first_stream).await; + write_empty_http_response(&mut first_stream).await; + + tokio::select! { + request = read_http_request(&mut first_stream) => { + write_empty_http_response(&mut first_stream).await; + !request.is_empty() + } + accepted = listener.accept() => { + let (mut second_stream, _) = accepted.unwrap(); + read_http_request(&mut second_stream).await; + write_empty_http_response(&mut second_stream).await; + false + } + _ = tokio::time::sleep(Duration::from_millis(500)) => false, + } + }); + + let url = format!("http://{addr}/reuse"); + let first_client = client::build_reqwest_client(false).unwrap(); + first_client + .get(&url) + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + let second_client = client::build_reqwest_client(false).unwrap(); + second_client + .get(&url) + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + + assert!(server.await.unwrap()); +} + +fn read_header_names(request: &str) -> Vec { + request + .lines() + .skip(1) + .take_while(|line| !line.is_empty()) + .filter_map(|line| { + line.split_once(':') + .map(|(name, _)| name.to_ascii_lowercase()) + }) + .collect() +} + +fn assert_header_subsequence(actual: &[String], expected: &[&str]) { + let mut offset = 0; + for expected_name in expected { + let Some(position) = actual[offset..] + .iter() + .position(|actual_name| actual_name == expected_name) + else { + panic!("missing header {expected_name}; actual order: {actual:?}"); + }; + offset += position + 1; + } +} + +async fn read_http_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() +} + +fn websocket_opening_header(request: &str, name: &str) -> Option { + request.lines().find_map(|line| { + let (header_name, value) = line.split_once(':')?; + header_name + .eq_ignore_ascii_case(name) + .then(|| value.trim().to_string()) + }) +} + +async fn write_empty_http_response(stream: &mut TcpStream) { + stream + .write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n") + .await + .unwrap(); +} + +async fn write_completed_sse_response(stream: &mut TcpStream) { + let body = concat!( + "event: response.completed\n", + "data: {\"response\":{\"id\":\"resp_order\",\"status\":\"completed\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n", + "\n", + ); + stream + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .await + .unwrap(); +} + +fn request_context<'a>( + request_id: &'a str, + account_id: Option<&'a str>, +) -> client::CodexRequestContext<'a> { + client::CodexRequestContext { + access_token: "access-token", + account_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, + installation_id: None, + session_id: None, + } +} + +fn prepared_websocket_request( + base_url: &str, +) -> codex_proxy_adapters::codex::websocket::connect::CodexWebSocketRequest { + let request = codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request( + base_url, + "dGhlIHNhbXBsZSBub25jZQ==", + vec![("authorization".to_string(), "Bearer access-token".to_string())], + &request, + ) + .expect("payload should serialize") +} + +fn pooled_websocket_request( + conversation_id: &str, +) -> codex_proxy_core::protocol::codex::responses::CodexResponsesRequest { + let mut request = + codex_proxy_core::protocol::codex::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + request.prompt_cache_key = Some(conversation_id.to_string()); + request +} + +fn completed_websocket_response( + response_id: &str, + input_tokens: u64, + output_tokens: u64, +) -> String { + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "object": "response", + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens + } + } + }) + .to_string() +} + +fn assert_substrings_appear_in_order(haystack: &str, needles: &[&str]) { + let mut cursor = 0; + for needle in needles { + let Some(offset) = haystack[cursor..].find(needle) else { + panic!("expected substring {needle:?} after byte {cursor} in:\n{haystack}"); + }; + cursor += offset + needle.len(); + } +} + +fn websocket_pool_config_for_tests( + maintenance_interval: Option, + ping_interval: Option, + liveness_timeout: Option, +) -> CodexWebSocketPoolConfig { + CodexWebSocketPoolConfig { + enabled: true, + max_age: Duration::from_secs(60), + max_per_account: 8, + maintenance_interval, + ping_interval, + ping_timeout: Duration::from_secs(1), + liveness_timeout, + } +} + +async fn write_compact_json_response(stream: &mut TcpStream) { + stream + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 15\r\n\r\n{\"id\":\"resp_1\"}", + ) + .await + .unwrap(); +} diff --git a/crates/adapters/tests/cookies.rs b/crates/adapters/tests/cookies.rs new file mode 100644 index 000000000..5d6521e00 --- /dev/null +++ b/crates/adapters/tests/cookies.rs @@ -0,0 +1,93 @@ +use chrono::{Duration, Utc}; +use codex_proxy_adapters::sqlite::{ + accounts::{NewAccount, SqliteAccountStore}, + cookies::SqliteCookieStore, +}; +use codex_proxy_core::accounts::model::AccountStatus; +use codex_proxy_platform::{crypto::SecretBox, storage::connect_sqlite}; +use secrecy::SecretString; +use sqlx::SqlitePool; + +#[tokio::test] +async fn cookie_store_should_capture_cf_clearance_and_replay_by_domain_and_path() { + let (pool, secret_box, _dir) = sqlite_cookie_store_parts("captured-cookies.sqlite").await; + seed_account(&pool, &secret_box, "acct_a").await; + let store = SqliteCookieStore::new(pool, secret_box); + + store + .capture_set_cookie("acct_a", "cf_clearance=root; Domain=.chatgpt.com; Path=/") + .await + .unwrap(); + store + .capture_set_cookie( + "acct_a", + "cf_clearance=codex; Domain=.chatgpt.com; Path=/codex", + ) + .await + .unwrap(); + + assert_eq!( + store + .cookie_header_for_request("acct_a", "chatgpt.com", "/codex/responses") + .await + .unwrap(), + Some("cf_clearance=codex; cf_clearance=root".to_string()) + ); + assert_eq!( + store + .cookie_header_for_request("acct_a", "chatgpt.com", "/api/codex/usage") + .await + .unwrap(), + Some("cf_clearance=root".to_string()) + ); +} + +#[tokio::test] +async fn cookie_store_should_ignore_non_capturable_set_cookie_headers() { + let (pool, secret_box, _dir) = sqlite_cookie_store_parts("ignored-cookies.sqlite").await; + seed_account(&pool, &secret_box, "acct_a").await; + let store = SqliteCookieStore::new(pool.clone(), secret_box); + + store + .capture_set_cookie("acct_a", "__cf_bm=bot-session; Domain=.chatgpt.com; Path=/") + .await + .unwrap(); + + let stored_count: (i64,) = + sqlx::query_as("select count(*) from account_cookies where account_id = ?") + .bind("acct_a") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(stored_count.0, 0); + assert_eq!( + store.cookie_header("acct_a", "chatgpt.com").await.unwrap(), + None + ); +} + +async fn sqlite_cookie_store_parts(name: &str) -> (SqlitePool, SecretBox, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join(name); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + (pool, SecretBox::new([21u8; 32]), dir) +} + +async fn seed_account(pool: &SqlitePool, secret_box: &SecretBox, id: &str) { + SqliteAccountStore::new(pool.clone(), secret_box.clone()) + .insert(NewAccount { + id: id.to_string(), + email: Some(format!("{id}@example.com")), + account_id: Some(format!("chatgpt-{id}")), + user_id: Some(format!("user-{id}")), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(format!("access-{id}").into()), + refresh_token: None, + access_token_expires_at: Some(Utc::now() + Duration::hours(1)), + status: AccountStatus::Active, + }) + .await + .unwrap(); +} diff --git a/crates/adapters/tests/events.rs b/crates/adapters/tests/events.rs new file mode 100644 index 000000000..8eef1f89d --- /dev/null +++ b/crates/adapters/tests/events.rs @@ -0,0 +1,71 @@ +use codex_proxy_adapters::sqlite::events::{EventLogFilter, SqliteEventLogStore}; +use codex_proxy_core::events::model::{EventLevel, EventLog}; +use codex_proxy_platform::storage::connect_sqlite; + +#[tokio::test] +async fn event_log_store_should_filter_and_cursor_page_events() { + let (store, _dir) = event_log_store("event-logs.sqlite").await; + let mut matching = EventLog::new("request", EventLevel::Error, "upstream timeout"); + matching.id = "log_matching".to_string(); + matching.request_id = Some("req_1".to_string()); + matching.route = Some("/v1/responses".to_string()); + store.append(&matching).await.unwrap(); + store + .append(&EventLog::new( + "request", + EventLevel::Info, + "upstream timeout", + )) + .await + .unwrap(); + store + .append(&EventLog::new( + "account", + EventLevel::Error, + "upstream timeout", + )) + .await + .unwrap(); + + let page = store + .list( + EventLogFilter { + kind: Some("request".to_string()), + level: Some(EventLevel::Error), + search: Some("timeout".to_string()), + ..EventLogFilter::default() + }, + None, + 1, + ) + .await + .unwrap(); + + assert_eq!(page.items.len(), 1); + assert_eq!(page.items[0].id, "log_matching"); + assert_eq!(page.items[0].request_id.as_deref(), Some("req_1")); + assert_eq!(page.items[0].route.as_deref(), Some("/v1/responses")); +} + +#[tokio::test] +async fn event_log_store_should_get_count_and_clear_events() { + let (store, _dir) = event_log_store("event-logs-clear.sqlite").await; + let mut event = EventLog::new("request", EventLevel::Warn, "detail"); + event.id = "log_detail".to_string(); + store.append(&event).await.unwrap(); + + let loaded = store.get("log_detail").await.unwrap().unwrap(); + assert_eq!(loaded.id, "log_detail"); + assert_eq!(store.count().await.unwrap(), 1); + assert_eq!(store.clear().await.unwrap(), 1); + assert_eq!(store.count().await.unwrap(), 0); +} + +async fn event_log_store(db_name: &str) -> (SqliteEventLogStore, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join(db_name); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .unwrap(); + (SqliteEventLogStore::new(pool), dir) +} diff --git a/crates/adapters/tests/models.rs b/crates/adapters/tests/models.rs new file mode 100644 index 000000000..a50910308 --- /dev/null +++ b/crates/adapters/tests/models.rs @@ -0,0 +1,93 @@ +use chrono::Utc; +use codex_proxy_adapters::sqlite::{accounts::SqliteAccountStore, models::ModelSnapshotRepository}; +use codex_proxy_core::accounts::{model::AccountStatus, ports::AccountStore}; +use codex_proxy_core::models::model::{BackendModelEntry, ModelConfig, ModelPlanSnapshot}; +use codex_proxy_platform::crypto::SecretBox; +use secrecy::SecretString; + +#[tokio::test] +async fn model_snapshot_repository_should_replace_and_load_plan_snapshots() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("models.sqlite"); + let pool = codex_proxy_platform::storage::connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .unwrap(); + let repo = ModelSnapshotRepository::new(pool); + let snapshot = ModelPlanSnapshot::from_backend_entries( + "team", + vec![BackendModelEntry { + id: Some("gpt-team".to_string()), + name: Some("GPT Team".to_string()), + ..BackendModelEntry::default() + }], + ); + + repo.replace_plan_snapshot(&snapshot).await.unwrap(); + let loaded = repo.list_plan_snapshots().await.unwrap(); + + assert_eq!(loaded.len(), 1); + assert_eq!(loaded[0].plan_type, "team"); + assert_eq!(loaded[0].models[0].id, "gpt-team"); + assert_eq!(loaded[0].models[0].source, "backend"); +} + +#[test] +fn model_config_should_be_usable_from_core_models() { + let config = ModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: None, + service_tier: None, + aliases: Default::default(), + }; + + assert_eq!(config.default_model, "gpt-5.5"); +} + +#[tokio::test] +async fn sqlite_account_store_should_list_pool_accounts() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("accounts.sqlite"); + let pool = codex_proxy_platform::storage::connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let secret_box = SecretBox::new([41u8; 32]); + let now = Utc::now().to_rfc3339(); + let access_token_cipher = secret_box + .encrypt(&SecretString::new("access-token".to_string().into())) + .expect("encrypt access token"); + let refresh_token_cipher = secret_box + .encrypt(&SecretString::new("refresh-token".to_string().into())) + .expect("encrypt refresh token"); + sqlx::query( + "insert into accounts ( + id, email, account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, + access_token_expires_at, status, added_at, updated_at + ) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind("acct_1") + .bind("user@example.com") + .bind("chatgpt-account") + .bind(Option::::None) + .bind("primary") + .bind("plus") + .bind(access_token_cipher) + .bind(Some(refresh_token_cipher)) + .bind(Option::::None) + .bind(match AccountStatus::Active { + AccountStatus::Active => "active", + _ => unreachable!(), + }) + .bind(&now) + .bind(&now) + .execute(&pool) + .await + .expect("insert account"); + + let store = SqliteAccountStore::new(pool, secret_box); + let accounts = store.list_pool_accounts().await.expect("list accounts"); + + assert_eq!(accounts.len(), 1); + assert_eq!(accounts[0].id, "acct_1"); + assert_eq!(accounts[0].access_token, "access-token"); + assert_eq!(accounts[0].plan_type.as_deref(), Some("plus")); +} diff --git a/tests/codex_gateway/oauth_refresh.rs b/crates/adapters/tests/oauth.rs similarity index 84% rename from tests/codex_gateway/oauth_refresh.rs rename to crates/adapters/tests/oauth.rs index 5d1142cb6..05f6ff9c0 100644 --- a/tests/codex_gateway/oauth_refresh.rs +++ b/crates/adapters/tests/oauth.rs @@ -3,12 +3,14 @@ use wiremock::{ Mock, MockServer, ResponseTemplate, }; -use codex_proxy_rs::codex::gateway::oauth::{ - OAuthClient, OAuthConfig, OAuthError, OpenAiOAuthRefresher, TokenRefresher, +use codex_proxy_adapters::oauth::openai::OpenAiOAuthClient; +use codex_proxy_core::auth::{ + oauth::{OAuthConfig, OAuthError}, + ports::{OAuthClient, TokenRefresher}, }; #[tokio::test] -async fn openai_oauth_refresher_should_exchange_refresh_token_with_form_body() { +async fn openai_oauth_client_should_exchange_refresh_token_with_form_body() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/oauth/token")) @@ -18,17 +20,9 @@ async fn openai_oauth_refresher_should_exchange_refresh_token_with_form_body() { }))) .mount(&server) .await; - let refresher = OpenAiOAuthRefresher::new( - reqwest::Client::builder().no_proxy().build().unwrap(), - OAuthConfig { - client_id: "codex-client".to_string(), - auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), - device_code_endpoint: "https://auth.openai.com/oauth/device/code".to_string(), - token_endpoint: format!("{}/oauth/token", server.uri()), - }, - ); + let client = test_oauth_client(&server); - let tokens = refresher.refresh("refresh-secret").await.unwrap(); + let tokens = client.refresh("refresh-secret").await.unwrap(); assert_eq!(tokens.access_token, "new-access"); assert_eq!(tokens.refresh_token.as_deref(), Some("new-refresh")); @@ -123,8 +117,8 @@ async fn openai_oauth_client_should_map_device_poll_pending_errors() { assert!(body.contains("client_id=codex-client")); } -fn test_oauth_client(server: &MockServer) -> OpenAiOAuthRefresher { - OpenAiOAuthRefresher::new( +fn test_oauth_client(server: &MockServer) -> OpenAiOAuthClient { + OpenAiOAuthClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), OAuthConfig { client_id: "codex-client".to_string(), diff --git a/crates/adapters/tests/refresh_leases.rs b/crates/adapters/tests/refresh_leases.rs new file mode 100644 index 000000000..255260936 --- /dev/null +++ b/crates/adapters/tests/refresh_leases.rs @@ -0,0 +1,68 @@ +use chrono::{Duration, TimeZone, Utc}; +use codex_proxy_adapters::sqlite::{ + accounts::{NewAccount, SqliteAccountStore}, + refresh_leases::SqliteRefreshLeaseStore, +}; +use codex_proxy_core::accounts::model::AccountStatus; +use codex_proxy_platform::{crypto::SecretBox, storage::connect_sqlite}; +use secrecy::SecretString; + +#[tokio::test] +async fn refresh_lease_store_should_acquire_release_and_respect_expiry() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("refresh-leases.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let secret_box = SecretBox::new([31u8; 32]); + let accounts = SqliteAccountStore::new(pool.clone(), secret_box); + accounts + .insert(NewAccount { + id: "acct-lease".to_string(), + email: None, + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access".to_string().into()), + refresh_token: Some(SecretString::new("refresh".to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let leases = SqliteRefreshLeaseStore::new(pool); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 8, 0, 0).unwrap(); + let first_expires_at = now + Duration::minutes(5); + let second_expires_at = now + Duration::minutes(10); + + assert!(leases + .try_acquire("acct-lease", "owner-a", first_expires_at, now) + .await + .expect("first owner should acquire")); + assert!(!leases + .try_acquire("acct-lease", "owner-b", second_expires_at, now) + .await + .expect("second owner should be blocked")); + assert!(leases + .try_acquire( + "acct-lease", + "owner-b", + second_expires_at, + first_expires_at + Duration::seconds(1), + ) + .await + .expect("second owner should acquire expired lease")); + assert!(!leases + .release("acct-lease", "owner-a") + .await + .expect("wrong owner should not release")); + assert!(leases + .release("acct-lease", "owner-b") + .await + .expect("current owner should release")); + assert!(leases + .try_acquire("acct-lease", "owner-c", second_expires_at, now) + .await + .expect("lease should be acquirable after release")); +} diff --git a/crates/adapters/tests/session_affinity.rs b/crates/adapters/tests/session_affinity.rs new file mode 100644 index 000000000..beb9c5741 --- /dev/null +++ b/crates/adapters/tests/session_affinity.rs @@ -0,0 +1,102 @@ +use chrono::{Duration, TimeZone, Utc}; +use codex_proxy_adapters::sqlite::session_affinity::SqliteSessionAffinityStore; +use codex_proxy_core::serving::affinity::SessionAffinityEntry; +use codex_proxy_platform::storage::connect_sqlite; + +#[tokio::test] +async fn session_affinity_store_should_upsert_and_list_active_records() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("session-affinity.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteSessionAffinityStore::new(pool); + let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(); + + store + .upsert( + "resp_1", + &SessionAffinityEntry { + account_id: "acct_1".to_string(), + conversation_id: "conv_1".to_string(), + turn_state: Some("turn_1".to_string()), + instructions_hash: Some("hash_1".to_string()), + input_tokens: Some(128), + function_call_ids: vec!["call_a".to_string(), "call_b".to_string()], + variant_hash: Some("variant_1".to_string()), + created_at: now, + }, + Duration::hours(4), + ) + .await + .expect("affinity should be stored"); + + let records = store + .list_active(now + Duration::minutes(1)) + .await + .expect("active affinities should load"); + + assert_eq!(records.len(), 1); + assert_eq!(records[0].response_id, "resp_1"); + assert_eq!(records[0].entry.account_id, "acct_1"); + assert_eq!(records[0].entry.function_call_ids, vec!["call_a", "call_b"]); +} + +#[tokio::test] +async fn session_affinity_store_should_delete_expired_records() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("session-affinity-cleanup.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteSessionAffinityStore::new(pool); + let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(); + + store + .upsert( + "resp_expired", + &SessionAffinityEntry { + account_id: "acct_1".to_string(), + conversation_id: "conv_1".to_string(), + turn_state: None, + instructions_hash: None, + input_tokens: None, + function_call_ids: Vec::new(), + variant_hash: None, + created_at: now - Duration::hours(5), + }, + Duration::hours(1), + ) + .await + .expect("expired affinity should be stored"); + store + .upsert( + "resp_active", + &SessionAffinityEntry { + account_id: "acct_2".to_string(), + conversation_id: "conv_2".to_string(), + turn_state: None, + instructions_hash: None, + input_tokens: None, + function_call_ids: Vec::new(), + variant_hash: None, + created_at: now, + }, + Duration::hours(1), + ) + .await + .expect("active affinity should be stored"); + + let deleted = store + .delete_expired(now) + .await + .expect("expired affinities should be deleted"); + let records = store + .list_active(now) + .await + .expect("active affinities should load"); + + assert_eq!(deleted, 1); + assert_eq!(records.len(), 1); + assert_eq!(records[0].response_id, "resp_active"); +} diff --git a/crates/assets/Cargo.toml b/crates/assets/Cargo.toml new file mode 100644 index 000000000..5e791e88c --- /dev/null +++ b/crates/assets/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "codex-proxy-assets" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +rust-version.workspace = true + +[dependencies] +axum.workspace = true +tokio.workspace = true + +[dev-dependencies] +tempfile = "3.27.0" +tokio.workspace = true +tower.workspace = true + +[lints] +workspace = true diff --git a/crates/assets/src/headers.rs b/crates/assets/src/headers.rs new file mode 100644 index 000000000..93d995343 --- /dev/null +++ b/crates/assets/src/headers.rs @@ -0,0 +1,60 @@ +//! 静态资源安全头。 + +use axum::http::{header, HeaderMap, HeaderValue}; + +const SPA_CACHE_CONTROL: &str = "no-cache"; +const ASSET_CACHE_CONTROL: &str = "public, max-age=31536000, immutable"; +const CONTENT_SECURITY_POLICY: &str = concat!( + "default-src 'self'; ", + "script-src 'self'; ", + "style-src 'self' 'unsafe-inline'; ", + "img-src 'self' data:; ", + "font-src 'self' data:; ", + "connect-src 'self'; ", + "base-uri 'self'; ", + "frame-ancestors 'none'" +); + +/// 返回给定请求路径的缓存策略。 +pub fn cache_control_for_path(path: &str) -> &'static str { + if path.starts_with("/assets/") { + ASSET_CACHE_CONTROL + } else { + SPA_CACHE_CONTROL + } +} + +/// 给静态资源响应写入缓存和安全头。 +pub fn apply_static_headers(headers: &mut HeaderMap, request_path: &str) { + headers.insert( + header::CACHE_CONTROL, + HeaderValue::from_static(cache_control_for_path(request_path)), + ); + headers.insert( + header::X_CONTENT_TYPE_OPTIONS, + HeaderValue::from_static("nosniff"), + ); + headers.insert("x-frame-options", HeaderValue::from_static("DENY")); + headers.insert("referrer-policy", HeaderValue::from_static("no-referrer")); + headers.insert( + header::CONTENT_SECURITY_POLICY, + HeaderValue::from_static(CONTENT_SECURITY_POLICY), + ); +} + +/// 根据文件扩展名返回常见前端资源的 content type。 +pub fn content_type_for_path(path: &str) -> HeaderValue { + if path.ends_with(".html") || path == "/" { + HeaderValue::from_static("text/html; charset=utf-8") + } else if path.ends_with(".js") { + HeaderValue::from_static("text/javascript; charset=utf-8") + } else if path.ends_with(".css") { + HeaderValue::from_static("text/css; charset=utf-8") + } else if path.ends_with(".json") { + HeaderValue::from_static("application/json") + } else if path.ends_with(".svg") { + HeaderValue::from_static("image/svg+xml") + } else { + HeaderValue::from_static("application/octet-stream") + } +} diff --git a/crates/assets/src/lib.rs b/crates/assets/src/lib.rs new file mode 100644 index 000000000..72929b6be --- /dev/null +++ b/crates/assets/src/lib.rs @@ -0,0 +1,4 @@ +//! 前端静态资源层，提供 SPA 回退和静态文件服务辅助。 + +pub mod headers; +pub mod router; diff --git a/crates/assets/src/router.rs b/crates/assets/src/router.rs new file mode 100644 index 000000000..2b61e8b5c --- /dev/null +++ b/crates/assets/src/router.rs @@ -0,0 +1,87 @@ +//! 资产路由。 + +use std::{ + path::{Component, Path, PathBuf}, + sync::Arc, +}; + +use axum::{ + body::Body, + extract::{Path as AxumPath, State}, + http::{header, StatusCode}, + response::{IntoResponse, Response}, + routing::get, + Router, +}; + +use crate::headers::{apply_static_headers, content_type_for_path}; + +#[derive(Clone)] +struct AssetState { + dist_dir: Arc, +} + +/// 构造前端静态资源路由。 +pub fn spa_router(dist_dir: impl AsRef) -> Router { + let state = AssetState { + dist_dir: Arc::new(dist_dir.as_ref().to_path_buf()), + }; + + Router::new() + .route("/", get(serve_index)) + .route("/assets/{*path}", get(serve_asset)) + .fallback(serve_spa_fallback) + .with_state(state) +} + +async fn serve_index(State(state): State) -> Response { + serve_file(&state.dist_dir.join("index.html"), "/").await +} + +async fn serve_spa_fallback(State(state): State) -> Response { + serve_file(&state.dist_dir.join("index.html"), "/").await +} + +async fn serve_asset( + State(state): State, + AxumPath(path): AxumPath, +) -> Response { + let Some(relative_path) = safe_relative_path(&path) else { + return StatusCode::NOT_FOUND.into_response(); + }; + let request_path = format!("/assets/{path}"); + + serve_file( + &state.dist_dir.join("assets").join(relative_path), + &request_path, + ) + .await +} + +async fn serve_file(path: &Path, request_path: &str) -> Response { + let Ok(bytes) = tokio::fs::read(path).await else { + return StatusCode::NOT_FOUND.into_response(); + }; + + let mut response = Response::new(Body::from(bytes)); + let path_for_content_type = path.to_string_lossy(); + let headers = response.headers_mut(); + headers.insert( + header::CONTENT_TYPE, + content_type_for_path(&path_for_content_type), + ); + apply_static_headers(headers, request_path); + response +} + +fn safe_relative_path(path: &str) -> Option { + let mut sanitized = PathBuf::new(); + for component in Path::new(path).components() { + match component { + Component::Normal(segment) => sanitized.push(segment), + _ => return None, + } + } + + (!sanitized.as_os_str().is_empty()).then_some(sanitized) +} diff --git a/crates/assets/tests/assets.rs b/crates/assets/tests/assets.rs new file mode 100644 index 000000000..20fe9a7ca --- /dev/null +++ b/crates/assets/tests/assets.rs @@ -0,0 +1,94 @@ +use axum::{ + body::{to_bytes, Body}, + http::{header, Request, StatusCode}, +}; +use tower::ServiceExt; + +#[test] +fn asset_headers_should_distinguish_spa_and_fingerprinted_assets() { + assert_eq!( + codex_proxy_assets::headers::cache_control_for_path("/"), + "no-cache" + ); + assert_eq!( + codex_proxy_assets::headers::cache_control_for_path("/assets/app.abc123.js"), + "public, max-age=31536000, immutable" + ); +} + +#[tokio::test] +async fn asset_router_should_serve_index_and_spa_fallback() { + let dir = tempfile::tempdir().expect("temp dir"); + std::fs::write( + dir.path().join("index.html"), + "
Codex Proxy
", + ) + .expect("index should be written"); + std::fs::create_dir(dir.path().join("assets")).expect("assets dir should be created"); + + let router = codex_proxy_assets::router::spa_router(dir.path()); + + let index = router + .clone() + .oneshot(request("/")) + .await + .expect("index response"); + let fallback = router + .oneshot(request("/admin/accounts")) + .await + .expect("fallback response"); + + assert_eq!(index.status(), StatusCode::OK); + assert_eq!(fallback.status(), StatusCode::OK); + assert_body_contains(index, "Codex Proxy").await; + assert_body_contains(fallback, "Codex Proxy").await; +} + +#[tokio::test] +async fn asset_router_should_serve_assets_with_cache_and_security_headers() { + let dir = tempfile::tempdir().expect("temp dir"); + std::fs::write(dir.path().join("index.html"), "").expect("write index"); + std::fs::create_dir(dir.path().join("assets")).expect("assets dir should be created"); + std::fs::write( + dir.path().join("assets").join("app.abc123.js"), + "window.__codexProxy = true;", + ) + .expect("asset should be written"); + + let response = codex_proxy_assets::router::spa_router(dir.path()) + .oneshot(request("/assets/app.abc123.js")) + .await + .expect("asset response"); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.headers().get(header::CACHE_CONTROL), + Some(&header::HeaderValue::from_static( + "public, max-age=31536000, immutable" + )) + ); + assert_eq!( + response.headers().get(header::X_CONTENT_TYPE_OPTIONS), + Some(&header::HeaderValue::from_static("nosniff")) + ); + assert_body_contains(response, "__codexProxy").await; +} + +fn request(uri: &str) -> Request { + Request::builder() + .uri(uri) + .body(Body::empty()) + .expect("request should build") +} + +async fn assert_body_contains(response: axum::response::Response, expected: &str) { + let body = to_bytes(response.into_body(), usize::MAX) + .await + .expect("body should collect"); + let text = std::str::from_utf8(&body).expect("body should be utf8"); + + assert!( + text.contains(expected), + "body should contain `{expected}`, got `{text}`" + ); +} diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml new file mode 100644 index 000000000..81e86a82e --- /dev/null +++ b/crates/core/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "codex-proxy-core" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +rust-version.workspace = true + +[dependencies] +async-trait.workspace = true +base64.workspace = true +chrono.workspace = true +hex.workspace = true +indexmap.workspace = true +rand.workspace = true +serde.workspace = true +serde_json.workspace = true +sha2.workspace = true +thiserror.workspace = true +tokio.workspace = true +tracing.workspace = true +uuid.workspace = true + +[lints] +workspace = true + +[dev-dependencies] +tokio.workspace = true diff --git a/crates/core/src/accounts/cloudflare.rs b/crates/core/src/accounts/cloudflare.rs new file mode 100644 index 000000000..7da32a882 --- /dev/null +++ b/crates/core/src/accounts/cloudflare.rs @@ -0,0 +1,127 @@ +//! Cloudflare recovery state for accounts. + +use std::{collections::HashMap, sync::Arc}; + +use chrono::{DateTime, Duration, Utc}; +use tokio::sync::RwLock; + +const PATH_BLOCK_THRESHOLD: u32 = 3; +const PATH_BLOCK_STALE_AFTER: Duration = Duration::hours(1); + +#[derive(Debug, Clone, Copy)] +struct PathBlockState { + count: u32, + last_at: DateTime, +} + +/// Tracks per-account Cloudflare path-block failures. +#[derive(Debug, Clone, Default)] +pub struct CloudflarePathBlockTracker { + counts: Arc>>, +} + +impl CloudflarePathBlockTracker { + /// Creates an empty path-block tracker. + pub fn new() -> Self { + Self::default() + } + + /// Records one path-block failure and returns the current non-stale count. + pub async fn record_path_block(&self, account_id: &str, now: DateTime) -> u32 { + let mut counts = self.counts.write().await; + let count = counts + .get(account_id) + .filter(|state| now.signed_duration_since(state.last_at) <= PATH_BLOCK_STALE_AFTER) + .map(|state| state.count.saturating_add(1)) + .unwrap_or(1); + counts.insert( + account_id.to_string(), + PathBlockState { + count, + last_at: now, + }, + ); + count + } + + /// Clears any tracked path-block failures for an account. + pub async fn reset(&self, account_id: &str) { + self.counts.write().await.remove(account_id); + } + + /// Returns the current non-stale path-block count for an account. + pub async fn count(&self, account_id: &str, now: DateTime) -> u32 { + self.counts + .read() + .await + .get(account_id) + .filter(|state| now.signed_duration_since(state.last_at) <= PATH_BLOCK_STALE_AFTER) + .map(|state| state.count) + .unwrap_or_default() + } + + /// Returns whether the account should be disabled at the current count. + pub async fn should_disable(&self, account_id: &str, now: DateTime) -> bool { + self.count(account_id, now).await >= PATH_BLOCK_THRESHOLD + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn record_path_block_should_disable_after_three_recent_failures() { + let tracker = CloudflarePathBlockTracker::new(); + let now = Utc::now(); + + assert_eq!(tracker.record_path_block("acct_1", now).await, 1); + assert_eq!( + tracker + .record_path_block("acct_1", now + Duration::seconds(1)) + .await, + 2 + ); + assert_eq!( + tracker + .record_path_block("acct_1", now + Duration::seconds(2)) + .await, + 3 + ); + + assert!( + tracker + .should_disable("acct_1", now + Duration::seconds(2)) + .await + ); + } + + #[tokio::test] + async fn record_path_block_should_ignore_stale_failures() { + let tracker = CloudflarePathBlockTracker::new(); + let now = Utc::now(); + + assert_eq!(tracker.record_path_block("acct_1", now).await, 1); + + assert_eq!( + tracker + .record_path_block("acct_1", now + Duration::hours(2)) + .await, + 1 + ); + } + + #[tokio::test] + async fn reset_should_clear_account_count() { + let tracker = CloudflarePathBlockTracker::new(); + let now = Utc::now(); + tracker.record_path_block("acct_1", now).await; + tracker + .record_path_block("acct_1", now + Duration::seconds(1)) + .await; + + tracker.reset("acct_1").await; + + assert_eq!(tracker.count("acct_1", now + Duration::seconds(2)).await, 0); + } +} diff --git a/crates/core/src/accounts/cookies.rs b/crates/core/src/accounts/cookies.rs new file mode 100644 index 000000000..b2f39a372 --- /dev/null +++ b/crates/core/src/accounts/cookies.rs @@ -0,0 +1,21 @@ +//! Cookie 捕获与重放策略。 + +/// 单个 Cookie 条目。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CookieEntry { + /// Cookie 名称。 + pub name: String, + /// Cookie 值。 + pub value: String, +} + +/// 将 Cookie 条目组装为 HTTP Cookie 头。 +pub fn cookie_header(cookies: &[CookieEntry]) -> Option { + let header = cookies + .iter() + .filter(|cookie| !cookie.name.trim().is_empty()) + .map(|cookie| format!("{}={}", cookie.name.trim(), cookie.value)) + .collect::>() + .join("; "); + (!header.is_empty()).then_some(header) +} diff --git a/src/codex/accounts/jwt.rs b/crates/core/src/accounts/jwt.rs similarity index 71% rename from src/codex/accounts/jwt.rs rename to crates/core/src/accounts/jwt.rs index c90200e49..0307fe7bb 100644 --- a/src/codex/accounts/jwt.rs +++ b/crates/core/src/accounts/jwt.rs @@ -1,14 +1,19 @@ use base64::Engine; -use chrono::{DateTime, Utc}; +use chrono::{DateTime, TimeZone, Utc}; use serde_json::Value; +/// JWT 过期状态。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum JwtExpiry { + /// token 已过期。 Expired, + /// token 仍然有效。 Valid, + /// token 缺失、格式错误或不包含可解析的 exp。 MissingOrInvalid, } +/// 按给定时间点判断 JWT 的 `exp` 是否已过期。 pub fn jwt_expiry(token: &str, now: DateTime) -> JwtExpiry { let Some(exp) = jwt_exp(token) else { return JwtExpiry::MissingOrInvalid; @@ -20,6 +25,11 @@ pub fn jwt_expiry(token: &str, now: DateTime) -> JwtExpiry { } } +/// 读取 JWT `exp` 并转换成 UTC 时间。 +pub fn jwt_expiration(token: &str) -> Option> { + Utc.timestamp_opt(jwt_exp(token)?, 0).single() +} + fn jwt_exp(token: &str) -> Option { let mut parts = token.split('.'); let _header = parts.next()?; @@ -42,7 +52,7 @@ mod tests { use chrono::{TimeZone, Utc}; use serde_json::json; - use super::{jwt_expiry, JwtExpiry}; + use super::{jwt_expiration, jwt_expiry, JwtExpiry}; #[test] fn jwt_expiry_should_detect_expired_valid_and_invalid_tokens() { @@ -59,6 +69,16 @@ mod tests { assert_eq!(jwt_expiry("not-a-jwt", now), JwtExpiry::MissingOrInvalid); } + #[test] + fn jwt_expiration_should_return_exp_as_utc_datetime() { + let expires_at = Utc.with_ymd_and_hms(2026, 6, 14, 8, 0, 0).unwrap(); + + assert_eq!( + jwt_expiration(&test_jwt(expires_at.timestamp())), + Some(expires_at) + ); + } + fn test_jwt(exp: i64) -> String { let header = json!({"alg": "none", "typ": "JWT"}); let payload = json!({ "exp": exp }); diff --git a/crates/core/src/accounts/lifecycle.rs b/crates/core/src/accounts/lifecycle.rs new file mode 100644 index 000000000..b9d66c6b3 --- /dev/null +++ b/crates/core/src/accounts/lifecycle.rs @@ -0,0 +1,20 @@ +//! 账号生命周期状态转换规则。 + +use chrono::{DateTime, Utc}; + +use crate::accounts::model::{Account, AccountStatus}; + +/// 根据 token 过期时间推导账号运行时状态。 +pub fn effective_status(account: &Account, now: DateTime) -> AccountStatus { + if account.status != AccountStatus::Active { + return account.status; + } + if account + .access_token_expires_at + .is_some_and(|expires_at| expires_at <= now) + { + AccountStatus::Expired + } else { + AccountStatus::Active + } +} diff --git a/crates/core/src/accounts/mod.rs b/crates/core/src/accounts/mod.rs new file mode 100644 index 000000000..b1a3adccc --- /dev/null +++ b/crates/core/src/accounts/mod.rs @@ -0,0 +1,20 @@ +//! 账号领域中的纯逻辑辅助。 + +/// Cloudflare 阻断恢复状态。 +pub mod cloudflare; +/// Cookie 捕获与重放策略。 +pub mod cookies; +/// JWT 过期判断。 +pub mod jwt; +/// 账号生命周期状态转换。 +pub mod lifecycle; +/// 账号领域模型。 +pub mod model; +/// 账号池调度策略。 +pub mod pool; +/// 账号领域端口。 +pub mod ports; +/// 账号管理用例。 +pub mod service; +/// 账号用量累积策略。 +pub mod usage; diff --git a/src/codex/accounts/model.rs b/crates/core/src/accounts/model.rs similarity index 65% rename from src/codex/accounts/model.rs rename to crates/core/src/accounts/model.rs index 9ecb0b82b..7034c4568 100644 --- a/src/codex/accounts/model.rs +++ b/crates/core/src/accounts/model.rs @@ -1,54 +1,98 @@ +//! 账号领域模型。 + use chrono::{DateTime, Duration, Utc}; use serde::{Deserialize, Serialize}; +/// 账号当前状态。 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] pub enum AccountStatus { + /// 账号可正常对外服务。 Active, + /// 访问令牌已经过期，需要重新刷新。 Expired, + /// 当前配额已经耗尽。 QuotaExhausted, + /// 账号正处于刷新流程中。 Refreshing, + /// 账号被显式禁用。 Disabled, + /// 账号被上游封禁。 Banned, } +/// 账号聚合根。 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Account { + /// 账号主键。 pub id: String, + /// 展示用邮箱地址。 pub email: Option, + /// 上游账号标识。 pub account_id: Option, + /// 上游用户标识。 pub user_id: Option, + /// 后台自定义标签。 pub label: Option, + /// 订阅计划类型。 pub plan_type: Option, + /// 当前访问令牌。 pub access_token: String, + /// 刷新令牌。 pub refresh_token: Option, + /// 访问令牌过期时间。 pub access_token_expires_at: Option>, + /// 当前账号状态。 pub status: AccountStatus, + /// 是否已经触发配额封顶。 pub quota_limit_reached: bool, + /// 是否需要执行额外配额校验。 pub quota_verify_required: bool, + /// 配额冷却结束时间。 pub quota_cooldown_until: Option>, + /// Cloudflare 冷却结束时间。 pub cloudflare_cooldown_until: Option>, + /// 历史总请求数。 pub request_count: u64, + /// 历史空响应次数。 pub empty_response_count: u64, + /// 历史图片输入 token 数。 pub image_input_tokens: u64, + /// 历史图片输出 token 数。 pub image_output_tokens: u64, + /// 历史图片请求数。 pub image_request_count: u64, + /// 历史图片请求失败数。 pub image_request_failed_count: u64, + /// 当前窗口请求数。 pub window_request_count: u64, + /// 当前窗口输入 token 数。 pub window_input_tokens: u64, + /// 当前窗口输出 token 数。 pub window_output_tokens: u64, + /// 当前窗口缓存 token 数。 pub window_cached_tokens: u64, + /// 当前窗口图片输入 token 数。 pub window_image_input_tokens: u64, + /// 当前窗口图片输出 token 数。 pub window_image_output_tokens: u64, + /// 当前窗口图片请求数。 pub window_image_request_count: u64, + /// 当前窗口图片请求失败数。 pub window_image_request_failed_count: u64, + /// 当前统计窗口起始时间。 pub window_started_at: Option>, + /// 当前统计窗口重置时间。 pub window_reset_at: Option>, + /// 当前限流窗口大小（秒）。 pub limit_window_seconds: Option, + /// 首次加入时间。 pub added_at: String, + /// 最近一次使用时间。 pub last_used_at: Option, } impl Account { + /// 构造测试默认账号。 pub fn test(id: &str, status: AccountStatus) -> Self { Self { id: id.to_string(), diff --git a/src/codex/accounts/pool.rs b/crates/core/src/accounts/pool.rs similarity index 87% rename from src/codex/accounts/pool.rs rename to crates/core/src/accounts/pool.rs index a03f8ef04..6fd0c907f 100644 --- a/src/codex/accounts/pool.rs +++ b/crates/core/src/accounts/pool.rs @@ -1,3 +1,5 @@ +//! 账号池调度策略。 + use std::{ cmp::Ordering, collections::{BTreeMap, VecDeque}, @@ -5,25 +7,37 @@ use std::{ use chrono::{DateTime, Duration, Utc}; -use crate::codex::accounts::{ +use crate::accounts::{ jwt::{jwt_expiry, JwtExpiry}, model::{Account, AccountStatus}, + service::AccountService, }; +/// 账号轮转策略。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RotationStrategy { + /// 优先选择当前窗口用量最低的账号。 LeastUsed, + /// 按候选账号顺序循环选择。 RoundRobin, + /// 优先选择最近使用过的账号。 Sticky, } +/// 账号池配置。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct AccountPoolOptions { + /// 单账号允许的最大并发请求数。 pub max_concurrent_per_account: usize, + /// 在途槽位未释放时的过期时间。 pub stale_slot_ttl: Duration, + /// 账号选择轮转策略。 pub rotation_strategy: RotationStrategy, + /// 是否跳过已标记配额耗尽的账号。 pub skip_quota_limited: bool, + /// 订阅层级优先级，越靠前优先级越高。 pub tier_priority: Vec, + /// 模型到允许订阅计划的映射。 pub model_plan_allowlist: BTreeMap>, } @@ -40,15 +54,21 @@ impl Default for AccountPoolOptions { } } +/// 账号获取请求。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct AccountAcquireRequest { + /// 本次请求使用的模型。 pub model: String, + /// 本次请求需要排除的账号 ID。 pub exclude_account_ids: Vec, + /// 会话亲和性建议的优先账号 ID。 pub preferred_account_id: Option, + /// 本次调度使用的当前时间。 pub now: DateTime, } impl AccountAcquireRequest { + /// 构造账号获取请求。 pub fn new(model: impl Into, now: DateTime) -> Self { Self { model: model.into(), @@ -58,6 +78,7 @@ impl AccountAcquireRequest { } } + /// 设置需要排除的账号 ID。 pub fn with_exclude_account_ids( mut self, account_ids: impl IntoIterator>, @@ -66,49 +87,76 @@ impl AccountAcquireRequest { self } + /// 设置会话亲和性建议的优先账号 ID。 pub fn with_preferred_account_id(mut self, account_id: impl Into) -> Self { self.preferred_account_id = Some(account_id.into()); self } } +/// 成功获取到的账号及调度元数据。 #[derive(Debug, Clone)] pub struct AcquiredAccount { + /// 被选中的账号快照。 pub account: Account, + /// 同一账号上一个在途槽位的创建时间。 pub previous_slot_at: Option>, } +/// 账号池容量摘要。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct AccountCapacitySummary { + /// 单账号最大并发请求数。 pub max_concurrent_per_account: usize, + /// 所有可用账号合计槽位数。 pub total_slots: usize, + /// 当前已占用槽位数。 pub used_slots: usize, + /// 当前可用槽位数。 pub available_slots: usize, } +/// 运行时窗口用量增量。 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct AccountWindowUsageDelta { + /// 输入 token 增量。 pub input_tokens: u64, + /// 输出 token 增量。 pub output_tokens: u64, + /// 缓存 token 增量。 pub cached_tokens: u64, + /// 图片输入 token 增量。 pub image_input_tokens: u64, + /// 图片输出 token 增量。 pub image_output_tokens: u64, + /// 是否记录一次成功图片请求。 pub image_request_succeeded: bool, + /// 是否记录一次失败图片请求。 pub image_request_failed: bool, } +/// 账号池状态摘要。 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct AccountPoolStatusSummary { + /// 账号总数。 pub total: usize, + /// 可调度账号数。 pub active: usize, + /// token 过期账号数。 pub expired: usize, + /// 配额耗尽账号数。 pub quota_exhausted: usize, + /// 处于限流冷却的账号数。 pub rate_limited: usize, + /// 正在刷新账号数。 pub refreshing: usize, + /// 禁用账号数。 pub disabled: usize, + /// 封禁账号数。 pub banned: usize, } +/// 纯内存账号池，负责账号调度和运行时状态维护。 #[derive(Debug)] pub struct AccountPool { accounts: BTreeMap, @@ -125,6 +173,7 @@ impl Default for AccountPool { } impl AccountPool { + /// 使用指定配置创建账号池。 pub fn with_options(options: AccountPoolOptions) -> Self { Self { accounts: BTreeMap::new(), @@ -135,16 +184,19 @@ impl AccountPool { } } + /// 插入或替换账号。 pub fn insert(&mut self, account: Account) { self.accounts.insert(account.id.clone(), account); } + /// 移除账号及其在途槽位。 pub fn remove(&mut self, account_id: &str) -> bool { let removed = self.accounts.remove(account_id).is_some(); self.slots.remove(account_id); removed } + /// 清空账号池。 pub fn clear(&mut self) { self.accounts.clear(); self.slots.clear(); @@ -152,10 +204,12 @@ impl AccountPool { self.least_used_cursor = 0; } + /// 替换模型计划 allowlist。 pub fn set_model_plan_allowlist(&mut self, allowlist: BTreeMap>) { self.options.model_plan_allowlist = allowlist; } + /// 更新账号标签。 pub fn set_label(&mut self, account_id: &str, label: Option) -> bool { let Some(account) = self.accounts.get_mut(account_id) else { return false; @@ -164,6 +218,7 @@ impl AccountPool { true } + /// 更新账号状态。 pub fn set_status(&mut self, account_id: &str, status: AccountStatus) -> bool { let Some(account) = self.accounts.get_mut(account_id) else { return false; @@ -175,6 +230,7 @@ impl AccountPool { true } + /// 重置账号累计和窗口用量。 pub fn reset_usage(&mut self, account_id: &str) -> bool { let Some(account) = self.accounts.get_mut(account_id) else { return false; @@ -198,6 +254,7 @@ impl AccountPool { true } + /// 标记账号配额限流直到指定时间。 pub fn mark_quota_limited_until( &mut self, account_id: &str, @@ -226,6 +283,7 @@ impl AccountPool { true } + /// 设置账号是否需要额外配额校验。 pub fn set_quota_verify_required(&mut self, account_id: &str, required: bool) -> bool { let Some(account) = self.accounts.get_mut(account_id) else { return false; @@ -234,6 +292,7 @@ impl AccountPool { true } + /// 同步上游返回的限流窗口信息。 pub fn sync_rate_limit_window( &mut self, account_id: &str, @@ -243,6 +302,7 @@ impl AccountPool { self.sync_rate_limit_window_at(account_id, new_reset_at, limit_window_seconds, Utc::now()) } + /// 记录账号运行时窗口 token 用量。 pub fn record_window_token_usage( &mut self, account_id: &str, @@ -290,6 +350,7 @@ impl AccountPool { true } + /// 标记账号处于 Cloudflare 冷却期。 pub fn set_cloudflare_cooldown_until( &mut self, account_id: &str, @@ -303,11 +364,13 @@ impl AccountPool { true } + /// 使用当前时间为指定模型获取账号。 pub fn acquire(&mut self, model: &str) -> Option { self.acquire_with(AccountAcquireRequest::new(model, Utc::now())) .map(|acquired| acquired.account) } + /// 使用完整调度请求获取账号。 pub fn acquire_with(&mut self, request: AccountAcquireRequest) -> Option { self.cleanup_stale_slots(request.now); self.refresh_account_statuses(request.now); @@ -337,6 +400,7 @@ impl AccountPool { }) } + /// 释放指定账号的一个在途槽位。 pub fn release(&mut self, account_id: &str) { let Some(slots) = self.slots.get_mut(account_id) else { return; @@ -347,6 +411,7 @@ impl AccountPool { } } + /// 计算账号池容量摘要。 pub fn capacity_summary(&mut self, now: DateTime) -> AccountCapacitySummary { self.cleanup_stale_slots(now); self.refresh_account_statuses(now); @@ -354,7 +419,12 @@ impl AccountPool { .accounts .values() .filter(|account| { - account.status == AccountStatus::Active && self.is_quota_available(account, now) + account.status == AccountStatus::Active + && AccountService::quota_available_at( + account, + now, + self.options.skip_quota_limited, + ) }) .count(); let total_slots = active_accounts * self.options.max_concurrent_per_account; @@ -377,6 +447,7 @@ impl AccountPool { } } + /// 计算账号池状态摘要。 pub fn status_summary(&mut self, now: DateTime) -> AccountPoolStatusSummary { self.cleanup_stale_slots(now); self.refresh_account_statuses(now); @@ -386,7 +457,13 @@ impl AccountPool { }; for account in self.accounts.values() { match account.status { - AccountStatus::Active if !self.is_quota_available(account, now) => { + AccountStatus::Active + if !AccountService::quota_available_at( + account, + now, + self.options.skip_quota_limited, + ) => + { summary.rate_limited += 1; } AccountStatus::Active => summary.active += 1, @@ -400,7 +477,7 @@ impl AccountPool { summary } - /// 获取所有配额锁定的账户ID列表（用于主动配额刷新） + /// 获取所有处于配额锁定状态的账号 ID。 pub fn list_quota_locked_accounts(&self) -> Vec { self.accounts .values() @@ -458,24 +535,17 @@ impl AccountPool { fn is_base_available(&self, account: &Account, request: &AccountAcquireRequest) -> bool { account.status == AccountStatus::Active && self.slot_count(&account.id) < self.options.max_concurrent_per_account - && self.is_quota_available(account, request.now) + && AccountService::quota_available_at( + account, + request.now, + self.options.skip_quota_limited, + ) && self.is_model_allowed(account, &request.model) && !request .exclude_account_ids .iter() .any(|account_id| account_id == &account.id) - && account - .cloudflare_cooldown_until - .is_none_or(|cooldown_until| request.now >= cooldown_until) - } - - fn is_quota_available(&self, account: &Account, now: DateTime) -> bool { - if !self.options.skip_quota_limited || !account.quota_limit_reached { - return true; - } - account - .quota_cooldown_until - .is_some_and(|cooldown_until| now >= cooldown_until) + && AccountService::cloudflare_available_at(account, request.now) } fn slot_count(&self, account_id: &str) -> usize { diff --git a/crates/core/src/accounts/ports.rs b/crates/core/src/accounts/ports.rs new file mode 100644 index 000000000..5c4c073a6 --- /dev/null +++ b/crates/core/src/accounts/ports.rs @@ -0,0 +1,100 @@ +//! 账号领域端口。 + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use thiserror::Error; + +use crate::accounts::{ + model::{Account, AccountStatus}, + usage::AccountUsageDelta, +}; + +/// 账号存储错误。 +#[derive(Debug, Error)] +pub enum AccountStoreError { + /// 底层存储失败。 + #[error("account store operation failed: {message}")] + OperationFailed { + /// 错误说明。 + message: String, + }, +} + +/// 账号存储结果类型。 +pub type AccountStoreResult = Result; + +/// 提供运行时账号列表的端口。 +#[async_trait] +pub trait AccountStore: Send + Sync + 'static { + /// 列出当前账号池可见的账号。 + async fn list_pool_accounts(&self) -> AccountStoreResult>; + + /// 读取单个账号池账号快照。 + async fn get_pool_account(&self, account_id: &str) -> AccountStoreResult> { + let accounts = self.list_pool_accounts().await?; + Ok(accounts + .into_iter() + .find(|account| account.id == account_id)) + } + + /// 标记账号进入配额冷却期。 + async fn mark_quota_limited_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> AccountStoreResult; + + /// 标记账号进入 Cloudflare 冷却期。 + async fn set_cloudflare_cooldown_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> AccountStoreResult; + + /// 更新账号状态。 + async fn set_status(&self, account_id: &str, status: AccountStatus) + -> AccountStoreResult; + + /// 记录账号用量增量。 + async fn record_usage_delta( + &self, + account_id: &str, + usage: AccountUsageDelta, + ) -> AccountStoreResult<()>; + + /// 读取账号当前配额 JSON。 + async fn get_quota_json(&self, _account_id: &str) -> AccountStoreResult> { + Ok(None) + } + + /// 更新账号当前配额 JSON。 + async fn update_quota_json( + &self, + _account_id: &str, + _quota_json: &str, + ) -> AccountStoreResult { + Ok(false) + } + + /// 同步账号当前 rate-limit 统计窗口。 + async fn sync_rate_limit_window( + &self, + _account_id: &str, + _reset_at: DateTime, + _limit_window_seconds: Option, + ) -> AccountStoreResult<()> { + Ok(()) + } + + /// 记录账号被用于一次外部请求。 + async fn record_request(&self, account_id: &str) -> AccountStoreResult<()> { + self.record_usage_delta( + account_id, + AccountUsageDelta { + requests: 1, + ..AccountUsageDelta::default() + }, + ) + .await + } +} diff --git a/crates/core/src/accounts/service.rs b/crates/core/src/accounts/service.rs new file mode 100644 index 000000000..5f232199b --- /dev/null +++ b/crates/core/src/accounts/service.rs @@ -0,0 +1,44 @@ +//! 账号管理用例。 + +use chrono::{DateTime, Utc}; + +use crate::accounts::model::{Account, AccountStatus}; + +/// 账号服务。 +#[derive(Debug, Clone, Default)] +pub struct AccountService; + +impl AccountService { + /// 判断账号是否可参与请求调度。 + pub fn can_serve(account: &Account) -> bool { + account.status == AccountStatus::Active && !account.quota_limit_reached + } + + /// 判断账号在指定时间是否处于可服务的基础状态。 + pub fn can_serve_at(account: &Account, now: DateTime, skip_quota_limited: bool) -> bool { + account.status == AccountStatus::Active + && Self::quota_available_at(account, now, skip_quota_limited) + && Self::cloudflare_available_at(account, now) + } + + /// 判断账号配额锁定是否已经释放。 + pub fn quota_available_at( + account: &Account, + now: DateTime, + skip_quota_limited: bool, + ) -> bool { + if !skip_quota_limited || !account.quota_limit_reached { + return true; + } + account + .quota_cooldown_until + .is_some_and(|cooldown_until| now >= cooldown_until) + } + + /// 判断 Cloudflare 冷却是否已经释放。 + pub fn cloudflare_available_at(account: &Account, now: DateTime) -> bool { + account + .cloudflare_cooldown_until + .is_none_or(|cooldown_until| now >= cooldown_until) + } +} diff --git a/crates/core/src/accounts/usage.rs b/crates/core/src/accounts/usage.rs new file mode 100644 index 000000000..d4056cb51 --- /dev/null +++ b/crates/core/src/accounts/usage.rs @@ -0,0 +1,41 @@ +//! 账号用量累积策略。 + +/// 账号用量增量。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct AccountUsageDelta { + /// 请求数增量。 + pub requests: u64, + /// 输入 token 增量。 + pub input_tokens: u64, + /// 输出 token 增量。 + pub output_tokens: u64, + /// 缓存 token 增量。 + pub cached_tokens: u64, + /// 空响应数增量。 + pub empty_responses: u64, + /// 图片工具输入 token 增量。 + pub image_input_tokens: u64, + /// 图片工具输出 token 增量。 + pub image_output_tokens: u64, + /// 图片工具成功请求数增量。 + pub image_requests: u64, + /// 图片工具失败请求数增量。 + pub image_request_failures: u64, +} + +impl AccountUsageDelta { + /// 合并两个用量增量。 + pub fn merged(self, other: Self) -> Self { + Self { + requests: self.requests + other.requests, + input_tokens: self.input_tokens + other.input_tokens, + output_tokens: self.output_tokens + other.output_tokens, + cached_tokens: self.cached_tokens + other.cached_tokens, + empty_responses: self.empty_responses + other.empty_responses, + image_input_tokens: self.image_input_tokens + other.image_input_tokens, + image_output_tokens: self.image_output_tokens + other.image_output_tokens, + image_requests: self.image_requests + other.image_requests, + image_request_failures: self.image_request_failures + other.image_request_failures, + } + } +} diff --git a/crates/core/src/admin/auth.rs b/crates/core/src/admin/auth.rs new file mode 100644 index 000000000..e8bf725e2 --- /dev/null +++ b/crates/core/src/admin/auth.rs @@ -0,0 +1,30 @@ +//! 管理员认证领域服务。 + +/// 管理员登录请求。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminLoginRequest { + /// 用户名。 + pub username: String, + /// 明文密码。 + pub password: String, +} + +/// 管理员认证服务。 +#[derive(Debug, Clone)] +pub struct AdminAuthService { + default_username: String, +} + +impl AdminAuthService { + /// 构造管理员认证服务。 + pub fn new(default_username: impl Into) -> Self { + Self { + default_username: default_username.into(), + } + } + + /// 判断用户名是否匹配配置中的管理员用户名。 + pub fn username_matches(&self, username: &str) -> bool { + self.default_username == username + } +} diff --git a/crates/core/src/admin/client_keys.rs b/crates/core/src/admin/client_keys.rs new file mode 100644 index 000000000..644cd7d9d --- /dev/null +++ b/crates/core/src/admin/client_keys.rs @@ -0,0 +1,26 @@ +//! 客户端 API Key 业务服务。 + +use std::sync::Arc; + +use crate::admin::ports::{ClientKeyStore, ClientKeyStoreResult}; + +/// 客户端 API Key 服务。 +#[derive(Clone)] +pub struct ClientKeyService { + store: Arc, +} + +impl ClientKeyService { + /// 构造服务。 + pub fn new(store: Arc) -> Self { + Self { store } + } + + /// 验证客户端 API Key。 + pub async fn verify(&self, plaintext: &str) -> ClientKeyStoreResult { + if !plaintext.starts_with("cpr_") { + return Ok(false); + } + self.store.verify_and_touch(plaintext).await + } +} diff --git a/crates/core/src/admin/mod.rs b/crates/core/src/admin/mod.rs new file mode 100644 index 000000000..ac2a81f5e --- /dev/null +++ b/crates/core/src/admin/mod.rs @@ -0,0 +1,6 @@ +//! 管理员领域用例和端口。 + +pub mod auth; +pub mod client_keys; +pub mod ports; +pub mod settings; diff --git a/crates/core/src/admin/ports.rs b/crates/core/src/admin/ports.rs new file mode 100644 index 000000000..c3bd21002 --- /dev/null +++ b/crates/core/src/admin/ports.rs @@ -0,0 +1,25 @@ +//! 管理员领域端口。 + +use async_trait::async_trait; +use thiserror::Error; + +/// 客户端 API Key 存储错误。 +#[derive(Debug, Error)] +pub enum ClientKeyStoreError { + /// 底层存储失败。 + #[error("client key store operation failed: {message}")] + OperationFailed { + /// 错误说明。 + message: String, + }, +} + +/// 客户端 API Key 存储结果类型。 +pub type ClientKeyStoreResult = Result; + +/// 提供客户端 API Key 验证能力的端口。 +#[async_trait] +pub trait ClientKeyStore: Send + Sync + 'static { + /// 验证明文客户端 API Key，并在成功时记录使用时间。 + async fn verify_and_touch(&self, plaintext: &str) -> ClientKeyStoreResult; +} diff --git a/crates/core/src/admin/settings.rs b/crates/core/src/admin/settings.rs new file mode 100644 index 000000000..08a612dc2 --- /dev/null +++ b/crates/core/src/admin/settings.rs @@ -0,0 +1,339 @@ +//! 管理端设置领域逻辑。 + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +const ROTATION_STRATEGIES: [&str; 3] = ["least_used", "round_robin", "sticky"]; + +/// 管理端可变设置。 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct AdminSettings { + /// 默认模型。 + pub default_model: String, + /// 默认 reasoning effort。 + pub default_reasoning_effort: Option, + /// 默认服务层级。 + pub service_tier: Option, + /// 模型别名。 + pub model_aliases: BTreeMap, + /// 是否启用 token 刷新。 + pub refresh_enabled: bool, + /// token 刷新提前秒数。 + pub refresh_margin_seconds: u64, + /// token 刷新并发度。 + 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: AdminQuotaWarningThresholds, + /// 是否跳过配额耗尽账号。 + 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 Default for AdminSettings { + fn default() -> Self { + Self { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: None, + service_tier: None, + model_aliases: BTreeMap::new(), + refresh_enabled: true, + refresh_margin_seconds: 300, + refresh_concurrency: 2, + max_concurrent_per_account: 3, + request_interval_ms: 50, + rotation_strategy: "least_used".to_string(), + tier_priority: Vec::new(), + quota_refresh_interval_minutes: 5, + quota_warning_thresholds: AdminQuotaWarningThresholds::default(), + quota_skip_exhausted: true, + logs_enabled: false, + logs_capacity: 2_000, + logs_capture_body: false, + usage_history_retention_days: None, + } + } +} + +/// 配额预警阈值集合。 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminQuotaWarningThresholds { + /// 主阈值列表。 + pub primary: Vec, + /// 次阈值列表。 + pub secondary: Vec, +} + +impl Default for AdminQuotaWarningThresholds { + fn default() -> Self { + Self { + primary: vec![80, 90], + secondary: vec![80, 90], + } + } +} + +/// 管理端设置补丁。 +#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminSettingsPatch { + /// 默认模型。 + pub default_model: Option, + /// 默认 reasoning effort。 + pub default_reasoning_effort: Option, + /// 默认服务层级。 + pub service_tier: Option, + /// 模型别名。 + pub model_aliases: Option>, + /// 是否启用 token 刷新。 + pub refresh_enabled: Option, + /// token 刷新提前秒数。 + pub refresh_margin_seconds: Option, + /// token 刷新并发度。 + pub refresh_concurrency: Option, + /// 单账号最大并发请求数。 + pub max_concurrent_per_account: Option, + /// 单账号请求间隔毫秒。 + pub request_interval_ms: Option, + /// 账号轮换策略。 + pub rotation_strategy: Option, + /// 订阅层级优先级。 + pub tier_priority: Option>, + /// 配额刷新间隔分钟。 + pub quota_refresh_interval_minutes: Option, + /// 配额预警阈值。 + pub quota_warning_thresholds: Option, + /// 是否跳过配额耗尽账号。 + pub quota_skip_exhausted: Option, + /// 是否启用日志。 + pub logs_enabled: Option, + /// 日志容量。 + pub logs_capacity: Option, + /// 是否捕获请求/响应体。 + pub logs_capture_body: Option, + /// 用量历史保留天数。 + pub usage_history_retention_days: Option, +} + +/// 设置领域服务。 +#[derive(Debug, Clone, Default)] +pub struct SettingsService; + +impl SettingsService { + /// 合并并验证补丁设置。 + pub fn apply_patch( + current: &mut AdminSettings, + patch: AdminSettingsPatch, + ) -> Result<(), SettingsServiceError> { + if let Some(default_model) = patch.default_model { + current.default_model = non_empty_string("defaultModel", default_model)?; + } + if let Some(default_reasoning_effort) = patch.default_reasoning_effort { + current.default_reasoning_effort = Some(non_empty_string( + "defaultReasoningEffort", + default_reasoning_effort, + )?); + } + if let Some(service_tier) = patch.service_tier { + current.service_tier = Some(non_empty_string("serviceTier", service_tier)?); + } + if let Some(model_aliases) = patch.model_aliases { + current.model_aliases = validate_aliases(model_aliases)?; + } + if let Some(refresh_enabled) = patch.refresh_enabled { + current.refresh_enabled = refresh_enabled; + } + if let Some(refresh_margin_seconds) = patch.refresh_margin_seconds { + current.refresh_margin_seconds = refresh_margin_seconds; + } + if let Some(refresh_concurrency) = patch.refresh_concurrency { + current.refresh_concurrency = positive_u32("refreshConcurrency", refresh_concurrency)?; + } + if let Some(max_concurrent_per_account) = patch.max_concurrent_per_account { + current.max_concurrent_per_account = + positive_usize("maxConcurrentPerAccount", max_concurrent_per_account)?; + } + if let Some(request_interval_ms) = patch.request_interval_ms { + current.request_interval_ms = request_interval_ms; + } + if let Some(rotation_strategy) = patch.rotation_strategy { + current.rotation_strategy = validate_rotation_strategy(rotation_strategy)?; + } + if let Some(tier_priority) = patch.tier_priority { + current.tier_priority = validate_tier_priority(tier_priority)?; + } + if let Some(quota_refresh_interval_minutes) = patch.quota_refresh_interval_minutes { + current.quota_refresh_interval_minutes = positive_u64( + "quotaRefreshIntervalMinutes", + quota_refresh_interval_minutes, + )?; + } + if let Some(quota_warning_thresholds) = patch.quota_warning_thresholds { + validate_thresholds("a_warning_thresholds)?; + current.quota_warning_thresholds = quota_warning_thresholds; + } + if let Some(quota_skip_exhausted) = patch.quota_skip_exhausted { + current.quota_skip_exhausted = quota_skip_exhausted; + } + if let Some(logs_enabled) = patch.logs_enabled { + current.logs_enabled = logs_enabled; + } + if let Some(logs_capacity) = patch.logs_capacity { + current.logs_capacity = positive_u32("logsCapacity", logs_capacity)?; + } + if let Some(logs_capture_body) = patch.logs_capture_body { + current.logs_capture_body = logs_capture_body; + } + if let Some(usage_history_retention_days) = patch.usage_history_retention_days { + current.usage_history_retention_days = Some(positive_u64( + "usageHistoryRetentionDays", + usage_history_retention_days, + )?); + } + + Ok(()) + } +} + +/// 设置领域错误。 +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum SettingsServiceError { + /// 设置字段非法。 + #[error("invalid setting `{field}`: {message}")] + InvalidField { + /// 字段名。 + field: String, + /// 错误说明。 + message: String, + }, +} + +impl SettingsServiceError { + /// 返回出错字段名。 + pub fn field(&self) -> &str { + match self { + Self::InvalidField { field, .. } => field, + } + } + + /// 返回字段错误说明。 + pub fn message(&self) -> &str { + match self { + Self::InvalidField { message, .. } => message, + } + } +} + +fn validate_aliases( + aliases: BTreeMap, +) -> Result, SettingsServiceError> { + let mut validated = BTreeMap::new(); + for (alias, target) in aliases { + let alias = non_empty_string("modelAliases", alias)?; + let target = non_empty_string("modelAliases", target)?; + validated.insert(alias, target); + } + Ok(validated) +} + +fn validate_tier_priority(tiers: Vec) -> Result, SettingsServiceError> { + let mut validated = Vec::with_capacity(tiers.len()); + for tier in tiers { + validated.push(non_empty_string("tierPriority", tier)?); + } + Ok(validated) +} + +fn validate_rotation_strategy(strategy: String) -> Result { + let strategy = non_empty_string("rotationStrategy", strategy)?; + if ROTATION_STRATEGIES.contains(&strategy.as_str()) { + Ok(strategy) + } else { + Err(invalid_field( + "rotationStrategy", + "must be one of least_used, round_robin, sticky", + )) + } +} + +fn validate_thresholds( + thresholds: &AdminQuotaWarningThresholds, +) -> Result<(), SettingsServiceError> { + if thresholds.primary.is_empty() && thresholds.secondary.is_empty() { + return Err(invalid_field( + "quotaWarningThresholds", + "must include at least one threshold", + )); + } + if thresholds + .primary + .iter() + .chain(thresholds.secondary.iter()) + .any(|threshold| *threshold > 100) + { + return Err(invalid_field( + "quotaWarningThresholds", + "thresholds must be between 0 and 100", + )); + } + Ok(()) +} + +fn non_empty_string(field: &str, value: String) -> Result { + let value = value.trim().to_string(); + if value.is_empty() { + Err(invalid_field(field, "must not be empty")) + } else { + Ok(value) + } +} + +fn positive_u32(field: &str, value: u32) -> Result { + if value == 0 { + Err(invalid_field(field, "must be greater than 0")) + } else { + Ok(value) + } +} + +fn positive_u64(field: &str, value: u64) -> Result { + if value == 0 { + Err(invalid_field(field, "must be greater than 0")) + } else { + Ok(value) + } +} + +fn positive_usize(field: &str, value: usize) -> Result { + if value == 0 { + Err(invalid_field(field, "must be greater than 0")) + } else { + Ok(value) + } +} + +fn invalid_field(field: &str, message: impl Into) -> SettingsServiceError { + SettingsServiceError::InvalidField { + field: field.to_string(), + message: message.into(), + } +} diff --git a/crates/core/src/auth/mod.rs b/crates/core/src/auth/mod.rs new file mode 100644 index 000000000..f2d6f342d --- /dev/null +++ b/crates/core/src/auth/mod.rs @@ -0,0 +1,8 @@ +//! 认证领域中的纯逻辑与端口。 + +/// OAuth token 刷新策略。 +pub mod oauth; +/// 认证上游端口定义。 +pub mod ports; +/// 管理员会话领域逻辑。 +pub mod session; diff --git a/crates/core/src/auth/oauth.rs b/crates/core/src/auth/oauth.rs new file mode 100644 index 000000000..bcd6dfab2 --- /dev/null +++ b/crates/core/src/auth/oauth.rs @@ -0,0 +1,417 @@ +//! OAuth 领域逻辑。 + +use std::{collections::BTreeMap, sync::Arc}; + +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use chrono::{DateTime, Duration, Utc}; +use rand::Rng; +use serde::Deserialize; +use sha2::{Digest, Sha256}; +use thiserror::Error; +use tokio::sync::Semaphore; + +use crate::accounts::model::{Account, AccountStatus}; +use crate::auth::ports::TokenRefresher; + +/// 一次刷新返回的 token 对。 +#[derive(Debug, Clone)] +pub struct TokenPair { + /// 新的访问令牌。 + pub access_token: String, + /// 可选的新刷新令牌。 + pub refresh_token: Option, +} + +/// OAuth 客户端静态配置。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OAuthConfig { + /// OAuth 客户端 ID。 + pub client_id: String, + /// 浏览器授权入口。 + pub auth_endpoint: String, + /// 设备码申请入口。 + pub device_code_endpoint: String, + /// Token 交换入口。 + pub token_endpoint: String, +} + +/// 设备码登录响应。 +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +pub struct DeviceCode { + /// 设备码。 + pub device_code: String, + /// 用户输入码。 + pub user_code: String, + /// 激活入口。 + pub verification_uri: String, + /// 带用户码的激活入口。 + pub verification_uri_complete: String, + /// 过期秒数。 + pub expires_in: u64, + /// 建议轮询间隔秒数。 + pub interval: u64, +} + +/// OAuth 调用错误。 +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum OAuthError { + /// 设备授权尚未完成。 + #[error("authorization pending")] + AuthorizationPending, + /// 上游要求放慢轮询速度。 + #[error("slow down")] + SlowDown, + /// 上游显式拒绝请求。 + #[error("OAuth request rejected: {0}")] + Rejected(String), + /// HTTP 或解析层失败。 + #[error("OAuth transport failed")] + Transport, +} + +impl OAuthError { + /// 返回可用于设备码轮询响应的标准错误码。 + pub fn pending_code(&self) -> Option<&'static str> { + match self { + Self::AuthorizationPending => Some("authorization_pending"), + Self::SlowDown => Some("slow_down"), + Self::Rejected(_) | Self::Transport => None, + } + } +} + +/// 一次 PKCE 登录开始后返回给前端的数据。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PkceLogin { + /// 引导用户跳转的授权地址。 + pub auth_url: String, + /// 回调校验状态值。 + pub state: String, +} + +/// 存储在服务端的 PKCE 会话。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PkceSession { + /// PKCE code verifier。 + pub code_verifier: String, + /// OAuth 回调地址。 + pub redirect_uri: String, + /// 登录完成后跳回的 host。 + pub return_host: String, +} + +#[derive(Debug, Clone)] +struct PendingPkceSession { + session: PkceSession, + created_at: DateTime, + exchanging: bool, +} + +/// 管理 PKCE 授权回调状态的内存会话仓。 +#[derive(Debug)] +pub struct PkceSessionStore { + pending: BTreeMap, + completed: BTreeMap>, + ttl: Duration, +} + +impl Default for PkceSessionStore { + fn default() -> Self { + Self { + pending: BTreeMap::new(), + completed: BTreeMap::new(), + ttl: Duration::minutes(5), + } + } +} + +impl PkceSessionStore { + /// 开始一次新的 PKCE 登录流程。 + pub fn start_login(&mut self, return_host: &str, config: &OAuthConfig) -> PkceLogin { + self.cleanup(Utc::now()); + let (code_verifier, code_challenge) = generate_pkce_pair(); + let state = random_state(); + let redirect_uri = "http://localhost:1455/auth/callback".to_string(); + let auth_url = build_auth_url(config, &redirect_uri, &state, &code_challenge); + + self.pending.insert( + state.clone(), + PendingPkceSession { + session: PkceSession { + code_verifier, + redirect_uri, + return_host: return_host.to_string(), + }, + created_at: Utc::now(), + exchanging: false, + }, + ); + + PkceLogin { auth_url, state } + } + + /// 尝试独占获取一条待完成的 PKCE 会话。 + pub fn try_acquire(&mut self, state: &str) -> Option { + self.cleanup(Utc::now()); + let pending = self.pending.get_mut(state)?; + if pending.exchanging { + return None; + } + pending.exchanging = true; + Some(pending.session.clone()) + } + + /// 释放一次失败的会话占用。 + pub fn release(&mut self, state: &str) { + if let Some(pending) = self.pending.get_mut(state) { + pending.exchanging = false; + } + } + + /// 将给定 state 标记为已完成。 + pub fn complete(&mut self, state: &str) { + self.pending.remove(state); + self.completed.insert(state.to_string(), Utc::now()); + } + + /// 判断 state 是否已经完成或正在处理中。 + pub fn is_completed_or_exchanging(&mut self, state: &str) -> bool { + self.cleanup(Utc::now()); + self.completed.contains_key(state) + || self + .pending + .get(state) + .is_some_and(|pending| pending.exchanging) + } + + fn cleanup(&mut self, now: DateTime) { + self.pending + .retain(|_, pending| now - pending.created_at <= self.ttl); + self.completed + .retain(|_, completed_at| now - *completed_at <= self.ttl); + } +} + +/// 刷新任务的调度策略。 +#[derive(Debug, Clone, Copy)] +pub struct RefreshPolicy { + /// 提前多久开始刷新访问令牌。 + pub refresh_margin_seconds: u64, + /// 允许并发执行的刷新任务数。 + pub refresh_concurrency: u32, +} + +/// 触发刷新动作的原因。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RefreshTrigger { + /// 在访问令牌即将过期前触发刷新。 + BeforeExpiry, + /// 在上游返回未授权后立即刷新。 + Unauthorized, +} + +/// 上游刷新失败后的领域结果。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +pub enum RefreshFailure { + /// 刷新令牌无效或已过期。 + #[error("refresh token is invalid or expired")] + InvalidGrant, + /// 账号配额耗尽。 + #[error("account quota is exhausted")] + QuotaExhausted, + /// 账号被上游封禁。 + #[error("account is banned")] + Banned, + /// 账号被显式禁用。 + #[error("account is disabled")] + Disabled, + /// 刷新请求在传输层失败。 + #[error("refresh transport failed")] + Transport, +} + +/// 调度器自身的执行错误。 +#[derive(Debug, Error)] +pub enum RefreshError { + /// 并发限制信号量已关闭。 + #[error("refresh task semaphore closed")] + ConcurrencyClosed, + /// 刷新请求在传输层失败。 + #[error("refresh transport failed")] + Transport, +} + +/// 负责执行单账号刷新策略的调度器。 +#[derive(Clone)] +pub struct RefreshScheduler { + policy: RefreshPolicy, + client: Arc, + semaphore: Arc, +} + +impl RefreshScheduler +where + C: TokenRefresher, +{ + /// 使用策略和上游刷新端口构造调度器。 + pub fn new(policy: RefreshPolicy, client: C) -> Self { + let concurrency = policy.refresh_concurrency.max(1) as usize; + Self { + policy, + client: Arc::new(client), + semaphore: Arc::new(Semaphore::new(concurrency)), + } + } + + /// 在给定时间点按触发原因刷新账号。 + /// + /// 当账号不需要刷新时返回原账号快照；当刷新失败但属于可映射的领域错误时， + /// 返回更新过状态的账号。 + /// + /// # Errors + /// + /// 当并发控制已关闭或刷新传输失败时返回 [`RefreshError`]。 + pub async fn refresh_account_at( + &self, + account: &Account, + trigger: RefreshTrigger, + now: DateTime, + ) -> Result { + if !self.should_refresh(account, trigger, now) { + return Ok(account.clone()); + } + + let Some(refresh_token) = account.refresh_token.as_deref() else { + let mut expired = account.clone(); + expired.status = AccountStatus::Expired; + return Ok(expired); + }; + + let _permit = self + .semaphore + .acquire() + .await + .map_err(|_| RefreshError::ConcurrencyClosed)?; + + match self.client.refresh(refresh_token).await { + Ok(token_pair) => Ok(apply_token_pair(account, token_pair)), + Err(RefreshFailure::Transport) => Err(RefreshError::Transport), + Err(error) => Ok(apply_refresh_failure(account, error)), + } + } + + /// 判断账号在给定触发原因下是否需要刷新。 + pub fn should_refresh_account_at( + &self, + account: &Account, + trigger: RefreshTrigger, + now: DateTime, + ) -> bool { + self.should_refresh(account, trigger, now) + } + + fn should_refresh( + &self, + account: &Account, + trigger: RefreshTrigger, + now: DateTime, + ) -> bool { + if account.status != AccountStatus::Active { + return false; + } + + match trigger { + RefreshTrigger::Unauthorized => true, + RefreshTrigger::BeforeExpiry => account + .access_token_expires_at + .is_some_and(|expires_at| expires_at <= now + self.refresh_margin()), + } + } + + fn refresh_margin(&self) -> Duration { + let seconds = self.policy.refresh_margin_seconds.min(86_400 * 7) as i64; + Duration::seconds(seconds) + } +} + +/// 将新的 token 对应用到账号快照上。 +pub fn apply_token_pair(account: &Account, token_pair: TokenPair) -> Account { + let mut refreshed = account.clone(); + refreshed.access_token = token_pair.access_token; + + // 刷新响应不返回 refresh_token 时，继续保留旧值，避免永久失去刷新能力。 + if let Some(refresh_token) = token_pair.refresh_token { + refreshed.refresh_token = Some(refresh_token); + } + + refreshed.status = AccountStatus::Active; + refreshed +} + +/// 将刷新失败映射为账号状态变更。 +pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Account { + let mut updated = account.clone(); + updated.status = match failure { + RefreshFailure::InvalidGrant => AccountStatus::Expired, + RefreshFailure::QuotaExhausted => AccountStatus::QuotaExhausted, + RefreshFailure::Banned => AccountStatus::Banned, + RefreshFailure::Disabled => AccountStatus::Disabled, + RefreshFailure::Transport => AccountStatus::Active, + }; + updated +} + +fn generate_pkce_pair() -> (String, String) { + let mut bytes = [0u8; 32]; + rand::rng().fill_bytes(&mut bytes); + let code_verifier = URL_SAFE_NO_PAD.encode(bytes); + let digest = Sha256::digest(code_verifier.as_bytes()); + let code_challenge = URL_SAFE_NO_PAD.encode(digest); + (code_verifier, code_challenge) +} + +fn random_state() -> String { + let mut bytes = [0u8; 16]; + rand::rng().fill_bytes(&mut bytes); + bytes.iter().map(|byte| format!("{byte:02x}")).collect() +} + +fn build_auth_url( + config: &OAuthConfig, + redirect_uri: &str, + state: &str, + code_challenge: &str, +) -> String { + let params = [ + ("response_type", "code"), + ("client_id", config.client_id.as_str()), + ("redirect_uri", redirect_uri), + ("scope", "openid profile email offline_access"), + ("code_challenge", code_challenge), + ("code_challenge_method", "S256"), + ("id_token_add_organizations", "true"), + ("codex_cli_simplified_flow", "true"), + ("state", state), + ("originator", "codex_cli_rs"), + ]; + let query = params + .iter() + .map(|(key, value)| format!("{}={}", percent_encode(key), percent_encode(value))) + .collect::>() + .join("&"); + format!("{}?{query}", config.auth_endpoint) +} + +fn percent_encode(value: &str) -> String { + let mut encoded = String::new(); + for byte in value.bytes() { + match byte { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' => { + encoded.push(char::from(byte)); + } + b' ' => encoded.push_str("%20"), + _ => encoded.push_str(&format!("%{byte:02X}")), + } + } + encoded +} diff --git a/crates/core/src/auth/ports.rs b/crates/core/src/auth/ports.rs new file mode 100644 index 000000000..5f4feeee9 --- /dev/null +++ b/crates/core/src/auth/ports.rs @@ -0,0 +1,30 @@ +//! 认证相关的上游端口。 + +use async_trait::async_trait; + +use crate::auth::oauth::{DeviceCode, OAuthError, RefreshFailure, TokenPair}; + +/// 刷新令牌的上游端口。 +#[async_trait] +pub trait TokenRefresher: Send + Sync + 'static { + /// 使用给定刷新令牌换取新的 token 对。 + async fn refresh(&self, refresh_token: &str) -> Result; +} + +/// OAuth 交互上游端口。 +#[async_trait] +pub trait OAuthClient: Send + Sync + 'static { + /// 使用 PKCE 授权码交换 token 对。 + async fn exchange_code( + &self, + code: &str, + code_verifier: &str, + redirect_uri: &str, + ) -> Result; + + /// 请求设备码登录信息。 + async fn request_device_code(&self) -> Result; + + /// 轮询设备码换取 token 对。 + async fn poll_device_token(&self, device_code: &str) -> Result; +} diff --git a/crates/core/src/auth/session.rs b/crates/core/src/auth/session.rs new file mode 100644 index 000000000..f7f48f311 --- /dev/null +++ b/crates/core/src/auth/session.rs @@ -0,0 +1,24 @@ +//! 管理员会话领域逻辑。 + +use chrono::{DateTime, Duration, Utc}; + +/// 管理员会话。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminSession { + /// 会话 ID。 + pub id: String, + /// 过期时间。 + pub expires_at: DateTime, +} + +impl AdminSession { + /// 判断会话在给定时间是否已经过期。 + pub fn is_expired_at(&self, now: DateTime) -> bool { + now >= self.expires_at + } +} + +/// 根据 TTL 计算会话过期时间。 +pub fn session_expiry(now: DateTime, ttl_minutes: i64) -> DateTime { + now + Duration::minutes(ttl_minutes.max(0)) +} diff --git a/crates/core/src/error.rs b/crates/core/src/error.rs new file mode 100644 index 000000000..01d448a45 --- /dev/null +++ b/crates/core/src/error.rs @@ -0,0 +1,23 @@ +//! 核心领域通用错误。 + +use thiserror::Error; + +/// 核心领域错误。 +#[derive(Debug, Error)] +pub enum CoreError { + /// 输入不符合领域约束。 + #[error("invalid domain input: {message}")] + InvalidInput { + /// 错误说明。 + message: String, + }, + /// 依赖端口调用失败。 + #[error("port operation failed: {message}")] + Port { + /// 错误说明。 + message: String, + }, +} + +/// 核心领域结果。 +pub type CoreResult = Result; diff --git a/crates/core/src/events/mod.rs b/crates/core/src/events/mod.rs new file mode 100644 index 000000000..ddc9a5c13 --- /dev/null +++ b/crates/core/src/events/mod.rs @@ -0,0 +1,8 @@ +//! 事件日志领域。 + +/// 事件日志模型。 +pub mod model; +/// 事件日志端口。 +pub mod ports; +/// 事件日志策略服务。 +pub mod service; diff --git a/src/codex/events/event.rs b/crates/core/src/events/model.rs similarity index 56% rename from src/codex/events/event.rs rename to crates/core/src/events/model.rs index 5c62848a9..0eaa0a56c 100644 --- a/src/codex/events/event.rs +++ b/crates/core/src/events/model.rs @@ -1,49 +1,59 @@ -use chrono::Utc; +//! 事件日志模型。 + +use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; -use serde_json::{json, Value}; +use serde_json::Value; use uuid::Uuid; +/// 事件等级。 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum EventLevel { + /// 调试。 Debug, + /// 信息。 Info, + /// 警告。 Warn, + /// 错误。 Error, } -impl EventLevel { - pub fn as_str(self) -> &'static str { - match self { - Self::Debug => "debug", - Self::Info => "info", - Self::Warn => "warn", - Self::Error => "error", - } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] +/// 结构化事件日志。 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "camelCase")] pub struct EventLog { + /// 事件 ID。 pub id: String, + /// 请求 ID。 pub request_id: Option, + /// 事件类别。 pub kind: String, + /// 事件等级。 pub level: EventLevel, + /// 账号 ID。 pub account_id: Option, + /// HTTP 路由。 pub route: Option, + /// 模型名。 pub model: Option, + /// HTTP 状态码。 pub status_code: Option, + /// 请求耗时毫秒。 pub latency_ms: Option, + /// 展示消息。 pub message: String, + /// 结构化元数据。 pub metadata: Value, - pub created_at: String, + /// 创建时间。 + pub created_at: DateTime, } impl EventLog { + /// 构造事件日志。 pub fn new(kind: impl Into, level: EventLevel, message: impl Into) -> Self { Self { - id: Uuid::new_v4().to_string(), + id: format!("log_{}", Uuid::new_v4().simple()), request_id: None, kind: kind.into(), level, @@ -53,8 +63,8 @@ impl EventLog { status_code: None, latency_ms: None, message: message.into(), - metadata: json!({}), - created_at: Utc::now().to_rfc3339(), + metadata: Value::Object(Default::default()), + created_at: Utc::now(), } } } diff --git a/crates/core/src/events/ports.rs b/crates/core/src/events/ports.rs new file mode 100644 index 000000000..71cdac291 --- /dev/null +++ b/crates/core/src/events/ports.rs @@ -0,0 +1,27 @@ +//! 事件日志端口。 + +use async_trait::async_trait; +use thiserror::Error; + +use crate::events::model::EventLog; + +/// 事件日志存储错误。 +#[derive(Debug, Error)] +pub enum EventLogStoreError { + /// 底层存储失败。 + #[error("event log store operation failed: {message}")] + OperationFailed { + /// 错误说明。 + message: String, + }, +} + +/// 事件日志存储结果。 +pub type EventLogStoreResult = Result; + +/// 事件日志存储端口。 +#[async_trait] +pub trait EventLogStore: Send + Sync + 'static { + /// 写入事件日志。 + async fn append(&self, event: &EventLog) -> EventLogStoreResult<()>; +} diff --git a/crates/core/src/events/service.rs b/crates/core/src/events/service.rs new file mode 100644 index 000000000..9501c131a --- /dev/null +++ b/crates/core/src/events/service.rs @@ -0,0 +1,21 @@ +//! 事件日志策略服务。 + +use crate::events::model::{EventLevel, EventLog}; + +/// 事件日志服务。 +#[derive(Debug, Clone)] +pub struct EventLogService { + enabled: bool, +} + +impl EventLogService { + /// 构造事件日志服务。 + pub fn new(enabled: bool) -> Self { + Self { enabled } + } + + /// 判断事件是否应该记录。 + pub fn should_record(&self, event: &EventLog) -> bool { + self.enabled || event.level == EventLevel::Error + } +} diff --git a/src/codex/gateway/conversation_identity.rs b/crates/core/src/gateway/conversation.rs similarity index 97% rename from src/codex/gateway/conversation_identity.rs rename to crates/core/src/gateway/conversation.rs index 1d74f4b0f..d051e511c 100644 --- a/src/codex/gateway/conversation_identity.rs +++ b/crates/core/src/gateway/conversation.rs @@ -7,7 +7,7 @@ use serde_json::Value; use sha2::{Digest, Sha256}; use uuid::Uuid; -use crate::codex::gateway::transport::types::CodexResponsesRequest; +use crate::protocol::codex::responses::CodexResponsesRequest; const LEADING_SYSTEM_REMINDER_OPEN: &str = ""; const LEADING_SYSTEM_REMINDER_CLOSE: &str = ""; @@ -142,10 +142,10 @@ fn normalize_conversation_anchor_text(text: &str) -> String { rest.to_string() } -/// 构建账号作用域的身份哈希 +/// 构建账号作用域的身份哈希。 /// /// 格式：`{prefix}_{sha256(kind + "\0" + account_scope + "\0" + client_value)[..32]}` -/// 其中 prefix 为 "cp"（conversation）或 "cw"（window） +/// 其中 prefix 为 `cp`（conversation）或 `cw`（window）。 fn build_account_scoped_identity(kind: &str, account_scope: &str, client_value: &str) -> String { let mut hasher = Sha256::new(); hasher.update(kind.as_bytes()); diff --git a/src/codex/gateway/fingerprint/model.rs b/crates/core/src/gateway/fingerprint.rs similarity index 86% rename from src/codex/gateway/fingerprint/model.rs rename to crates/core/src/gateway/fingerprint.rs index 150725fab..236c73f91 100644 --- a/src/codex/gateway/fingerprint/model.rs +++ b/crates/core/src/gateway/fingerprint.rs @@ -1,20 +1,31 @@ use indexmap::IndexMap; use serde::{Deserialize, Serialize}; +/// 上游请求指纹。 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct Fingerprint { + /// 客户端来源名。 pub originator: String, + /// 应用版本。 pub app_version: String, + /// 构建号。 pub build_number: String, + /// 平台名。 pub platform: String, + /// 架构名。 pub arch: String, + /// Chromium 主版本。 pub chromium_version: String, + /// User-Agent 模板。 pub user_agent_template: String, + /// 默认请求头。 pub default_headers: IndexMap, + /// 请求头顺序。 pub header_order: Vec, } impl Fingerprint { + /// 返回默认的 Codex Desktop 指纹。 pub fn default_codex_desktop() -> Self { Self { originator: "Codex Desktop".to_string(), @@ -29,10 +40,12 @@ impl Fingerprint { } } + /// 返回测试用默认指纹。 pub fn default_for_tests() -> Self { Self::default_codex_desktop() } + /// 返回默认请求头集合。 pub fn default_headers() -> IndexMap { let mut headers = IndexMap::new(); headers.insert( @@ -48,6 +61,7 @@ impl Fingerprint { headers } + /// 返回默认请求头顺序。 pub fn default_header_order() -> Vec { vec![ "authorization".to_string(), @@ -73,6 +87,7 @@ impl Fingerprint { ] } + /// 根据模板展开最终 User-Agent。 pub fn user_agent(&self) -> String { self.user_agent_template .replace("{version}", &self.app_version) @@ -80,6 +95,7 @@ impl Fingerprint { .replace("{arch}", &self.arch) } + /// 生成 `sec-ch-ua` 头值。 pub fn sec_ch_ua(&self) -> String { format!( "\"Chromium\";v=\"{}\", \"Not:A-Brand\";v=\"24\"", diff --git a/crates/core/src/gateway/installation.rs b/crates/core/src/gateway/installation.rs new file mode 100644 index 000000000..33eb49eae --- /dev/null +++ b/crates/core/src/gateway/installation.rs @@ -0,0 +1,50 @@ +//! Installation ID 的纯规则。 +//! +//! 该模块只负责 UUID 格式校验和生成。`~/.codex/installation_id` 与数据目录 +//! 持久化属于平台文件系统职责，放在 `platform::storage` 中。 + +use uuid::Uuid; + +/// 生成新的 installation ID。 +pub fn generate_installation_id() -> String { + Uuid::new_v4().to_string() +} + +/// 解析并规范化 installation ID。 +/// +/// 输入会先去除首尾空白；非 UUID 字符串返回 `None`。 +pub fn parse_installation_id(raw: &str) -> Option { + let trimmed = raw.trim(); + Uuid::parse_str(trimmed).ok()?; + Some(trimmed.to_string()) +} + +#[cfg(test)] +mod tests { + use uuid::Uuid; + + use super::*; + + #[test] + fn parse_installation_id_should_trim_and_accept_uuid() { + let id = Uuid::new_v4().to_string(); + + let parsed = parse_installation_id(&format!(" {id}\n")); + + assert_eq!(parsed, Some(id)); + } + + #[test] + fn parse_installation_id_should_reject_non_uuid() { + let parsed = parse_installation_id("not-a-uuid"); + + assert_eq!(parsed, None); + } + + #[test] + fn generate_installation_id_should_return_uuid() { + let generated = generate_installation_id(); + + assert!(Uuid::parse_str(&generated).is_ok()); + } +} diff --git a/crates/core/src/gateway/mod.rs b/crates/core/src/gateway/mod.rs new file mode 100644 index 000000000..e7ea9c3cc --- /dev/null +++ b/crates/core/src/gateway/mod.rs @@ -0,0 +1,10 @@ +//! 上游网关领域中的纯逻辑辅助。 + +/// 会话与缓存键派生。 +pub mod conversation; +/// 上游指纹。 +pub mod fingerprint; +/// 安装 ID 生成。 +pub mod installation; +/// 上游网关端口。 +pub mod ports; diff --git a/crates/core/src/gateway/ports.rs b/crates/core/src/gateway/ports.rs new file mode 100644 index 000000000..64318a23b --- /dev/null +++ b/crates/core/src/gateway/ports.rs @@ -0,0 +1,42 @@ +//! 上游 Codex 网关端口。 + +use async_trait::async_trait; +use thiserror::Error; + +use crate::models::model::BackendModelEntry; + +/// 拉取上游模型目录时的请求上下文。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CodexModelCatalogRequest<'a> { + /// 当前账号访问令牌。 + pub access_token: &'a str, + /// 上游账号 ID。 + pub account_id: Option<&'a str>, + /// 请求 ID。 + pub request_id: &'a str, + /// Codex installation id。 + pub installation_id: Option<&'a str>, + /// 订阅计划类型。 + pub plan_type: &'a str, +} + +/// 上游模型目录客户端错误。 +#[derive(Debug, Error)] +pub enum CodexModelCatalogClientError { + /// 上游请求失败。 + #[error("model catalog request failed: {message}")] + RequestFailed { + /// 错误说明。 + message: String, + }, +} + +/// 上游模型目录客户端。 +#[async_trait] +pub trait CodexModelCatalogClient: Send + Sync + 'static { + /// 读取当前账号可见的上游模型目录。 + async fn fetch_models( + &self, + request: &CodexModelCatalogRequest<'_>, + ) -> Result, CodexModelCatalogClientError>; +} diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs new file mode 100644 index 000000000..165cb93e0 --- /dev/null +++ b/crates/core/src/lib.rs @@ -0,0 +1,14 @@ +#![deny(missing_docs)] + +//! 核心领域层，承载纯业务模型、协议语义和策略。 + +pub mod accounts; +pub mod admin; +pub mod auth; +pub mod error; +pub mod events; +pub mod gateway; +pub mod models; +pub mod protocol; +pub mod serving; +pub mod usage; diff --git a/src/codex/models/catalog.rs b/crates/core/src/models/catalog.rs similarity index 73% rename from src/codex/models/catalog.rs rename to crates/core/src/models/catalog.rs index 7d2db12b2..13e7cf0b3 100644 --- a/src/codex/models/catalog.rs +++ b/crates/core/src/models/catalog.rs @@ -1,67 +1,16 @@ +//! 模型目录聚合与纯解析规则。 + use std::collections::{BTreeMap, BTreeSet}; -use serde::{Deserialize, Serialize}; - -use crate::config::ModelConfig; - +use crate::models::model::{ + BackendModelEntry, BackendReasoningEffort, CodexModelInfo, ModelConfig, ModelDebugEntry, + ModelPlanSnapshot, ModelStoreDebug, ParsedModelName, ReasoningEffortInfo, +}; const SERVICE_TIER_SUFFIXES: [&str; 2] = ["fast", "flex"]; const REASONING_EFFORT_SUFFIXES: [&str; 6] = ["none", "minimal", "low", "medium", "high", "xhigh"]; const STATIC_REASONING_EFFORTS: [&str; 4] = ["low", "medium", "high", "xhigh"]; -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct ReasoningEffortInfo { - pub reasoning_effort: String, - pub description: String, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct CodexModelInfo { - pub id: String, - pub display_name: String, - pub description: String, - pub is_default: bool, - pub supported_reasoning_efforts: Vec, - pub default_reasoning_effort: String, - pub input_modalities: Vec, - pub output_modalities: Vec, - pub supports_personality: bool, - pub upgrade: Option, - pub source: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub context_window: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub max_context_window: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub max_output_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub truncation_policy_limit: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ParsedModelName { - pub model_id: String, - pub reasoning_effort: Option, - pub service_tier: Option, -} - -#[derive(Debug, Clone, Serialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct ModelStoreDebug { - pub total_models: usize, - pub static_models: usize, - pub alias_count: usize, - pub models: Vec, -} - -#[derive(Debug, Clone, Serialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct ModelDebugEntry { - pub id: String, - pub source: String, -} - +/// 模型目录聚合。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct ModelCatalog { models: Vec, @@ -70,57 +19,8 @@ pub struct ModelCatalog { model_plan_index: BTreeMap>, } -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] -pub struct BackendReasoningEffort { - #[serde(default, rename = "reasoning_effort", alias = "reasoningEffort")] - pub reasoning_effort: Option, - #[serde(default)] - pub effort: Option, - #[serde(default)] - pub description: Option, -} - -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] -pub struct BackendTruncationPolicy { - pub limit: Option, -} - -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] -pub struct BackendModelEntry { - pub slug: Option, - pub id: Option, - pub name: Option, - pub display_name: Option, - pub description: Option, - pub is_default: Option, - pub default_reasoning_effort: Option, - pub default_reasoning_level: Option, - #[serde(default)] - pub supported_reasoning_efforts: Vec, - #[serde(default)] - pub supported_reasoning_levels: Vec, - pub input_modalities: Option>, - pub output_modalities: Option>, - pub supports_personality: Option, - pub upgrade: Option, - #[serde(alias = "contextWindow")] - pub context_window: Option, - #[serde(alias = "maxContextWindow")] - pub max_context_window: Option, - #[serde(alias = "maxOutputTokens")] - pub max_output_tokens: Option, - #[serde(alias = "truncationPolicy")] - pub truncation_policy: Option, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct ModelPlanSnapshot { - pub plan_type: String, - pub models: Vec, -} - impl ModelPlanSnapshot { + /// 将上游模型条目标准化为可持久化快照。 pub fn from_backend_entries( plan_type: impl Into, entries: Vec, @@ -133,6 +33,7 @@ impl ModelPlanSnapshot { } impl ModelCatalog { + /// 从静态配置构造模型目录。 pub fn from_config(config: &ModelConfig) -> Self { let default_model = config.default_model.trim().to_string(); let aliases = normalize_aliases(&config.aliases); @@ -144,6 +45,7 @@ impl ModelCatalog { } } + /// 从静态配置和后端快照构造模型目录。 pub fn from_config_and_snapshots( config: &ModelConfig, snapshots: &[ModelPlanSnapshot], @@ -159,7 +61,7 @@ impl ModelCatalog { for snapshot in snapshots { for model in &snapshot.models { - // 后端模型是按账号 plan 返回的；catalog 需要去重展示，调度层仍要保留 model -> plans。 + // 后端模型按账号计划返回；目录对外去重，但仍保留 model -> plans 映射。 models_by_id.insert(model.id.clone(), model.clone()); let plans = model_plan_index.entry(model.id.clone()).or_default(); if !plans.contains(&snapshot.plan_type) { @@ -181,10 +83,12 @@ impl ModelCatalog { } } + /// 返回对外可见的模型列表副本。 pub fn models(&self) -> Vec { self.models.clone() } + /// 按模型 ID 查询模型信息。 pub fn model_info(&self, model_id: &str) -> Option { self.models .iter() @@ -192,6 +96,7 @@ impl ModelCatalog { .cloned() } + /// 返回模型目录调试视图。 pub fn debug(&self) -> ModelStoreDebug { ModelStoreDebug { total_models: self.models.len(), @@ -208,10 +113,12 @@ impl ModelCatalog { } } + /// 返回 model -> plans allowlist。 pub fn model_plan_allowlist(&self) -> BTreeMap> { self.model_plan_index.clone() } + /// 判断输入模型名是否在目录中可识别。 pub fn is_recognized_model_name(&self, input: &str) -> bool { let trimmed = input.trim(); if trimmed.is_empty() { @@ -231,6 +138,7 @@ impl ModelCatalog { || self.model_info(&stripped.model_name).is_some() } + /// 解析外部传入的模型名，提取别名、推理强度和服务层级后缀。 pub fn parse_model_name(&self, input: &str) -> ParsedModelName { let trimmed = input.trim(); if self.aliases.contains_key(trimmed) || self.model_info(trimmed).is_some() { @@ -249,6 +157,7 @@ impl ModelCatalog { } } + /// 将标准化模型名重新拼成展示名。 pub fn build_display_model_name(parsed: &ParsedModelName) -> String { let mut name = parsed.model_id.clone(); if let Some(reasoning_effort) = &parsed.reasoning_effort { @@ -361,9 +270,13 @@ fn normalize_backend_model(raw: BackendModelEntry) -> CodexModelInfo { let id = first_non_empty([raw.slug.as_deref(), raw.id.as_deref(), raw.name.as_deref()]) .unwrap_or("unknown") .to_string(); - let display_name = first_non_empty([raw.display_name.as_deref(), raw.name.as_deref()]) - .unwrap_or(&id) - .to_string(); + let display_name = first_non_empty([ + raw.display_name.as_deref(), + raw.title.as_deref(), + raw.name.as_deref(), + ]) + .unwrap_or(&id) + .to_string(); let reasoning_entries = if raw.supported_reasoning_efforts.is_empty() { raw.supported_reasoning_levels } else { diff --git a/crates/core/src/models/mod.rs b/crates/core/src/models/mod.rs new file mode 100644 index 000000000..4af25742e --- /dev/null +++ b/crates/core/src/models/mod.rs @@ -0,0 +1,10 @@ +//! 模型目录领域。 + +/// 模型目录聚合与纯解析规则。 +pub mod catalog; +/// 模型目录配置与纯数据类型。 +pub mod model; +/// 模型快照存储端口。 +pub mod ports; +/// 模型目录服务用例。 +pub mod service; diff --git a/crates/core/src/models/model.rs b/crates/core/src/models/model.rs new file mode 100644 index 000000000..7dbc3a77e --- /dev/null +++ b/crates/core/src/models/model.rs @@ -0,0 +1,183 @@ +//! 模型目录领域中的纯数据类型。 + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +/// 模型目录默认值与别名配置。 +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +pub struct ModelConfig { + /// 默认模型名。 + pub default_model: String, + /// 默认推理强度。 + pub default_reasoning_effort: Option, + /// 默认服务层级。 + pub service_tier: Option, + /// 模型别名映射。 + pub aliases: BTreeMap, +} + +/// 单个 reasoning effort 的标准化展示信息。 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct ReasoningEffortInfo { + /// 推理强度名。 + pub reasoning_effort: String, + /// 展示描述。 + pub description: String, +} + +/// 对外暴露的单个模型目录条目。 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct CodexModelInfo { + /// 模型唯一 ID。 + pub id: String, + /// 展示名。 + pub display_name: String, + /// 描述文本。 + pub description: String, + /// 是否默认模型。 + pub is_default: bool, + /// 支持的推理强度列表。 + pub supported_reasoning_efforts: Vec, + /// 默认推理强度。 + pub default_reasoning_effort: String, + /// 输入模态。 + pub input_modalities: Vec, + /// 输出模态。 + pub output_modalities: Vec, + /// 是否支持 personality。 + pub supports_personality: bool, + /// 升级提示。 + pub upgrade: Option, + /// 来源标记。 + pub source: String, + /// 当前上下文窗口。 + #[serde(skip_serializing_if = "Option::is_none")] + pub context_window: Option, + /// 最大上下文窗口。 + #[serde(skip_serializing_if = "Option::is_none")] + pub max_context_window: Option, + /// 最大输出 token 数。 + #[serde(skip_serializing_if = "Option::is_none")] + pub max_output_tokens: Option, + /// 截断策略限制。 + #[serde(skip_serializing_if = "Option::is_none")] + pub truncation_policy_limit: Option, +} + +/// 解析模型名后得到的标准化结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ParsedModelName { + /// 最终模型 ID。 + pub model_id: String, + /// 推理强度后缀。 + pub reasoning_effort: Option, + /// 服务层级后缀。 + pub service_tier: Option, +} + +/// 模型目录调试视图。 +#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct ModelStoreDebug { + /// 模型总数。 + pub total_models: usize, + /// 静态模型数。 + pub static_models: usize, + /// 别名数量。 + pub alias_count: usize, + /// 调试模型条目。 + pub models: Vec, +} + +/// 单个调试模型条目。 +#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct ModelDebugEntry { + /// 模型 ID。 + pub id: String, + /// 来源标记。 + pub source: String, +} + +/// 按计划类型持久化的模型快照。 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct ModelPlanSnapshot { + /// 订阅计划类型。 + pub plan_type: String, + /// 该计划可见的模型列表。 + pub models: Vec, +} + +/// 上游声明的 reasoning effort 条目。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct BackendReasoningEffort { + /// reasoning_effort 字段。 + #[serde(default, rename = "reasoning_effort", alias = "reasoningEffort")] + pub reasoning_effort: Option, + /// effort 字段的兼容别名。 + #[serde(default)] + pub effort: Option, + /// 展示描述。 + #[serde(default)] + pub description: Option, +} + +/// 上游截断策略定义。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct BackendTruncationPolicy { + /// 上下文截断限制。 + pub limit: Option, +} + +/// 上游 `/codex/models` 返回的原始模型条目。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct BackendModelEntry { + /// slug。 + pub slug: Option, + /// id。 + pub id: Option, + /// name。 + pub name: Option, + /// display_name。 + pub display_name: Option, + /// title。 + pub title: Option, + /// description。 + pub description: Option, + /// 是否默认模型。 + pub is_default: Option, + /// 默认 reasoning effort。 + pub default_reasoning_effort: Option, + /// 默认 reasoning level。 + pub default_reasoning_level: Option, + /// 支持的 reasoning effort 列表。 + #[serde(default)] + pub supported_reasoning_efforts: Vec, + /// 支持的 reasoning level 列表。 + #[serde(default)] + pub supported_reasoning_levels: Vec, + /// 输入模态。 + pub input_modalities: Option>, + /// 输出模态。 + pub output_modalities: Option>, + /// 是否支持 personality。 + pub supports_personality: Option, + /// 升级提示。 + pub upgrade: Option, + /// contextWindow 别名。 + #[serde(alias = "contextWindow")] + pub context_window: Option, + /// maxContextWindow 别名。 + #[serde(alias = "maxContextWindow")] + pub max_context_window: Option, + /// maxOutputTokens 别名。 + #[serde(alias = "maxOutputTokens")] + pub max_output_tokens: Option, + /// truncationPolicy 别名。 + #[serde(alias = "truncationPolicy")] + pub truncation_policy: Option, +} diff --git a/crates/core/src/models/ports.rs b/crates/core/src/models/ports.rs new file mode 100644 index 000000000..37f25a666 --- /dev/null +++ b/crates/core/src/models/ports.rs @@ -0,0 +1,33 @@ +//! 模型快照存储端口。 + +use async_trait::async_trait; +use thiserror::Error; + +use crate::models::model::ModelPlanSnapshot; + +/// 模型快照存储错误。 +#[derive(Debug, Error)] +pub enum ModelSnapshotStoreError { + /// 底层存储失败。 + #[error("model snapshot store operation failed: {message}")] + OperationFailed { + /// 错误说明。 + message: String, + }, +} + +/// 模型快照存储结果类型。 +pub type ModelSnapshotStoreResult = Result; + +/// 模型快照存储端口。 +#[async_trait] +pub trait ModelSnapshotStore: Send + Sync + 'static { + /// 用单个计划快照替换同名快照。 + async fn replace_plan_snapshot( + &self, + snapshot: &ModelPlanSnapshot, + ) -> ModelSnapshotStoreResult<()>; + + /// 列出所有计划快照。 + async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult>; +} diff --git a/crates/core/src/models/service.rs b/crates/core/src/models/service.rs new file mode 100644 index 000000000..cb465a214 --- /dev/null +++ b/crates/core/src/models/service.rs @@ -0,0 +1,233 @@ +//! 模型目录服务用例。 + +use std::{collections::BTreeMap, sync::Arc}; + +use thiserror::Error; + +use crate::{ + accounts::model::{Account, AccountStatus}, + gateway::ports::{CodexModelCatalogClient, CodexModelCatalogRequest}, + models::{ + catalog::ModelCatalog, + model::{ModelConfig, ModelPlanSnapshot}, + ports::{ModelSnapshotStore, ModelSnapshotStoreError}, + }, +}; + +/// 模型刷新摘要。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ModelRefreshResult { + /// 成功刷新并写入的计划数。 + pub refreshed_plans: usize, + /// 本次成功写入的模型数。 + pub model_count: usize, + /// 刷新失败的计划数。 + pub failed_plans: usize, +} + +/// 模型服务错误。 +#[derive(Debug, Error, PartialEq, Eq)] +pub enum ModelServiceError { + /// 没有注入快照存储。 + #[error("model snapshot store is unavailable")] + SnapshotStoreUnavailable, + /// 没有注入上游模型客户端。 + #[error("model catalog client is unavailable")] + UpstreamClientUnavailable, + /// 没有可用账号。 + #[error("no active accounts available for model refresh")] + NoAccounts, + /// 快照写入失败。 + #[error("failed to store model snapshot")] + StoreSnapshot, + /// 刷新后重新读取快照失败。 + #[error("failed to load model snapshots")] + LoadSnapshots, + /// 所有计划都刷新失败。 + #[error("all model refresh plans failed")] + AllPlansFailed(ModelRefreshResult), +} + +/// 模型到可用计划的映射。 +pub type ModelPlanAllowlist = BTreeMap>; + +/// 可共享更新的模型计划映射缓存。 +pub type SharedModelPlanAllowlist = Arc>; + +/// 模型目录服务。 +#[derive(Clone)] +pub struct ModelService { + config: ModelConfig, + snapshot_store: Option>, + upstream_client: Option>, + model_plan_allowlist: Option, +} + +impl ModelService { + /// 构造模型服务。 + pub fn new( + config: ModelConfig, + snapshot_store: Option>, + upstream_client: Option>, + model_plan_allowlist: Option, + ) -> Self { + Self { + config, + snapshot_store, + upstream_client, + model_plan_allowlist, + } + } + + /// 返回模型目录的静态配置。 + pub fn config(&self) -> &ModelConfig { + &self.config + } + + /// 构造当前对外暴露的模型目录。 + pub async fn catalog(&self) -> ModelCatalog { + let Some(snapshot_store) = self.snapshot_store.as_ref() else { + return ModelCatalog::from_config(&self.config); + }; + + match snapshot_store.list_plan_snapshots().await { + Ok(snapshots) if !snapshots.is_empty() => { + ModelCatalog::from_config_and_snapshots(&self.config, &snapshots) + } + Ok(_) => ModelCatalog::from_config(&self.config), + Err(error) => { + tracing::warn!(error = %error, "加载模型快照失败，回退到静态目录"); + ModelCatalog::from_config(&self.config) + } + } + } + + /// 刷新活跃账号对应的后端模型目录。 + pub async fn refresh_backend_models( + &self, + accounts: &[Account], + request_id: &str, + ) -> Result { + self.refresh_backend_models_with_installation_id(accounts, request_id, None) + .await + } + + /// 使用运行时 installation id 刷新活跃账号对应的后端模型目录。 + pub async fn refresh_backend_models_with_installation_id( + &self, + accounts: &[Account], + request_id: &str, + installation_id: Option<&str>, + ) -> Result { + let snapshot_store = self + .snapshot_store + .as_ref() + .ok_or(ModelServiceError::SnapshotStoreUnavailable)?; + let upstream_client = self + .upstream_client + .as_ref() + .ok_or(ModelServiceError::UpstreamClientUnavailable)?; + + let plan_accounts = distinct_active_plan_accounts(accounts); + if plan_accounts.is_empty() { + return Err(ModelServiceError::NoAccounts); + } + + let mut result = ModelRefreshResult { + refreshed_plans: 0, + model_count: 0, + failed_plans: 0, + }; + + for (plan_type, account) in plan_accounts { + let request = CodexModelCatalogRequest { + access_token: &account.access_token, + account_id: account.account_id.as_deref(), + request_id, + installation_id, + plan_type: &plan_type, + }; + + let entries = match upstream_client.fetch_models(&request).await { + Ok(entries) if !entries.is_empty() => entries, + Ok(_) => { + result.failed_plans += 1; + continue; + } + Err(error) => { + tracing::warn!(error = %error, plan_type, "刷新后端模型失败"); + result.failed_plans += 1; + continue; + } + }; + + let snapshot = ModelPlanSnapshot::from_backend_entries(plan_type, entries); + result.model_count += snapshot.models.len(); + snapshot_store + .replace_plan_snapshot(&snapshot) + .await + .map_err(map_store_snapshot_error)?; + result.refreshed_plans += 1; + } + + if result.refreshed_plans == 0 { + return Err(ModelServiceError::AllPlansFailed(result)); + } + + let allowlist = self.model_plan_allowlist_from_store().await?; + if let Some(shared_allowlist) = &self.model_plan_allowlist { + *shared_allowlist.lock().await = allowlist; + } + + Ok(result) + } + + /// 读取当前缓存的 model -> plans allowlist。 + pub async fn model_plan_allowlist(&self) -> Result { + self.model_plan_allowlist_from_store().await + } + + async fn model_plan_allowlist_from_store( + &self, + ) -> Result { + let snapshot_store = self + .snapshot_store + .as_ref() + .ok_or(ModelServiceError::SnapshotStoreUnavailable)?; + let snapshots = snapshot_store + .list_plan_snapshots() + .await + .map_err(map_load_snapshots_error)?; + Ok( + ModelCatalog::from_config_and_snapshots(&self.config, &snapshots) + .model_plan_allowlist(), + ) + } +} + +fn distinct_active_plan_accounts(accounts: &[Account]) -> Vec<(String, Account)> { + let mut by_plan = 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_with(|| account.clone()); + } + + by_plan.into_iter().collect() +} + +fn map_store_snapshot_error(_: ModelSnapshotStoreError) -> ModelServiceError { + ModelServiceError::StoreSnapshot +} + +fn map_load_snapshots_error(_: ModelSnapshotStoreError) -> ModelServiceError { + ModelServiceError::LoadSnapshots +} diff --git a/crates/core/src/protocol/codex/chat.rs b/crates/core/src/protocol/codex/chat.rs new file mode 100644 index 000000000..03a6878e9 --- /dev/null +++ b/crates/core/src/protocol/codex/chat.rs @@ -0,0 +1,11 @@ +//! Codex Chat 扩展。 + +use serde::{Deserialize, Serialize}; + +/// Codex chat 侧的扩展选项。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct CodexChatOptions { + /// 是否包含 reasoning 摘要。 + pub include_reasoning: bool, +} diff --git a/src/codex/gateway/transport/rate_limits.rs b/crates/core/src/protocol/codex/events.rs similarity index 55% rename from src/codex/gateway/transport/rate_limits.rs rename to crates/core/src/protocol/codex/events.rs index 96085edca..57bd43b52 100644 --- a/src/codex/gateway/transport/rate_limits.rs +++ b/crates/core/src/protocol/codex/events.rs @@ -1,60 +1,234 @@ +//! Codex 事件语义、用量提取与限流解析。 + use std::{ collections::BTreeMap, time::{SystemTime, UNIX_EPOCH}, }; use chrono::{DateTime, Duration, Utc}; +use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; +use crate::protocol::codex::sse::{parse_sse_events, SseError}; + +/// 从 Codex/OpenAI usage 结构中提取出的标准化 token 用量。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct TokenUsage { + /// 输入 token 数。 + pub input_tokens: u64, + /// 输出 token 数。 + pub output_tokens: u64, + /// 命中缓存的输入 token 数。 + pub cached_tokens: u64, + /// 图片工具输入 token 数。 + pub image_input_tokens: u64, + /// 图片工具输出 token 数。 + pub image_output_tokens: u64, + /// 主模型总 token 数，不包含图片工具 token。 + pub total_tokens: u64, +} + +impl TokenUsage { + /// 构造仅包含主模型 token 的用量快照。 + pub fn new(input_tokens: u64, output_tokens: u64, cached_tokens: u64) -> Self { + Self { + input_tokens, + output_tokens, + cached_tokens, + image_input_tokens: 0, + image_output_tokens: 0, + total_tokens: input_tokens + output_tokens, + } + } + + /// 将两份用量按字段累加。 + pub fn merged(self, other: Self) -> Self { + Self { + input_tokens: self.input_tokens + other.input_tokens, + output_tokens: self.output_tokens + other.output_tokens, + cached_tokens: self.cached_tokens + other.cached_tokens, + image_input_tokens: self.image_input_tokens + other.image_input_tokens, + image_output_tokens: self.image_output_tokens + other.image_output_tokens, + total_tokens: self.total_tokens + other.total_tokens, + } + } +} + +/// 从单个 JSON 响应体中提取用量。 +/// +/// 该函数同时兼容 Codex usage 结构和 OpenAI usage 结构。 +pub fn extract_usage(body: &Value) -> Option { + let usage = body.get("usage").unwrap_or(body); + if !usage.is_object() { + return None; + } + + let input_tokens = number_field(usage, "input_tokens") + .or_else(|| number_field(usage, "prompt_tokens")) + .unwrap_or_default(); + let output_tokens = number_field(usage, "output_tokens") + .or_else(|| number_field(usage, "completion_tokens")) + .unwrap_or_default(); + let cached_tokens = nested_number_field(usage, &["input_tokens_details", "cached_tokens"]) + .or_else(|| nested_number_field(usage, &["prompt_tokens_details", "cached_tokens"])) + .or_else(|| number_field(usage, "cached_tokens")) + .unwrap_or_default(); + let image_input_tokens = + nested_number_field(body, &["tool_usage", "image_gen", "input_tokens"]).unwrap_or_default(); + let image_output_tokens = + nested_number_field(body, &["tool_usage", "image_gen", "output_tokens"]) + .unwrap_or_default(); + let total_tokens = number_field(usage, "total_tokens").unwrap_or(input_tokens + output_tokens); + + let has_usage = [ + "input_tokens", + "output_tokens", + "prompt_tokens", + "completion_tokens", + "cached_tokens", + "total_tokens", + ] + .iter() + .any(|field| usage.get(*field).is_some()) + || usage.get("input_tokens_details").is_some() + || usage.get("prompt_tokens_details").is_some() + || image_input_tokens > 0 + || image_output_tokens > 0; + + has_usage.then_some(TokenUsage { + input_tokens, + output_tokens, + cached_tokens, + image_input_tokens, + image_output_tokens, + total_tokens, + }) +} + +/// 从完整 SSE 文本中提取最终可见用量。 +/// +/// 如果存在 `response.completed` 的 usage，则优先返回它；否则回退到最后一条 +/// 可见 usage 事件。 +/// +/// # Errors +/// +/// 当输入不是合法 SSE 流时，返回 [`SseError`]。 +pub fn extract_sse_usage(body: &str) -> Result, SseError> { + let events = parse_sse_events(body)?; + let mut fallback_usage: Option = None; + + for event in events { + if event.data == "[DONE]" { + continue; + } + let Ok(value) = serde_json::from_str::(&event.data) else { + continue; + }; + let event_usage = value + .get("response") + .and_then(extract_usage) + .or_else(|| extract_usage(&value)); + if let Some(event_usage) = event_usage { + if event.event.as_deref() == Some("response.completed") { + return Ok(Some(event_usage)); + } + fallback_usage = Some(event_usage); + } + } + + Ok(fallback_usage) +} + +/// 从上游错误响应体中提取 retry-after 秒数。 +/// +/// 支持结构化的 `resets_in_seconds` / `resets_at`，也支持官方错误消息里的 +/// `try again in 11.054s` / `try again in 28ms` 文本。 +pub fn retry_after_seconds_from_body(body: &str) -> Option { + let value = serde_json::from_str::(body).ok()?; + let error = value + .pointer("/response/error") + .or_else(|| value.get("error")) + .unwrap_or(&value); + if let Some(seconds) = error + .get("resets_in_seconds") + .and_then(Value::as_u64) + .filter(|seconds| *seconds > 0) + { + return Some(seconds); + } + retry_after_seconds_from_resets_at(error) + .or_else(|| retry_after_seconds_from_rate_limit_message(error)) +} + +/// 标准化的单个限流窗口。 #[derive(Debug, Clone, Copy, PartialEq)] pub struct RateLimitWindow { + /// 已使用百分比。 pub used_percent: f64, + /// 窗口分钟数。 pub window_minutes: Option, + /// 重置时间戳。 pub reset_at: Option, } impl RateLimitWindow { + /// 以秒返回限流窗口大小。 pub fn limit_window_seconds(self) -> Option { self.window_minutes .and_then(|minutes| minutes.checked_mul(60)) } + /// 将重置时间戳转换为 UTC 时间。 pub fn reset_at_datetime(self) -> Option> { let reset_at = self.reset_at?; DateTime::::from_timestamp(reset_at, 0) } } +/// 复合限流信息，通常用于 code review 配额。 #[derive(Debug, Clone, Copy, PartialEq)] pub struct RateLimitDetails { + /// 当前请求是否被允许。 pub allowed: Option, + /// 当前窗口是否已经触顶。 pub limit_reached: Option, + /// 主限流窗口。 pub primary: Option, + /// 次级限流窗口。 pub secondary: Option, } +/// 从 header 或内部事件中解析出的完整限流状态。 #[derive(Debug, Clone, Copy, PartialEq)] pub struct ParsedRateLimits { + /// 主限流窗口。 pub primary: Option, + /// 次级限流窗口。 pub secondary: Option, + /// code review 限流窗口。 pub code_review: Option, } impl ParsedRateLimits { + /// 返回主窗口的 UTC 重置时间。 pub fn primary_reset_at(self) -> Option> { self.primary.and_then(RateLimitWindow::reset_at_datetime) } + /// 返回主窗口的秒级窗口长度。 pub fn primary_limit_window_seconds(self) -> Option { self.primary.and_then(RateLimitWindow::limit_window_seconds) } + /// 返回主窗口是否已经触顶。 pub fn primary_limit_reached(self) -> bool { self.primary .is_some_and(|window| window.used_percent >= 100.0) } } +/// 从响应头对中解析限流信息。 pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option { let mut normalized = BTreeMap::new(); for (name, value) in headers { @@ -78,6 +252,7 @@ pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option Option { if value .get("type") @@ -117,12 +292,14 @@ pub fn parse_rate_limits_event(value: &Value) -> Option { }) } +/// 从原始 JSON 文本中解析内部 `codex.rate_limits` 事件。 pub fn parse_rate_limits_event_raw(raw: &str) -> Option { serde_json::from_str::(raw) .ok() .and_then(|value| parse_rate_limits_event(&value)) } +/// 将解析后的限流状态转换为配额响应体。 pub fn rate_limit_quota( rate_limits: &ParsedRateLimits, plan_type: Option<&str>, @@ -142,6 +319,7 @@ pub fn rate_limit_quota( quota } +/// 将限流状态转换回 HTTP 头键值对。 pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(String, String)> { let mut headers = Vec::new(); push_window_headers(&mut headers, "x-codex-primary", rate_limits.primary); @@ -161,6 +339,7 @@ pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(Strin headers } +/// 为退避秒数增加抖动，避免同一时间窗口内集中重试。 pub fn cooldown_with_jitter(seconds: u64, factor_bps: u16) -> Duration { let factor_bps = u64::from(factor_bps); let variance = seconds.saturating_mul(factor_bps) / 10_000; @@ -176,6 +355,94 @@ pub fn cooldown_with_jitter(seconds: u64, factor_bps: u16) -> Duration { Duration::seconds(lower.saturating_add(offset).min(i64::MAX as u64) as i64) } +fn number_field(value: &Value, field: &str) -> Option { + value.get(field)?.as_u64() +} + +fn nested_number_field(value: &Value, path: &[&str]) -> Option { + let mut current = value; + for segment in path { + current = current.get(*segment)?; + } + current.as_u64() +} + +fn retry_after_seconds_from_resets_at(error: &Value) -> Option { + let resets_at = error.get("resets_at").and_then(Value::as_u64)?; + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .ok()? + .as_secs(); + (resets_at > now).then_some(resets_at - now) +} + +fn retry_after_seconds_from_rate_limit_message(error: &Value) -> Option { + let code = error + .get("code") + .or_else(|| error.get("type")) + .and_then(Value::as_str)?; + if code != "rate_limit_exceeded" { + return None; + } + let message = error.get("message").and_then(Value::as_str)?; + parse_try_again_delay_seconds(message) +} + +fn parse_try_again_delay_seconds(message: &str) -> Option { + let lower = message.to_ascii_lowercase(); + let marker = "try again in"; + let offset = lower.find(marker)? + marker.len(); + let remainder = message.get(offset..)?.trim_start(); + let number_end = number_prefix_len(remainder)?; + let value = remainder.get(..number_end)?.parse::().ok()?; + if !value.is_finite() || value <= 0.0 { + return None; + } + let unit_text = remainder + .get(number_end..)? + .trim_start() + .to_ascii_lowercase(); + let unit = unit_token(&unit_text)?; + match unit { + "ms" => positive_seconds_ceil(value / 1000.0), + "s" | "second" | "seconds" => positive_seconds_ceil(value), + _ => None, + } +} + +fn number_prefix_len(input: &str) -> Option { + let mut seen_digit = false; + let mut seen_dot = false; + let mut end = 0; + for (index, ch) in input.char_indices() { + if ch.is_ascii_digit() { + seen_digit = true; + end = index + ch.len_utf8(); + } else if ch == '.' && !seen_dot { + seen_dot = true; + end = index + ch.len_utf8(); + } else { + break; + } + } + seen_digit.then_some(end) +} + +fn unit_token(input: &str) -> Option<&str> { + let end = input + .char_indices() + .find_map(|(index, ch)| (!ch.is_ascii_alphabetic()).then_some(index)) + .unwrap_or(input.len()); + (end > 0).then_some(&input[..end]) +} + +fn positive_seconds_ceil(seconds: f64) -> Option { + if !seconds.is_finite() || seconds <= 0.0 || seconds > u64::MAX as f64 { + return None; + } + Some(seconds.ceil() as u64) +} + fn push_window_headers( headers: &mut Vec<(String, String)>, prefix: &str, @@ -353,115 +620,3 @@ fn is_review_limit_name(value: Option<&str>) -> bool { "review" | "code_review" | "codex_review" | "codex_code_review" ) } - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::{ - parse_rate_limit_headers, parse_rate_limits_event, rate_limit_quota, RateLimitWindow, - }; - - #[test] - fn parse_rate_limit_headers_should_extract_primary_secondary_and_review_windows() { - let headers = vec![ - ( - "x-codex-primary-used-percent".to_string(), - "100".to_string(), - ), - ( - "x-codex-primary-window-minutes".to_string(), - "5".to_string(), - ), - ( - "x-codex-primary-reset-at".to_string(), - "1893456300".to_string(), - ), - ( - "x-codex-secondary-used-percent".to_string(), - "42.5".to_string(), - ), - ( - "x-codex-secondary-window-minutes".to_string(), - "10080".to_string(), - ), - ( - "x-codex-code-review-primary-used-percent".to_string(), - "80".to_string(), - ), - ( - "x-codex-code-review-primary-reset-at".to_string(), - "1893456600".to_string(), - ), - ]; - - let parsed = parse_rate_limit_headers(&headers).unwrap(); - - assert_eq!( - parsed.primary, - Some(RateLimitWindow { - used_percent: 100.0, - window_minutes: Some(5), - reset_at: Some(1_893_456_300), - }) - ); - assert_eq!(parsed.secondary.unwrap().window_minutes, Some(10080)); - assert_eq!( - parsed.code_review.unwrap().primary.unwrap().reset_at, - Some(1_893_456_600) - ); - } - - #[test] - fn parse_rate_limits_event_should_extract_internal_websocket_rate_limits() { - let event = json!({ - "type": "codex.rate_limits", - "rate_limits": { - "primary": { - "used_percent": 99.5, - "window_minutes": 300, - "reset_at": 1893456300 - }, - "secondary": { - "used_percent": 10, - "window_minutes": 10080, - "reset_at": 1894056000 - } - } - }); - - let parsed = parse_rate_limits_event(&event).unwrap(); - - assert_eq!(parsed.primary.unwrap().used_percent, 99.5); - assert_eq!(parsed.secondary.unwrap().reset_at, Some(1_894_056_000)); - } - - #[test] - fn rate_limit_quota_should_preserve_existing_credits_when_passive_data_lacks_credits() { - let headers = vec![ - ("x-codex-primary-used-percent".to_string(), "25".to_string()), - ( - "x-codex-primary-window-minutes".to_string(), - "5".to_string(), - ), - ( - "x-codex-primary-reset-at".to_string(), - "1893456300".to_string(), - ), - ]; - let existing = json!({ - "credits": { - "has_credits": true, - "unlimited": false, - "balance": 12 - } - }); - let parsed = parse_rate_limit_headers(&headers).unwrap(); - - let quota = rate_limit_quota(&parsed, Some("plus"), Some(&existing)); - - assert_eq!(quota["plan_type"], "plus"); - assert_eq!(quota["rate_limit"]["remaining_percent"], 75); - assert_eq!(quota["credits"]["balance"], 12); - } -} diff --git a/crates/core/src/protocol/codex/mod.rs b/crates/core/src/protocol/codex/mod.rs new file mode 100644 index 000000000..06dd44d50 --- /dev/null +++ b/crates/core/src/protocol/codex/mod.rs @@ -0,0 +1,14 @@ +//! Codex 协议类型。 + +/// Codex Chat 扩展字段。 +pub mod chat; +/// Codex 事件语义与用量提取。 +pub mod events; +/// Codex Responses / Compact 请求体。 +pub mod responses; +/// Codex JSON schema 处理。 +pub(crate) mod schema; +/// SSE 事件解析。 +pub mod sse; +/// WebSocket 帧类型与纯编解码。 +pub mod websocket; diff --git a/src/codex/gateway/transport/types.rs b/crates/core/src/protocol/codex/responses.rs similarity index 74% rename from src/codex/gateway/transport/types.rs rename to crates/core/src/protocol/codex/responses.rs index 26f9477cb..9c4518880 100644 --- a/src/codex/gateway/transport/types.rs +++ b/crates/core/src/protocol/codex/responses.rs @@ -1,64 +1,95 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; +/// Codex Responses 上游请求体。 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CodexResponsesRequest { + /// 模型名。 pub model: String, + /// 指令文本。 pub instructions: String, + /// 输入消息与结构化条目。 pub input: Vec, + /// 是否流式返回。 pub stream: bool, + /// 是否要求上游存储响应。 pub store: bool, #[serde(skip_serializing_if = "Option::is_none")] + /// reasoning 配置。 pub reasoning: Option, #[serde(skip_serializing_if = "Option::is_none")] + /// 工具定义。 pub tools: Option>, #[serde(skip_serializing_if = "Option::is_none")] + /// 工具选择策略。 pub tool_choice: Option, #[serde(skip_serializing_if = "Option::is_none")] + /// 是否允许并行工具调用。 pub parallel_tool_calls: Option, #[serde(skip_serializing_if = "Option::is_none")] + /// 输出文本格式配置。 pub text: Option, #[serde(skip_serializing_if = "Option::is_none")] + /// 图片生成开关。 pub generate: Option, #[serde(skip)] + /// tuple schema 原始定义，仅供响应重构时使用。 pub tuple_schema: Option, #[serde(skip_serializing_if = "Option::is_none")] + /// service tier。 pub service_tier: Option, #[serde(skip_serializing_if = "Option::is_none")] + /// 用于显式续链的前一个 response ID。 pub previous_response_id: Option, #[serde(skip_serializing_if = "Option::is_none")] + /// 提示缓存键。 pub prompt_cache_key: Option, #[serde(skip)] + /// 是否由客户端显式提供了 prompt cache key。 pub explicit_prompt_cache_key: bool, #[serde(skip)] + /// 客户端会话 ID。 pub client_conversation_id: Option, #[serde(skip)] + /// 变体身份键。 pub variant_identity: Option, #[serde(skip_serializing_if = "Option::is_none")] + /// include 列表。 pub include: Option>, #[serde(skip_serializing_if = "Option::is_none")] + /// 传给上游的 client metadata。 pub client_metadata: Option, #[serde(skip)] + /// 是否偏好 WebSocket 传输。 pub use_websocket: bool, #[serde(skip)] + /// 是否强制 HTTP SSE。 pub force_http_sse: bool, #[serde(skip)] + /// turn state 透传头。 pub turn_state: Option, #[serde(skip)] + /// turn metadata 透传头。 pub turn_metadata: Option, #[serde(skip)] + /// beta features 透传头。 pub beta_features: Option, #[serde(skip)] + /// 客户端版本头。 pub version: Option, #[serde(skip)] + /// timing metrics 透传头。 pub include_timing_metrics: Option, #[serde(skip)] + /// codex window id。 pub codex_window_id: Option, #[serde(skip)] + /// 父线程 ID。 pub parent_thread_id: Option, } impl CodexResponsesRequest { + /// 构造默认的 HTTP SSE 请求。 pub fn new_http_sse( model: impl Into, instructions: impl Into, @@ -97,6 +128,7 @@ impl CodexResponsesRequest { } } + /// 判断请求是否声明了图片生成工具。 pub fn expects_image_generation(&self) -> bool { self.tools.as_deref().is_some_and(|tools| { tools @@ -106,17 +138,25 @@ impl CodexResponsesRequest { } } +/// Codex compact 端点请求体。 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CodexCompactRequest { + /// 模型名。 pub model: String, + /// 输入消息与结构化条目。 pub input: Vec, + /// 指令文本。 pub instructions: String, #[serde(skip_serializing_if = "Option::is_none")] + /// 工具定义。 pub tools: Option>, #[serde(skip_serializing_if = "Option::is_none")] + /// 是否允许并行工具调用。 pub parallel_tool_calls: Option, #[serde(skip_serializing_if = "Option::is_none")] + /// reasoning 配置。 pub reasoning: Option, #[serde(skip_serializing_if = "Option::is_none")] + /// 输出文本格式配置。 pub text: Option, } diff --git a/src/codex/gateway/protocol/tuple_schema.rs b/crates/core/src/protocol/codex/schema.rs similarity index 99% rename from src/codex/gateway/protocol/tuple_schema.rs rename to crates/core/src/protocol/codex/schema.rs index 8e4f3f20b..321286131 100644 --- a/src/codex/gateway/protocol/tuple_schema.rs +++ b/crates/core/src/protocol/codex/schema.rs @@ -1,3 +1,5 @@ +//! Codex 协议所需的 JSON schema 纯处理逻辑。 + use serde_json::{Map, Value}; pub(crate) struct PreparedSchema { @@ -18,6 +20,7 @@ pub(crate) fn prepare_schema(schema: Value) -> PreparedSchema { } } +#[allow(dead_code)] pub(crate) fn reconvert_tuple_values(data: Value, schema: &Value) -> Value { reconvert_tuple_values_with_root(data, schema, schema) } @@ -116,6 +119,7 @@ fn has_tuple_schemas(value: &Value) -> bool { { return true; } + object_values_have_tuple_schema(object, "properties") || object.get("items").is_some_and(has_tuple_schemas) || schema_array_has_tuple_schema(object, "oneOf") diff --git a/src/codex/gateway/transport/sse.rs b/crates/core/src/protocol/codex/sse.rs similarity index 92% rename from src/codex/gateway/transport/sse.rs rename to crates/core/src/protocol/codex/sse.rs index e295c0347..915d4442c 100644 --- a/src/codex/gateway/transport/sse.rs +++ b/crates/core/src/protocol/codex/sse.rs @@ -1,21 +1,35 @@ +//! SSE 事件解析与编码。 + use thiserror::Error; -pub const MAX_SSE_EVENT_BUFFER_BYTES: usize = 64 * 1024 * 1024; - +/// 单条 SSE 事件。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct SseEvent { + /// 事件名。 pub event: Option, + /// 数据体。 pub data: String, + /// 可选 ID。 pub id: Option, + /// 可选 retry。 pub retry: Option, } +/// 单事件缓冲上限。 +pub const MAX_SSE_EVENT_BUFFER_BYTES: usize = 64 * 1024 * 1024; + +/// SSE 解析错误。 #[derive(Debug, Error, PartialEq, Eq)] pub enum SseError { + /// retry 字段不是合法整数。 #[error("invalid SSE retry value: {0}")] InvalidRetry(String), + /// 单个事件缓冲超过上限。 #[error("SSE buffer exceeded {max_bytes} bytes — aborting stream")] - BufferExceeded { max_bytes: usize }, + BufferExceeded { + /// 上限字节数。 + max_bytes: usize, + }, } #[derive(Debug, Default)] @@ -60,6 +74,7 @@ impl EventBuilder { } } +/// 解析 SSE 事件流。 pub fn parse_sse_events(input: &str) -> Result, SseError> { let mut events = Vec::new(); let mut builder = EventBuilder::default(); @@ -110,6 +125,23 @@ pub fn parse_sse_events(input: &str) -> Result, SseError> { Ok(events) } +/// 编码单条 SSE 事件。 +pub fn encode_sse_event(event: &str, data: &str) -> String { + let mut frame = String::new(); + if !event.is_empty() { + frame.push_str("event: "); + frame.push_str(event); + frame.push('\n'); + } + for line in data.split('\n') { + frame.push_str("data: "); + frame.push_str(line); + frame.push('\n'); + } + frame.push('\n'); + frame +} + fn track_event_buffer_bytes(current_bytes: &mut usize, line: &str) -> Result<(), SseError> { let line_separator_bytes = usize::from(*current_bytes != 0); *current_bytes = current_bytes @@ -166,22 +198,6 @@ fn non_sse_error_message(raw: &str) -> String { .to_string() } -pub fn encode_sse_event(event: &str, data: &str) -> String { - let mut frame = String::new(); - if !event.is_empty() { - frame.push_str("event: "); - frame.push_str(event); - frame.push('\n'); - } - for line in data.split('\n') { - frame.push_str("data: "); - frame.push_str(line); - frame.push('\n'); - } - frame.push('\n'); - frame -} - fn split_sse_field(line: &str) -> (&str, &str) { let Some((field, value)) = line.split_once(':') else { return (line, ""); diff --git a/src/codex/gateway/transport/websocket/codec.rs b/crates/core/src/protocol/codex/websocket.rs similarity index 50% rename from src/codex/gateway/transport/websocket/codec.rs rename to crates/core/src/protocol/codex/websocket.rs index 41c50e038..2a9031a7a 100644 --- a/src/codex/gateway/transport/websocket/codec.rs +++ b/crates/core/src/protocol/codex/websocket.rs @@ -1,279 +1,260 @@ -use reqwest::StatusCode; -use serde::{Deserialize, Serialize}; -use serde_json::{json, Map, Value}; -use tokio_tungstenite::tungstenite::{http::HeaderMap as WsHeaderMap, Error as WsError, Message}; +//! Codex WebSocket 帧类型与纯编解码。 -use crate::codex::gateway::transport::{ - retry_after::retry_after_seconds_from_body, sse::encode_sse_event, types::CodexResponsesRequest, -}; +use serde::{ser::SerializeMap, Deserialize, Serialize}; +use serde_json::{Map, Value}; -use super::CodexWebSocketError; +use crate::protocol::codex::responses::CodexResponsesRequest; +use crate::protocol::codex::sse::encode_sse_event; +use crate::serving::responses::{http_sse_fallback_allowed, transport_for_request, CodexTransport}; const REDACTED_PAYLOAD_VALUE: &str = ""; -pub(super) struct ClassifiedWebSocketError { - pub(super) status: StatusCode, +/// WebSocket 消息包装。 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct WsMessage { + /// 消息类型。 + #[serde(rename = "type")] + pub message_type: String, + /// 消息负载。 + pub payload: T, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum WebSocketErrorClassificationProfile { - OneShot, - Pooled, +/// Codex WebSocket 聚合响应。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] +pub struct CodexWebSocketResponse { + /// 聚合后的事件。 + pub events: Vec, } -#[derive(Debug, Clone, PartialEq, serde::Serialize)] +/// Codex WebSocket 流响应。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] +pub struct CodexWebSocketStreamResponse { + /// 聚合响应。 + pub response: CodexWebSocketResponse, +} + +/// WebSocket 握手审计快照。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct OpeningAuditSnapshot { + /// 请求行。 + pub request_line: String, + /// 请求头顺序。 + pub header_order: Vec, + /// 红action后的请求头。 + pub headers: Vec, +} + +/// WebSocket 握手审计请求头。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct OpeningAuditHeader { + /// 请求头名。 + pub name: String, + /// 请求头值。 + pub value: String, +} + +/// WebSocket payload 审计快照。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] pub struct PayloadAuditSnapshot { + /// 按构造顺序记录的顶层字段。 pub top_level_keys: Vec, + /// 红action后的 JSON payload。 pub body: Value, } +/// WebSocket 一致性差异。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct WsParityDiff { + /// 差异说明。 + pub differences: Vec, +} + +/// WebSocket audit artifact. +#[derive(Debug, Clone, Serialize, PartialEq)] +pub struct WebSocketAuditArtifact { + /// 实际选择的传输模式。 + pub transport_mode: String, + /// 当前请求是否允许 HTTP/SSE fallback。 + pub fallback_allowed: bool, + /// 打开握手快照。 + pub opening: Option, + /// 首个 `response.create` payload 快照。 + pub payload: Option, + /// 失败快照。 + pub error: Option, +} + +/// WebSocket audit error snapshot. +#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +pub struct WebSocketAuditErrorSnapshot { + /// 错误分类。 + pub classification: String, + /// 错误消息。 + pub message: String, +} + +/// Structured parity diff between two WebSocket audit artifacts. +#[derive(Debug, Clone, Serialize, PartialEq)] +pub struct WebSocketParityDiff { + /// 逐字段差异。 + pub differences: Vec, +} + +/// Single WebSocket parity difference. +#[derive(Debug, Clone, Serialize, PartialEq)] +pub struct WebSocketParityDifference { + /// 字段路径。 + pub path: String, + /// 当前值。 + pub current: Value, + /// 参考值。 + pub reference: Value, +} + +/// WebSocket 错误帧分类模式。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WebSocketErrorClassificationProfile { + /// 单次 WebSocket 请求。 + OneShot, + /// 连接池复用请求。 + Pooled, +} + +/// 从 WebSocket 错误帧推导出的上游错误分类。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ClassifiedWebSocketError { + /// 应映射给 HTTP 层的状态码。 + pub status_code: u16, +} + #[derive(Debug, Deserialize)] -#[allow(dead_code)] struct ResponseCompleted { - id: String, - #[serde(default)] - usage: Option, - #[serde(default)] - end_turn: Option, + #[serde(rename = "id")] + _id: String, + #[serde(default, rename = "usage")] + _usage: Option, + #[serde(default, rename = "end_turn")] + _end_turn: Option, } #[derive(Debug, Deserialize)] -#[allow(dead_code)] struct ResponseCompletedUsage { - input_tokens: i64, - input_tokens_details: Option, - output_tokens: i64, - output_tokens_details: Option, - total_tokens: i64, + #[serde(rename = "input_tokens")] + _input_tokens: i64, + #[serde(default, rename = "input_tokens_details")] + _input_tokens_details: Option, + #[serde(rename = "output_tokens")] + _output_tokens: i64, + #[serde(default, rename = "output_tokens_details")] + _output_tokens_details: Option, + #[serde(rename = "total_tokens")] + _total_tokens: i64, } #[derive(Debug, Deserialize)] -#[allow(dead_code)] struct ResponseCompletedInputTokensDetails { - cached_tokens: i64, + #[serde(rename = "cached_tokens")] + _cached_tokens: i64, } #[derive(Debug, Deserialize)] -#[allow(dead_code)] struct ResponseCompletedOutputTokensDetails { - reasoning_tokens: i64, + #[serde(rename = "reasoning_tokens")] + _reasoning_tokens: i64, } #[derive(Debug, Deserialize)] -#[allow(dead_code)] struct ResponsesStreamEventShape { #[serde(rename = "type")] - kind: String, - headers: Option, - metadata: Option, - response: Option, - item: Option, - item_id: Option, - call_id: Option, - delta: Option, - summary_index: Option, - content_index: Option, + _kind: String, + #[serde(default, rename = "headers")] + _headers: Option, + #[serde(default, rename = "metadata")] + _metadata: Option, + #[serde(default, rename = "response")] + _response: Option, + #[serde(default, rename = "item")] + _item: Option, + #[serde(default, rename = "item_id")] + _item_id: Option, + #[serde(default, rename = "call_id")] + _call_id: Option, + #[serde(default, rename = "delta")] + _delta: Option, + #[serde(default, rename = "summary_index")] + _summary_index: Option, + #[serde(default, rename = "content_index")] + _content_index: Option, } -#[derive(Serialize)] -#[serde(tag = "type")] -enum ResponsesWsRequest<'a> { - #[serde(rename = "response.create")] - ResponseCreate(ResponseCreateWsPayload<'a>), -} - -#[derive(Serialize)] -struct ResponseCreateWsPayload<'a> { - model: &'a str, - #[serde(skip_serializing_if = "str::is_empty")] - instructions: &'a str, - #[serde(skip_serializing_if = "Option::is_none")] - previous_response_id: Option<&'a str>, - input: &'a [Value], - tools: Vec, - tool_choice: Value, - parallel_tool_calls: bool, - reasoning: Option<&'a Value>, - store: bool, - stream: bool, - include: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - service_tier: Option<&'a str>, - #[serde(skip_serializing_if = "Option::is_none")] - prompt_cache_key: Option<&'a str>, - #[serde(skip_serializing_if = "Option::is_none")] - text: Option<&'a Value>, - #[serde(skip_serializing_if = "Option::is_none")] - generate: Option, - #[serde(skip_serializing_if = "Option::is_none")] - client_metadata: Option<&'a Value>, -} - -impl<'a> ResponsesWsRequest<'a> { - fn from_request(request: &'a CodexResponsesRequest) -> Self { - Self::ResponseCreate(ResponseCreateWsPayload { - model: &request.model, - instructions: &request.instructions, - previous_response_id: request.previous_response_id.as_deref(), - input: &request.input, - tools: request.tools.clone().unwrap_or_default(), - tool_choice: request.tool_choice.clone().unwrap_or_else(|| json!("auto")), - parallel_tool_calls: request.parallel_tool_calls.unwrap_or(true), - reasoning: request.reasoning.as_ref(), - store: request.store, - stream: request.stream, - include: request.include.clone().unwrap_or_default(), - service_tier: request.service_tier.as_deref(), - prompt_cache_key: request.prompt_cache_key.as_deref(), - text: request.text.as_ref(), - generate: request.generate, - client_metadata: request.client_metadata.as_ref(), - }) - } -} - -pub(super) fn websocket_request_text( - request: &CodexResponsesRequest, +/// 将单条 WebSocket 消息编码为文本帧。 +pub fn encode_ws_message( + message: &WsMessage, ) -> Result { - serde_json::to_string(&ResponsesWsRequest::from_request(request)) + serde_json::to_string(message) } -pub(super) fn websocket_request_body(request: &CodexResponsesRequest) -> Value { - serde_json::to_value(ResponsesWsRequest::from_request(request)).unwrap_or(Value::Null) +/// 从文本帧解码单条 WebSocket 消息。 +pub fn decode_ws_message Deserialize<'de>>( + frame: &str, +) -> Result, serde_json::Error> { + serde_json::from_str(frame) } -pub fn websocket_payload_audit_snapshot(request: &CodexResponsesRequest) -> PayloadAuditSnapshot { - let body = websocket_request_body(request); - websocket_payload_audit_snapshot_from_body_and_keys(&body, websocket_payload_keys(request)) +/// 将一条公开 WebSocket JSON 事件编码为 SSE 帧。 +pub fn websocket_event_to_sse_frame(raw: &str) -> Option { + let event = websocket_event_type(raw)?; + if is_internal_websocket_event(&event) { + return None; + } + if websocket_event_should_skip(raw) { + return None; + } + Some(encode_sse_event(&event, raw)) } -pub(super) fn websocket_payload_audit_snapshot_from_request_body( - request: &CodexResponsesRequest, - body: &Value, -) -> PayloadAuditSnapshot { - websocket_payload_audit_snapshot_from_body_and_keys(body, websocket_payload_keys(request)) -} - -fn websocket_payload_audit_snapshot_from_body_and_keys( - body: &Value, - top_level_keys: Vec, -) -> PayloadAuditSnapshot { - PayloadAuditSnapshot { - top_level_keys, - body: redact_payload_body(body.clone()), - } -} - -fn websocket_payload_keys(request: &CodexResponsesRequest) -> Vec { - let mut keys = Vec::from(["type".to_string(), "model".to_string()]); - if !request.instructions.is_empty() { - keys.push("instructions".to_string()); - } - if request.previous_response_id.is_some() { - keys.push("previous_response_id".to_string()); - } - keys.extend( - [ - "input", - "tools", - "tool_choice", - "parallel_tool_calls", - "reasoning", - "store", - "stream", - "include", - ] - .into_iter() - .map(ToString::to_string), - ); - if request.service_tier.is_some() { - keys.push("service_tier".to_string()); - } - if request.prompt_cache_key.is_some() { - keys.push("prompt_cache_key".to_string()); - } - if request.text.is_some() { - keys.push("text".to_string()); - } - if request.generate.is_some() { - keys.push("generate".to_string()); - } - if request.client_metadata.is_some() { - keys.push("client_metadata".to_string()); - } - keys -} - -fn redact_payload_body(body: Value) -> Value { - let Value::Object(body) = body else { - return body; - }; - - Value::Object( - body.into_iter() - .map(|(key, value)| { - let value = if is_sensitive_payload_key(&key) { - Value::String(REDACTED_PAYLOAD_VALUE.to_string()) - } else { - redact_nested_payload_value(value) - }; - (key, value) - }) - .collect::>(), - ) -} - -fn redact_nested_payload_value(value: Value) -> Value { - match value { - Value::Array(values) => Value::Array( - values - .into_iter() - .map(redact_nested_payload_value) - .collect(), - ), - Value::Object(object) => Value::Object( - object - .into_iter() - .map(|(key, value)| (key, redact_nested_payload_value(value))) - .collect(), - ), - value => value, - } -} - -fn is_sensitive_payload_key(key: &str) -> bool { - matches!( - key, - "instructions" - | "input" - | "previous_response_id" - | "prompt_cache_key" - | "client_metadata" - | "tools" - ) -} - -pub(super) fn websocket_message_text( - message: Message, -) -> Result, CodexWebSocketError> { - match message { - Message::Text(text) => Ok(Some(text.to_string())), - Message::Binary(_) => Err(CodexWebSocketError::UnexpectedBinaryEvent), - Message::Close(_) => Err(CodexWebSocketError::ClosedByServerBeforeCompleted), - _ => Ok(None), - } -} - -pub(super) fn websocket_event_type(raw: &str) -> Option { +fn websocket_event_type(raw: &str) -> Option { serde_json::from_str::(raw).ok().and_then(|value| { value .get("type") .and_then(Value::as_str) - .map(str::to_string) + .map(ToString::to_string) }) } -pub(super) fn metadata_turn_state(raw: &str) -> Option { +fn is_internal_websocket_event(event: &str) -> bool { + event == "codex.rate_limits" +} + +fn websocket_event_should_skip(raw: &str) -> bool { + websocket_metadata_turn_state(raw).is_some() + || websocket_event_shape_parse_error(raw) + || websocket_response_completed_missing_response(raw) + || websocket_response_created_missing_response(raw) + || websocket_response_output_text_delta_missing_delta(raw) + || websocket_delta_event_missing_official_required_fields(raw) + || websocket_output_item_event_missing_item(raw) + || websocket_output_item_event_non_object_item(raw) + || websocket_output_item_event_invalid_item_type_tag(raw) + || websocket_output_item_event_invalid_metadata(raw) + || websocket_message_output_item_event_invalid_required_fields(raw) + || websocket_agent_message_output_item_event_invalid_required_fields(raw) + || websocket_reasoning_output_item_event_invalid_required_fields(raw) + || websocket_function_call_output_item_event_invalid_required_fields(raw) + || websocket_function_call_output_payload_item_event_invalid_required_fields(raw) + || websocket_custom_tool_call_output_item_event_invalid_required_fields(raw) + || websocket_custom_tool_call_output_payload_item_event_invalid_required_fields(raw) + || websocket_tool_search_call_output_item_event_invalid_required_fields(raw) + || websocket_tool_search_output_item_event_invalid_required_fields(raw) + || websocket_local_shell_call_output_item_event_invalid_required_fields(raw) + || websocket_web_search_call_output_item_event_invalid_required_fields(raw) + || websocket_image_generation_call_output_item_event_invalid_required_fields(raw) + || websocket_compaction_output_item_event_invalid_required_fields(raw) + || websocket_reasoning_summary_part_added_missing_summary_index(raw) +} + +/// 从 `response.metadata` 帧中提取 `x-codex-turn-state`。 +pub fn websocket_metadata_turn_state(raw: &str) -> Option { let value = serde_json::from_str::(raw).ok()?; if value.get("type").and_then(Value::as_str) != Some("response.metadata") { return None; @@ -300,25 +281,8 @@ fn json_value_as_string(value: &Value) -> Option { } } -fn json_field_absent_or_null(value: &Value, field: &str) -> bool { - matches!(value.get(field), None | Some(Value::Null)) -} - -pub(super) fn incomplete_response_reason(raw: &str) -> Option { - let value = serde_json::from_str::(raw).ok()?; - if value.get("type").and_then(Value::as_str) != Some("response.incomplete") { - return None; - } - Some( - value - .pointer("/response/incomplete_details/reason") - .and_then(Value::as_str) - .unwrap_or("unknown") - .to_string(), - ) -} - -pub(super) fn response_completed_parse_error(raw: &str) -> Option { +/// 校验 `response.completed` 的官方响应形状。 +pub fn websocket_response_completed_parse_error(raw: &str) -> Option { let value = serde_json::from_str::(raw).ok()?; if value.get("type").and_then(Value::as_str) != Some("response.completed") { return None; @@ -333,11 +297,32 @@ pub(super) fn response_completed_parse_error(raw: &str) -> Option { } } -pub(super) fn responses_stream_event_shape_parse_error(raw: &str) -> bool { +/// 从 `response.incomplete` 中提取 incomplete reason。 +pub fn websocket_incomplete_response_reason(raw: &str) -> Option { + let value = serde_json::from_str::(raw).ok()?; + if value.get("type").and_then(Value::as_str) != Some("response.incomplete") { + return None; + } + Some( + value + .pointer("/response/incomplete_details/reason") + .and_then(Value::as_str) + .unwrap_or("unknown") + .to_string(), + ) +} + +/// 判断事件是否不符合官方流事件的基本字段类型。 +pub fn websocket_event_shape_parse_error(raw: &str) -> bool { serde_json::from_str::(raw).is_err() } -pub(super) fn response_completed_missing_response(raw: &str) -> bool { +fn json_field_absent_or_null(value: &Value, field: &str) -> bool { + matches!(value.get(field), None | Some(Value::Null)) +} + +/// 判断 `response.completed` 是否缺少 `response`。 +pub fn websocket_response_completed_missing_response(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; @@ -345,7 +330,8 @@ pub(super) fn response_completed_missing_response(raw: &str) -> bool { && json_field_absent_or_null(&value, "response") } -pub(super) fn response_created_missing_response(raw: &str) -> bool { +/// 判断 `response.created` 是否缺少 `response`。 +pub fn websocket_response_created_missing_response(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; @@ -353,7 +339,8 @@ pub(super) fn response_created_missing_response(raw: &str) -> bool { && json_field_absent_or_null(&value, "response") } -pub(super) fn response_output_text_delta_missing_delta(raw: &str) -> bool { +/// 判断 `response.output_text.delta` 是否缺少 `delta`。 +pub fn websocket_response_output_text_delta_missing_delta(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; @@ -361,7 +348,8 @@ pub(super) fn response_output_text_delta_missing_delta(raw: &str) -> bool { && json_field_absent_or_null(&value, "delta") } -pub(super) fn delta_event_missing_official_required_fields(raw: &str) -> bool { +/// 判断 delta 事件是否缺少官方必需字段。 +pub fn websocket_delta_event_missing_official_required_fields(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; @@ -383,36 +371,31 @@ pub(super) fn delta_event_missing_official_required_fields(raw: &str) -> bool { } } -pub(super) fn output_item_event_missing_item(raw: &str) -> bool { +/// 判断 output item 事件是否缺少 `item`。 +pub fn websocket_output_item_event_missing_item(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) && json_field_absent_or_null(&value, "item") + is_output_item_event(&value) && json_field_absent_or_null(&value, "item") } -pub(super) fn output_item_event_non_object_item(raw: &str) -> bool { +/// 判断 output item 事件的 `item` 是否不是对象。 +pub fn websocket_output_item_event_non_object_item(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) && value - .get("item") - .is_some_and(|item| !item.is_null() && !item.is_object()) + is_output_item_event(&value) + && value + .get("item") + .is_some_and(|item| !item.is_null() && !item.is_object()) } -pub(super) fn output_item_event_invalid_item_type_tag(raw: &str) -> bool { +/// 判断 output item 事件的 `item.type` 是否缺失。 +pub fn websocket_output_item_event_invalid_item_type_tag(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -421,14 +404,12 @@ pub(super) fn output_item_event_invalid_item_type_tag(raw: &str) -> bool { .is_some_and(|item| item.get("type").and_then(Value::as_str).is_none()) } -pub(super) fn output_item_event_invalid_metadata(raw: &str) -> bool { +/// 判断 output item 事件的 metadata 是否无效。 +pub fn websocket_output_item_event_invalid_metadata(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -439,21 +420,25 @@ pub(super) fn output_item_event_invalid_metadata(raw: &str) -> bool { if metadata.is_null() { return false; } - metadata .as_object() .is_none_or(|metadata| optional_string_field_invalid(metadata, "turn_id")) }) } -pub(super) fn message_output_item_event_invalid_required_fields(raw: &str) -> bool { +fn is_output_item_event(value: &Value) -> bool { + matches!( + value.get("type").and_then(Value::as_str), + Some("response.output_item.done" | "response.output_item.added") + ) +} + +/// 判断 message output item 是否缺少官方必需字段。 +pub fn websocket_message_output_item_event_invalid_required_fields(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -473,7 +458,12 @@ pub(super) fn message_output_item_event_invalid_required_fields(raw: &str) -> bo }) } -fn optional_message_phase_invalid(item: &serde_json::Map) -> bool { +fn optional_string_field_invalid(item: &Map, field: &str) -> bool { + item.get(field) + .is_some_and(|value| !value.is_null() && !value.is_string()) +} + +fn optional_message_phase_invalid(item: &Map) -> bool { item.get("phase").is_some_and(|phase| { !phase.is_null() && !matches!(phase.as_str(), Some("commentary" | "final_answer")) }) @@ -505,14 +495,12 @@ fn valid_image_detail(detail: &Value) -> bool { matches!(detail.as_str(), Some("auto" | "low" | "high" | "original")) } -pub(super) fn agent_message_output_item_event_invalid_required_fields(raw: &str) -> bool { +/// 判断 agent_message output item 是否缺少官方必需字段。 +pub fn websocket_agent_message_output_item_event_invalid_required_fields(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -548,14 +536,12 @@ fn agent_message_content_item_invalid_required_fields(content_item: &Value) -> b } } -pub(super) fn reasoning_output_item_event_invalid_required_fields(raw: &str) -> bool { +/// 判断 reasoning output item 是否缺少官方必需字段。 +pub fn websocket_reasoning_output_item_event_invalid_required_fields(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -611,14 +597,12 @@ fn reasoning_content_item_invalid_required_fields(content_item: &Value) -> bool } } -pub(super) fn function_call_output_item_event_invalid_required_fields(raw: &str) -> bool { +/// 判断 function_call output item 是否缺少官方必需字段。 +pub fn websocket_function_call_output_item_event_invalid_required_fields(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -634,14 +618,14 @@ pub(super) fn function_call_output_item_event_invalid_required_fields(raw: &str) }) } -pub(super) fn function_call_output_payload_item_event_invalid_required_fields(raw: &str) -> bool { +/// 判断 function_call_output output item 是否缺少官方必需字段。 +pub fn websocket_function_call_output_payload_item_event_invalid_required_fields( + raw: &str, +) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -656,54 +640,6 @@ pub(super) fn function_call_output_payload_item_event_invalid_required_fields(ra }) } -pub(super) fn custom_tool_call_output_item_event_invalid_required_fields(raw: &str) -> bool { - let Ok(value) = serde_json::from_str::(raw) else { - return false; - }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { - return false; - } - value - .get("item") - .and_then(Value::as_object) - .is_some_and(|item| { - item.get("type").and_then(Value::as_str) == Some("custom_tool_call") - && (optional_string_field_invalid(item, "id") - || optional_string_field_invalid(item, "status") - || item.get("call_id").and_then(Value::as_str).is_none() - || item.get("name").and_then(Value::as_str).is_none() - || item.get("input").and_then(Value::as_str).is_none()) - }) -} - -pub(super) fn custom_tool_call_output_payload_item_event_invalid_required_fields( - raw: &str, -) -> bool { - let Ok(value) = serde_json::from_str::(raw) else { - return false; - }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { - return false; - } - value - .get("item") - .and_then(Value::as_object) - .is_some_and(|item| { - item.get("type").and_then(Value::as_str) == Some("custom_tool_call_output") - && (optional_string_field_invalid(item, "name") - || item.get("call_id").and_then(Value::as_str).is_none() - || item - .get("output") - .is_none_or(function_output_payload_invalid_required_fields)) - }) -} - fn function_output_payload_invalid_required_fields(output: &Value) -> bool { if output.is_string() { return false; @@ -742,14 +678,56 @@ fn function_output_content_item_invalid_required_fields(content_item: &Value) -> } } -pub(super) fn tool_search_call_output_item_event_invalid_required_fields(raw: &str) -> bool { +/// 判断 custom_tool_call output item 是否缺少官方必需字段。 +pub fn websocket_custom_tool_call_output_item_event_invalid_required_fields(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { + return false; + } + value + .get("item") + .and_then(Value::as_object) + .is_some_and(|item| { + item.get("type").and_then(Value::as_str) == Some("custom_tool_call") + && (optional_string_field_invalid(item, "id") + || optional_string_field_invalid(item, "status") + || item.get("call_id").and_then(Value::as_str).is_none() + || item.get("name").and_then(Value::as_str).is_none() + || item.get("input").and_then(Value::as_str).is_none()) + }) +} + +/// 判断 custom_tool_call_output output item 是否缺少官方必需字段。 +pub fn websocket_custom_tool_call_output_payload_item_event_invalid_required_fields( + raw: &str, +) -> bool { + let Ok(value) = serde_json::from_str::(raw) else { + return false; + }; + if !is_output_item_event(&value) { + return false; + } + value + .get("item") + .and_then(Value::as_object) + .is_some_and(|item| { + item.get("type").and_then(Value::as_str) == Some("custom_tool_call_output") + && (optional_string_field_invalid(item, "name") + || item.get("call_id").and_then(Value::as_str).is_none() + || item + .get("output") + .is_none_or(function_output_payload_invalid_required_fields)) + }) +} + +/// 判断 tool_search_call output item 是否缺少官方必需字段。 +pub fn websocket_tool_search_call_output_item_event_invalid_required_fields(raw: &str) -> bool { + let Ok(value) = serde_json::from_str::(raw) else { + return false; + }; + if !is_output_item_event(&value) { return false; } value @@ -765,14 +743,12 @@ pub(super) fn tool_search_call_output_item_event_invalid_required_fields(raw: &s }) } -pub(super) fn tool_search_output_item_event_invalid_required_fields(raw: &str) -> bool { +/// 判断 tool_search_output output item 是否缺少官方必需字段。 +pub fn websocket_tool_search_output_item_event_invalid_required_fields(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -787,14 +763,12 @@ pub(super) fn tool_search_output_item_event_invalid_required_fields(raw: &str) - }) } -pub(super) fn local_shell_call_output_item_event_invalid_required_fields(raw: &str) -> bool { +/// 判断 local_shell_call output item 是否缺少官方必需字段。 +pub fn websocket_local_shell_call_output_item_event_invalid_required_fields(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -812,14 +786,45 @@ pub(super) fn local_shell_call_output_item_event_invalid_required_fields(raw: &s }) } -pub(super) fn web_search_call_output_item_event_invalid_required_fields(raw: &str) -> bool { +fn valid_local_shell_status(status: Option<&Value>) -> bool { + matches!( + status.and_then(Value::as_str), + Some("completed" | "in_progress" | "incomplete") + ) +} + +fn local_shell_action_invalid_required_fields(action: &Map) -> bool { + action.get("type").and_then(Value::as_str) != Some("exec") + || action + .get("command") + .and_then(Value::as_array) + .is_none_or(|command| command.iter().any(|part| !part.is_string())) + || optional_u64_field_invalid(action, "timeout_ms") + || optional_string_field_invalid(action, "working_directory") + || optional_string_field_invalid(action, "user") + || optional_string_map_field_invalid(action, "env") +} + +fn optional_u64_field_invalid(item: &Map, field: &str) -> bool { + item.get(field) + .is_some_and(|value| !value.is_null() && value.as_u64().is_none()) +} + +fn optional_string_map_field_invalid(item: &Map, field: &str) -> bool { + item.get(field).is_some_and(|value| { + !value.is_null() + && value + .as_object() + .is_none_or(|object| object.values().any(|value| !value.is_string())) + }) +} + +/// 判断 web_search_call output item 是否缺少官方必需字段。 +pub fn websocket_web_search_call_output_item_event_invalid_required_fields(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -859,33 +864,23 @@ fn web_search_action_invalid_required_fields(action: &Value) -> bool { } } -fn valid_local_shell_status(status: Option<&Value>) -> bool { - matches!( - status.and_then(Value::as_str), - Some("completed" | "in_progress" | "incomplete") - ) +fn optional_string_array_field_invalid(item: &Map, field: &str) -> bool { + item.get(field).is_some_and(|value| { + !value.is_null() + && value + .as_array() + .is_none_or(|items| items.iter().any(|item| !item.is_string())) + }) } -fn local_shell_action_invalid_required_fields(action: &serde_json::Map) -> bool { - action.get("type").and_then(Value::as_str) != Some("exec") - || action - .get("command") - .and_then(Value::as_array) - .is_none_or(|command| command.iter().any(|part| !part.is_string())) - || optional_u64_field_invalid(action, "timeout_ms") - || optional_string_field_invalid(action, "working_directory") - || optional_string_field_invalid(action, "user") - || optional_string_map_field_invalid(action, "env") -} - -pub(super) fn image_generation_call_output_item_event_invalid_required_fields(raw: &str) -> bool { +/// 判断 image_generation_call output item 是否缺少官方必需字段。 +pub fn websocket_image_generation_call_output_item_event_invalid_required_fields( + raw: &str, +) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -900,42 +895,12 @@ pub(super) fn image_generation_call_output_item_event_invalid_required_fields(ra }) } -fn optional_string_field_invalid(item: &serde_json::Map, field: &str) -> bool { - item.get(field) - .is_some_and(|value| !value.is_null() && !value.is_string()) -} - -fn optional_string_array_field_invalid(item: &serde_json::Map, field: &str) -> bool { - item.get(field).is_some_and(|value| { - !value.is_null() - && value - .as_array() - .is_none_or(|items| items.iter().any(|item| !item.is_string())) - }) -} - -fn optional_u64_field_invalid(item: &serde_json::Map, field: &str) -> bool { - item.get(field) - .is_some_and(|value| !value.is_null() && value.as_u64().is_none()) -} - -fn optional_string_map_field_invalid(item: &serde_json::Map, field: &str) -> bool { - item.get(field).is_some_and(|value| { - !value.is_null() - && value - .as_object() - .is_none_or(|object| object.values().any(|value| !value.is_string())) - }) -} - -pub(super) fn compaction_output_item_event_invalid_required_fields(raw: &str) -> bool { +/// 判断 compaction output item 是否缺少官方必需字段。 +pub fn websocket_compaction_output_item_event_invalid_required_fields(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; - if !matches!( - value.get("type").and_then(Value::as_str), - Some("response.output_item.done" | "response.output_item.added") - ) { + if !is_output_item_event(&value) { return false; } value @@ -951,7 +916,8 @@ pub(super) fn compaction_output_item_event_invalid_required_fields(raw: &str) -> }) } -pub(super) fn reasoning_summary_part_added_missing_summary_index(raw: &str) -> bool { +/// 判断 reasoning summary part added 事件是否缺少 summary_index。 +pub fn websocket_reasoning_summary_part_added_missing_summary_index(raw: &str) -> bool { let Ok(value) = serde_json::from_str::(raw) else { return false; }; @@ -959,19 +925,13 @@ pub(super) fn reasoning_summary_part_added_missing_summary_index(raw: &str) -> b && value.get("summary_index").and_then(Value::as_i64).is_none() } -pub(super) fn websocket_sse_chunk(raw: &str, event: &str) -> String { - encode_sse_event(event, raw) +/// 判断 WebSocket 事件是否会结束当前响应流。 +pub fn is_terminal_websocket_event(event: &str) -> bool { + matches!(event, "response.completed" | "response.failed" | "error") } -pub(super) fn is_internal_websocket_event(raw: &str) -> bool { - websocket_event_type(raw).as_deref() == Some("codex.rate_limits") -} - -pub(super) fn is_terminal_websocket_event(event: &str) -> bool { - event == "response.completed" || event == "response.failed" || event == "error" -} - -pub(super) fn classify_ws_error_frame( +/// 将 WebSocket 错误帧映射为上游 HTTP 状态分类。 +pub fn classify_websocket_error_frame( raw: &str, profile: WebSocketErrorClassificationProfile, ) -> Option { @@ -980,92 +940,79 @@ pub(super) fn classify_ws_error_frame( if event_type != "error" && event_type != "response.failed" { return None; } - let code = value + + let code = websocket_error_code(&value); + if code.as_deref() == Some("websocket_connection_limit_reached") { + return Some(ClassifiedWebSocketError { status_code: 503 }); + } + + if event_type == "error" { + if let Some(status_code) = explicit_error_status_code(&value) { + if (200..=299).contains(&status_code) { + return None; + } + return Some(ClassifiedWebSocketError { status_code }); + } + } + + if let Some(code) = code { + if let Some(status_code) = rotatable_error_status_code(&code, profile) { + return Some(ClassifiedWebSocketError { status_code }); + } + } + + if event_type == "response.failed" { + return Some(ClassifiedWebSocketError { status_code: 503 }); + } + None +} + +fn websocket_error_code(value: &Value) -> Option { + value .pointer("/response/error/code") .or_else(|| value.pointer("/response/error/type")) .or_else(|| value.pointer("/error/code")) .or_else(|| value.pointer("/error/type")) .and_then(Value::as_str) - .map(str::to_ascii_lowercase); - if code.as_deref() == Some("websocket_connection_limit_reached") { - return Some(ClassifiedWebSocketError { - status: StatusCode::SERVICE_UNAVAILABLE, - }); - } - if event_type == "error" { - if let Some(status) = value - .get("status") - .or_else(|| value.get("status_code")) - .and_then(Value::as_u64) - .and_then(|status| u16::try_from(status).ok()) - .and_then(|status| StatusCode::from_u16(status).ok()) + .map(str::to_ascii_lowercase) +} + +fn explicit_error_status_code(value: &Value) -> Option { + value + .get("status") + .or_else(|| value.get("status_code")) + .and_then(Value::as_u64) + .and_then(|status| u16::try_from(status).ok()) +} + +fn rotatable_error_status_code( + code: &str, + profile: WebSocketErrorClassificationProfile, +) -> Option { + match code { + "usage_limit_reached" | "rate_limit_exceeded" | "rate_limit_reached" => Some(429), + "quota_exhausted" | "quota_exceeded" | "insufficient_quota" | "payment_required" => { + Some(402) + } + "usage_not_included" => Some(429), + "unauthorized" | "token_invalid" | "token_expired" | "account_deactivated" => Some(401), + "forbidden" | "account_banned" | "banned" => Some(403), + "context_length_exceeded" | "invalid_prompt" | "cyber_policy" | "invalid_request" => { + Some(400) + } + "previous_response_not_found" => Some(400), + "server_is_overloaded" | "slow_down" => Some(503), + "websocket_connection_limit_reached" + if profile == WebSocketErrorClassificationProfile::Pooled => { - if status.is_success() { - return None; - } - return Some(ClassifiedWebSocketError { status }); + Some(503) } - } - if let Some(code) = code { - if let Some(status) = rotatable_error_status(&code, profile) { - return Some(ClassifiedWebSocketError { status }); - } - } - if event_type == "response.failed" { - return Some(ClassifiedWebSocketError { - status: StatusCode::SERVICE_UNAVAILABLE, - }); - } - None -} - -pub(super) fn codex_websocket_transport_error(error: WsError) -> CodexWebSocketError { - match error { - WsError::Http(response) => { - let (parts, body) = (*response).into_parts(); - let body = body - .and_then(|bytes| String::from_utf8(bytes).ok()) - .unwrap_or_default(); - CodexWebSocketError::Upstream { - status: parts.status, - retry_after_seconds: retry_after_seconds(&parts.headers) - .or_else(|| retry_after_seconds_from_body(&body)), - body, - } - } - error => CodexWebSocketError::Transport(error), + _ => None, } } -pub(super) fn turn_state(headers: &WsHeaderMap) -> Option { - headers - .get("x-codex-turn-state") - .and_then(|value| value.to_str().ok()) - .map(ToString::to_string) -} - -pub(super) fn set_cookie_headers(headers: &WsHeaderMap) -> Vec { - headers - .get_all("set-cookie") - .iter() - .filter_map(|value| value.to_str().ok().map(ToString::to_string)) - .collect() -} - -pub(super) fn rate_limit_headers(headers: &WsHeaderMap) -> Vec<(String, String)> { - headers - .iter() - .filter(|(name, _)| is_rate_limit_header(name.as_str())) - .filter_map(|(name, value)| { - value - .to_str() - .ok() - .map(|value| (name.as_str().to_string(), value.to_string())) - }) - .collect() -} - -pub(super) fn retry_after_seconds_from_wrapped_error_headers(raw: &str) -> Option { +/// 从包裹在 WebSocket `error.headers` 中的 retry-after 值提取秒数。 +pub fn retry_after_seconds_from_wrapped_error_headers(raw: &str) -> Option { let value = serde_json::from_str::(raw).ok()?; if value.get("type").and_then(Value::as_str) != Some("error") { return None; @@ -1096,52 +1043,366 @@ fn retry_after_seconds_from_json_header_value(value: &Value) -> Option { (seconds > 0).then_some(seconds) } -fn rotatable_error_status( - code: &str, - profile: WebSocketErrorClassificationProfile, -) -> Option { - match code { - "usage_limit_reached" | "rate_limit_exceeded" | "rate_limit_reached" => { - Some(StatusCode::TOO_MANY_REQUESTS) - } - "quota_exhausted" | "quota_exceeded" | "insufficient_quota" | "payment_required" => { - Some(StatusCode::PAYMENT_REQUIRED) - } - "usage_not_included" => Some(StatusCode::TOO_MANY_REQUESTS), - "unauthorized" | "token_invalid" | "token_expired" | "account_deactivated" => { - Some(StatusCode::UNAUTHORIZED) - } - "forbidden" | "account_banned" | "banned" => Some(StatusCode::FORBIDDEN), - "context_length_exceeded" | "invalid_prompt" | "cyber_policy" | "invalid_request" => { - Some(StatusCode::BAD_REQUEST) - } - "previous_response_not_found" => Some(StatusCode::BAD_REQUEST), - "server_is_overloaded" | "slow_down" => Some(StatusCode::SERVICE_UNAVAILABLE), - "websocket_connection_limit_reached" - if profile == WebSocketErrorClassificationProfile::Pooled => - { - Some(StatusCode::SERVICE_UNAVAILABLE) - } - _ => None, +/// 生成 Responses WebSocket payload 审计快照。 +pub fn websocket_payload_audit_snapshot(request: &CodexResponsesRequest) -> PayloadAuditSnapshot { + let body = websocket_response_create_payload(request); + PayloadAuditSnapshot { + top_level_keys: websocket_payload_keys(request), + body: redact_payload_body(body), } } -fn is_rate_limit_header(name: &str) -> bool { - let name = name.to_ascii_lowercase(); - name == "retry-after" - || name.contains("ratelimit") - || name.contains("rate-limit") - || name.starts_with("x-codex-primary-") - || name.starts_with("x-codex-secondary-") - || name.starts_with("x-codex-code-review-") - || name.starts_with("x-codex-review-") - || name.starts_with("x-code-review-") +/// Build a WebSocket audit artifact for a single opening attempt. +pub fn websocket_audit_artifact_from_attempt( + request: &CodexResponsesRequest, + opening: OpeningAuditSnapshot, + payload: PayloadAuditSnapshot, +) -> WebSocketAuditArtifact { + WebSocketAuditArtifact { + transport_mode: websocket_transport_mode_name(request).to_string(), + fallback_allowed: http_sse_fallback_allowed(request), + opening: Some(opening), + payload: Some(payload), + error: None, + } } -fn retry_after_seconds(headers: &WsHeaderMap) -> Option { - headers - .get("retry-after") - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.trim().parse::().ok()) - .filter(|seconds| *seconds > 0) +/// Diff two WebSocket audit artifacts by the fields used for parity checks. +pub fn websocket_parity_diff( + current: &WebSocketAuditArtifact, + reference: &WebSocketAuditArtifact, +) -> WebSocketParityDiff { + let mut differences = Vec::new(); + push_websocket_parity_difference( + &mut differences, + "transport_mode", + serde_json::json!(current.transport_mode), + serde_json::json!(reference.transport_mode), + ); + push_websocket_parity_difference( + &mut differences, + "fallback_allowed", + serde_json::json!(current.fallback_allowed), + serde_json::json!(reference.fallback_allowed), + ); + push_websocket_parity_difference( + &mut differences, + "opening.request_line", + opening_request_line(¤t.opening), + opening_request_line(&reference.opening), + ); + push_websocket_parity_difference( + &mut differences, + "opening.header_order", + opening_header_order(¤t.opening), + opening_header_order(&reference.opening), + ); + push_websocket_parity_difference( + &mut differences, + "opening.sec_websocket_extensions", + opening_header_value(¤t.opening, "Sec-WebSocket-Extensions"), + opening_header_value(&reference.opening, "Sec-WebSocket-Extensions"), + ); + push_websocket_parity_difference( + &mut differences, + "payload.top_level_keys", + payload_top_level_keys(¤t.payload), + payload_top_level_keys(&reference.payload), + ); + push_websocket_parity_difference( + &mut differences, + "error.classification", + error_classification(¤t.error), + error_classification(&reference.error), + ); + WebSocketParityDiff { differences } +} + +fn websocket_transport_mode_name(request: &CodexResponsesRequest) -> &'static str { + match transport_for_request(request) { + CodexTransport::HttpSse => "http_sse", + CodexTransport::WebSocketPreferred => "websocket_preferred", + CodexTransport::WebSocketRequired => "websocket_required", + } +} + +fn push_websocket_parity_difference( + differences: &mut Vec, + path: &str, + current: Value, + reference: Value, +) { + if current == reference { + return; + } + differences.push(WebSocketParityDifference { + path: path.to_string(), + current, + reference, + }); +} + +fn opening_request_line(opening: &Option) -> Value { + opening + .as_ref() + .map(|opening| serde_json::json!(opening.request_line)) + .unwrap_or(Value::Null) +} + +fn opening_header_order(opening: &Option) -> Value { + opening + .as_ref() + .map(|opening| serde_json::json!(opening.header_order)) + .unwrap_or(Value::Null) +} + +fn opening_header_value(opening: &Option, name: &str) -> Value { + opening + .as_ref() + .and_then(|opening| { + opening + .headers + .iter() + .find(|header| header.name.eq_ignore_ascii_case(name)) + }) + .map(|header| serde_json::json!(header.value)) + .unwrap_or(Value::Null) +} + +fn payload_top_level_keys(payload: &Option) -> Value { + payload + .as_ref() + .map(|payload| serde_json::json!(payload.top_level_keys)) + .unwrap_or(Value::Null) +} + +fn error_classification(error: &Option) -> Value { + error + .as_ref() + .map(|error| serde_json::json!(error.classification)) + .unwrap_or(Value::Null) +} + +/// 生成 Responses WebSocket `response.create` payload。 +pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Value { + let mut body = Map::new(); + insert_value( + &mut body, + "type", + Value::String("response.create".to_string()), + ); + insert_value(&mut body, "model", Value::String(request.model.clone())); + if !request.instructions.is_empty() { + insert_value( + &mut body, + "instructions", + Value::String(request.instructions.clone()), + ); + } + if let Some(previous_response_id) = &request.previous_response_id { + insert_value( + &mut body, + "previous_response_id", + Value::String(previous_response_id.clone()), + ); + } + insert_value(&mut body, "input", Value::Array(request.input.clone())); + insert_value( + &mut body, + "tools", + Value::Array(request.tools.clone().unwrap_or_default()), + ); + insert_value( + &mut body, + "tool_choice", + request + .tool_choice + .clone() + .unwrap_or(Value::String("auto".to_string())), + ); + insert_value( + &mut body, + "parallel_tool_calls", + Value::Bool(request.parallel_tool_calls.unwrap_or(true)), + ); + insert_value( + &mut body, + "reasoning", + request.reasoning.clone().unwrap_or(Value::Null), + ); + insert_value(&mut body, "store", Value::Bool(request.store)); + insert_value(&mut body, "stream", Value::Bool(request.stream)); + insert_value( + &mut body, + "include", + request + .include + .clone() + .unwrap_or_default() + .into_iter() + .map(Value::String) + .collect::>() + .into(), + ); + if let Some(service_tier) = &request.service_tier { + insert_value( + &mut body, + "service_tier", + Value::String(service_tier.clone()), + ); + } + if let Some(prompt_cache_key) = &request.prompt_cache_key { + insert_value( + &mut body, + "prompt_cache_key", + Value::String(prompt_cache_key.clone()), + ); + } + if let Some(text) = &request.text { + insert_value(&mut body, "text", text.clone()); + } + if let Some(generate) = request.generate { + insert_value(&mut body, "generate", Value::Bool(generate)); + } + if let Some(client_metadata) = &request.client_metadata { + insert_value(&mut body, "client_metadata", client_metadata.clone()); + } + Value::Object(body) +} + +/// 生成 Responses WebSocket `response.create` 文本帧内容。 +pub fn websocket_response_create_payload_text( + request: &CodexResponsesRequest, +) -> Result { + serde_json::to_string(&OrderedResponseCreatePayload(request)) +} + +struct OrderedResponseCreatePayload<'a>(&'a CodexResponsesRequest); + +impl Serialize for OrderedResponseCreatePayload<'_> { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + let request = self.0; + let mut map = serializer.serialize_map(Some(websocket_payload_keys(request).len()))?; + map.serialize_entry("type", "response.create")?; + map.serialize_entry("model", &request.model)?; + if !request.instructions.is_empty() { + map.serialize_entry("instructions", &request.instructions)?; + } + if let Some(previous_response_id) = &request.previous_response_id { + map.serialize_entry("previous_response_id", previous_response_id)?; + } + map.serialize_entry("input", &request.input)?; + map.serialize_entry("tools", request.tools.as_deref().unwrap_or(&[]))?; + if let Some(tool_choice) = &request.tool_choice { + map.serialize_entry("tool_choice", tool_choice)?; + } else { + map.serialize_entry("tool_choice", "auto")?; + } + map.serialize_entry( + "parallel_tool_calls", + &request.parallel_tool_calls.unwrap_or(true), + )?; + if let Some(reasoning) = &request.reasoning { + map.serialize_entry("reasoning", reasoning)?; + } else { + map.serialize_entry("reasoning", &Value::Null)?; + } + map.serialize_entry("store", &request.store)?; + map.serialize_entry("stream", &request.stream)?; + map.serialize_entry("include", request.include.as_deref().unwrap_or(&[]))?; + if let Some(service_tier) = &request.service_tier { + map.serialize_entry("service_tier", service_tier)?; + } + if let Some(prompt_cache_key) = &request.prompt_cache_key { + map.serialize_entry("prompt_cache_key", prompt_cache_key)?; + } + if let Some(text) = &request.text { + map.serialize_entry("text", text)?; + } + if let Some(generate) = request.generate { + map.serialize_entry("generate", &generate)?; + } + if let Some(client_metadata) = &request.client_metadata { + map.serialize_entry("client_metadata", client_metadata)?; + } + map.end() + } +} + +fn insert_value(body: &mut Map, key: &str, value: Value) { + body.insert(key.to_string(), value); +} + +fn websocket_payload_keys(request: &CodexResponsesRequest) -> Vec { + let mut keys = vec!["type".to_string(), "model".to_string()]; + if !request.instructions.is_empty() { + keys.push("instructions".to_string()); + } + if request.previous_response_id.is_some() { + keys.push("previous_response_id".to_string()); + } + keys.extend( + [ + "input", + "tools", + "tool_choice", + "parallel_tool_calls", + "reasoning", + "store", + "stream", + "include", + ] + .into_iter() + .map(ToString::to_string), + ); + if request.service_tier.is_some() { + keys.push("service_tier".to_string()); + } + if request.prompt_cache_key.is_some() { + keys.push("prompt_cache_key".to_string()); + } + if request.text.is_some() { + keys.push("text".to_string()); + } + if request.generate.is_some() { + keys.push("generate".to_string()); + } + if request.client_metadata.is_some() { + keys.push("client_metadata".to_string()); + } + keys +} + +fn redact_payload_body(body: Value) -> Value { + let Value::Object(body) = body else { + return body; + }; + + Value::Object( + body.into_iter() + .map(|(key, value)| { + let value = if is_sensitive_payload_key(&key) { + Value::String(REDACTED_PAYLOAD_VALUE.to_string()) + } else { + value + }; + (key, value) + }) + .collect(), + ) +} + +fn is_sensitive_payload_key(key: &str) -> bool { + matches!( + key, + "instructions" + | "input" + | "previous_response_id" + | "prompt_cache_key" + | "client_metadata" + | "tools" + ) } diff --git a/crates/core/src/protocol/mod.rs b/crates/core/src/protocol/mod.rs new file mode 100644 index 000000000..e32969f07 --- /dev/null +++ b/crates/core/src/protocol/mod.rs @@ -0,0 +1,4 @@ +//! 协议语义与纯数据类型。 + +pub mod codex; +pub mod openai; diff --git a/crates/core/src/protocol/openai/chat.rs b/crates/core/src/protocol/openai/chat.rs new file mode 100644 index 000000000..cd049a951 --- /dev/null +++ b/crates/core/src/protocol/openai/chat.rs @@ -0,0 +1,620 @@ +//! OpenAI Chat Completions 请求到 Codex Responses 请求的纯转换。 + +use std::collections::{BTreeMap, BTreeSet}; + +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use thiserror::Error; +use uuid::Uuid; + +use crate::protocol::codex::{ + responses::CodexResponsesRequest, + schema::{prepare_schema, reconvert_tuple_values}, + sse::{parse_sse_events, SseError}, +}; + +/// OpenAI Chat Completions 请求体。 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ChatCompletionRequest { + /// 客户端请求的模型名。 + pub model: String, + /// 是否流式返回。 + #[serde(default)] + pub stream: bool, + /// 按 OpenAI 语义组织的消息列表。 + pub messages: Vec, + /// 可选的 reasoning effort。 + #[serde(default)] + pub reasoning_effort: Option, + /// 可选的 service tier。 + #[serde(default)] + pub service_tier: Option, + /// 原生 tools 字段。 + #[serde(default)] + pub tools: Option>, + /// tool choice 配置。 + #[serde(default)] + pub tool_choice: Option, + /// parallel tool calls 配置。 + #[serde(default)] + pub parallel_tool_calls: Option, + /// 旧式 functions 字段，会转换为 Codex tools。 + #[serde(default)] + pub functions: Option>, + /// 响应格式配置。 + #[serde(default)] + pub response_format: Option, + /// 可选的 user 字段，用作稳定会话锚点。 + #[serde(default)] + pub user: Option, +} + +/// Chat 请求中的单条消息。 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ChatMessage { + /// 角色名。 + pub role: String, + /// 消息内容，既可能是字符串也可能是结构化数组。 + #[serde(default)] + pub content: Option, + /// function/tool 相关的名称。 + #[serde(default)] + pub name: Option, + /// assistant 消息内的 tool_calls。 + #[serde(default)] + pub tool_calls: Option>, + /// tool 响应对应的调用 ID。 + #[serde(default)] + pub tool_call_id: Option, + /// 旧式 function_call 字段。 + #[serde(default)] + pub function_call: Option, +} + +/// Chat 转译失败的错误。 +#[derive(Debug, Error, PartialEq, Eq)] +pub enum ChatTranslationError { + /// 请求没有任何消息，无法生成有效上游输入。 + #[error("messages must not be empty")] + EmptyMessages, +} + +/// 将 OpenAI Chat Completions 请求纯转换为 Codex Responses 请求。 +pub fn translate_chat_to_codex( + request: ChatCompletionRequest, +) -> Result { + if request.messages.is_empty() { + return Err(ChatTranslationError::EmptyMessages); + } + + let client_conversation_id = request + .user + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string); + let instructions = chat_instructions(&request.messages); + let mut input = chat_input(request.messages); + if input.is_empty() { + input.push(json!({"role": "user", "content": ""})); + } + + let mut codex_request = CodexResponsesRequest::new_http_sse(request.model, instructions, input); + codex_request.force_http_sse = true; + codex_request.tools = codex_tools(request.tools, request.functions); + codex_request.tool_choice = request.tool_choice; + codex_request.parallel_tool_calls = request.parallel_tool_calls; + + let response_format = response_format_text(request.response_format); + codex_request.text = response_format.text; + codex_request.tuple_schema = response_format.tuple_schema; + codex_request.service_tier = request.service_tier; + codex_request.prompt_cache_key = client_conversation_id.clone(); + codex_request.client_conversation_id = client_conversation_id; + + if let Some(effort) = request.reasoning_effort { + codex_request.reasoning = Some(json!({"effort": effort, "summary": "auto"})); + } + + Ok(codex_request) +} + +/// 将 Codex SSE 完成响应转换为 OpenAI Chat Completions 响应体。 +pub fn chat_completion_from_codex_sse( + body: &str, + model: &str, + include_reasoning: bool, + tuple_schema: Option<&Value>, +) -> Result, SseError> { + let events = parse_sse_events(body)?; + let mut content = String::new(); + let mut reasoning = String::new(); + let mut function_call_items = BTreeMap::new(); + let mut finished_call_ids = BTreeSet::new(); + let mut tool_calls = Vec::new(); + let mut completed_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) { + 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.output_item.added") => { + if let Some((item_id, call_id, name)) = function_call_item(value.get("item")) { + function_call_items.insert(item_id, FunctionCallInfo { call_id, name }); + } + } + Some("response.function_call_arguments.done") => { + if let Some(tool_call) = + tool_call_from_done_event(&value, &function_call_items, &mut finished_call_ids) + { + tool_calls.push(tool_call); + } + } + Some("response.output_item.done") => { + if let Some(tool_call) = tool_call_from_output_item( + value.get("item"), + &function_call_items, + &mut finished_call_ids, + ) { + tool_calls.push(tool_call); + } + } + Some("response.completed") => { + completed_response = value.get("response").cloned(); + } + _ => {} + } + } + + let Some(response) = completed_response else { + return Ok(None); + }; + if content.is_empty() { + content = output_text_from_response(&response); + } + if let Some(tuple_schema) = tuple_schema { + content = reconvert_tuple_text(&content, tuple_schema).unwrap_or(content); + } + let mut message = json!({ + "role": "assistant", + "content": if content.is_empty() && !tool_calls.is_empty() { + Value::Null + } else { + Value::String(content) + }, + }); + if include_reasoning && !reasoning.is_empty() { + message["reasoning_content"] = Value::String(reasoning); + } + let has_tool_calls = !tool_calls.is_empty(); + if has_tool_calls { + message["tool_calls"] = Value::Array(tool_calls); + } + + Ok(Some(json!({ + "id": format!("chatcmpl-{}", Uuid::new_v4().simple()), + "object": "chat.completion", + "created": Utc::now().timestamp(), + "model": model, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" }, + }], + "usage": openai_usage(response.get("usage")), + }))) +} + +fn chat_instructions(messages: &[ChatMessage]) -> String { + let instructions = messages + .iter() + .filter(|message| message.role == "system" || message.role == "developer") + .map(|message| extract_text(&message.content)) + .filter(|text| !text.is_empty()) + .collect::>() + .join("\n\n"); + + if instructions.is_empty() { + "You are a helpful assistant.".to_string() + } else { + instructions + } +} + +fn chat_input(messages: Vec) -> Vec { + let mut input = Vec::new(); + + for message in messages { + match message.role.as_str() { + "system" | "developer" => {} + "assistant" => push_assistant_message(&mut input, message), + "tool" => input.push(json!({ + "type": "function_call_output", + "call_id": message.tool_call_id.unwrap_or_else(|| "unknown".to_string()), + "output": extract_text(&message.content), + })), + "function" => input.push(json!({ + "type": "function_call_output", + "call_id": format!("fc_{}", message.name.unwrap_or_else(|| "unknown".to_string())), + "output": extract_text(&message.content), + })), + _ => input.push(json!({ + "role": "user", + "content": extract_content(&message.content), + })), + } + } + + input +} + +fn push_assistant_message(input: &mut Vec, message: ChatMessage) { + let text = extract_text(&message.content); + let has_tool_calls = message + .tool_calls + .as_ref() + .is_some_and(|calls| !calls.is_empty()); + + if !text.is_empty() || (!has_tool_calls && message.function_call.is_none()) { + input.push(json!({"role": "assistant", "content": text})); + } + + if let Some(tool_calls) = message.tool_calls { + for tool_call in tool_calls { + let call_id = tool_call + .get("id") + .and_then(Value::as_str) + .unwrap_or("unknown"); + let function = tool_call.get("function").unwrap_or(&Value::Null); + let name = function + .get("name") + .and_then(Value::as_str) + .unwrap_or("unknown"); + let arguments = function + .get("arguments") + .and_then(Value::as_str) + .unwrap_or(""); + input.push(json!({ + "type": "function_call", + "call_id": call_id, + "name": name, + "arguments": arguments, + })); + } + } + + if let Some(function_call) = message.function_call { + let name = function_call + .get("name") + .and_then(Value::as_str) + .unwrap_or("unknown"); + let arguments = function_call + .get("arguments") + .and_then(Value::as_str) + .unwrap_or(""); + input.push(json!({ + "type": "function_call", + "call_id": format!("fc_{name}"), + "name": name, + "arguments": arguments, + })); + } +} + +fn extract_text(content: &Option) -> String { + let Some(content) = content else { + return String::new(); + }; + if let Some(text) = content.as_str() { + return text.to_string(); + } + let Some(parts) = content.as_array() else { + return String::new(); + }; + + parts + .iter() + .filter_map(|part| { + (part.get("type").and_then(Value::as_str) == Some("text")) + .then(|| part.get("text").and_then(Value::as_str)) + .flatten() + }) + .collect::>() + .join("\n") +} + +fn extract_content(content: &Option) -> Value { + let Some(content) = content else { + return Value::String(String::new()); + }; + if content.as_str().is_some() { + return content.clone(); + } + + let Some(parts) = content.as_array() else { + return Value::String(String::new()); + }; + let has_image = parts + .iter() + .any(|part| part.get("type").and_then(Value::as_str) == Some("image_url")); + if !has_image { + return Value::String(extract_text(&Some(content.clone()))); + } + + let codex_parts = parts + .iter() + .filter_map(codex_content_part) + .collect::>(); + if codex_parts.is_empty() { + Value::String(String::new()) + } else { + Value::Array(codex_parts) + } +} + +fn codex_content_part(part: &Value) -> Option { + match part.get("type").and_then(Value::as_str)? { + "text" => part + .get("text") + .and_then(Value::as_str) + .map(|text| json!({"type": "input_text", "text": text})), + "image_url" => image_url(part).map(|url| json!({"type": "input_image", "image_url": url})), + _ => None, + } +} + +fn image_url(part: &Value) -> Option<&str> { + let image = part.get("image_url")?; + image + .as_str() + .or_else(|| image.get("url").and_then(Value::as_str)) +} + +fn reconvert_tuple_text(text: &str, tuple_schema: &Value) -> Option { + let parsed = serde_json::from_str::(text).ok()?; + Some(reconvert_tuple_values(parsed, tuple_schema).to_string()) +} + +fn codex_tools(tools: Option>, functions: Option>) -> Option> { + if let Some(tools) = tools.filter(|tools| !tools.is_empty()) { + return Some(tools); + } + + functions + .filter(|functions| !functions.is_empty()) + .map(|functions| { + functions + .into_iter() + .map(|function| json!({"type": "function", "function": function})) + .collect() + }) +} + +struct PreparedResponseFormat { + text: Option, + tuple_schema: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct FunctionCallInfo { + call_id: String, + name: String, +} + +fn response_format_text(response_format: Option) -> PreparedResponseFormat { + let Some(format) = response_format else { + return PreparedResponseFormat { + text: None, + tuple_schema: None, + }; + }; + let Some(kind) = format.get("type").and_then(Value::as_str) else { + return PreparedResponseFormat { + text: None, + tuple_schema: None, + }; + }; + + match kind { + "json_object" => PreparedResponseFormat { + text: Some(json!({"format": {"type": "json_object"}})), + tuple_schema: None, + }, + "json_schema" => { + let Some(schema) = format.get("json_schema") else { + return PreparedResponseFormat { + text: None, + tuple_schema: None, + }; + }; + let name = schema + .get("name") + .and_then(Value::as_str) + .unwrap_or("response"); + let prepared_schema = + prepare_schema(schema.get("schema").cloned().unwrap_or_else(|| json!({}))); + let mut codex_format = json!({ + "type": "json_schema", + "name": name, + "schema": prepared_schema.schema, + }); + if let Some(strict) = schema.get("strict").and_then(Value::as_bool) { + codex_format["strict"] = Value::Bool(strict); + } + + PreparedResponseFormat { + text: Some(json!({"format": codex_format})), + tuple_schema: prepared_schema.original_schema, + } + } + _ => PreparedResponseFormat { + text: None, + tuple_schema: None, + }, + } +} + +fn output_text_from_response(response: &Value) -> String { + response + .get("output_text") + .and_then(Value::as_str) + .map(ToString::to_string) + .unwrap_or_else(|| { + response + .get("output") + .and_then(Value::as_array) + .map(|items| { + items + .iter() + .filter_map(output_text_from_item) + .collect::>() + .join("\n\n") + }) + .unwrap_or_default() + }) +} + +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 function_call_item(item: Option<&Value>) -> Option<(String, String, String)> { + let item = item?; + if item.get("type").and_then(Value::as_str) != Some("function_call") { + return None; + } + let item_id = item.get("id").and_then(Value::as_str)?; + let call_id = item.get("call_id").and_then(Value::as_str)?; + let name = item.get("name").and_then(Value::as_str)?; + Some((item_id.to_string(), call_id.to_string(), name.to_string())) +} + +fn tool_call_from_done_event( + event: &Value, + function_call_items: &BTreeMap, + finished_call_ids: &mut BTreeSet, +) -> Option { + let event_id = event + .get("call_id") + .or_else(|| event.get("item_id")) + .and_then(Value::as_str)?; + let info = function_call_items.get(event_id); + let call_id = info + .map(|info| info.call_id.as_str()) + .unwrap_or(event_id) + .to_string(); + if !finished_call_ids.insert(call_id.clone()) { + return None; + } + let name = event + .get("name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()) + .or_else(|| info.map(|info| info.name.as_str())) + .unwrap_or("unknown"); + let arguments = event + .get("arguments") + .and_then(Value::as_str) + .unwrap_or_default(); + Some(openai_tool_call(call_id, name, arguments)) +} + +fn tool_call_from_output_item( + item: Option<&Value>, + function_call_items: &BTreeMap, + finished_call_ids: &mut BTreeSet, +) -> Option { + let item = item?; + if item.get("type").and_then(Value::as_str) != Some("function_call") { + return None; + } + let item_id = item.get("id").and_then(Value::as_str); + let info = item_id.and_then(|item_id| function_call_items.get(item_id)); + let call_id = item + .get("call_id") + .and_then(Value::as_str) + .or_else(|| info.map(|info| info.call_id.as_str())) + .or(item_id)?; + if !finished_call_ids.insert(call_id.to_string()) { + return None; + } + let name = item + .get("name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()) + .or_else(|| info.map(|info| info.name.as_str())) + .unwrap_or("unknown"); + let arguments = item + .get("arguments") + .and_then(Value::as_str) + .unwrap_or_default(); + Some(openai_tool_call(call_id.to_string(), name, arguments)) +} + +fn openai_tool_call(call_id: String, name: &str, arguments: &str) -> Value { + json!({ + "id": call_id, + "type": "function", + "function": { + "name": name, + "arguments": arguments, + }, + }) +} + +fn openai_usage(usage: Option<&Value>) -> Value { + let prompt_tokens = usage + .and_then(|usage| usage.get("input_tokens")) + .and_then(Value::as_u64) + .unwrap_or_default(); + let completion_tokens = usage + .and_then(|usage| usage.get("output_tokens")) + .and_then(Value::as_u64) + .unwrap_or_default(); + let cached_tokens = usage + .and_then(|usage| usage.get("input_tokens_details")) + .and_then(|details| details.get("cached_tokens")) + .and_then(Value::as_u64) + .unwrap_or_default(); + let reasoning_tokens = usage + .and_then(|usage| usage.get("output_tokens_details")) + .and_then(|details| details.get("reasoning_tokens")) + .and_then(Value::as_u64); + let mut openai_usage = json!({ + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens.saturating_add(completion_tokens), + "prompt_tokens_details": { + "cached_tokens": cached_tokens, + }, + }); + if let Some(reasoning_tokens) = reasoning_tokens { + openai_usage["completion_tokens_details"] = json!({ + "reasoning_tokens": reasoning_tokens, + }); + } + openai_usage +} diff --git a/crates/core/src/protocol/openai/errors.rs b/crates/core/src/protocol/openai/errors.rs new file mode 100644 index 000000000..5f7929048 --- /dev/null +++ b/crates/core/src/protocol/openai/errors.rs @@ -0,0 +1,15 @@ +//! OpenAI 错误响应类型。 + +use serde::{Deserialize, Serialize}; + +/// OpenAI 错误对象。 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct OpenAiError { + /// 错误消息。 + pub message: String, + /// 错误类型。 + #[serde(rename = "type")] + pub error_type: String, + /// 错误代码。 + pub code: Option, +} diff --git a/crates/core/src/protocol/openai/mod.rs b/crates/core/src/protocol/openai/mod.rs new file mode 100644 index 000000000..cad3dd3bf --- /dev/null +++ b/crates/core/src/protocol/openai/mod.rs @@ -0,0 +1,6 @@ +//! OpenAI 协议输入类型与纯转换逻辑。 + +pub mod chat; +pub mod errors; +pub mod models; +pub mod responses; diff --git a/crates/core/src/protocol/openai/models.rs b/crates/core/src/protocol/openai/models.rs new file mode 100644 index 000000000..a7c8b358e --- /dev/null +++ b/crates/core/src/protocol/openai/models.rs @@ -0,0 +1,25 @@ +//! OpenAI 模型列表响应类型。 + +use serde::{Deserialize, Serialize}; + +/// OpenAI 模型对象。 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct OpenAiModel { + /// 模型 ID。 + pub id: String, + /// 对象类型。 + pub object: String, + /// 创建时间戳。 + pub created: i64, + /// 所有者。 + pub owned_by: String, +} + +/// OpenAI 模型列表。 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct OpenAiModelList { + /// 对象类型。 + pub object: String, + /// 模型数据。 + pub data: Vec, +} diff --git a/crates/core/src/protocol/openai/responses.rs b/crates/core/src/protocol/openai/responses.rs new file mode 100644 index 000000000..5bde9597a --- /dev/null +++ b/crates/core/src/protocol/openai/responses.rs @@ -0,0 +1,745 @@ +//! OpenAI Responses API 类型。 + +use std::collections::BTreeSet; + +use serde::{Deserialize, Serialize}; +use serde_json::{json, Map, Value}; + +use crate::protocol::codex::{ + responses::{CodexCompactRequest, CodexResponsesRequest}, + schema::{prepare_schema, reconvert_tuple_values}, + sse::{parse_sse_events, SseError}, +}; + +/// OpenAI 请求声明的响应格式类别。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ResponseFormat { + /// 普通文本输出。 + Text, + /// JSON object 输出。 + JsonObject, + /// JSON schema 输出。 + JsonSchema, +} + +/// Responses API 请求体。 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct OpenAiResponsesRequest { + /// 模型名。 + pub model: String, + /// 输入内容。 + #[serde(default = "default_responses_input")] + pub input: Value, + /// 指令文本。 + #[serde(default)] + pub instructions: Option, + /// reasoning 配置。 + #[serde(default)] + pub reasoning: Option, + /// 工具定义。 + #[serde(default)] + pub tools: Option, + /// service tier。 + #[serde(default)] + pub service_tier: Option, + /// 工具选择策略。 + #[serde(default)] + pub tool_choice: Option, + /// 是否允许并行工具调用。 + #[serde(default)] + pub parallel_tool_calls: Option, + /// 输出文本格式配置。 + #[serde(default)] + pub text: Option, + /// 图片生成开关。 + #[serde(default)] + pub generate: Option, + /// 提示缓存键。 + #[serde(default)] + pub prompt_cache_key: Option, + /// include 列表。 + #[serde(default)] + pub include: Option, + /// 传给上游的 client metadata。 + #[serde(default)] + pub client_metadata: Option, + /// 前一个响应 ID。 + #[serde(default)] + pub previous_response_id: Option, + /// Codex turn state 透传头。 + #[serde(default, rename = "turnState", alias = "turn_state")] + pub turn_state: Option, + /// Codex turn metadata 透传头。 + #[serde(default, rename = "turnMetadata", alias = "turn_metadata")] + pub turn_metadata: Option, + /// beta features 透传头。 + #[serde(default, rename = "betaFeatures", alias = "beta_features")] + pub beta_features: Option, + /// 客户端版本头。 + #[serde(default)] + pub version: Option, + /// timing metrics 透传头。 + #[serde( + default, + rename = "includeTimingMetrics", + alias = "include_timing_metrics" + )] + pub include_timing_metrics: Option, + /// codex window id。 + #[serde(default, rename = "codexWindowId", alias = "codex_window_id")] + pub codex_window_id: Option, + /// 父线程 ID。 + #[serde(default, rename = "parentThreadId", alias = "parent_thread_id")] + pub parent_thread_id: Option, + /// 是否偏好 WebSocket 传输。 + #[serde(default)] + pub use_websocket: Option, + /// 是否流式返回。 + #[serde(default = "default_responses_stream")] + pub stream: bool, +} + +fn default_responses_input() -> Value { + Value::Array(Vec::new()) +} + +fn default_responses_stream() -> bool { + true +} + +/// 将 OpenAI Responses 请求转换为 Codex Responses 请求。 +pub fn translate_response_to_codex(request: OpenAiResponsesRequest) -> CodexResponsesRequest { + let prepared_text = prepare_text_format(request.text, true); + let mut codex_request = CodexResponsesRequest::new_http_sse( + request.model, + request.instructions.unwrap_or_default(), + sanitize_responses_input(request.input), + ); + codex_request.previous_response_id = request.previous_response_id; + codex_request.turn_state = request.turn_state; + codex_request.turn_metadata = non_empty_string(request.turn_metadata); + codex_request.beta_features = non_empty_string(request.beta_features); + codex_request.include_timing_metrics = non_empty_string(request.include_timing_metrics); + codex_request.version = non_empty_string(request.version); + codex_request.codex_window_id = non_empty_string(request.codex_window_id); + codex_request.parent_thread_id = non_empty_string(request.parent_thread_id); + codex_request.reasoning = responses_reasoning(request.reasoning); + codex_request.tools = non_empty_array(request.tools); + codex_request.tool_choice = request.tool_choice; + codex_request.parallel_tool_calls = request.parallel_tool_calls; + codex_request.text = prepared_text.text; + codex_request.tuple_schema = prepared_text.tuple_schema; + codex_request.generate = request.generate; + codex_request.service_tier = non_empty_string(request.service_tier); + codex_request.prompt_cache_key = non_empty_string(request.prompt_cache_key); + codex_request.explicit_prompt_cache_key = codex_request.prompt_cache_key.is_some(); + codex_request.include = string_array(request.include); + codex_request.client_metadata = sanitize_client_metadata(request.client_metadata); + ensure_reasoning_include(&mut codex_request); + match request.use_websocket { + Some(true) => codex_request.use_websocket = true, + Some(false) => codex_request.force_http_sse = true, + None => {} + } + codex_request +} + +/// 将 OpenAI Responses 请求转换为 Codex compact 请求。 +pub fn translate_response_to_compact(request: OpenAiResponsesRequest) -> CodexCompactRequest { + CodexCompactRequest { + model: request.model, + input: sanitize_responses_input(request.input), + instructions: request.instructions.unwrap_or_default(), + tools: non_empty_array(request.tools), + parallel_tool_calls: request.parallel_tool_calls, + reasoning: compact_reasoning(request.reasoning), + text: prepare_text_format(request.text, false).text, + } +} + +struct PreparedTextFormat { + text: Option, + tuple_schema: Option, +} + +fn prepare_text_format(text: Option, prepare_tuple_schema: bool) -> PreparedTextFormat { + let Some(Value::Object(text)) = text else { + return PreparedTextFormat { + text: None, + tuple_schema: None, + }; + }; + let Some(Value::Object(format)) = text.get("format") else { + return PreparedTextFormat { + text: None, + tuple_schema: None, + }; + }; + let Some(format_type) = format.get("type").and_then(Value::as_str) else { + return PreparedTextFormat { + text: None, + tuple_schema: None, + }; + }; + + let mut sanitized_format = Map::new(); + sanitized_format.insert("type".to_string(), Value::String(format_type.to_string())); + if let Some(name) = format.get("name").and_then(Value::as_str) { + sanitized_format.insert("name".to_string(), Value::String(name.to_string())); + } + + let mut tuple_schema = None; + if let Some(Value::Object(schema)) = format.get("schema") { + let schema = Value::Object(schema.clone()); + let schema = if prepare_tuple_schema { + let prepared = prepare_schema(schema); + tuple_schema = prepared.original_schema; + prepared.schema + } else { + schema + }; + sanitized_format.insert("schema".to_string(), schema); + } + if let Some(strict) = format.get("strict").and_then(Value::as_bool) { + sanitized_format.insert("strict".to_string(), Value::Bool(strict)); + } + + PreparedTextFormat { + text: Some(json!({"format": sanitized_format})), + tuple_schema, + } +} + +fn sanitize_responses_input(input: Value) -> Vec { + match input { + Value::Array(items) => sanitize_codex_input_items(items), + Value::Null => Vec::new(), + value => vec![value], + } +} + +fn sanitize_codex_input_items(input: Vec) -> Vec { + input + .into_iter() + .filter_map(|item| { + let Value::Object(object) = item else { + return Some(item); + }; + match object.get("type").and_then(Value::as_str) { + Some("reasoning") => sanitize_reasoning_item(&object), + Some("compaction") => sanitize_compaction_item(&object), + _ => Some(Value::Object(object)), + } + }) + .collect() +} + +fn sanitize_reasoning_item(item: &Map) -> Option { + let id = non_empty_str(item.get("id"))?; + let summary = sanitize_summary(item.get("summary"))?; + let mut sanitized = Map::new(); + sanitized.insert("type".to_string(), Value::String("reasoning".to_string())); + sanitized.insert("id".to_string(), Value::String(id.to_string())); + sanitized.insert("summary".to_string(), Value::Array(summary)); + if let Some(status) = item + .get("status") + .and_then(Value::as_str) + .filter(|status| matches!(*status, "in_progress" | "completed" | "incomplete")) + { + sanitized.insert("status".to_string(), Value::String(status.to_string())); + } + if let Some(encrypted_content) = non_empty_str(item.get("encrypted_content")) { + sanitized.insert( + "encrypted_content".to_string(), + Value::String(encrypted_content.to_string()), + ); + } + if let Some(content) = sanitize_reasoning_content(item.get("content")) { + sanitized.insert("content".to_string(), Value::Array(content)); + } + Some(Value::Object(sanitized)) +} + +fn sanitize_summary(value: Option<&Value>) -> Option> { + let Value::Array(parts) = value? else { + return None; + }; + Some( + parts + .iter() + .filter_map(|part| { + let Value::Object(part) = part else { + return None; + }; + if part.get("type").and_then(Value::as_str) != Some("summary_text") { + return None; + } + let text = part.get("text").and_then(Value::as_str)?; + Some(json!({"type": "summary_text", "text": text})) + }) + .collect(), + ) +} + +fn sanitize_reasoning_content(value: Option<&Value>) -> Option> { + let Value::Array(parts) = value? else { + return None; + }; + let content = parts + .iter() + .filter_map(|part| { + let Value::Object(part) = part else { + return None; + }; + if part.get("type").and_then(Value::as_str) != Some("reasoning_text") { + return None; + } + let text = part.get("text").and_then(Value::as_str)?; + Some(json!({"type": "reasoning_text", "text": text})) + }) + .collect::>(); + (!content.is_empty()).then_some(content) +} + +fn sanitize_compaction_item(item: &Map) -> Option { + let encrypted_content = non_empty_str(item.get("encrypted_content"))?; + let mut sanitized = Map::new(); + sanitized.insert("type".to_string(), Value::String("compaction".to_string())); + sanitized.insert( + "encrypted_content".to_string(), + Value::String(encrypted_content.to_string()), + ); + if let Some(id) = non_empty_str(item.get("id")) { + sanitized.insert("id".to_string(), Value::String(id.to_string())); + } + Some(Value::Object(sanitized)) +} + +fn responses_reasoning(reasoning: Option) -> Option { + let Value::Object(input) = reasoning? else { + return None; + }; + let effort = input.get("effort").and_then(Value::as_str); + let summary = input + .get("summary") + .and_then(Value::as_str) + .unwrap_or("auto"); + let mut output = Map::new(); + output.insert("summary".to_string(), Value::String(summary.to_string())); + if let Some(effort) = effort { + output.insert("effort".to_string(), Value::String(effort.to_string())); + } + Some(Value::Object(output)) +} + +fn compact_reasoning(reasoning: Option) -> Option { + let Value::Object(input) = reasoning? else { + return None; + }; + let mut output = Map::new(); + if let Some(effort) = input.get("effort").and_then(Value::as_str) { + output.insert("effort".to_string(), Value::String(effort.to_string())); + } + if let Some(summary) = input.get("summary").and_then(Value::as_str) { + output.insert("summary".to_string(), Value::String(summary.to_string())); + } + (!output.is_empty()).then_some(Value::Object(output)) +} + +fn ensure_reasoning_include(request: &mut CodexResponsesRequest) { + if request.reasoning.is_none() { + return; + } + if request + .include + .as_ref() + .is_some_and(|include| !include.is_empty()) + { + return; + } + request.include = Some(vec!["reasoning.encrypted_content".to_string()]); +} + +fn sanitize_client_metadata(client_metadata: Option) -> Option { + let Value::Object(input) = client_metadata? else { + return None; + }; + let metadata = input + .into_iter() + .filter_map(|(key, value)| { + value + .as_str() + .map(|value| (key, Value::String(value.to_string()))) + }) + .collect::>(); + (!metadata.is_empty()).then_some(Value::Object(metadata)) +} + +fn non_empty_array(value: Option) -> Option> { + let Value::Array(values) = value? else { + return None; + }; + (!values.is_empty()).then_some(values) +} + +fn string_array(value: Option) -> Option> { + let Value::Array(values) = value? else { + return None; + }; + values + .into_iter() + .map(|value| match value { + Value::String(value) => Some(value), + _ => None, + }) + .collect() +} + +fn non_empty_string(value: Option) -> Option { + value.and_then(|value| { + let trimmed = value.trim(); + (!trimmed.is_empty()).then(|| trimmed.to_string()) + }) +} + +fn non_empty_str(value: Option<&Value>) -> Option<&str> { + let value = value?.as_str()?; + (!value.trim().is_empty()).then_some(value) +} + +/// 从 Codex SSE 收集出的 Responses API 非流式结果。 +#[derive(Debug, Clone, PartialEq)] +pub enum CollectedResponse { + /// 上游返回了完成响应。 + Completed(Value), + /// 上游返回了失败事件。 + Failed(ResponsesSseFailure), + /// 上游没有返回 `response.completed`。 + MissingCompleted, + /// 上游完成响应没有可见输出。 + Empty, +} + +/// Responses SSE 失败事件。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ResponsesSseFailure { + /// SSE 事件名。 + pub event: String, + /// 错误消息。 + pub message: String, + /// 上游错误代码。 + pub 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), + } + } +} + +/// 从完成 SSE 中提取会话亲和性和 replay 元数据。 +#[derive(Debug, Clone, PartialEq)] +pub struct CompletedResponseMetadata { + /// 响应 ID。 + pub response_id: String, + /// 完成响应中包含的函数调用 ID。 + pub function_call_ids: Vec, + /// 可用于 reasoning replay 的响应条目。 + pub replay_items: Vec, +} + +/// 将 Codex Responses SSE 完成响应收集为非流式 Responses API JSON。 +/// +/// # Errors +/// +/// 当输入不是合法 SSE 流时，返回 [`SseError`]。 +pub fn response_from_codex_sse( + body: &str, + tuple_schema: Option<&Value>, +) -> 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") => { + completed_response = value.get("response").cloned(); + } + 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); + }; + if is_empty_response(&response, &output_text, &output_items) { + return Ok(CollectedResponse::Empty); + } + + ensure_completed_response_output(&mut response, &output_items, &output_text); + reconvert_completed_response_tuple_values(&mut response, tuple_schema); + sync_output_text_from_output(&mut response); + Ok(CollectedResponse::Completed(response)) +} + +/// 从 Codex Responses SSE 中提取完成响应元数据。 +/// +/// # Errors +/// +/// 当输入不是合法 SSE 流时，返回 [`SseError`]。 +pub fn completed_response_metadata( + body: &str, +) -> Result, SseError> { + let events = parse_sse_events(body)?; + let mut response_id = None; + let mut function_call_ids = BTreeSet::new(); + let mut replay_items = Vec::new(); + + for event in events { + let Ok(value) = serde_json::from_str::(&event.data) else { + continue; + }; + match event.event.as_deref() { + Some("response.output_item.done") => { + if let Some(item) = value.get("item") { + collect_response_replay_items(item, &mut replay_items); + } + if let Some(call_id) = value.pointer("/item/call_id").and_then(Value::as_str) { + function_call_ids.insert(call_id.to_string()); + } + } + Some("response.completed") => { + response_id = value + .pointer("/response/id") + .and_then(Value::as_str) + .map(ToString::to_string); + collect_response_function_call_ids(&value, &mut function_call_ids); + if let Some(output) = value.pointer("/response/output").and_then(Value::as_array) { + for item in output { + collect_response_replay_items(item, &mut replay_items); + } + } + } + _ => {} + } + } + + Ok(response_id.map(|response_id| CompletedResponseMetadata { + response_id, + function_call_ids: function_call_ids.into_iter().collect(), + replay_items, + })) +} + +/// 对单个 Responses SSE 事件的数据执行 tuple schema 回转换。 +pub fn reconvert_responses_sse_event_tuple_values( + event_name: Option<&str>, + mut data: Value, + tuple_schema: &Value, +) -> Value { + match responses_event_type(event_name, &data) { + Some("response.output_text.delta") => { + reconvert_output_text_delta_tuple_values(&mut data, tuple_schema); + } + Some("response.output_item.done") => { + if let Some(item) = data.get_mut("item") { + reconvert_output_item_tuple_values(item, tuple_schema); + } + } + Some("response.completed") => { + if let Some(response) = data.get_mut("response") { + reconvert_completed_response_tuple_values(response, Some(tuple_schema)); + sync_output_text_from_output(response); + } + } + _ => {} + } + data +} + +fn responses_event_type<'a>(event_name: Option<&'a str>, data: &'a Value) -> Option<&'a str> { + event_name.or_else(|| data.get("type").and_then(Value::as_str)) +} + +fn reconvert_output_text_delta_tuple_values(data: &mut Value, tuple_schema: &Value) { + let Some(delta) = data.get("delta").and_then(Value::as_str) else { + return; + }; + let Ok(parsed) = serde_json::from_str::(delta) else { + return; + }; + let reconverted = reconvert_tuple_values(parsed, tuple_schema); + data["delta"] = Value::String(reconverted.to_string()); +} + +fn is_empty_response(response: &Value, output_text: &str, output_items: &[Value]) -> bool { + if !output_text.trim().is_empty() || !output_items.is_empty() { + return false; + } + + response + .pointer("/usage/output_tokens") + .and_then(Value::as_u64) + .unwrap_or_default() + == 0 +} + +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 reconvert_completed_response_tuple_values(response: &mut Value, tuple_schema: Option<&Value>) { + let Some(tuple_schema) = tuple_schema else { + return; + }; + let Some(items) = response.get_mut("output").and_then(Value::as_array_mut) else { + return; + }; + for item in items { + reconvert_output_item_tuple_values(item, tuple_schema); + } +} + +fn reconvert_output_item_tuple_values(item: &mut Value, tuple_schema: &Value) { + let Some(content) = item.get_mut("content").and_then(Value::as_array_mut) else { + return; + }; + for part in content { + let part_type = part.get("type").and_then(Value::as_str); + if !matches!(part_type, Some("output_text" | "text")) { + continue; + } + let Some(text) = part.get("text").and_then(Value::as_str) else { + continue; + }; + let Ok(parsed) = serde_json::from_str::(text) else { + continue; + }; + let reconverted = reconvert_tuple_values(parsed, tuple_schema); + part["text"] = Value::String(reconverted.to_string()); + } +} + +fn sync_output_text_from_output(response: &mut Value) { + 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 collect_response_function_call_ids(value: &Value, function_call_ids: &mut BTreeSet) { + let Some(output) = value.pointer("/response/output").and_then(Value::as_array) else { + return; + }; + for item in output { + if let Some(call_id) = item.get("call_id").and_then(Value::as_str) { + function_call_ids.insert(call_id.to_string()); + } + } +} + +fn collect_response_replay_items(item: &Value, replay_items: &mut Vec) { + if matches!( + item.get("type").and_then(Value::as_str), + Some("reasoning" | "function_call") + ) { + replay_items.push(item.clone()); + } +} + +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 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) +} diff --git a/crates/core/src/serving/affinity.rs b/crates/core/src/serving/affinity.rs new file mode 100644 index 000000000..a029a51b6 --- /dev/null +++ b/crates/core/src/serving/affinity.rs @@ -0,0 +1,238 @@ +//! 会话亲和性策略。 + +use std::collections::BTreeMap; + +use chrono::{DateTime, Duration, Utc}; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +use crate::{ + gateway::conversation::derive_stable_conversation_key, + protocol::codex::responses::CodexResponsesRequest, +}; + +/// 会话亲和性键。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SessionAffinityKey { + /// 客户端会话 ID。 + pub conversation_id: String, +} + +/// 会话亲和性条目。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SessionAffinityEntry { + /// 创建该响应的账号 ID。 + pub account_id: String, + /// 对话 ID。 + pub conversation_id: String, + /// 上游 turn state。 + pub turn_state: Option, + /// 指令哈希。 + pub instructions_hash: Option, + /// 输入 token 数。 + pub input_tokens: Option, + /// 函数调用 ID 列表。 + pub function_call_ids: Vec, + /// 用于区分并发对话分支的变体哈希。 + pub variant_hash: Option, + /// 创建时间。 + pub created_at: DateTime, +} + +/// 持久化的会话亲和性条目。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StoredSessionAffinity { + /// 响应 ID。 + pub response_id: String, + /// 亲和性条目。 + pub entry: SessionAffinityEntry, + /// 持久化过期时间。 + pub expires_at: DateTime, +} + +/// 纯内存会话亲和性映射。 +#[derive(Debug, Clone)] +pub struct SessionAffinityMap { + entries: BTreeMap, + ttl: Duration, +} + +impl SessionAffinityMap { + /// 构造指定 TTL 的会话亲和性映射。 + pub fn new(ttl: Duration) -> Self { + Self { + entries: BTreeMap::new(), + ttl, + } + } + + /// 记录响应 ID 到亲和性条目的映射。 + pub fn record(&mut self, response_id: String, entry: SessionAffinityEntry) { + self.entries.insert(response_id, entry); + } + + /// 从持久化记录恢复未过期条目。 + pub fn restore(&mut self, records: Vec, now: DateTime) -> usize { + let mut restored = 0usize; + for record in records { + if self.stored_record_is_active(&record, now) { + self.entries.insert(record.response_id, record.entry); + restored = restored.saturating_add(1); + } + } + restored + } + + /// 根据响应 ID 查找账号 ID。 + pub fn lookup_account(&self, response_id: &str, now: DateTime) -> Option { + self.active_entry(response_id, now) + .map(|entry| entry.account_id.clone()) + } + + /// 根据响应 ID 查找对话 ID。 + pub fn lookup_conversation_id(&self, response_id: &str, now: DateTime) -> Option { + self.active_entry(response_id, now) + .map(|entry| entry.conversation_id.clone()) + } + + /// 根据响应 ID 查找 turn state。 + pub fn lookup_turn_state(&self, response_id: &str, now: DateTime) -> Option { + self.active_entry(response_id, now) + .and_then(|entry| entry.turn_state.clone()) + } + + /// 根据响应 ID 查找指令哈希。 + pub fn lookup_instructions_hash( + &self, + response_id: &str, + now: DateTime, + ) -> Option { + self.active_entry(response_id, now) + .and_then(|entry| entry.instructions_hash.clone()) + } + + /// 根据响应 ID 查找函数调用 ID 列表。 + pub fn lookup_function_call_ids(&self, response_id: &str, now: DateTime) -> Vec { + self.active_entry(response_id, now) + .map(|entry| entry.function_call_ids.clone()) + .unwrap_or_default() + } + + /// 查找指定对话和可选变体下最新的响应 ID。 + pub fn lookup_latest_response_by_conversation( + &self, + conversation_id: &str, + max_age: Option, + variant_hash: Option<&str>, + now: DateTime, + ) -> Option { + self.entries + .iter() + .filter(|(_, entry)| entry.conversation_id == conversation_id) + .filter(|(_, entry)| self.entry_is_active(entry, now)) + .filter(|(_, entry)| { + variant_hash + .is_none_or(|variant_hash| entry.variant_hash.as_deref() == Some(variant_hash)) + }) + .filter(|(_, entry)| { + max_age.is_none_or(|max_age| now.signed_duration_since(entry.created_at) <= max_age) + }) + .max_by_key(|(_, entry)| entry.created_at) + .map(|(response_id, _)| response_id.clone()) + } + + /// 删除响应 ID 映射。 + pub fn forget(&mut self, response_id: &str) -> bool { + self.entries.remove(response_id).is_some() + } + + /// 返回当前内存条目数。 + pub fn len(&self) -> usize { + self.entries.len() + } + + /// 返回是否没有内存条目。 + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + fn active_entry(&self, response_id: &str, now: DateTime) -> Option<&SessionAffinityEntry> { + self.entries + .get(response_id) + .filter(|entry| self.entry_is_active(entry, now)) + } + + fn stored_record_is_active(&self, record: &StoredSessionAffinity, now: DateTime) -> bool { + record.expires_at > now && self.entry_is_active(&record.entry, now) + } + + fn entry_is_active(&self, entry: &SessionAffinityEntry, now: DateTime) -> bool { + entry + .created_at + .checked_add_signed(self.ttl) + .is_some_and(|expires_at| expires_at > now) + } +} + +/// 准备用于区分并发分支的变体身份。 +pub fn prepare_variant_identity(request: &mut CodexResponsesRequest) { + request.variant_identity = build_variant_identity(request); +} + +/// 计算请求变体哈希。 +pub fn compute_variant_hash(request: &CodexResponsesRequest) -> String { + compute_variant_hash_with_identity(request, request.variant_identity.as_deref()) +} + +fn compute_variant_hash_with_identity( + request: &CodexResponsesRequest, + identity: Option<&str>, +) -> String { + let tools_json = request + .tools + .as_ref() + .map_or_else(|| "[]".to_string(), |tools| tools_json(tools)); + let mut hasher = Sha256::new(); + hasher.update(request.instructions.as_bytes()); + hasher.update(b"\0"); + hasher.update(tools_json.as_bytes()); + if let Some(identity) = identity + .map(str::trim) + .filter(|identity| !identity.is_empty()) + { + hasher.update(b"\0"); + hasher.update(identity.as_bytes()); + } + hex::encode(hasher.finalize()).chars().take(12).collect() +} + +fn tools_json(tools: &[Value]) -> String { + serde_json::to_string(tools).unwrap_or_else(|_| "[]".to_string()) +} + +fn build_variant_identity(request: &CodexResponsesRequest) -> Option { + let mut parts = Vec::with_capacity(2); + if let Some(window_id) = non_empty_str(request.codex_window_id.as_deref()) { + parts.push(format!("window:{window_id}")); + } + if request.explicit_prompt_cache_key + || non_empty_str(request.client_conversation_id.as_deref()).is_some() + { + if let Some(anchor) = derive_stable_conversation_key(request) { + parts.push(format!("anchor:{anchor}")); + } + } + + (!parts.is_empty()).then(|| parts.join("\0")) +} + +fn non_empty_str(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +/// 计算 instructions 哈希。 +pub fn hash_instructions(instructions: Option<&str>) -> String { + let mut hasher = Sha256::new(); + hasher.update(instructions.unwrap_or_default().as_bytes()); + hex::encode(hasher.finalize()) +} diff --git a/crates/core/src/serving/chat.rs b/crates/core/src/serving/chat.rs new file mode 100644 index 000000000..ca9a603c4 --- /dev/null +++ b/crates/core/src/serving/chat.rs @@ -0,0 +1,13 @@ +//! Chat 补全编排。 + +use crate::protocol::{ + codex::responses::CodexResponsesRequest, + openai::chat::{translate_chat_to_codex, ChatCompletionRequest, ChatTranslationError}, +}; + +/// 将 Chat 请求准备为上游 Responses 请求。 +pub fn prepare_chat_request( + request: ChatCompletionRequest, +) -> Result { + translate_chat_to_codex(request) +} diff --git a/crates/core/src/serving/errors.rs b/crates/core/src/serving/errors.rs new file mode 100644 index 000000000..b0d0d8c51 --- /dev/null +++ b/crates/core/src/serving/errors.rs @@ -0,0 +1,17 @@ +//! 服务编排错误。 + +use thiserror::Error; + +/// 服务编排错误。 +#[derive(Debug, Error, PartialEq, Eq)] +pub enum ServingError { + /// 没有可用账号。 + #[error("no available account")] + NoAvailableAccount, + /// 上游请求失败。 + #[error("upstream request failed: {message}")] + Upstream { + /// 错误说明。 + message: String, + }, +} diff --git a/crates/core/src/serving/fallback.rs b/crates/core/src/serving/fallback.rs new file mode 100644 index 000000000..a41cf032c --- /dev/null +++ b/crates/core/src/serving/fallback.rs @@ -0,0 +1,16 @@ +//! 上游回退策略。 + +/// 判断状态码是否表示上游限流。 +pub fn status_code_is_rate_limited(status_code: u16) -> bool { + status_code == 429 +} + +/// 判断状态码是否表示账号配额耗尽。 +pub fn status_code_is_quota_exhausted(status_code: u16) -> bool { + status_code == 402 +} + +/// 判断状态码是否表示可换账号重试的临时上游错误。 +pub fn status_code_is_transient_upstream(status_code: u16) -> bool { + matches!(status_code, 500..=599) +} diff --git a/src/codex/serving/dispatch/implicit_resume.rs b/crates/core/src/serving/implicit_resume.rs similarity index 60% rename from src/codex/serving/dispatch/implicit_resume.rs rename to crates/core/src/serving/implicit_resume.rs index 1386554b0..d55954d18 100644 --- a/src/codex/serving/dispatch/implicit_resume.rs +++ b/crates/core/src/serving/implicit_resume.rs @@ -1,20 +1,27 @@ +//! Responses 隐式续接策略。 + use std::collections::HashSet; use serde_json::Value; -use crate::codex::gateway::transport::types::CodexResponsesRequest; - -#[derive(Debug, Clone)] -pub(crate) struct ImplicitResumeSnapshot { - input: Vec, - previous_response_id: Option, - turn_state: Option, - use_websocket: bool, - force_http_sse: bool, +/// 隐式续接前的可恢复请求状态。 +#[derive(Debug, Clone, PartialEq)] +pub struct ImplicitResumeSnapshot { + /// 完整输入历史。 + pub input: Vec, + /// 原始 previous_response_id。 + pub previous_response_id: Option, + /// 原始 turn state。 + pub turn_state: Option, + /// 原始 WebSocket 偏好。 + pub use_websocket: bool, + /// 原始 HTTP SSE 强制标记。 + pub force_http_sse: bool, } impl ImplicitResumeSnapshot { - pub(super) fn capture(request: &CodexResponsesRequest) -> Self { + /// 捕获当前请求中隐式续接会改写的字段。 + pub fn capture(request: &crate::protocol::codex::responses::CodexResponsesRequest) -> Self { Self { input: request.input.clone(), previous_response_id: request.previous_response_id.clone(), @@ -24,7 +31,8 @@ impl ImplicitResumeSnapshot { } } - pub(crate) fn restore(self, request: &mut CodexResponsesRequest) { + /// 恢复隐式续接前的请求字段。 + pub fn restore(self, request: &mut crate::protocol::codex::responses::CodexResponsesRequest) { request.input = self.input; request.previous_response_id = self.previous_response_id; request.turn_state = self.turn_state; @@ -33,7 +41,8 @@ impl ImplicitResumeSnapshot { } } -pub(super) fn continuation_input_start(input: &[Value]) -> usize { +/// 返回续接输入在完整输入中的起始位置。 +pub fn continuation_input_start(input: &[Value]) -> usize { let mut last_model_output_index = None; for (index, item) in input.iter().enumerate() { if item.get("role").is_some() { @@ -49,7 +58,8 @@ pub(super) fn continuation_input_start(input: &[Value]) -> usize { last_model_output_index.map_or(0, |index| index.saturating_add(1)) } -pub(super) fn implicit_resume_allowed( +/// 判断完整输入历史是否可以隐式续接到已记录响应。 +pub fn implicit_resume_allowed( continuation_input: &[Value], full_input: &[Value], stored_function_call_ids: &[String], @@ -98,3 +108,36 @@ fn inline_function_call_ids(input: &[Value]) -> HashSet { .map(ToString::to_string) .collect() } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn continuation_input_start_should_follow_last_assistant_message() { + let input = vec![ + json!({"role": "user", "content": "first"}), + json!({"role": "assistant", "content": "cached"}), + json!({"role": "user", "content": "continue"}), + ]; + + assert_eq!(continuation_input_start(&input), 2); + } + + #[test] + fn implicit_resume_allowed_should_reject_unmatched_function_call_output() { + let input = vec![json!({ + "type": "function_call_output", + "call_id": "missing", + "output": "tool output" + })]; + + assert!(!implicit_resume_allowed( + &input, + &input, + &["expected".to_string()] + )); + } +} diff --git a/crates/core/src/serving/mod.rs b/crates/core/src/serving/mod.rs new file mode 100644 index 000000000..7bef82b24 --- /dev/null +++ b/crates/core/src/serving/mod.rs @@ -0,0 +1,26 @@ +//! 对外服务编排策略。 + +/// 亲和性策略。 +pub mod affinity; +/// Chat 补全编排。 +pub mod chat; +/// 服务错误。 +pub mod errors; +/// 回退策略。 +pub mod fallback; +/// Responses 隐式续接策略。 +pub mod implicit_resume; +/// 配额策略。 +pub mod quota; +/// Responses reasoning replay 缓存策略。 +pub mod reasoning_replay; +/// 错误恢复策略。 +pub mod recovery; +/// Responses 编排。 +pub mod responses; +/// 路由策略。 +pub mod routing; +/// 流生命周期策略。 +pub mod stream; +/// 用量跟踪策略。 +pub mod usage; diff --git a/src/codex/accounts/service/quota.rs b/crates/core/src/serving/quota.rs similarity index 53% rename from src/codex/accounts/service/quota.rs rename to crates/core/src/serving/quota.rs index 69d8fcd80..c2038b951 100644 --- a/src/codex/accounts/service/quota.rs +++ b/crates/core/src/serving/quota.rs @@ -1,162 +1,21 @@ -use crate::codex::accounts::model::AccountStatus; - -use super::{ - health::{apply_codex_account_error, fetch_account_usage, public_codex_error}, - AccountQuotaError, AccountQuotaResult, AccountQuotaWarning, AccountQuotaWarnings, - AccountService, AccountServiceError, QuotaWarningLevel, QuotaWarningWindow, -}; +//! 配额检查策略。 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 async fn quota_warnings(&self) -> Result { - let snapshots = self - .repository()? - .list_quota_snapshots() - .await - .map_err(|_| AccountServiceError::QuotaWarnings)?; - let primary_thresholds = sorted_thresholds(&self.config.quota.warning_thresholds.primary); - let secondary_thresholds = - sorted_thresholds(&self.config.quota.warning_thresholds.secondary); - let mut warnings = Vec::new(); - let mut updated_at = None; - - for snapshot in snapshots { - let Ok(quota) = serde_json::from_str::(&snapshot.quota_json) else { - continue; - }; - let before_len = warnings.len(); - if let Some(warning) = warning_from_quota_window( - &snapshot.account_id, - snapshot.email.as_deref(), - "a, - "rate_limit", - QuotaWarningWindow::Primary, - &primary_thresholds, - ) { - warnings.push(warning); - } - if let Some(warning) = warning_from_quota_window( - &snapshot.account_id, - snapshot.email.as_deref(), - "a, - "secondary_rate_limit", - QuotaWarningWindow::Secondary, - &secondary_thresholds, - ) { - warnings.push(warning); - } - if warnings.len() > before_len { - updated_at = max_optional_datetime(updated_at, snapshot.quota_fetched_at); - } - } - - Ok(AccountQuotaWarnings { - warnings, - updated_at, - }) - } +/// 判断配额百分比是否达到阈值。 +pub fn quota_reached(used_percent: f64, threshold: f64) -> bool { + used_percent >= threshold } -fn warning_from_quota_window( - account_id: &str, - email: Option<&str>, - quota: &Value, - field: &str, - window: QuotaWarningWindow, - thresholds: &[u8], -) -> Option { - let quota_window = quota.get(field).filter(|value| !value.is_null())?; - let used_percent = quota_window - .get("used_percent") - .and_then(Value::as_f64) - .filter(|value| value.is_finite())?; - let level = warning_level(used_percent, thresholds)?; - - Some(AccountQuotaWarning { - account_id: account_id.to_string(), - email: email.map(str::to_string), - window, - level, - used_percent, - reset_at: quota_window.get("reset_at").and_then(Value::as_i64), - }) -} - -fn warning_level(used_percent: f64, thresholds: &[u8]) -> Option { - let matched_index = thresholds - .iter() - .rposition(|threshold| used_percent >= f64::from(*threshold))?; - if matched_index + 1 == thresholds.len() { - Some(QuotaWarningLevel::Critical) - } else { - Some(QuotaWarningLevel::Warning) - } -} - -fn sorted_thresholds(thresholds: &[u8]) -> Vec { - let mut thresholds = thresholds.to_vec(); - thresholds.sort_unstable(); - thresholds.dedup(); - thresholds -} - -fn max_optional_datetime( - current: Option>, - candidate: Option>, -) -> Option> { - match (current, candidate) { - (Some(current), Some(candidate)) => Some(current.max(candidate)), - (Some(current), None) => Some(current), - (None, Some(candidate)) => Some(candidate), - (None, None) => None, - } -} - -pub(crate) fn quota_from_usage(usage: &Value) -> Value { +/// 从 Codex usage 响应中提取持久化配额快照。 +pub 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") @@ -185,6 +44,7 @@ pub(crate) fn quota_from_usage(usage: &Value) -> Value { }), ); } + 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)) diff --git a/src/codex/serving/dispatch/reasoning_replay.rs b/crates/core/src/serving/reasoning_replay.rs similarity index 62% rename from src/codex/serving/dispatch/reasoning_replay.rs rename to crates/core/src/serving/reasoning_replay.rs index d8d98b7b2..88e59f5f5 100644 --- a/src/codex/serving/dispatch/reasoning_replay.rs +++ b/crates/core/src/serving/reasoning_replay.rs @@ -1,15 +1,16 @@ -use std::{collections::HashMap, time::Duration}; +//! Responses reasoning replay 缓存策略。 -use chrono::{DateTime, Utc}; +use std::collections::{BTreeMap, HashSet}; + +use chrono::{DateTime, Duration, Utc}; use serde_json::{json, Map, Value}; -use tokio::sync::RwLock; -const REASONING_REPLAY_TTL: Duration = Duration::from_secs(55 * 60); const MAX_ENTRIES: usize = 512; const MAX_ENTRY_BYTES: usize = 256 * 1024; const MAX_TOTAL_BYTES: usize = 4 * 1024 * 1024; -#[derive(Debug, Clone)] +/// 单条 reasoning replay 缓存记录。 +#[derive(Debug, Clone, PartialEq)] struct ReasoningReplayEntry { account_id: String, conversation_id: String, @@ -20,31 +21,39 @@ struct ReasoningReplayEntry { sequence: u64, } -#[derive(Debug, Default)] -struct ReasoningReplayState { - entries: HashMap, +/// 纯内存 reasoning replay 缓存。 +#[derive(Debug, Clone)] +pub struct ReasoningReplayCache { + entries: BTreeMap, total_bytes: usize, next_sequence: u64, -} - -#[derive(Debug, Default)] -pub(super) struct ReasoningReplayCache { - state: RwLock, + ttl: Duration, } impl ReasoningReplayCache { - pub(super) async fn record( - &self, - response_id: &str, + /// 构造指定 TTL 的缓存。 + pub fn new(ttl: Duration) -> Self { + Self { + entries: BTreeMap::new(), + total_bytes: 0, + next_sequence: 0, + ttl, + } + } + + /// 记录响应对应的 replay 条目，返回实际保存的条目数。 + pub fn record( + &mut self, + response_id: String, account_id: &str, conversation_id: &str, variant_hash: &str, items: &[Value], + now: DateTime, ) -> usize { let sanitized = sanitize_reasoning_replay_items(items); - let mut state = self.state.write().await; - cleanup_expired(&mut state, Utc::now()); - delete_entry(&mut state, response_id); + self.cleanup_expired(now); + self.delete_entry(&response_id); if sanitized.is_empty() { return 0; @@ -54,38 +63,39 @@ impl ReasoningReplayCache { return 0; } - let sequence = state.next_sequence; - state.next_sequence = state.next_sequence.saturating_add(1); - state.total_bytes = state.total_bytes.saturating_add(byte_size); - state.entries.insert( - response_id.to_string(), + let sequence = self.next_sequence; + self.next_sequence = self.next_sequence.saturating_add(1); + self.total_bytes = self.total_bytes.saturating_add(byte_size); + self.entries.insert( + response_id, ReasoningReplayEntry { account_id: account_id.to_string(), conversation_id: conversation_id.to_string(), variant_hash: variant_hash.to_string(), items: sanitized, byte_size, - created_at: Utc::now(), + created_at: now, sequence, }, ); - evict_overflow(&mut state); - state - .entries - .get(response_id) + self.evict_overflow(); + self.entries + .values() + .max_by_key(|entry| entry.sequence) .map_or(0, |entry| entry.items.len()) } - pub(super) async fn lookup( - &self, + /// 查找指定响应与身份对应的 replay 条目。 + pub fn lookup( + &mut self, response_id: &str, account_id: &str, conversation_id: &str, variant_hash: &str, + now: DateTime, ) -> Vec { - let mut state = self.state.write().await; - cleanup_expired(&mut state, Utc::now()); - let Some(entry) = state.entries.get(response_id) else { + self.cleanup_expired(now); + let Some(entry) = self.entries.get(response_id) else { return Vec::new(); }; if entry.account_id != account_id @@ -97,15 +107,16 @@ impl ReasoningReplayCache { entry.items.clone() } - pub(super) async fn evict_by_identity( - &self, + /// 按账号、对话、变体驱逐 replay 条目。 + pub fn evict_by_identity( + &mut self, account_id: &str, conversation_id: &str, variant_hash: &str, + now: DateTime, ) -> usize { - let mut state = self.state.write().await; - cleanup_expired(&mut state, Utc::now()); - let response_ids = state + self.cleanup_expired(now); + let response_ids = self .entries .iter() .filter(|(_, entry)| { @@ -117,14 +128,50 @@ impl ReasoningReplayCache { .collect::>(); let evicted = response_ids.len(); for response_id in response_ids { - delete_entry(&mut state, &response_id); + self.delete_entry(&response_id); } evicted } + + fn cleanup_expired(&mut self, now: DateTime) { + let ttl = self.ttl; + let response_ids = self + .entries + .iter() + .filter(|(_, entry)| entry_expired(entry.created_at, ttl, now)) + .map(|(response_id, _)| response_id.clone()) + .collect::>(); + for response_id in response_ids { + self.delete_entry(&response_id); + } + } + + fn evict_overflow(&mut self) { + while self.entries.len() > MAX_ENTRIES || self.total_bytes > MAX_TOTAL_BYTES { + let Some(response_id) = self.oldest_response_id() else { + return; + }; + self.delete_entry(&response_id); + } + } + + fn oldest_response_id(&self) -> Option { + self.entries + .iter() + .min_by_key(|(_, entry)| (entry.created_at, entry.sequence)) + .map(|(response_id, _)| response_id.clone()) + } + + fn delete_entry(&mut self, response_id: &str) { + let Some(entry) = self.entries.remove(response_id) else { + return; + }; + self.total_bytes = self.total_bytes.saturating_sub(entry.byte_size); + } } fn sanitize_reasoning_replay_items(items: &[Value]) -> Vec { - let mut seen = std::collections::HashSet::new(); + let mut seen = HashSet::new(); let mut sanitized = Vec::new(); for item in items { let Some(replay_item) = @@ -245,44 +292,39 @@ fn estimate_items_bytes(items: &[Value]) -> usize { serde_json::to_vec(items).map_or(usize::MAX, |bytes| bytes.len()) } -fn cleanup_expired(state: &mut ReasoningReplayState, now: DateTime) { - let response_ids = state - .entries - .iter() - .filter(|(_, entry)| entry_expired(entry.created_at, now)) - .map(|(response_id, _)| response_id.clone()) - .collect::>(); - for response_id in response_ids { - delete_entry(state, &response_id); +fn entry_expired(created_at: DateTime, ttl: Duration, now: DateTime) -> bool { + now.signed_duration_since(created_at) >= ttl +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn reasoning_replay_cache_should_sanitize_reasoning_items() { + let mut cache = ReasoningReplayCache::new(Duration::minutes(55)); + let saved = cache.record( + "resp_1".to_string(), + "acct_1", + "conv_1", + "variant_1", + &[json!({ + "type": "reasoning", + "id": "rs_1", + "status": "completed", + "summary": [{"type": "summary_text", "text": "kept"}], + "encrypted_content": "enc", + "extra": "drop" + })], + Utc::now(), + ); + + let replay = cache.lookup("resp_1", "acct_1", "conv_1", "variant_1", Utc::now()); + + assert_eq!(saved, 1); + assert_eq!(replay[0].get("extra"), None); + assert_eq!(replay[0]["encrypted_content"], "enc"); } } - -fn entry_expired(created_at: DateTime, now: DateTime) -> bool { - now.signed_duration_since(created_at) - .to_std() - .is_ok_and(|age| age > REASONING_REPLAY_TTL) -} - -fn evict_overflow(state: &mut ReasoningReplayState) { - while state.entries.len() > MAX_ENTRIES || state.total_bytes > MAX_TOTAL_BYTES { - let Some(response_id) = oldest_response_id(state) else { - return; - }; - delete_entry(state, &response_id); - } -} - -fn oldest_response_id(state: &ReasoningReplayState) -> Option { - state - .entries - .iter() - .min_by_key(|(_, entry)| (entry.created_at, entry.sequence)) - .map(|(response_id, _)| response_id.clone()) -} - -fn delete_entry(state: &mut ReasoningReplayState, response_id: &str) { - let Some(entry) = state.entries.remove(response_id) else { - return; - }; - state.total_bytes = state.total_bytes.saturating_sub(entry.byte_size); -} diff --git a/crates/core/src/serving/recovery.rs b/crates/core/src/serving/recovery.rs new file mode 100644 index 000000000..9d8cb86d8 --- /dev/null +++ b/crates/core/src/serving/recovery.rs @@ -0,0 +1,6 @@ +//! 错误恢复规则。 + +/// 判断状态码是否允许在同一账号上重试。 +pub fn status_code_allows_same_account_retry(status_code: u16) -> bool { + matches!(status_code, 500..=599) +} diff --git a/crates/core/src/serving/responses.rs b/crates/core/src/serving/responses.rs new file mode 100644 index 000000000..2a04b03fd --- /dev/null +++ b/crates/core/src/serving/responses.rs @@ -0,0 +1,129 @@ +//! Responses 创建编排。 + +use serde_json::{Map, Value}; + +use crate::{ + models::model::{ModelConfig, ParsedModelName}, + protocol::codex::responses::CodexResponsesRequest, +}; + +/// Codex Responses 请求的上游传输决策。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CodexTransport { + /// 使用 HTTP SSE。 + HttpSse, + /// 优先使用 WebSocket，失败时允许回退 HTTP SSE。 + WebSocketPreferred, + /// 必须使用 WebSocket，不允许 HTTP SSE 回退。 + WebSocketRequired, +} + +/// 判断 Responses 请求应使用哪种上游传输。 +pub fn transport_for_request(request: &CodexResponsesRequest) -> CodexTransport { + if request.force_http_sse { + return CodexTransport::HttpSse; + } + + if request.previous_response_id.is_some() { + return CodexTransport::WebSocketRequired; + } + + CodexTransport::WebSocketPreferred +} + +/// 判断请求在 WebSocket 失败后是否允许 HTTP SSE 回退。 +pub fn http_sse_fallback_allowed(request: &CodexResponsesRequest) -> bool { + !matches!( + transport_for_request(request), + CodexTransport::WebSocketRequired + ) +} + +/// 应用模型名后缀和模型默认值对应的 Responses 上游请求选项。 +pub fn apply_response_model_options( + request: &mut CodexResponsesRequest, + parsed_model: &ParsedModelName, + config: &ModelConfig, +) { + request.model = parsed_model.model_id.clone(); + apply_reasoning_options(request, parsed_model, config); + apply_service_tier_options(request, parsed_model, config); +} + +fn apply_reasoning_options( + request: &mut CodexResponsesRequest, + parsed_model: &ParsedModelName, + config: &ModelConfig, +) { + let existing_reasoning = request.reasoning.take(); + let existing_object = match existing_reasoning { + Some(Value::Object(object)) => object, + Some(_) | None => Map::new(), + }; + let effort = existing_object + .get("effort") + .and_then(Value::as_str) + .map(ToString::to_string) + .or_else(|| non_empty_string(parsed_model.reasoning_effort.as_deref())) + .or_else(|| non_empty_string(config.default_reasoning_effort.as_deref())); + if effort.is_none() && existing_object.is_empty() { + request.reasoning = None; + return; + } + + let summary = existing_object + .get("summary") + .and_then(Value::as_str) + .filter(|value| !value.trim().is_empty()) + .unwrap_or("auto"); + let mut reasoning = Map::new(); + reasoning.insert("summary".to_string(), Value::String(summary.to_string())); + if let Some(effort) = effort { + reasoning.insert("effort".to_string(), Value::String(effort)); + } + request.reasoning = Some(Value::Object(reasoning)); + ensure_reasoning_include(request); +} + +fn apply_service_tier_options( + request: &mut CodexResponsesRequest, + parsed_model: &ParsedModelName, + config: &ModelConfig, +) { + request.service_tier = request + .service_tier + .take() + .and_then(|value| non_empty_string(Some(&value))) + .or_else(|| non_empty_string(parsed_model.service_tier.as_deref())) + .or_else(|| non_empty_string(config.service_tier.as_deref())) + .map(normalize_service_tier_for_upstream); +} + +fn ensure_reasoning_include(request: &mut CodexResponsesRequest) { + if request.reasoning.is_none() { + return; + } + if request + .include + .as_ref() + .is_some_and(|include| !include.is_empty()) + { + return; + } + request.include = Some(vec!["reasoning.encrypted_content".to_string()]); +} + +fn normalize_service_tier_for_upstream(service_tier: String) -> String { + if service_tier == "fast" { + "priority".to_string() + } else { + service_tier + } +} + +fn non_empty_string(value: Option<&str>) -> Option { + value.and_then(|value| { + let trimmed = value.trim(); + (!trimmed.is_empty()).then(|| trimmed.to_string()) + }) +} diff --git a/crates/core/src/serving/routing.rs b/crates/core/src/serving/routing.rs new file mode 100644 index 000000000..b4a3a3408 --- /dev/null +++ b/crates/core/src/serving/routing.rs @@ -0,0 +1,8 @@ +//! 账号选择路由策略。 + +/// 路由决策。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RoutingDecision { + /// 选中的账号 ID。 + pub account_id: String, +} diff --git a/crates/core/src/serving/stream.rs b/crates/core/src/serving/stream.rs new file mode 100644 index 000000000..fddebf6cb --- /dev/null +++ b/crates/core/src/serving/stream.rs @@ -0,0 +1,10 @@ +//! 流生命周期策略。 + +/// 流结束原因。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StreamFinishReason { + /// 正常完成。 + Completed, + /// 上游中断。 + Interrupted, +} diff --git a/crates/core/src/serving/usage.rs b/crates/core/src/serving/usage.rs new file mode 100644 index 000000000..77edd411f --- /dev/null +++ b/crates/core/src/serving/usage.rs @@ -0,0 +1,8 @@ +//! 服务用量跟踪策略。 + +use crate::protocol::codex::events::TokenUsage; + +/// 判断用量是否为空。 +pub fn has_billable_usage(usage: TokenUsage) -> bool { + usage.total_tokens > 0 || usage.image_input_tokens > 0 || usage.image_output_tokens > 0 +} diff --git a/crates/core/src/usage/mod.rs b/crates/core/src/usage/mod.rs new file mode 100644 index 000000000..081b5a747 --- /dev/null +++ b/crates/core/src/usage/mod.rs @@ -0,0 +1,8 @@ +//! 用量聚合领域。 + +/// 用量模型。 +pub mod model; +/// 用量存储端口。 +pub mod ports; +/// 用量聚合服务。 +pub mod service; diff --git a/crates/core/src/usage/model.rs b/crates/core/src/usage/model.rs new file mode 100644 index 000000000..5b11ab3b8 --- /dev/null +++ b/crates/core/src/usage/model.rs @@ -0,0 +1,28 @@ +//! 用量聚合模型。 + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +/// 用量聚合窗口。 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct UsageWindow { + /// 窗口开始时间。 + pub started_at: DateTime, + /// 请求数。 + pub request_count: u64, + /// 输入 token 数。 + pub input_tokens: u64, + /// 输出 token 数。 + pub output_tokens: u64, +} + +/// 用量快照。 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct UsageSnapshot { + /// 账号 ID。 + pub account_id: String, + /// 当前窗口。 + pub window: UsageWindow, +} diff --git a/crates/core/src/usage/ports.rs b/crates/core/src/usage/ports.rs new file mode 100644 index 000000000..b07eb0498 --- /dev/null +++ b/crates/core/src/usage/ports.rs @@ -0,0 +1,27 @@ +//! 用量存储端口。 + +use async_trait::async_trait; +use thiserror::Error; + +use crate::usage::model::UsageSnapshot; + +/// 用量存储错误。 +#[derive(Debug, Error)] +pub enum UsageStoreError { + /// 底层存储失败。 + #[error("usage store operation failed: {message}")] + OperationFailed { + /// 错误说明。 + message: String, + }, +} + +/// 用量存储结果。 +pub type UsageStoreResult = Result; + +/// 用量存储端口。 +#[async_trait] +pub trait UsageStore: Send + Sync + 'static { + /// 写入用量快照。 + async fn record_snapshot(&self, snapshot: &UsageSnapshot) -> UsageStoreResult<()>; +} diff --git a/crates/core/src/usage/service.rs b/crates/core/src/usage/service.rs new file mode 100644 index 000000000..9157154be --- /dev/null +++ b/crates/core/src/usage/service.rs @@ -0,0 +1,34 @@ +//! 用量聚合策略。 + +use crate::{ + accounts::usage::AccountUsageDelta, protocol::codex::events::TokenUsage, + usage::model::UsageWindow, +}; + +/// 用量聚合服务。 +#[derive(Debug, Clone, Default)] +pub struct UsageService; + +impl UsageService { + /// 将 token 增量累加到窗口。 + pub fn add_tokens(window: &mut UsageWindow, input_tokens: u64, output_tokens: u64) { + window.request_count += 1; + window.input_tokens += input_tokens; + window.output_tokens += output_tokens; + } + + /// 将标准化 token 用量转换为账号持久化用量增量。 + pub fn account_delta_from_token_usage(usage: TokenUsage) -> AccountUsageDelta { + AccountUsageDelta { + requests: 0, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + empty_responses: 0, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_requests: 0, + image_request_failures: 0, + } + } +} diff --git a/tests/codex_accounts/pool_scheduling.rs b/crates/core/tests/account_pool.rs similarity index 86% rename from tests/codex_accounts/pool_scheduling.rs rename to crates/core/tests/account_pool.rs index f741ecad0..09be14821 100644 --- a/tests/codex_accounts/pool_scheduling.rs +++ b/crates/core/tests/account_pool.rs @@ -1,15 +1,15 @@ use std::collections::BTreeMap; +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use chrono::{Duration, TimeZone, Utc}; -use codex_proxy_rs::codex::accounts::{ +use codex_proxy_core::accounts::{ model::{Account, AccountStatus}, pool::{ AccountAcquireRequest, AccountPool, AccountPoolOptions, AccountWindowUsageDelta, RotationStrategy, }, }; - -use crate::support::admin_accounts::test_jwt; +use serde_json::{json, Value}; #[test] fn account_pool_should_respect_max_concurrent_slots_per_account() { @@ -98,6 +98,58 @@ fn account_pool_should_not_shorten_existing_quota_cooldown() { .is_none()); } +#[test] +fn reset_usage_should_clear_runtime_counters_and_preserve_window_reset() { + let now = fixed_time(); + let mut account = Account::test("acct_a", AccountStatus::Active); + account.request_count = 7; + account.empty_response_count = 2; + account.image_input_tokens = 11; + account.image_output_tokens = 13; + account.image_request_count = 3; + account.image_request_failed_count = 1; + account.window_request_count = 5; + account.window_input_tokens = 19; + account.window_output_tokens = 23; + account.window_cached_tokens = 29; + account.window_image_input_tokens = 31; + account.window_image_output_tokens = 37; + account.window_image_request_count = 2; + account.window_image_request_failed_count = 1; + account.window_reset_at = Some(now + Duration::seconds(300)); + account.limit_window_seconds = Some(300); + account.last_used_at = Some("2026-06-12T12:00:00Z".to_string()); + let mut pool = AccountPool::default(); + pool.insert(account); + + assert!(pool.reset_usage("acct_a")); + assert!(!pool.reset_usage("missing")); + + let acquired = pool + .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .unwrap() + .account; + assert_eq!(acquired.request_count, 1); + assert_eq!(acquired.empty_response_count, 0); + assert_eq!(acquired.image_input_tokens, 0); + assert_eq!(acquired.image_output_tokens, 0); + assert_eq!(acquired.image_request_count, 0); + assert_eq!(acquired.image_request_failed_count, 0); + assert_eq!(acquired.window_request_count, 1); + assert_eq!(acquired.window_input_tokens, 0); + assert_eq!(acquired.window_output_tokens, 0); + assert_eq!(acquired.window_cached_tokens, 0); + assert_eq!(acquired.window_image_input_tokens, 0); + assert_eq!(acquired.window_image_output_tokens, 0); + assert_eq!(acquired.window_image_request_count, 0); + assert_eq!(acquired.window_image_request_failed_count, 0); + assert_eq!(acquired.window_reset_at, Some(now + Duration::seconds(300))); + assert_ne!( + acquired.last_used_at.as_deref(), + Some("2026-06-12T12:00:00Z") + ); +} + #[test] fn account_pool_should_not_replace_known_window_length_with_cooldown_seconds() { let now = Utc::now(); @@ -134,7 +186,7 @@ fn acquire_should_skip_accounts_with_expired_token_metadata() { fn acquire_should_skip_accounts_with_expired_jwt_when_metadata_is_missing() { let now = Utc::now(); let mut expired = Account::test("expired", AccountStatus::Active); - expired.access_token = test_jwt(None, None, None, None, -60); + expired.access_token = test_jwt(-60); expired.access_token_expires_at = None; let mut pool = AccountPool::default(); pool.insert(expired); @@ -497,3 +549,15 @@ fn record_window_token_usage_should_accumulate_runtime_window_tokens() { fn fixed_time() -> chrono::DateTime { Utc.with_ymd_and_hms(2026, 6, 11, 8, 0, 0).unwrap() } + +fn test_jwt(exp_offset_seconds: i64) -> String { + let payload = json!({ + "exp": Utc::now().timestamp() + exp_offset_seconds, + }); + 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/crates/core/tests/admin_settings.rs b/crates/core/tests/admin_settings.rs new file mode 100644 index 000000000..2028bd4f2 --- /dev/null +++ b/crates/core/tests/admin_settings.rs @@ -0,0 +1,85 @@ +use std::collections::BTreeMap; + +use codex_proxy_core::admin::settings::{ + AdminQuotaWarningThresholds, AdminSettings, AdminSettingsPatch, SettingsService, +}; + +#[test] +fn settings_service_should_apply_retained_settings_patch() { + let mut current = AdminSettings { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: Some("high".to_string()), + service_tier: Some("flex".to_string()), + model_aliases: BTreeMap::from([("codex-fast".to_string(), "gpt-5.5".to_string())]), + refresh_enabled: true, + refresh_margin_seconds: 240, + refresh_concurrency: 2, + max_concurrent_per_account: 4, + request_interval_ms: 50, + rotation_strategy: "least_used".to_string(), + tier_priority: vec!["team".to_string(), "plus".to_string()], + quota_refresh_interval_minutes: 5, + quota_warning_thresholds: AdminQuotaWarningThresholds { + primary: vec![80, 90], + secondary: vec![70, 95], + }, + quota_skip_exhausted: true, + logs_enabled: true, + logs_capacity: 2_000, + logs_capture_body: false, + usage_history_retention_days: Some(30), + }; + + SettingsService::apply_patch( + &mut current, + AdminSettingsPatch { + default_model: Some(" gpt-6 ".to_string()), + default_reasoning_effort: Some("medium".to_string()), + service_tier: Some("priority".to_string()), + model_aliases: Some(BTreeMap::from([( + "fast".to_string(), + "gpt-6-fast".to_string(), + )])), + refresh_enabled: Some(false), + refresh_margin_seconds: Some(180), + refresh_concurrency: Some(3), + max_concurrent_per_account: Some(5), + request_interval_ms: Some(125), + rotation_strategy: Some("round_robin".to_string()), + tier_priority: Some(vec!["pro".to_string(), "plus".to_string()]), + quota_refresh_interval_minutes: Some(15), + quota_warning_thresholds: Some(AdminQuotaWarningThresholds { + primary: vec![75, 90], + secondary: vec![65, 95], + }), + quota_skip_exhausted: Some(false), + logs_enabled: Some(false), + logs_capacity: Some(3_000), + logs_capture_body: Some(true), + usage_history_retention_days: Some(60), + }, + ) + .expect("settings patch should be valid"); + + assert_eq!(current.default_model, "gpt-6"); + assert_eq!(current.rotation_strategy, "round_robin"); + assert_eq!(current.quota_warning_thresholds.primary, vec![75, 90]); + assert_eq!(current.usage_history_retention_days, Some(60)); + assert!(!current.refresh_enabled); + assert!(current.logs_capture_body); +} + +#[test] +fn settings_service_should_reject_invalid_patch_values() { + let mut current = AdminSettings::default(); + let error = SettingsService::apply_patch( + &mut current, + AdminSettingsPatch { + rotation_strategy: Some("random".to_string()), + ..AdminSettingsPatch::default() + }, + ) + .expect_err("invalid rotation strategy should be rejected"); + + assert_eq!(error.field(), "rotationStrategy"); +} diff --git a/crates/core/tests/auth.rs b/crates/core/tests/auth.rs new file mode 100644 index 000000000..0dfbfdb07 --- /dev/null +++ b/crates/core/tests/auth.rs @@ -0,0 +1,59 @@ +use codex_proxy_core::auth::oauth::{OAuthConfig, PkceSessionStore}; + +#[test] +fn pkce_session_store_should_build_login_url_and_acquire_session() { + let mut store = PkceSessionStore::default(); + let config = OAuthConfig { + client_id: "codex-client".to_string(), + auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), + device_code_endpoint: "https://auth.openai.com/oauth/device/code".to_string(), + token_endpoint: "https://auth.openai.com/oauth/token".to_string(), + }; + + let login = store.start_login("console.example.com", &config); + let session = store + .try_acquire(&login.state) + .expect("pkce session should be available"); + + assert!(login + .auth_url + .starts_with("https://auth.openai.com/oauth/authorize?")); + assert!(login.auth_url.contains("response_type=code")); + assert!(login.auth_url.contains("client_id=codex-client")); + assert!(login + .auth_url + .contains("redirect_uri=http%3A%2F%2Flocalhost%3A1455%2Fauth%2Fcallback")); + assert!(login + .auth_url + .contains("scope=openid%20profile%20email%20offline_access")); + assert!(login.auth_url.contains(&format!("state={}", login.state))); + assert_eq!(session.redirect_uri, "http://localhost:1455/auth/callback"); + assert_eq!(session.return_host, "console.example.com"); + assert!(!session.code_verifier.is_empty()); +} + +#[test] +fn pkce_session_store_should_treat_acquired_and_completed_states_as_unavailable() { + let mut store = PkceSessionStore::default(); + let config = OAuthConfig { + client_id: "codex-client".to_string(), + auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), + device_code_endpoint: "https://auth.openai.com/oauth/device/code".to_string(), + token_endpoint: "https://auth.openai.com/oauth/token".to_string(), + }; + + let login = store.start_login("console.example.com", &config); + let _session = store + .try_acquire(&login.state) + .expect("first acquire should succeed"); + + assert!(store.try_acquire(&login.state).is_none()); + assert!(store.is_completed_or_exchanging(&login.state)); + + store.release(&login.state); + assert!(store.try_acquire(&login.state).is_some()); + + store.complete(&login.state); + assert!(store.try_acquire(&login.state).is_none()); + assert!(store.is_completed_or_exchanging(&login.state)); +} diff --git a/crates/core/tests/models.rs b/crates/core/tests/models.rs new file mode 100644 index 000000000..34346ecca --- /dev/null +++ b/crates/core/tests/models.rs @@ -0,0 +1,302 @@ +use std::{collections::BTreeMap, sync::Arc}; + +use async_trait::async_trait; + +use codex_proxy_core::accounts::model::{Account, AccountStatus}; +use codex_proxy_core::models::{ + catalog::ModelCatalog, + model::{ + BackendModelEntry, BackendReasoningEffort, BackendTruncationPolicy, ModelConfig, + ModelPlanSnapshot, ParsedModelName, + }, + ports::{ModelSnapshotStore, ModelSnapshotStoreResult}, + service::{ModelRefreshResult, ModelService, ModelServiceError}, +}; + +use codex_proxy_core::gateway::ports::{ + CodexModelCatalogClient, CodexModelCatalogClientError, CodexModelCatalogRequest, +}; + +#[test] +fn model_catalog_should_parse_alias_reasoning_and_service_tier_suffixes() { + let mut aliases = BTreeMap::new(); + aliases.insert("codex-fast".to_string(), "gpt-5.5".to_string()); + let catalog = ModelCatalog::from_config(&ModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: None, + service_tier: None, + aliases, + }); + + let parsed = catalog.parse_model_name("codex-fast-high-flex"); + + assert_eq!( + parsed, + ParsedModelName { + model_id: "gpt-5.5".to_string(), + reasoning_effort: Some("high".to_string()), + service_tier: Some("flex".to_string()) + } + ); +} + +#[test] +fn model_catalog_should_fall_back_to_default_for_unknown_model_names() { + let catalog = ModelCatalog::from_config(&ModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: None, + service_tier: None, + aliases: BTreeMap::new(), + }); + + let parsed = catalog.parse_model_name("unknown-medium"); + + assert_eq!(parsed.model_id, "gpt-5.5"); + assert_eq!(parsed.reasoning_effort, Some("medium".to_string())); +} + +#[test] +fn model_catalog_should_parse_none_and_minimal_reasoning_suffixes() { + let catalog = ModelCatalog::from_config(&ModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: None, + service_tier: None, + aliases: BTreeMap::new(), + }); + + let none = catalog.parse_model_name("gpt-5.5-none"); + let minimal = catalog.parse_model_name("gpt-5.5-minimal-fast"); + + assert_eq!(none.reasoning_effort, Some("none".to_string())); + assert_eq!(minimal.reasoning_effort, Some("minimal".to_string())); + assert_eq!(minimal.service_tier, Some("fast".to_string())); +} + +#[test] +fn model_catalog_should_merge_backend_snapshots_and_build_plan_allowlist() { + let snapshot = ModelPlanSnapshot::from_backend_entries( + "plus", + vec![BackendModelEntry { + slug: Some("gpt-6".to_string()), + display_name: Some("GPT-6".to_string()), + description: Some("Backend model".to_string()), + is_default: Some(true), + default_reasoning_level: Some("minimal".to_string()), + supported_reasoning_levels: vec![BackendReasoningEffort { + effort: Some("minimal".to_string()), + description: Some("Minimal".to_string()), + ..BackendReasoningEffort::default() + }], + input_modalities: Some(vec!["text".to_string(), "image".to_string()]), + output_modalities: Some(vec!["text".to_string()]), + supports_personality: Some(true), + context_window: Some(200_000), + max_output_tokens: Some(16_384), + truncation_policy: Some(BackendTruncationPolicy { + limit: Some(131_072), + }), + ..BackendModelEntry::default() + }], + ); + let catalog = ModelCatalog::from_config_and_snapshots( + &ModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: None, + service_tier: None, + aliases: BTreeMap::new(), + }, + &[snapshot], + ); + + let info = catalog.model_info("gpt-6").unwrap(); + assert_eq!(info.display_name, "GPT-6"); + assert_eq!(info.default_reasoning_effort, "minimal"); + assert_eq!( + info.supported_reasoning_efforts[0].reasoning_effort, + "minimal" + ); + assert_eq!(info.context_window, Some(200_000)); + assert_eq!(info.truncation_policy_limit, Some(131_072)); + assert_eq!( + catalog.model_plan_allowlist().get("gpt-6").unwrap(), + &vec!["plus".to_string()] + ); +} + +#[tokio::test] +async fn model_service_should_fall_back_to_static_catalog_when_snapshot_store_is_missing() { + let service = ModelService::new(test_model_config(), None, None, None); + + let catalog = service.catalog().await; + + assert_eq!(catalog.models().len(), 1); + assert_eq!(catalog.models()[0].id, "gpt-5.5"); +} + +#[tokio::test] +async fn model_service_should_refresh_distinct_active_plans_and_build_allowlist() { + let snapshot_store = Arc::new(InMemorySnapshotStore::default()); + let upstream = Arc::new(FakeModelCatalogClient::with_models(BTreeMap::from([ + ("plus".to_string(), backend_models("gpt-6")), + ("team".to_string(), backend_models("gpt-7")), + ]))); + let service = ModelService::new( + test_model_config(), + Some(snapshot_store.clone()), + Some(upstream.clone()), + None, + ); + + let result = service + .refresh_backend_models( + &[ + active_account("acct-plus-1", "plus"), + active_account("acct-plus-2", "plus"), + active_account("acct-team", "team"), + inactive_account("acct-disabled", "plus"), + ], + "req-model-refresh", + ) + .await + .expect("refresh should succeed"); + + assert_eq!( + result, + ModelRefreshResult { + refreshed_plans: 2, + model_count: 2, + failed_plans: 0, + } + ); + assert_eq!(upstream.recorded_plan_types().await, vec!["plus", "team"]); + + let snapshots = snapshot_store.list_plan_snapshots().await.unwrap(); + assert_eq!(snapshots.len(), 2); + assert_eq!(snapshots[0].plan_type, "plus"); + assert_eq!(snapshots[1].plan_type, "team"); + + let allowlist = service + .model_plan_allowlist() + .await + .expect("allowlist should be available"); + assert_eq!(allowlist.get("gpt-6").unwrap(), &vec!["plus".to_string()]); + assert_eq!(allowlist.get("gpt-7").unwrap(), &vec!["team".to_string()]); +} + +#[tokio::test] +async fn model_service_should_return_no_accounts_when_no_active_plans_exist() { + let snapshot_store = Arc::new(InMemorySnapshotStore::default()); + let upstream = Arc::new(FakeModelCatalogClient::default()); + let service = ModelService::new( + test_model_config(), + Some(snapshot_store), + Some(upstream), + None, + ); + + let error = service + .refresh_backend_models( + &[inactive_account("acct-disabled", "plus")], + "req-model-refresh", + ) + .await + .expect_err("refresh should fail"); + + assert!(matches!(error, ModelServiceError::NoAccounts)); +} + +fn test_model_config() -> ModelConfig { + ModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: None, + service_tier: None, + aliases: BTreeMap::new(), + } +} + +fn backend_models(model_id: &str) -> Vec { + vec![BackendModelEntry { + slug: Some(model_id.to_string()), + display_name: Some(model_id.to_uppercase()), + default_reasoning_level: Some("minimal".to_string()), + supported_reasoning_levels: vec![BackendReasoningEffort { + effort: Some("minimal".to_string()), + description: Some("Minimal".to_string()), + ..BackendReasoningEffort::default() + }], + ..BackendModelEntry::default() + }] +} + +fn active_account(id: &str, plan_type: &str) -> Account { + let mut account = Account::test(id, AccountStatus::Active); + account.account_id = Some(format!("upstream-{id}")); + account.plan_type = Some(plan_type.to_string()); + account +} + +fn inactive_account(id: &str, plan_type: &str) -> Account { + let mut account = active_account(id, plan_type); + account.status = AccountStatus::Disabled; + account +} + +#[derive(Default)] +struct InMemorySnapshotStore { + snapshots: tokio::sync::Mutex>, +} + +#[async_trait] +impl ModelSnapshotStore for InMemorySnapshotStore { + async fn replace_plan_snapshot( + &self, + snapshot: &ModelPlanSnapshot, + ) -> ModelSnapshotStoreResult<()> { + self.snapshots + .lock() + .await + .insert(snapshot.plan_type.clone(), snapshot.clone()); + Ok(()) + } + + async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult> { + Ok(self.snapshots.lock().await.values().cloned().collect()) + } +} + +#[derive(Default)] +struct FakeModelCatalogClient { + models_by_plan: BTreeMap>, + requested_plan_types: tokio::sync::Mutex>, +} + +impl FakeModelCatalogClient { + fn with_models(models_by_plan: BTreeMap>) -> Self { + Self { + models_by_plan, + requested_plan_types: tokio::sync::Mutex::new(Vec::new()), + } + } + + async fn recorded_plan_types(&self) -> Vec { + self.requested_plan_types.lock().await.clone() + } +} + +#[async_trait] +impl CodexModelCatalogClient for FakeModelCatalogClient { + async fn fetch_models( + &self, + request: &CodexModelCatalogRequest<'_>, + ) -> Result, CodexModelCatalogClientError> { + self.requested_plan_types + .lock() + .await + .push(request.plan_type.to_string()); + Ok(self + .models_by_plan + .get(request.plan_type) + .cloned() + .unwrap_or_default()) + } +} diff --git a/crates/core/tests/protocol.rs b/crates/core/tests/protocol.rs new file mode 100644 index 000000000..cdf16d088 --- /dev/null +++ b/crates/core/tests/protocol.rs @@ -0,0 +1,1774 @@ +use codex_proxy_core::models::model::{ModelConfig, ParsedModelName}; +use codex_proxy_core::protocol::codex::events::{ + cooldown_with_jitter, extract_sse_usage, extract_usage, parse_rate_limit_headers, + parse_rate_limits_event, rate_limit_quota, retry_after_seconds_from_body, RateLimitWindow, + TokenUsage, +}; +use codex_proxy_core::protocol::codex::responses::CodexResponsesRequest; +use codex_proxy_core::protocol::codex::websocket::{ + classify_websocket_error_frame, is_terminal_websocket_event, + retry_after_seconds_from_wrapped_error_headers, + websocket_agent_message_output_item_event_invalid_required_fields, + websocket_audit_artifact_from_attempt, + websocket_compaction_output_item_event_invalid_required_fields, + websocket_custom_tool_call_output_item_event_invalid_required_fields, + websocket_custom_tool_call_output_payload_item_event_invalid_required_fields, + websocket_delta_event_missing_official_required_fields, websocket_event_shape_parse_error, + websocket_event_to_sse_frame, + websocket_function_call_output_item_event_invalid_required_fields, + websocket_function_call_output_payload_item_event_invalid_required_fields, + websocket_image_generation_call_output_item_event_invalid_required_fields, + websocket_incomplete_response_reason, + websocket_local_shell_call_output_item_event_invalid_required_fields, + websocket_message_output_item_event_invalid_required_fields, websocket_metadata_turn_state, + websocket_output_item_event_invalid_item_type_tag, + websocket_output_item_event_invalid_metadata, websocket_output_item_event_missing_item, + websocket_output_item_event_non_object_item, websocket_parity_diff, + websocket_payload_audit_snapshot, + websocket_reasoning_output_item_event_invalid_required_fields, + websocket_reasoning_summary_part_added_missing_summary_index, + websocket_response_completed_missing_response, websocket_response_completed_parse_error, + websocket_response_created_missing_response, + websocket_response_output_text_delta_missing_delta, + websocket_tool_search_call_output_item_event_invalid_required_fields, + websocket_tool_search_output_item_event_invalid_required_fields, + websocket_web_search_call_output_item_event_invalid_required_fields, OpeningAuditHeader, + OpeningAuditSnapshot, PayloadAuditSnapshot, WebSocketAuditErrorSnapshot, + WebSocketErrorClassificationProfile, +}; +use codex_proxy_core::protocol::openai::responses::{ + response_from_codex_sse, translate_response_to_codex, translate_response_to_compact, + CollectedResponse, OpenAiResponsesRequest, +}; +use codex_proxy_core::serving::responses::{ + apply_response_model_options, http_sse_fallback_allowed, transport_for_request, CodexTransport, +}; +use serde_json::json; + +#[test] +fn codex_responses_request_should_enable_http_sse_defaults() { + let request = CodexResponsesRequest::new_http_sse("gpt-5.5", "be brief", vec![json!({})]); + + assert!(request.stream); +} + +#[test] +fn codex_responses_transport_should_prefer_websocket_without_history() { + let request = CodexResponsesRequest::new_http_sse("gpt-5.5", "be brief", vec![json!({})]); + + assert_eq!( + transport_for_request(&request), + CodexTransport::WebSocketPreferred + ); + assert!(http_sse_fallback_allowed(&request)); +} + +#[test] +fn codex_responses_transport_should_require_websocket_for_previous_response_id() { + let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "be brief", vec![json!({})]); + request.previous_response_id = Some("resp_previous".to_string()); + + assert_eq!( + transport_for_request(&request), + CodexTransport::WebSocketRequired + ); + assert!(!http_sse_fallback_allowed(&request)); +} + +#[test] +fn codex_responses_transport_should_allow_forced_http_sse() { + let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "be brief", vec![json!({})]); + request.previous_response_id = Some("resp_previous".to_string()); + request.use_websocket = true; + request.force_http_sse = true; + + assert_eq!(transport_for_request(&request), CodexTransport::HttpSse); + assert!(http_sse_fallback_allowed(&request)); +} + +#[test] +fn codex_responses_transport_flags_should_not_serialize_to_upstream_json() { + let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "be brief", vec![json!({})]); + request.use_websocket = true; + request.force_http_sse = true; + + let body = serde_json::to_value(&request).expect("request should serialize"); + + assert!(body.get("use_websocket").is_none()); + assert!(body.get("useWebSocket").is_none()); + assert!(body.get("force_http_sse").is_none()); + assert!(body.get("forceHttpSse").is_none()); +} + +#[test] +fn codex_websocket_payload_audit_snapshot_should_redact_user_content() { + let mut request = CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "private instructions", + vec![json!({ + "role": "user", + "content": "private prompt", + })], + ); + request.previous_response_id = Some("resp_secret".to_string()); + request.reasoning = Some(json!({"effort": "medium"})); + request.tools = Some(vec![json!({ + "type": "function", + "name": "private_tool", + "description": "private schema", + })]); + request.text = Some(json!({"format": {"type": "text"}})); + request.service_tier = Some("flex".to_string()); + request.prompt_cache_key = Some("cache-secret".to_string()); + request.generate = Some(false); + request.include = Some(vec!["reasoning.encrypted_content".to_string()]); + request.client_metadata = Some(json!({ + "thread_id": "thread-secret", + "safe": "value", + })); + + let snapshot = websocket_payload_audit_snapshot(&request); + + assert_eq!( + snapshot.top_level_keys, + vec![ + "type", + "model", + "instructions", + "previous_response_id", + "input", + "tools", + "tool_choice", + "parallel_tool_calls", + "reasoning", + "store", + "stream", + "include", + "service_tier", + "prompt_cache_key", + "text", + "generate", + "client_metadata", + ] + ); + assert_eq!(snapshot.body["type"], "response.create"); + assert_eq!(snapshot.body["model"], "gpt-5.5"); + assert_eq!(snapshot.body["stream"], true); + assert_eq!(snapshot.body["instructions"], ""); + assert_eq!(snapshot.body["input"], ""); + assert_eq!(snapshot.body["previous_response_id"], ""); + assert_eq!(snapshot.body["prompt_cache_key"], ""); + assert_eq!(snapshot.body["client_metadata"], ""); + assert_eq!(snapshot.body["tools"], ""); +} + +#[test] +fn codex_websocket_response_create_payload_text_should_preserve_old_field_order() { + let mut request = CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "private capture instructions", + vec![json!({ + "role": "user", + "content": "private capture prompt", + })], + ); + request.prompt_cache_key = Some("session-1".to_string()); + request.generate = Some(false); + request.client_metadata = Some(json!({ + "thread_id": "capture-thread-secret", + "safe": "capture", + })); + + let payload = + codex_proxy_core::protocol::codex::websocket::websocket_response_create_payload_text( + &request, + ) + .expect("payload should serialize"); + + assert_substrings_appear_in_order( + &payload, + &[ + "\"type\":\"response.create\"", + "\"model\":\"gpt-5.5\"", + "\"instructions\":\"private capture instructions\"", + "\"input\":", + "\"tools\":[]", + "\"tool_choice\":\"auto\"", + "\"parallel_tool_calls\":true", + "\"reasoning\":null", + "\"store\":false", + "\"stream\":true", + "\"include\":[]", + "\"prompt_cache_key\":\"session-1\"", + "\"generate\":false", + "\"client_metadata\":", + ], + ); +} + +#[test] +fn codex_websocket_audit_artifact_should_record_transport_opening_and_payload() { + let mut request = CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "private instructions", + vec![json!({ + "role": "user", + "content": "private prompt", + })], + ); + request.previous_response_id = Some("resp_secret".to_string()); + let opening = OpeningAuditSnapshot { + request_line: "GET /backend-api/codex/responses HTTP/1.1".to_string(), + header_order: vec!["Host".to_string(), "authorization".to_string()], + headers: vec![ + OpeningAuditHeader { + name: "Host".to_string(), + value: "chatgpt.com".to_string(), + }, + OpeningAuditHeader { + name: "authorization".to_string(), + value: "".to_string(), + }, + ], + }; + let payload = websocket_payload_audit_snapshot(&request); + + let artifact = websocket_audit_artifact_from_attempt(&request, opening.clone(), payload); + + assert_eq!(artifact.transport_mode, "websocket_required"); + assert!(!artifact.fallback_allowed); + assert_eq!(artifact.opening, Some(opening)); + assert_eq!( + artifact.payload.expect("payload").body["input"], + "" + ); + assert_eq!(artifact.error, None); +} + +#[test] +fn codex_websocket_parity_diff_should_report_header_payload_and_error_differences() { + let current = codex_proxy_core::protocol::codex::websocket::WebSocketAuditArtifact { + transport_mode: "websocket_preferred".to_string(), + fallback_allowed: true, + opening: Some(OpeningAuditSnapshot { + request_line: "GET /codex/responses HTTP/1.1".to_string(), + header_order: vec!["Host".to_string(), "Connection".to_string()], + headers: vec![OpeningAuditHeader { + name: "Sec-WebSocket-Extensions".to_string(), + value: "permessage-deflate; client_max_window_bits".to_string(), + }], + }), + payload: Some(PayloadAuditSnapshot { + top_level_keys: vec!["type".to_string(), "model".to_string()], + body: json!({"type": "response.create"}), + }), + error: Some(WebSocketAuditErrorSnapshot { + classification: "opening_failed".to_string(), + message: "connection refused".to_string(), + }), + }; + let reference = codex_proxy_core::protocol::codex::websocket::WebSocketAuditArtifact { + transport_mode: "websocket_required".to_string(), + fallback_allowed: false, + opening: Some(OpeningAuditSnapshot { + request_line: "GET /codex/responses?source=audit HTTP/1.1".to_string(), + header_order: vec!["Host".to_string(), "Authorization".to_string()], + headers: vec![OpeningAuditHeader { + name: "Sec-WebSocket-Extensions".to_string(), + value: "identity".to_string(), + }], + }), + payload: Some(PayloadAuditSnapshot { + top_level_keys: vec!["type".to_string(), "previous_response_id".to_string()], + body: json!({"type": "response.create"}), + }), + error: Some(WebSocketAuditErrorSnapshot { + classification: "timeout".to_string(), + message: "timed out".to_string(), + }), + }; + + let diff = websocket_parity_diff(¤t, &reference); + + assert!(diff.differences.iter().any(|difference| { + difference.path == "opening.header_order" + && difference.current == json!(["Host", "Connection"]) + && difference.reference == json!(["Host", "Authorization"]) + })); + assert!(diff + .differences + .iter() + .any(|difference| difference.path == "payload.top_level_keys")); + assert!(diff + .differences + .iter() + .any(|difference| difference.path == "error.classification")); +} + +#[test] +fn codex_websocket_event_to_sse_frame_should_encode_public_events_only() { + let event = json!({ + "type": "response.completed", + "response": { + "id": "resp_ws", + "object": "response", + "usage": { + "input_tokens": 4, + "output_tokens": 2, + "total_tokens": 6 + } + } + }) + .to_string(); + + let frame = websocket_event_to_sse_frame(&event).expect("public event should encode"); + + assert_eq!( + frame, + format!("event: response.completed\ndata: {event}\n\n") + ); + assert!(websocket_event_to_sse_frame(r#"{"type":"codex.rate_limits"}"#).is_none()); + assert!(websocket_event_to_sse_frame(r#"{"response":{}}"#).is_none()); + assert!(websocket_event_to_sse_frame("not-json").is_none()); +} + +#[test] +fn codex_websocket_metadata_turn_state_should_extract_case_insensitive_header() { + let event = json!({ + "type": "response.metadata", + "headers": { + "X-Codex-Turn-State": ["turn-from-metadata"] + } + }) + .to_string(); + + assert_eq!( + websocket_metadata_turn_state(&event).as_deref(), + Some("turn-from-metadata") + ); +} + +#[test] +fn codex_websocket_terminal_event_should_match_completed_failed_and_error() { + assert!(is_terminal_websocket_event("response.completed")); + assert!(is_terminal_websocket_event("response.failed")); + assert!(is_terminal_websocket_event("error")); + assert!(!is_terminal_websocket_event("response.output_text.delta")); +} + +#[test] +fn codex_websocket_error_frame_should_classify_rotatable_codes() { + let frame = json!({ + "type": "response.failed", + "response": { + "error": { + "code": "usage_limit_reached", + "message": "Rate limit reached. Please try again in 11.054s." + } + } + }) + .to_string(); + + let error = + classify_websocket_error_frame(&frame, WebSocketErrorClassificationProfile::OneShot) + .expect("rate limit frame should classify"); + + assert_eq!(error.status_code, 429); +} + +#[test] +fn codex_websocket_error_frame_should_classify_old_gateway_special_codes() { + let cases = [ + ("quota_exhausted", 402), + ("token_expired", 401), + ("account_banned", 403), + ("previous_response_not_found", 400), + ("server_is_overloaded", 503), + ("usage_not_included", 429), + ]; + + for (code, expected_status) in cases { + let frame = json!({ + "type": "response.failed", + "response": { + "error": { + "code": code, + "message": "classified by old gateway parity" + } + } + }) + .to_string(); + + let error = + classify_websocket_error_frame(&frame, WebSocketErrorClassificationProfile::OneShot) + .expect("special gateway code should classify"); + + assert_eq!(error.status_code, expected_status, "code: {code}"); + } +} + +#[test] +fn codex_websocket_error_frame_should_honor_explicit_error_status() { + let frame = json!({ + "type": "error", + "status_code": 403, + "error": { + "code": "forbidden" + } + }) + .to_string(); + + let error = + classify_websocket_error_frame(&frame, WebSocketErrorClassificationProfile::OneShot) + .expect("explicit status should classify"); + + assert_eq!(error.status_code, 403); +} + +#[test] +fn codex_websocket_error_frame_should_ignore_success_status_and_unmapped_error() { + let success_status = json!({ + "type": "error", + "status": 200, + "error": { + "code": "informational" + } + }) + .to_string(); + let unmapped_error = json!({ + "type": "error", + "error": { + "code": "not_a_rotatable_error" + } + }) + .to_string(); + + assert!(classify_websocket_error_frame( + &success_status, + WebSocketErrorClassificationProfile::OneShot + ) + .is_none()); + assert!(classify_websocket_error_frame( + &unmapped_error, + WebSocketErrorClassificationProfile::OneShot + ) + .is_none()); +} + +#[test] +fn codex_websocket_wrapped_error_headers_should_extract_retry_after() { + let frame = json!({ + "type": "error", + "headers": { + "retry-after": ["37"] + }, + "error": { + "code": "rate_limit_exceeded" + } + }) + .to_string(); + + assert_eq!( + retry_after_seconds_from_wrapped_error_headers(&frame), + Some(37) + ); +} + +#[test] +fn codex_websocket_response_completed_parse_error_should_validate_shape() { + let frame = json!({ + "type": "response.completed", + "response": { + "id": "resp_bad", + "usage": { + "input_tokens": "bad", + "output_tokens": 1, + "total_tokens": 1 + } + } + }) + .to_string(); + + let error = websocket_response_completed_parse_error(&frame) + .expect("invalid response.completed should report a parse error"); + + assert!(error.contains("failed to parse ResponseCompleted")); +} + +#[test] +fn codex_websocket_response_completed_parse_error_should_reject_missing_id_and_incomplete_usage() { + let missing_id = json!({ + "type": "response.completed", + "response": { + "object": "response", + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2 + } + } + }) + .to_string(); + let incomplete_usage = json!({ + "type": "response.completed", + "response": { + "id": "resp_incomplete_usage", + "object": "response", + "usage": { + "input_tokens": 1, + "output_tokens": 1 + } + } + }) + .to_string(); + + assert!(websocket_response_completed_parse_error(&missing_id) + .expect("missing id should be rejected") + .contains("missing field")); + assert!(websocket_response_completed_parse_error(&incomplete_usage) + .expect("incomplete usage should be rejected") + .contains("missing field")); +} + +#[test] +fn codex_websocket_event_shape_parse_error_should_detect_field_type_mismatches() { + let frame = json!({ + "type": "response.output_text.delta", + "delta": 42 + }) + .to_string(); + + assert!(websocket_event_shape_parse_error(&frame)); +} + +#[test] +fn codex_websocket_response_events_should_detect_missing_required_fields() { + assert!(websocket_response_created_missing_response( + r#"{"type":"response.created"}"# + )); + assert!(websocket_response_output_text_delta_missing_delta( + r#"{"type":"response.output_text.delta"}"# + )); +} + +#[test] +fn codex_websocket_output_item_events_should_detect_invalid_item_shape() { + assert!(websocket_output_item_event_missing_item( + r#"{"type":"response.output_item.added"}"# + )); + assert!(websocket_output_item_event_non_object_item( + r#"{"type":"response.output_item.done","item":"bad"}"# + )); + assert!(websocket_output_item_event_invalid_item_type_tag( + r#"{"type":"response.output_item.done","item":{}}"# + )); +} + +#[test] +fn codex_websocket_message_output_item_should_validate_required_fields() { + let frame = json!({ + "type": "response.output_item.done", + "item": { + "type": "message", + "role": "assistant", + "content": [{ + "type": "output_text" + }] + } + }) + .to_string(); + + assert!(websocket_message_output_item_event_invalid_required_fields( + &frame + )); +} + +#[test] +fn codex_websocket_optional_null_fields_should_match_missing_field_behavior() { + let message = json!({ + "type": "response.output_item.done", + "item": { + "type": "message", + "id": null, + "role": "assistant", + "phase": null, + "metadata": null, + "content": [{ + "type": "output_text", + "text": "hello" + }] + } + }) + .to_string(); + let custom_tool = json!({ + "type": "response.output_item.done", + "item": { + "type": "custom_tool_call", + "id": null, + "status": null, + "call_id": "call_1", + "name": "render", + "input": "{}" + } + }) + .to_string(); + let context_compaction = json!({ + "type": "response.output_item.done", + "item": { + "type": "context_compaction", + "encrypted_content": null + } + }) + .to_string(); + + assert!(websocket_event_to_sse_frame(&message).is_some()); + assert!(websocket_event_to_sse_frame(&custom_tool).is_some()); + assert!(websocket_event_to_sse_frame(&context_compaction).is_some()); +} + +#[test] +fn codex_websocket_response_control_events_should_detect_incomplete_and_missing_completed_response() +{ + let incomplete = json!({ + "type": "response.incomplete", + "response": { + "incomplete_details": { + "reason": "max_output_tokens" + } + } + }) + .to_string(); + + assert_eq!( + websocket_incomplete_response_reason(&incomplete).as_deref(), + Some("max_output_tokens") + ); + assert!(websocket_response_completed_missing_response( + r#"{"type":"response.completed"}"# + )); +} + +#[test] +fn codex_websocket_delta_events_should_validate_official_required_fields() { + assert!(websocket_delta_event_missing_official_required_fields( + r#"{"type":"response.custom_tool_call_input.delta","delta":"x"}"# + )); + assert!(websocket_delta_event_missing_official_required_fields( + r#"{"type":"response.reasoning_summary_text.delta","delta":"x"}"# + )); + assert!(websocket_delta_event_missing_official_required_fields( + r#"{"type":"response.reasoning_text.delta","delta":"x"}"# + )); +} + +#[test] +fn codex_websocket_output_item_metadata_and_agent_should_validate_required_fields() { + let invalid_metadata = json!({ + "type": "response.output_item.done", + "item": { + "type": "message", + "metadata": { + "turn_id": 7 + } + } + }) + .to_string(); + let invalid_agent = json!({ + "type": "response.output_item.done", + "item": { + "type": "agent_message", + "author": "assistant", + "content": [{ + "type": "input_text" + }] + } + }) + .to_string(); + + assert!(websocket_output_item_event_invalid_metadata( + &invalid_metadata + )); + assert!(websocket_agent_message_output_item_event_invalid_required_fields(&invalid_agent)); +} + +#[test] +fn codex_websocket_output_item_optional_fields_should_reject_old_gateway_edge_cases() { + let invalid_tool_search_call = json!({ + "type": "response.output_item.done", + "item": { + "type": "tool_search_call", + "call_id": {}, + "status": [], + "execution": "client", + "arguments": {} + } + }) + .to_string(); + let invalid_custom_tool_call = json!({ + "type": "response.output_item.done", + "item": { + "type": "custom_tool_call", + "id": {}, + "status": [], + "call_id": "call_1", + "name": "render", + "input": "{}" + } + }) + .to_string(); + let invalid_custom_tool_output = json!({ + "type": "response.output_item.done", + "item": { + "type": "custom_tool_call_output", + "call_id": "call_1", + "name": {}, + "output": "ok" + } + }) + .to_string(); + let invalid_tool_search_output = json!({ + "type": "response.output_item.done", + "item": { + "type": "tool_search_output", + "call_id": {}, + "status": "completed", + "execution": "server", + "tools": [] + } + }) + .to_string(); + let invalid_image = json!({ + "type": "response.output_item.done", + "item": { + "type": "image_generation_call", + "id": "img_1", + "status": "completed", + "revised_prompt": {}, + "result": "base64" + } + }) + .to_string(); + let invalid_context_compaction = json!({ + "type": "response.output_item.done", + "item": { + "type": "context_compaction", + "encrypted_content": {} + } + }) + .to_string(); + + assert!( + websocket_tool_search_call_output_item_event_invalid_required_fields( + &invalid_tool_search_call + ) + ); + assert!( + websocket_custom_tool_call_output_item_event_invalid_required_fields( + &invalid_custom_tool_call + ) + ); + assert!( + websocket_custom_tool_call_output_payload_item_event_invalid_required_fields( + &invalid_custom_tool_output + ) + ); + assert!( + websocket_tool_search_output_item_event_invalid_required_fields( + &invalid_tool_search_output + ) + ); + assert!( + websocket_image_generation_call_output_item_event_invalid_required_fields(&invalid_image) + ); + assert!( + websocket_compaction_output_item_event_invalid_required_fields(&invalid_context_compaction) + ); +} + +#[test] +fn codex_websocket_reasoning_and_function_items_should_validate_required_fields() { + let invalid_reasoning = json!({ + "type": "response.output_item.done", + "item": { + "type": "reasoning", + "summary": [{ + "type": "summary_text" + }] + } + }) + .to_string(); + let invalid_function_call = json!({ + "type": "response.output_item.done", + "item": { + "type": "function_call", + "name": "tool", + "call_id": "call_1" + } + }) + .to_string(); + let invalid_function_output = json!({ + "type": "response.output_item.done", + "item": { + "type": "function_call_output", + "call_id": "call_1", + "output": [{ + "type": "input_text" + }] + } + }) + .to_string(); + + assert!(websocket_reasoning_output_item_event_invalid_required_fields(&invalid_reasoning)); + assert!( + websocket_function_call_output_item_event_invalid_required_fields(&invalid_function_call) + ); + assert!( + websocket_function_call_output_payload_item_event_invalid_required_fields( + &invalid_function_output + ) + ); +} + +#[test] +fn codex_websocket_custom_and_tool_search_items_should_validate_required_fields() { + let invalid_custom_call = json!({ + "type": "response.output_item.done", + "item": { + "type": "custom_tool_call", + "call_id": "call_1", + "name": "custom" + } + }) + .to_string(); + let invalid_custom_output = json!({ + "type": "response.output_item.done", + "item": { + "type": "custom_tool_call_output", + "output": "ok" + } + }) + .to_string(); + let invalid_tool_search_call = json!({ + "type": "response.output_item.done", + "item": { + "type": "tool_search_call", + "arguments": {} + } + }) + .to_string(); + let invalid_tool_search_output = json!({ + "type": "response.output_item.done", + "item": { + "type": "tool_search_output", + "status": "completed", + "execution": "remote" + } + }) + .to_string(); + + assert!( + websocket_custom_tool_call_output_item_event_invalid_required_fields(&invalid_custom_call) + ); + assert!( + websocket_custom_tool_call_output_payload_item_event_invalid_required_fields( + &invalid_custom_output + ) + ); + assert!( + websocket_tool_search_call_output_item_event_invalid_required_fields( + &invalid_tool_search_call + ) + ); + assert!( + websocket_tool_search_output_item_event_invalid_required_fields( + &invalid_tool_search_output + ) + ); +} + +#[test] +fn codex_websocket_local_web_image_and_compaction_items_should_validate_required_fields() { + let invalid_local_shell = json!({ + "type": "response.output_item.done", + "item": { + "type": "local_shell_call", + "call_id": "call_1", + "status": "failed", + "action": { + "type": "exec", + "command": "ls" + } + } + }) + .to_string(); + let invalid_web_search = json!({ + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "status": "completed", + "action": {} + } + }) + .to_string(); + let invalid_image = json!({ + "type": "response.output_item.done", + "item": { + "type": "image_generation_call", + "id": "img_1", + "status": "completed" + } + }) + .to_string(); + let invalid_compaction = json!({ + "type": "response.output_item.done", + "item": { + "type": "compaction" + } + }) + .to_string(); + + assert!( + websocket_local_shell_call_output_item_event_invalid_required_fields(&invalid_local_shell) + ); + assert!( + websocket_web_search_call_output_item_event_invalid_required_fields(&invalid_web_search) + ); + assert!( + websocket_image_generation_call_output_item_event_invalid_required_fields(&invalid_image) + ); + assert!(websocket_compaction_output_item_event_invalid_required_fields(&invalid_compaction)); +} + +#[test] +fn codex_websocket_local_shell_and_web_search_should_reject_nested_shape_mismatches() { + let invalid_local_shell = json!({ + "type": "response.output_item.done", + "item": { + "type": "local_shell_call", + "call_id": "call_1", + "status": "completed", + "action": { + "type": "exec", + "command": "pwd", + "timeout_ms": "1000", + "env": { + "PATH": {} + } + } + } + }) + .to_string(); + let invalid_web_search = json!({ + "type": "response.output_item.done", + "item": { + "type": "web_search_call", + "status": "completed", + "action": { + "type": "search", + "query": {}, + "queries": ["valid", {}] + } + } + }) + .to_string(); + + assert!( + websocket_local_shell_call_output_item_event_invalid_required_fields(&invalid_local_shell) + ); + assert!( + websocket_web_search_call_output_item_event_invalid_required_fields(&invalid_web_search) + ); +} + +#[test] +fn codex_websocket_reasoning_summary_part_should_require_summary_index() { + assert!( + websocket_reasoning_summary_part_added_missing_summary_index( + r#"{"type":"response.reasoning_summary_part.added"}"# + ) + ); +} + +#[test] +fn codex_websocket_event_to_sse_frame_should_skip_invalid_json_and_shape_mismatches() { + let invalid_delta = json!({ + "type": "response.output_text.delta", + "delta": 42 + }) + .to_string(); + let non_object_item = json!({ + "type": "response.output_item.done", + "item": 123 + }) + .to_string(); + + assert!(websocket_event_to_sse_frame("not-json-from-upstream").is_none()); + assert!(websocket_event_to_sse_frame(&invalid_delta).is_none()); + assert!(websocket_event_to_sse_frame(&non_object_item).is_none()); +} + +#[test] +fn chat_completion_request_should_translate_to_codex_request() { + let request = codex_proxy_core::protocol::openai::chat::ChatCompletionRequest { + model: "gpt-5.5".to_string(), + stream: true, + messages: vec![ + codex_proxy_core::protocol::openai::chat::ChatMessage { + role: "system".to_string(), + content: Some(json!("be brief")), + name: None, + tool_calls: None, + tool_call_id: None, + function_call: None, + }, + codex_proxy_core::protocol::openai::chat::ChatMessage { + role: "user".to_string(), + content: Some(json!("hello")), + name: None, + tool_calls: None, + tool_call_id: None, + function_call: None, + }, + ], + reasoning_effort: Some("medium".to_string()), + service_tier: Some("auto".to_string()), + tools: None, + tool_choice: None, + parallel_tool_calls: None, + functions: None, + response_format: None, + user: Some(" client-123 ".to_string()), + }; + + let codex = codex_proxy_core::protocol::openai::chat::translate_chat_to_codex(request) + .expect("chat request should translate"); + + assert_eq!(codex.model, "gpt-5.5"); + assert!(codex.force_http_sse); + assert_eq!(codex.prompt_cache_key.as_deref(), Some("client-123")); + assert_eq!(codex.client_conversation_id.as_deref(), Some("client-123")); +} + +#[test] +fn sse_parser_should_combine_multiline_data_and_metadata() { + let events = codex_proxy_core::protocol::codex::sse::parse_sse_events( + "event: message\nid: evt_1\ndata: hello\ndata: world\nretry: 10\n\n", + ) + .expect("sse should parse"); + + assert_eq!(events[0].data, "hello\nworld"); +} + +#[test] +fn chat_completion_from_codex_sse_should_convert_completed_response() { + let body = concat!( + "event: response.output_text.delta\n", + "data: {\"delta\":\"hello\"}\n\n", + "event: response.completed\n", + "data: {\"response\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":3}}}\n\n", + ); + + let response = codex_proxy_core::protocol::openai::chat::chat_completion_from_codex_sse( + body, "gpt-5.5", false, None, + ) + .expect("conversion should succeed") + .expect("completed response"); + + assert_eq!(response["choices"][0]["message"]["content"], "hello"); +} + +#[test] +fn response_from_codex_sse_should_collect_completed_response_with_deltas() { + let body = concat!( + "event: response.output_text.delta\n", + "data: {\"delta\":\"hello\"}\n\n", + "event: response.completed\n", + "data: {\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"status\":\"completed\",\"usage\":{\"input_tokens\":2,\"output_tokens\":3}}}\n\n", + ); + + let response = response_from_codex_sse(body, None).expect("conversion should succeed"); + + assert_eq!( + response, + CollectedResponse::Completed(json!({ + "id": "resp_1", + "object": "response", + "status": "completed", + "usage": { + "input_tokens": 2, + "output_tokens": 3 + }, + "output": [{ + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{ + "type": "output_text", + "text": "hello", + "annotations": [] + }] + }], + "output_text": "hello" + })) + ); +} + +#[test] +fn openai_response_request_should_translate_to_codex_request() { + let request = serde_json::from_value::(json!({ + "model": "gpt-5.5", + "input": "hello", + "previous_response_id": "resp_previous", + "turnState": "turn_previous", + "stream": false + })) + .expect("responses request should deserialize"); + + let codex = translate_response_to_codex(request); + + assert_eq!(codex.model, "gpt-5.5"); + assert_eq!(codex.input, vec![json!("hello")]); + assert_eq!(codex.previous_response_id.as_deref(), Some("resp_previous")); + assert_eq!(codex.turn_state.as_deref(), Some("turn_previous")); + assert!(!codex.use_websocket); + assert!(!codex.force_http_sse); +} + +#[test] +fn openai_response_request_should_default_missing_stream_to_true() { + let request = serde_json::from_value::(json!({ + "model": "gpt-5.5", + "input": "hello" + })) + .expect("responses request should deserialize"); + + let codex = translate_response_to_codex(request); + + assert!(codex.stream); +} + +#[test] +fn openai_response_request_should_prepare_tuple_schema_before_upstream() { + let request = serde_json::from_value::(json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "text": { + "format": { + "type": "json_schema", + "name": "TupleAnswer", + "schema": { + "type": "object", + "properties": { + "point": { + "type": "array", + "prefixItems": [ + {"type": "number"}, + {"type": "number"} + ], + "items": false + } + }, + "required": ["point"] + }, + "strict": true + } + } + })) + .expect("responses request should deserialize"); + + let codex = translate_response_to_codex(request); + let schema = &codex.text.as_ref().unwrap()["format"]["schema"]; + + assert_eq!( + schema["properties"]["point"], + json!({ + "type": "object", + "properties": { + "0": {"type": "number"}, + "1": {"type": "number"} + }, + "required": ["0", "1"], + "additionalProperties": false + }) + ); + assert_eq!( + codex.tuple_schema.unwrap()["properties"]["point"]["type"], + "array" + ); +} + +#[test] +fn openai_response_request_should_forward_parity_fields_to_codex() { + let request = serde_json::from_value::(json!({ + "model": "gpt-5.5-fast", + "stream": false, + "input": [], + "instructions": "be terse", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "tool_choice": {"type": "function", "function": {"name": "lookup"}}, + "parallel_tool_calls": true, + "tools": [{"type": "function", "name": "lookup"}], + "prompt_cache_key": "pcache", + "include": ["reasoning.encrypted_content"], + "client_metadata": {"safe": "yes", "drop": 42}, + "turnState": "turn-body", + "turnMetadata": "meta-body", + "betaFeatures": "beta-body", + "includeTimingMetrics": "true", + "version": "2026-06-12", + "codexWindowId": "window-body", + "parentThreadId": "parent-body", + "use_websocket": false + })) + .expect("responses request should deserialize"); + + let codex = translate_response_to_codex(request); + + assert_eq!(codex.instructions, "be terse"); + assert_eq!( + codex.reasoning, + Some(json!({"effort": "high", "summary": "auto"})) + ); + assert_eq!(codex.service_tier.as_deref(), Some("priority")); + assert_eq!( + codex.tool_choice, + Some(json!({"type": "function", "function": {"name": "lookup"}})) + ); + assert_eq!(codex.parallel_tool_calls, Some(true)); + assert_eq!( + codex.tools, + Some(vec![json!({"type": "function", "name": "lookup"})]) + ); + assert_eq!(codex.prompt_cache_key.as_deref(), Some("pcache")); + assert!(codex.explicit_prompt_cache_key); + assert_eq!( + codex.include, + Some(vec!["reasoning.encrypted_content".to_string()]) + ); + assert_eq!(codex.client_metadata, Some(json!({"safe": "yes"}))); + assert_eq!(codex.turn_state.as_deref(), Some("turn-body")); + assert_eq!(codex.turn_metadata.as_deref(), Some("meta-body")); + assert_eq!(codex.beta_features.as_deref(), Some("beta-body")); + assert_eq!(codex.include_timing_metrics.as_deref(), Some("true")); + assert_eq!(codex.version.as_deref(), Some("2026-06-12")); + assert_eq!(codex.codex_window_id.as_deref(), Some("window-body")); + assert_eq!(codex.parent_thread_id.as_deref(), Some("parent-body")); + assert!(codex.force_http_sse); +} + +#[test] +fn codex_responses_model_options_should_apply_suffix_defaults_and_include_reasoning() { + let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5-high-fast", "", Vec::new()); + let parsed = ParsedModelName { + model_id: "gpt-5.5".to_string(), + reasoning_effort: Some("high".to_string()), + service_tier: Some("fast".to_string()), + }; + let config = ModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: Some("medium".to_string()), + service_tier: Some("flex".to_string()), + aliases: Default::default(), + }; + + apply_response_model_options(&mut request, &parsed, &config); + + assert_eq!(request.model, "gpt-5.5"); + assert_eq!( + request.reasoning, + Some(json!({"summary": "auto", "effort": "high"})) + ); + assert_eq!(request.service_tier.as_deref(), Some("priority")); + assert_eq!( + request.include, + Some(vec!["reasoning.encrypted_content".to_string()]) + ); +} + +#[test] +fn codex_responses_model_options_should_preserve_client_include_and_normalize_body_tier() { + let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new()); + request.reasoning = Some(json!({"summary": "detailed"})); + request.service_tier = Some("fast".to_string()); + request.include = Some(vec!["file_search_call.results".to_string()]); + let parsed = ParsedModelName { + model_id: "gpt-5.5".to_string(), + reasoning_effort: Some("low".to_string()), + service_tier: Some("flex".to_string()), + }; + let config = ModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: Some("medium".to_string()), + service_tier: Some("flex".to_string()), + aliases: Default::default(), + }; + + apply_response_model_options(&mut request, &parsed, &config); + + assert_eq!( + request.reasoning, + Some(json!({"summary": "detailed", "effort": "low"})) + ); + assert_eq!(request.service_tier.as_deref(), Some("priority")); + assert_eq!( + request.include, + Some(vec!["file_search_call.results".to_string()]) + ); +} + +#[test] +fn openai_compact_request_should_drop_responses_only_fields_and_sanitize_input() { + let request = serde_json::from_value::(json!({ + "model": "gpt-5.5-fast", + "instructions": "compress the session", + "input": [ + {"role": "user", "content": "hello"}, + { + "type": "reasoning", + "id": "rs_1", + "status": "completed", + "summary": [{"type": "summary_text", "text": "kept"}], + "ignored": "drop" + }, + {"type": "compaction", "encrypted_content": "enc_compact"}, + {"type": "compaction", "id": "drop_missing_encrypted"} + ], + "tools": [{"type": "function", "name": "lookup"}], + "parallel_tool_calls": false, + "reasoning": {"effort": "high", "summary": "auto", "extra": "drop"}, + "text": { + "format": { + "type": "json_schema", + "name": "Compact", + "schema": {"type": "object"}, + "strict": true + } + }, + "stream": true, + "store": true, + "prompt_cache_key": "must_not_forward" + })) + .expect("responses request should deserialize"); + + let compact = translate_response_to_compact(request); + let body = serde_json::to_value(&compact).expect("compact request should serialize"); + + assert_eq!(body["model"], "gpt-5.5-fast"); + assert_eq!(body["instructions"], "compress the session"); + assert_eq!(body["parallel_tool_calls"], false); + assert_eq!( + body["reasoning"], + json!({"effort": "high", "summary": "auto"}) + ); + assert_eq!( + body["tools"], + json!([{"type": "function", "name": "lookup"}]) + ); + assert_eq!(body["text"]["format"]["type"], "json_schema"); + assert!(body.get("stream").is_none()); + assert!(body.get("store").is_none()); + assert!(body.get("prompt_cache_key").is_none()); + assert_eq!(body["input"].as_array().unwrap().len(), 3); + assert!(body["input"][1].get("ignored").is_none()); + assert_eq!(body["input"][2]["encrypted_content"], "enc_compact"); +} + +#[test] +fn openai_streaming_response_with_previous_response_should_require_websocket() { + let request = serde_json::from_value::(json!({ + "model": "gpt-5.5", + "input": "hello", + "previous_response_id": "resp_previous", + "stream": true + })) + .expect("responses request should deserialize"); + + let codex = translate_response_to_codex(request); + + assert_eq!(codex.previous_response_id.as_deref(), Some("resp_previous")); + assert!(codex.stream); + assert!(!codex.force_http_sse); +} + +#[test] +fn extract_usage_should_read_codex_usage_shape() { + let body = json!({ + "usage": { + "input_tokens": 12, + "output_tokens": 5, + "input_tokens_details": { + "cached_tokens": 3 + } + } + }); + + let usage = extract_usage(&body).expect("usage should exist"); + + assert_eq!( + usage, + TokenUsage { + input_tokens: 12, + output_tokens: 5, + cached_tokens: 3, + image_input_tokens: 0, + image_output_tokens: 0, + total_tokens: 17, + } + ); +} + +#[test] +fn extract_usage_should_read_image_generation_tokens_separately() { + let body = json!({ + "usage": { + "input_tokens": 12, + "output_tokens": 5, + "input_tokens_details": { + "cached_tokens": 3 + } + }, + "tool_usage": { + "image_gen": { + "input_tokens": 31, + "output_tokens": 9 + } + } + }); + + let usage = extract_usage(&body).expect("usage should exist"); + + assert_eq!( + usage, + TokenUsage { + input_tokens: 12, + output_tokens: 5, + cached_tokens: 3, + image_input_tokens: 31, + image_output_tokens: 9, + total_tokens: 17, + } + ); +} + +#[test] +fn extract_usage_should_merge_openai_usage_shape() { + let first = extract_usage(&json!({ + "usage": { + "prompt_tokens": 8, + "completion_tokens": 4, + "prompt_tokens_details": { + "cached_tokens": 2 + } + } + })) + .expect("usage should exist"); + let second = TokenUsage { + input_tokens: 1, + output_tokens: 2, + cached_tokens: 0, + image_input_tokens: 0, + image_output_tokens: 0, + total_tokens: 3, + }; + + assert_eq!( + first.merged(second), + TokenUsage { + input_tokens: 9, + output_tokens: 6, + cached_tokens: 2, + image_input_tokens: 0, + image_output_tokens: 0, + total_tokens: 15, + } + ); +} + +#[test] +fn extract_sse_usage_should_prefer_completed_response_usage() { + let body = concat!( + "event: response.created\n", + "data: {\"response\":{\"id\":\"resp_1\",\"usage\":{\"input_tokens\":3,\"output_tokens\":5,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n", + "event: response.completed\n", + "data: {\"response\":{\"id\":\"resp_1\",\"usage\":{\"input_tokens\":3,\"output_tokens\":5,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n", + ); + + let usage = extract_sse_usage(body) + .expect("usage extraction should succeed") + .expect("usage should exist"); + + assert_eq!( + usage, + TokenUsage { + input_tokens: 3, + output_tokens: 5, + cached_tokens: 1, + image_input_tokens: 0, + image_output_tokens: 0, + total_tokens: 8, + } + ); +} + +#[test] +fn extract_sse_usage_should_read_completed_image_generation_tokens() { + let body = concat!( + "event: response.completed\n", + "data: {\"response\":{\"id\":\"resp_1\",\"usage\":{\"input_tokens\":12,\"output_tokens\":5,\"input_tokens_details\":{\"cached_tokens\":3}},\"tool_usage\":{\"image_gen\":{\"input_tokens\":31,\"output_tokens\":9}}}}\n\n", + ); + + let usage = extract_sse_usage(body) + .expect("usage extraction should succeed") + .expect("usage should exist"); + + assert_eq!( + usage, + TokenUsage { + input_tokens: 12, + output_tokens: 5, + cached_tokens: 3, + image_input_tokens: 31, + image_output_tokens: 9, + total_tokens: 17, + } + ); +} + +#[test] +fn retry_after_seconds_from_body_should_read_structured_retry_delay() { + let body = json!({ + "response": { + "error": { + "resets_in_seconds": 45 + } + } + }) + .to_string(); + + assert_eq!(retry_after_seconds_from_body(&body), Some(45)); +} + +#[test] +fn retry_after_seconds_from_body_should_parse_rate_limit_message_seconds() { + let body = json!({ + "response": { + "error": { + "code": "rate_limit_exceeded", + "message": "Rate limit reached. Please try again in 11.054s." + } + } + }) + .to_string(); + + assert_eq!(retry_after_seconds_from_body(&body), Some(12)); +} + +#[test] +fn retry_after_seconds_from_body_should_parse_rate_limit_message_milliseconds() { + let body = json!({ + "error": { + "code": "rate_limit_exceeded", + "message": "Rate limit reached. Please try again in 28ms." + } + }) + .to_string(); + + assert_eq!(retry_after_seconds_from_body(&body), Some(1)); +} + +#[test] +fn retry_after_seconds_from_body_should_ignore_retry_message_for_other_codes() { + let body = json!({ + "response": { + "error": { + "code": "upstream_transient_error", + "message": "Try again in 35 seconds." + } + } + }) + .to_string(); + + assert_eq!(retry_after_seconds_from_body(&body), None); +} + +#[test] +fn parse_rate_limit_headers_should_extract_primary_secondary_and_review_windows() { + let headers = vec![ + ( + "x-codex-primary-used-percent".to_string(), + "100".to_string(), + ), + ( + "x-codex-primary-window-minutes".to_string(), + "5".to_string(), + ), + ( + "x-codex-primary-reset-at".to_string(), + "1893456300".to_string(), + ), + ( + "x-codex-secondary-used-percent".to_string(), + "42.5".to_string(), + ), + ( + "x-codex-secondary-window-minutes".to_string(), + "10080".to_string(), + ), + ( + "x-codex-code-review-primary-used-percent".to_string(), + "80".to_string(), + ), + ( + "x-codex-code-review-primary-reset-at".to_string(), + "1893456600".to_string(), + ), + ]; + + let parsed = parse_rate_limit_headers(&headers).expect("rate limits should parse"); + + assert_eq!( + parsed.primary, + Some(RateLimitWindow { + used_percent: 100.0, + window_minutes: Some(5), + reset_at: Some(1_893_456_300), + }) + ); + assert_eq!( + parsed.secondary.expect("secondary window").window_minutes, + Some(10080) + ); + assert_eq!( + parsed + .code_review + .expect("review window") + .primary + .expect("review primary") + .reset_at, + Some(1_893_456_600) + ); +} + +#[test] +fn parse_rate_limits_event_should_extract_internal_websocket_rate_limits() { + let event = json!({ + "type": "codex.rate_limits", + "rate_limits": { + "primary": { + "used_percent": 99.5, + "window_minutes": 300, + "reset_at": 1893456300 + }, + "secondary": { + "used_percent": 10, + "window_minutes": 10080, + "reset_at": 1894056000 + } + } + }); + + let parsed = parse_rate_limits_event(&event).expect("event should parse"); + + assert_eq!(parsed.primary.expect("primary window").used_percent, 99.5); + assert_eq!( + parsed.secondary.expect("secondary window").reset_at, + Some(1_894_056_000) + ); +} + +#[test] +fn rate_limit_quota_should_preserve_existing_credits_when_passive_data_lacks_credits() { + let headers = vec![ + ("x-codex-primary-used-percent".to_string(), "25".to_string()), + ( + "x-codex-primary-window-minutes".to_string(), + "5".to_string(), + ), + ( + "x-codex-primary-reset-at".to_string(), + "1893456300".to_string(), + ), + ]; + let existing = json!({ + "credits": { + "has_credits": true, + "unlimited": false, + "balance": 12 + } + }); + let parsed = parse_rate_limit_headers(&headers).expect("rate limits should parse"); + + let quota = rate_limit_quota(&parsed, Some("plus"), Some(&existing)); + + assert_eq!(quota["plan_type"], "plus"); + assert_eq!(quota["rate_limit"]["remaining_percent"], 75); + assert_eq!(quota["credits"]["balance"], 12); +} + +#[test] +fn cooldown_with_jitter_should_return_positive_duration_within_expected_range() { + let duration = cooldown_with_jitter(60, 2_000); + + assert!((48..=72).contains(&duration.num_seconds())); +} + +#[tokio::test] +async fn refresh_scheduler_should_refresh_before_expiry_and_preserve_refresh_token() { + use async_trait::async_trait; + use chrono::{Duration as ChronoDuration, Utc}; + use codex_proxy_core::accounts::model::{Account, AccountStatus}; + use codex_proxy_core::auth::oauth::{ + RefreshFailure, RefreshPolicy, RefreshScheduler, RefreshTrigger, TokenPair, + }; + use codex_proxy_core::auth::ports::TokenRefresher; + + #[derive(Clone)] + struct StaticRefreshClient { + result: Result, + } + + #[async_trait] + impl TokenRefresher for StaticRefreshClient { + async fn refresh(&self, _refresh_token: &str) -> Result { + self.result.clone() + } + } + + let now = Utc::now(); + let mut account = Account::test("acct_1", AccountStatus::Active); + account.access_token_expires_at = Some(now + ChronoDuration::seconds(60)); + account.refresh_token = Some("rt_keep".to_string()); + let scheduler = RefreshScheduler::new( + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 2, + }, + StaticRefreshClient { + result: Ok(TokenPair { + access_token: "new-access".to_string(), + refresh_token: None, + }), + }, + ); + + let refreshed = scheduler + .refresh_account_at(&account, RefreshTrigger::BeforeExpiry, now) + .await + .expect("refresh should succeed"); + + assert_eq!(refreshed.access_token, "new-access"); + assert_eq!(refreshed.refresh_token.as_deref(), Some("rt_keep")); + assert_eq!(refreshed.status, AccountStatus::Active); +} + +fn assert_substrings_appear_in_order(haystack: &str, needles: &[&str]) { + let mut cursor = 0; + for needle in needles { + let Some(offset) = haystack[cursor..].find(needle) else { + panic!("expected substring {needle:?} after byte {cursor} in:\n{haystack}"); + }; + cursor += offset + needle.len(); + } +} diff --git a/crates/core/tests/session_affinity.rs b/crates/core/tests/session_affinity.rs new file mode 100644 index 000000000..be6d3fde5 --- /dev/null +++ b/crates/core/tests/session_affinity.rs @@ -0,0 +1,106 @@ +use chrono::{Duration, TimeZone, Utc}; +use codex_proxy_core::serving::affinity::{ + SessionAffinityEntry, SessionAffinityMap, StoredSessionAffinity, +}; + +#[test] +fn session_affinity_map_should_restore_active_records_and_skip_expired_records() { + let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(); + let mut map = SessionAffinityMap::new(Duration::hours(4)); + + let restored = map.restore( + vec![ + StoredSessionAffinity { + response_id: "resp_active".to_string(), + entry: SessionAffinityEntry { + account_id: "acct_active".to_string(), + conversation_id: "conv_active".to_string(), + turn_state: Some("turn_active".to_string()), + instructions_hash: Some("hash_active".to_string()), + input_tokens: Some(42), + function_call_ids: vec!["call_1".to_string()], + variant_hash: Some("variant_a".to_string()), + created_at: now - Duration::minutes(30), + }, + expires_at: now + Duration::hours(1), + }, + StoredSessionAffinity { + response_id: "resp_expired".to_string(), + entry: SessionAffinityEntry { + account_id: "acct_expired".to_string(), + conversation_id: "conv_expired".to_string(), + turn_state: None, + instructions_hash: None, + input_tokens: None, + function_call_ids: Vec::new(), + variant_hash: None, + created_at: now - Duration::hours(5), + }, + expires_at: now + Duration::hours(1), + }, + ], + now, + ); + + assert_eq!(restored, 1); + assert_eq!( + map.lookup_account("resp_active", now), + Some("acct_active".to_string()) + ); + assert_eq!(map.lookup_account("resp_expired", now), None); +} + +#[test] +fn session_affinity_map_should_find_latest_response_by_conversation_and_variant() { + let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(); + let mut map = SessionAffinityMap::new(Duration::hours(4)); + + map.record( + "resp_old".to_string(), + SessionAffinityEntry { + account_id: "acct_a".to_string(), + conversation_id: "conv".to_string(), + turn_state: Some("turn_old".to_string()), + instructions_hash: None, + input_tokens: None, + function_call_ids: Vec::new(), + variant_hash: Some("variant_a".to_string()), + created_at: now - Duration::minutes(20), + }, + ); + map.record( + "resp_new".to_string(), + SessionAffinityEntry { + account_id: "acct_b".to_string(), + conversation_id: "conv".to_string(), + turn_state: Some("turn_new".to_string()), + instructions_hash: None, + input_tokens: None, + function_call_ids: Vec::new(), + variant_hash: Some("variant_a".to_string()), + created_at: now - Duration::minutes(5), + }, + ); + map.record( + "resp_other_variant".to_string(), + SessionAffinityEntry { + account_id: "acct_c".to_string(), + conversation_id: "conv".to_string(), + turn_state: Some("turn_other".to_string()), + instructions_hash: None, + input_tokens: None, + function_call_ids: Vec::new(), + variant_hash: Some("variant_b".to_string()), + created_at: now - Duration::minutes(1), + }, + ); + + assert_eq!( + map.lookup_latest_response_by_conversation("conv", None, Some("variant_a"), now), + Some("resp_new".to_string()) + ); + assert_eq!( + map.lookup_turn_state("resp_new", now), + Some("turn_new".to_string()) + ); +} diff --git a/crates/platform/Cargo.toml b/crates/platform/Cargo.toml new file mode 100644 index 000000000..e3443c1c2 --- /dev/null +++ b/crates/platform/Cargo.toml @@ -0,0 +1,41 @@ +[package] +name = "codex-proxy-platform" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +rust-version.workspace = true + +[dependencies] +aes-gcm.workspace = true +argon2.workspace = true +base64.workspace = true +chrono.workspace = true +config.workspace = true +dirs.workspace = true +flate2.workspace = true +hex.workspace = true +hmac.workspace = true +rand.workspace = true +rustls.workspace = true +rustls-native-certs.workspace = true +rustls-pki-types.workspace = true +secrecy.workspace = true +serde.workspace = true +serde_json.workspace = true +serde_yml.workspace = true +sha2.workspace = true +sqlx.workspace = true +thiserror.workspace = true +tokio.workspace = true +tracing.workspace = true +tracing-appender.workspace = true +tracing-subscriber.workspace = true +uuid.workspace = true +zeroize.workspace = true + +[lints] +workspace = true + +[dev-dependencies] +tempfile = "3.27.0" diff --git a/crates/platform/src/config/loader.rs b/crates/platform/src/config/loader.rs new file mode 100644 index 000000000..3f19f4e21 --- /dev/null +++ b/crates/platform/src/config/loader.rs @@ -0,0 +1,69 @@ +use std::path::Path; + +use serde::Serialize; + +use super::{ + types::{AppConfig, AuthConfig, LoggingConfig, ModelConfig, QuotaConfig, UsageStatsConfig}, + ConfigResult, ConfigWriteError, ConfigWriteResult, +}; + +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() + } + + /// 写入管理端设置使用的本地配置覆盖文件。 + pub async fn write_settings_overlay( + &self, + local_config_path: impl AsRef, + ) -> ConfigWriteResult<()> { + let local_config_path = local_config_path.as_ref(); + if let Some(parent) = local_config_path + .parent() + .filter(|path| !path.as_os_str().is_empty()) + { + tokio::fs::create_dir_all(parent) + .await + .map_err(ConfigWriteError::CreateDirectory)?; + } + + let overlay = SettingsConfigOverlay::from(self); + let content = serde_yml::to_string(&overlay).map_err(ConfigWriteError::Serialize)?; + tokio::fs::write(local_config_path, content) + .await + .map_err(ConfigWriteError::Write) + } +} + +#[derive(Debug, Serialize)] +struct SettingsConfigOverlay { + model: ModelConfig, + auth: AuthConfig, + quota: QuotaConfig, + usage_stats: UsageStatsConfig, + logging: LoggingConfig, +} + +impl From<&AppConfig> for SettingsConfigOverlay { + fn from(config: &AppConfig) -> Self { + Self { + model: config.model.clone(), + auth: config.auth.clone(), + quota: config.quota.clone(), + usage_stats: config.usage_stats.clone(), + logging: config.logging.clone(), + } + } +} diff --git a/crates/platform/src/config/mod.rs b/crates/platform/src/config/mod.rs new file mode 100644 index 000000000..b622b68d4 --- /dev/null +++ b/crates/platform/src/config/mod.rs @@ -0,0 +1,28 @@ +//! 配置加载与配置类型。 + +mod loader; +mod types; + +use thiserror::Error; + +pub use types::*; + +/// 配置加载结果类型。 +pub type ConfigResult = Result; + +/// 配置写入结果类型。 +pub type ConfigWriteResult = Result; + +/// 配置写入错误。 +#[derive(Debug, Error)] +pub enum ConfigWriteError { + /// 序列化配置覆盖失败。 + #[error("serialize config overlay")] + Serialize(#[source] serde_yml::Error), + /// 创建配置目录失败。 + #[error("create config directory")] + CreateDirectory(#[source] std::io::Error), + /// 写入配置覆盖文件失败。 + #[error("write config overlay")] + Write(#[source] std::io::Error), +} diff --git a/src/config/types.rs b/crates/platform/src/config/types.rs similarity index 66% rename from src/config/types.rs rename to crates/platform/src/config/types.rs index ec3bb1722..bf615aad9 100644 --- a/src/config/types.rs +++ b/crates/platform/src/config/types.rs @@ -2,96 +2,150 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; +/// 应用运行总配置。 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] pub struct AppConfig { + /// HTTP 服务配置。 pub server: ServerConfig, + /// API 地址配置。 pub api: ApiConfig, + /// 模型默认值配置。 pub model: ModelConfig, + /// 认证与刷新配置。 pub auth: AuthConfig, + /// 配额配置。 pub quota: QuotaConfig, + /// 用量统计配置。 pub usage_stats: UsageStatsConfig, + /// 数据库配置。 pub database: DatabaseConfig, + /// 安全材料配置。 pub security: SecurityConfig, + /// TLS 偏好配置。 pub tls: TlsConfig, + /// WebSocket 连接池配置。 pub ws_pool: WebSocketPoolConfig, + /// 管理员初始化配置。 pub admin: AdminConfig, + /// 日志配置。 pub logging: LoggingConfig, } +/// HTTP 监听配置。 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] pub struct ServerConfig { + /// 监听主机。 pub host: String, + /// 监听端口。 pub port: u16, } +/// 上游 API 地址配置。 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] pub struct ApiConfig { + /// 上游基础 URL。 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, } +/// 认证、轮换与 OAuth 配置。 #[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, + /// OAuth 客户端 ID。 pub oauth_client_id: String, + /// OAuth 授权端点。 pub oauth_auth_endpoint: String, + /// OAuth 令牌端点。 pub oauth_token_endpoint: String, } +/// 配额刷新与跳过配置。 #[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 { + /// SQLite 连接 URL。 pub url: String, } +/// 本地安全文件路径配置。 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] pub struct SecurityConfig { + /// 主密钥文件路径。 pub master_key_file: String, + /// API Key pepper 文件路径。 pub api_key_pepper_file: String, } +/// TLS/HTTP 协议偏好配置。 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] pub struct TlsConfig { + /// 是否强制 HTTP/1.1。 pub force_http11: bool, } +/// WebSocket 连接池配置。 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] #[serde(deny_unknown_fields)] pub struct WebSocketPoolConfig { + /// 是否启用连接池。 pub enabled: bool, + /// 连接最大存活时长。 pub max_age_ms: u64, + /// 单账号最大连接数。 pub max_per_account: usize, } @@ -105,9 +159,12 @@ impl Default for WebSocketPoolConfig { } } +/// 管理员初始化与会话配置。 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] pub struct AdminConfig { + /// 会话有效期（分钟）。 pub session_ttl_minutes: u64, + /// 会话清理周期（秒）。 #[serde(default = "default_session_cleanup_interval")] pub session_cleanup_interval_secs: u64, /// 默认管理员用户名（首次启动时创建） @@ -130,12 +187,18 @@ fn default_admin_password() -> String { "admin".to_string() } +/// 日志持久化配置。 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] #[serde(deny_unknown_fields)] pub struct LoggingConfig { + /// 日志目录。 pub directory: String, + /// 保留天数。 pub retention_days: usize, + /// 是否启用日志。 pub enabled: bool, + /// 内存容量。 pub capacity: u32, + /// 是否捕获请求体。 pub capture_body: bool, } diff --git a/crates/platform/src/crypto/hash.rs b/crates/platform/src/crypto/hash.rs new file mode 100644 index 000000000..a3f0f20f4 --- /dev/null +++ b/crates/platform/src/crypto/hash.rs @@ -0,0 +1,25 @@ +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use hmac::{Hmac, KeyInit, Mac}; +use sha2::Sha256; + +/// HMAC-SHA256 类型别名。 +pub type HmacSha256 = Hmac; + +/// 计算 HMAC-SHA256 并以 URL-safe base64 输出。 +pub fn hmac_sha256_base64(key: &[u8], payload: &[u8]) -> Option { + let mut mac = HmacSha256::new_from_slice(key).ok()?; + mac.update(payload); + Some(URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes())) +} + +/// 常数时间比较两个字节切片。 +pub fn constant_time_eq(left: &[u8], right: &[u8]) -> bool { + if left.len() != right.len() { + return false; + } + + left.iter() + .zip(right) + .fold(0u8, |acc, (left, right)| acc | (left ^ right)) + == 0 +} diff --git a/crates/platform/src/crypto/mod.rs b/crates/platform/src/crypto/mod.rs new file mode 100644 index 000000000..703b5ae9b --- /dev/null +++ b/crates/platform/src/crypto/mod.rs @@ -0,0 +1,7 @@ +//! 加密原语。 + +/// HMAC 和常数时间比较辅助。 +pub mod hash; +mod secret_box; + +pub use secret_box::{CryptoError, CryptoResult, SecretBox}; diff --git a/src/platform/crypto/secret_box.rs b/crates/platform/src/crypto/secret_box.rs similarity index 85% rename from src/platform/crypto/secret_box.rs rename to crates/platform/src/crypto/secret_box.rs index a61ddab34..5b7140be4 100644 --- a/src/platform/crypto/secret_box.rs +++ b/crates/platform/src/crypto/secret_box.rs @@ -1,3 +1,5 @@ +//! AES-256-GCM 密钥盒。 + use std::{fs, path::Path}; use aes_gcm::{ @@ -8,40 +10,54 @@ use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use secrecy::{ExposeSecret, SecretString}; use thiserror::Error; +/// 密钥盒相关错误。 #[derive(Debug, Error)] pub enum CryptoError { + /// 密钥长度不正确。 #[error("invalid secret key length")] InvalidKeyLength, + /// 加密失败。 #[error("secret encryption failed")] Encrypt, + /// 解密失败。 #[error("secret decryption failed")] Decrypt, + /// nonce 长度不正确。 #[error("invalid nonce length")] InvalidNonceLength, + /// Base64 编码/解码失败。 #[error("invalid secret encoding: {0}")] Decode(#[from] base64::DecodeError), + /// 密文版本不支持。 #[error("unsupported secret version")] UnsupportedVersion, + /// 解密结果不是 UTF-8。 #[error("secret is not valid utf-8: {0}")] Utf8(#[from] std::string::FromUtf8Error), + /// 读写密钥文件失败。 #[error("secret key file io error: {0}")] Io(#[from] std::io::Error), + /// 持久化密钥长度不正确。 #[error("stored secret key must decode to 32 bytes")] InvalidStoredKeyLength, } +/// 密钥盒操作结果。 pub type CryptoResult = Result; +/// 用于加密和解密敏感字符串的密钥盒。 #[derive(Clone)] pub struct SecretBox { key: [u8; 32], } impl SecretBox { + /// 使用 32 字节密钥创建密钥盒。 pub fn new(key: [u8; 32]) -> Self { Self { key } } + /// 从文件加载或生成密钥。 pub fn load_or_create(path: impl AsRef) -> CryptoResult { let path = path.as_ref(); if path.exists() { @@ -62,6 +78,7 @@ impl SecretBox { Ok(Self::new(key)) } + /// 加密敏感明文。 pub fn encrypt(&self, plaintext: &SecretString) -> CryptoResult { let cipher = Aes256Gcm::new_from_slice(&self.key).map_err(|_| CryptoError::InvalidKeyLength)?; @@ -78,6 +95,7 @@ impl SecretBox { )) } + /// 解密密文。 pub fn decrypt(&self, encoded: &str) -> CryptoResult { let mut parts = encoded.split(':'); let version = parts.next().unwrap_or_default(); diff --git a/src/platform/identity/admin_session.rs b/crates/platform/src/identity/admin_password.rs similarity index 85% rename from src/platform/identity/admin_session.rs rename to crates/platform/src/identity/admin_password.rs index dee3361ac..d1743136a 100644 --- a/src/platform/identity/admin_session.rs +++ b/crates/platform/src/identity/admin_password.rs @@ -3,8 +3,9 @@ use argon2::{ Argon2, }; -use crate::platform::identity::error::AuthResult; +use crate::identity::AuthResult; +/// 生成管理员密码哈希。 pub fn hash_admin_password(password: &str) -> AuthResult { let salt = SaltString::generate(&mut OsRng); Ok(Argon2::default() @@ -12,6 +13,7 @@ pub fn hash_admin_password(password: &str) -> AuthResult { .to_string()) } +/// 验证管理员密码。 pub fn verify_admin_password(password: &str, hash: &str) -> AuthResult { let parsed = PasswordHash::new(hash)?; Ok(Argon2::default() diff --git a/src/platform/identity/client_key.rs b/crates/platform/src/identity/client_key.rs similarity index 86% rename from src/platform/identity/client_key.rs rename to crates/platform/src/identity/client_key.rs index 153522ebb..2a2b3f57a 100644 --- a/src/platform/identity/client_key.rs +++ b/crates/platform/src/identity/client_key.rs @@ -5,27 +5,34 @@ use hmac::{Hmac, KeyInit, Mac}; use rand::Rng; use sha2::Sha256; -use crate::platform::identity::error::{AuthError, AuthResult}; +use crate::identity::{AuthError, AuthResult}; type HmacSha256 = Hmac; +/// 客户端 API Key 派生结果。 #[derive(Debug, Clone)] pub struct GeneratedClientApiKey { + /// 仅显示一次的明文密钥。 pub plaintext: String, + /// 前缀缓存。 pub prefix: String, + /// 持久化哈希。 pub key_hash: String, } +/// API Key 哈希器。 #[derive(Debug, Clone)] pub struct ApiKeyHasher { pepper: [u8; 32], } impl ApiKeyHasher { + /// 使用给定 pepper 构造哈希器。 pub fn new(pepper: [u8; 32]) -> Self { Self { pepper } } + /// 从文件加载或生成 pepper。 pub fn load_or_create(path: impl AsRef) -> AuthResult { let path = path.as_ref(); if path.exists() { @@ -46,6 +53,7 @@ impl ApiKeyHasher { Ok(Self::new(pepper)) } + /// 从字节切片生成哈希器。 pub fn try_from_slice(pepper: &[u8]) -> AuthResult { let pepper: [u8; 32] = pepper .try_into() @@ -53,6 +61,7 @@ impl ApiKeyHasher { Ok(Self::new(pepper)) } + /// 生成新的客户端 API Key。 pub fn generate_client_api_key(&self, _name: &str) -> GeneratedClientApiKey { let mut bytes = [0u8; 32]; rand::rng().fill_bytes(&mut bytes); @@ -67,12 +76,14 @@ impl ApiKeyHasher { } } + /// 计算 API Key 的 HMAC。 pub fn hash_client_api_key(&self, plaintext: &str) -> String { let mut mac = self.new_mac(); mac.update(plaintext.as_bytes()); URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes()) } + /// 校验 API Key 是否有效。 pub fn verify_client_api_key(&self, plaintext: &str, key_hash: &str) -> AuthResult { let Some(suffix) = plaintext.strip_prefix("cpr_") else { return Ok(false); diff --git a/src/platform/identity/error.rs b/crates/platform/src/identity/mod.rs similarity index 60% rename from src/platform/identity/error.rs rename to crates/platform/src/identity/mod.rs index ca6752538..dd5c2074e 100644 --- a/src/platform/identity/error.rs +++ b/crates/platform/src/identity/mod.rs @@ -1,17 +1,31 @@ +//! 身份哈希与本地认证相关原语。 + +mod admin_password; +mod client_key; + +pub use admin_password::{hash_admin_password, verify_admin_password}; +pub use client_key::{ApiKeyHasher, GeneratedClientApiKey}; + use thiserror::Error; +/// 身份原语错误。 #[derive(Debug, Error)] pub enum AuthError { + /// 密码哈希错误。 #[error("password hash error: {0}")] PasswordHash(String), + /// API Key 编码错误。 #[error("invalid api key encoding: {0}")] ApiKeyEncoding(#[from] base64::DecodeError), + /// Pepper 长度错误。 #[error("invalid api key pepper length")] InvalidPepperLength, + /// Pepper 文件读写错误。 #[error("api key pepper file io error: {0}")] PepperIo(#[from] std::io::Error), } +/// 身份原语结果。 pub type AuthResult = Result; impl From for AuthError { diff --git a/src/utils/json.rs b/crates/platform/src/json/mod.rs similarity index 66% rename from src/utils/json.rs rename to crates/platform/src/json/mod.rs index 0f6519517..710b6055b 100644 --- a/src/utils/json.rs +++ b/crates/platform/src/json/mod.rs @@ -1,9 +1,17 @@ +//! 无业务含义的 JSON 与分页辅助。 + +mod pagination; + +pub use pagination::{clamp_limit, decode_cursor, encode_cursor, Page}; + 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 { diff --git a/src/utils/pagination.rs b/crates/platform/src/json/pagination.rs similarity index 82% rename from src/utils/pagination.rs rename to crates/platform/src/json/pagination.rs index 3efcd0ff3..e82d54b63 100644 --- a/src/utils/pagination.rs +++ b/crates/platform/src/json/pagination.rs @@ -1,17 +1,22 @@ use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use serde::Serialize; +/// 通用分页结果。 #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct Page { + /// 当前页条目。 pub items: Vec, + /// 下一页游标。 pub next_cursor: Option, } +/// 编码游标。 pub fn encode_cursor(created_at: &str, id: &str) -> String { URL_SAFE_NO_PAD.encode(format!("{created_at}|{id}")) } +/// 解码游标。 pub fn decode_cursor(cursor: &str) -> Option<(String, String)> { let raw = URL_SAFE_NO_PAD.decode(cursor).ok()?; let text = String::from_utf8(raw).ok()?; @@ -19,6 +24,7 @@ pub fn decode_cursor(cursor: &str) -> Option<(String, String)> { Some((created_at.to_string(), id.to_string())) } +/// 限制分页大小。 pub fn clamp_limit(limit: u32) -> u32 { limit.clamp(1, 200) } diff --git a/crates/platform/src/lib.rs b/crates/platform/src/lib.rs new file mode 100644 index 000000000..63e6af79d --- /dev/null +++ b/crates/platform/src/lib.rs @@ -0,0 +1,10 @@ +#![deny(missing_docs)] + +//! 平台基础设施层，承载配置、加密、身份、存储、日志和 JSON 原语。 + +pub mod config; +pub mod crypto; +pub mod identity; +pub mod json; +pub mod logging; +pub mod storage; diff --git a/crates/platform/src/logging/mod.rs b/crates/platform/src/logging/mod.rs new file mode 100644 index 000000000..634fa29c6 --- /dev/null +++ b/crates/platform/src/logging/mod.rs @@ -0,0 +1,5 @@ +//! 日志初始化与文件轮转。 + +mod rotation; + +pub use rotation::{build_file_appender, init_tracing, LogError, LogGuard, RotationConfig}; diff --git a/src/platform/logging/rotation.rs b/crates/platform/src/logging/rotation.rs similarity index 88% rename from src/platform/logging/rotation.rs rename to crates/platform/src/logging/rotation.rs index de4ac1d49..ccc6580d1 100644 --- a/src/platform/logging/rotation.rs +++ b/crates/platform/src/logging/rotation.rs @@ -13,13 +13,17 @@ use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, Env const LOG_PREFIX: &str = "codex-proxy-rs"; const LOG_SUFFIX: &str = "log"; +/// 日志轮转配置。 #[derive(Debug, Clone)] pub struct RotationConfig { + /// 日志目录。 pub directory: PathBuf, + /// 保留周期。 pub retention_days: usize, } impl RotationConfig { + /// 创建轮转配置。 pub fn new(directory: impl AsRef, retention_days: usize) -> Self { Self { directory: directory.as_ref().to_path_buf(), @@ -28,20 +32,26 @@ impl RotationConfig { } } +/// 日志初始化错误。 #[derive(Debug, Error)] pub enum LogError { + /// 日志 IO 错误。 #[error("log io error: {0}")] Io(#[from] io::Error), + /// appender 初始化错误。 #[error("log appender initialization failed: {0}")] Appender(#[from] tracing_appender::rolling::InitError), + /// tracing subscriber 已初始化。 #[error("global tracing subscriber is already initialized")] SubscriberAlreadyInitialized, } +/// 日志守卫。 pub struct LogGuard { _guard: WorkerGuard, } +/// 构造文件 appender。 pub fn build_file_appender(config: &RotationConfig) -> Result { fs::create_dir_all(&config.directory)?; @@ -58,6 +68,7 @@ pub fn build_file_appender(config: &RotationConfig) -> Result Result { let appender = build_file_appender(&config)?; let (writer, guard) = tracing_appender::non_blocking(appender); diff --git a/crates/platform/src/storage/mod.rs b/crates/platform/src/storage/mod.rs new file mode 100644 index 000000000..d7c3128a6 --- /dev/null +++ b/crates/platform/src/storage/mod.rs @@ -0,0 +1,11 @@ +//! SQLite 存储初始化与路径辅助。 + +mod paths; +/// SQLite 连接初始化。 +pub mod sqlite; + +pub use paths::{ + codex_desktop_installation_id_path, data_dir, data_installation_id_path, ensure_data_dir, + load_or_create_installation_id, load_or_create_installation_id_from_paths, +}; +pub use sqlite::{connect_sqlite, SqlitePool}; diff --git a/crates/platform/src/storage/paths.rs b/crates/platform/src/storage/paths.rs new file mode 100644 index 000000000..fc84d99d0 --- /dev/null +++ b/crates/platform/src/storage/paths.rs @@ -0,0 +1,158 @@ +use std::{ + fs, io, + path::{Path, PathBuf}, +}; + +use uuid::Uuid; + +const INSTALLATION_ID_FILE_NAME: &str = "installation_id"; + +/// 返回本地数据目录。 +pub fn data_dir() -> PathBuf { + dirs::data_local_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("codex-proxy-rs") + .join("data") +} + +/// 确保本地数据目录存在。 +pub fn ensure_data_dir() -> std::io::Result { + let dir = data_dir(); + std::fs::create_dir_all(&dir)?; + Ok(dir) +} + +/// 返回真实 Codex Desktop installation ID 文件路径。 +pub fn codex_desktop_installation_id_path() -> Option { + dirs::home_dir().map(|home| home.join(".codex").join(INSTALLATION_ID_FILE_NAME)) +} + +/// 返回本应用数据目录中的 installation ID 文件路径。 +pub fn data_installation_id_path(data_dir: &Path) -> PathBuf { + data_dir.join(INSTALLATION_ID_FILE_NAME) +} + +/// 按兼容顺序读取或生成 installation ID。 +/// +/// 顺序为真实 Codex Desktop 文件、本应用数据目录文件、生成并写入本应用数据目录。 +/// 当 `data_dir` 为 `None` 且没有可读旧文件时，返回一个不会持久化的新 UUID。 +pub fn load_or_create_installation_id(data_dir: Option<&Path>) -> io::Result { + let codex_path = codex_desktop_installation_id_path(); + let data_path = data_dir.map(data_installation_id_path); + load_or_create_installation_id_from_paths(codex_path.as_deref(), data_path.as_deref()) +} + +/// 从显式文件路径按兼容顺序读取或生成 installation ID。 +/// +/// 此函数用于测试和调用方已经完成路径解析的场景。 +pub fn load_or_create_installation_id_from_paths( + codex_desktop_path: Option<&Path>, + data_path: Option<&Path>, +) -> io::Result { + if let Some(path) = codex_desktop_path { + if let Some(id) = read_installation_id(path)? { + return Ok(id); + } + } + + if let Some(path) = data_path { + if let Some(id) = read_installation_id(path)? { + return Ok(id); + } + + let generated = generate_installation_id(); + persist_installation_id(path, &generated)?; + return Ok(generated); + } + + Ok(generate_installation_id()) +} + +fn read_installation_id(path: &Path) -> io::Result> { + match fs::read_to_string(path) { + Ok(content) => Ok(parse_installation_id(&content)), + Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(None), + Err(error) => Err(error), + } +} + +fn persist_installation_id(path: &Path, installation_id: &str) -> io::Result<()> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + fs::write(path, installation_id) +} + +fn parse_installation_id(raw: &str) -> Option { + let trimmed = raw.trim(); + Uuid::parse_str(trimmed).ok()?; + Some(trimmed.to_string()) +} + +fn generate_installation_id() -> String { + Uuid::new_v4().to_string() +} + +#[cfg(test)] +mod tests { + use std::fs; + + use tempfile::TempDir; + use uuid::Uuid; + + use super::*; + + #[test] + fn resolve_installation_id_should_prefer_codex_desktop_file() { + let temp = TempDir::new().expect("tempdir should be created"); + let codex_path = temp.path().join(".codex").join("installation_id"); + let data_path = temp.path().join("data").join("installation_id"); + fs::create_dir_all(codex_path.parent().expect("codex path should have parent")) + .expect("codex parent should be created"); + fs::create_dir_all(data_path.parent().expect("data path should have parent")) + .expect("data parent should be created"); + let codex_id = Uuid::new_v4().to_string(); + let data_id = Uuid::new_v4().to_string(); + fs::write(&codex_path, &codex_id).expect("codex id should be written"); + fs::write(&data_path, data_id).expect("data id should be written"); + + let resolved = + load_or_create_installation_id_from_paths(Some(&codex_path), Some(&data_path)) + .expect("installation id should resolve"); + + assert_eq!(resolved, codex_id); + } + + #[test] + fn resolve_installation_id_should_read_data_file_when_codex_file_is_absent() { + let temp = TempDir::new().expect("tempdir should be created"); + let codex_path = temp.path().join(".codex").join("installation_id"); + let data_path = temp.path().join("data").join("installation_id"); + fs::create_dir_all(data_path.parent().expect("data path should have parent")) + .expect("data parent should be created"); + let data_id = Uuid::new_v4().to_string(); + fs::write(&data_path, &data_id).expect("data id should be written"); + + let resolved = + load_or_create_installation_id_from_paths(Some(&codex_path), Some(&data_path)) + .expect("installation id should resolve"); + + assert_eq!(resolved, data_id); + } + + #[test] + fn resolve_installation_id_should_generate_and_persist_when_files_are_absent() { + let temp = TempDir::new().expect("tempdir should be created"); + let codex_path = temp.path().join(".codex").join("installation_id"); + let data_path = temp.path().join("data").join("installation_id"); + + let resolved = + load_or_create_installation_id_from_paths(Some(&codex_path), Some(&data_path)) + .expect("installation id should resolve"); + + assert_eq!( + fs::read_to_string(&data_path).expect("generated id should be persisted"), + resolved + ); + } +} diff --git a/src/platform/storage/schema.sql b/crates/platform/src/storage/schema.sql similarity index 100% rename from src/platform/storage/schema.sql rename to crates/platform/src/storage/schema.sql diff --git a/src/platform/storage/db.rs b/crates/platform/src/storage/sqlite.rs similarity index 93% rename from src/platform/storage/db.rs rename to crates/platform/src/storage/sqlite.rs index da24dd544..2b0b1782f 100644 --- a/src/platform/storage/db.rs +++ b/crates/platform/src/storage/sqlite.rs @@ -1,7 +1,9 @@ use std::{fs, path::Path, str::FromStr, time::Duration}; -use sqlx::{sqlite::SqliteConnectOptions, SqlitePool}; +use sqlx::sqlite::SqliteConnectOptions; +pub use sqlx::SqlitePool; +/// 连接 SQLite 并初始化 schema。 pub async fn connect_sqlite(database_url: &str) -> Result { ensure_sqlite_parent_dir(database_url)?; let options = SqliteConnectOptions::from_str(database_url)? diff --git a/tests/admin/session.rs b/crates/platform/tests/admin_password.rs similarity index 60% rename from tests/admin/session.rs rename to crates/platform/tests/admin_password.rs index 23ebb9b67..ed83ba6a4 100644 --- a/tests/admin/session.rs +++ b/crates/platform/tests/admin_password.rs @@ -1,10 +1,9 @@ -use codex_proxy_rs::platform::identity::admin_session::{ - hash_admin_password, verify_admin_password, -}; +use codex_proxy_platform::identity::{hash_admin_password, verify_admin_password}; #[test] -fn admin_password_hash_is_not_a_client_api_key() { +fn admin_password_hash_should_not_accept_client_api_key_string() { let hash = hash_admin_password("correct horse battery staple").unwrap(); + assert!(verify_admin_password("correct horse battery staple", &hash).unwrap()); assert!(!verify_admin_password("cpr_fake_client_key", &hash).unwrap()); } diff --git a/tests/platform/client_key_auth.rs b/crates/platform/tests/client_key_auth.rs similarity index 88% rename from tests/platform/client_key_auth.rs rename to crates/platform/tests/client_key_auth.rs index 810e09f7c..6763371b6 100644 --- a/tests/platform/client_key_auth.rs +++ b/crates/platform/tests/client_key_auth.rs @@ -1,4 +1,4 @@ -use codex_proxy_rs::platform::identity::client_key::ApiKeyHasher; +use codex_proxy_platform::identity::ApiKeyHasher; #[test] fn client_api_key_has_proxy_prefix_and_verifies_against_hash() { diff --git a/tests/config.rs b/crates/platform/tests/config.rs similarity index 77% rename from tests/config.rs rename to crates/platform/tests/config.rs index a6ad1f8d3..32c49eb9c 100644 --- a/tests/config.rs +++ b/crates/platform/tests/config.rs @@ -1,9 +1,9 @@ -use codex_proxy_rs::config::{AppConfig, LoggingConfig}; +use codex_proxy_platform::config::{AppConfig, LoggingConfig}; use std::fs; #[test] fn default_config_keeps_only_codex_backend() { - let cfg: AppConfig = serde_yml::from_str(include_str!("../config.yaml")).unwrap(); + let cfg: AppConfig = serde_yml::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"); @@ -18,13 +18,27 @@ fn default_config_keeps_only_codex_backend() { cfg.auth.oauth_token_endpoint, "https://auth.openai.com/oauth/token" ); - assert_eq!(cfg.database.url, "sqlite://data/codex-proxy-rs.sqlite"); - assert_eq!(cfg.security.master_key_file, "data/master.key"); assert!(cfg.ws_pool.enabled); assert_eq!(cfg.ws_pool.max_age_ms, 3_300_000); assert_eq!(cfg.ws_pool.max_per_account, 8); } +#[test] +fn default_config_keeps_runtime_artifacts_under_runtime_directory() { + let cfg: AppConfig = serde_yml::from_str(include_str!("../../../config.yaml")).unwrap(); + + assert_eq!( + cfg.database.url, + "sqlite://.runtime/data/codex-proxy-rs.sqlite" + ); + assert_eq!(cfg.security.master_key_file, ".runtime/data/master.key"); + assert_eq!( + cfg.security.api_key_pepper_file, + ".runtime/data/api-key-pepper.key" + ); + assert_eq!(cfg.logging.directory, ".runtime/logs"); +} + #[test] fn config_loader_merges_config_and_local_yaml() { let dir = tempfile::tempdir().unwrap(); @@ -65,10 +79,10 @@ quota: usage_stats: history_retention_days: null database: - url: sqlite://data/codex-proxy-rs.sqlite + url: sqlite://.runtime/data/codex-proxy-rs.sqlite security: - master_key_file: data/master.key - api_key_pepper_file: data/api-key-pepper.key + master_key_file: .runtime/data/master.key + api_key_pepper_file: .runtime/data/api-key-pepper.key tls: force_http11: false ws_pool: @@ -78,7 +92,7 @@ ws_pool: admin: session_ttl_minutes: 1440 logging: - directory: logs + directory: .runtime/logs retention_days: 14 enabled: false capacity: 2000 @@ -105,7 +119,10 @@ ws_pool: assert_eq!(cfg.server.host, "0.0.0.0"); assert_eq!(cfg.server.port, 8080); - assert_eq!(cfg.database.url, "sqlite://data/codex-proxy-rs.sqlite"); + assert_eq!( + cfg.database.url, + "sqlite://.runtime/data/codex-proxy-rs.sqlite" + ); assert_eq!(cfg.model.default_model, "gpt-5.5"); assert_eq!(cfg.auth.max_concurrent_per_account, 3); assert_eq!(cfg.auth.oauth_client_id, "app_test_client"); @@ -128,7 +145,7 @@ ws_pool: fn logging_config_rejects_unsupported_size_rotation_field() { let err = serde_yml::from_str::( r#" -directory: logs +directory: .runtime/logs max_file_bytes: 10485760 retention_days: 14 enabled: false diff --git a/tests/platform/crypto.rs b/crates/platform/tests/crypto.rs similarity index 95% rename from tests/platform/crypto.rs rename to crates/platform/tests/crypto.rs index 496949fb9..3d102bb12 100644 --- a/tests/platform/crypto.rs +++ b/crates/platform/tests/crypto.rs @@ -1,4 +1,4 @@ -use codex_proxy_rs::platform::crypto::SecretBox; +use codex_proxy_platform::crypto::SecretBox; use secrecy::{ExposeSecret, SecretString}; #[test] diff --git a/tests/platform/log_rotation.rs b/crates/platform/tests/log_rotation.rs similarity index 90% rename from tests/platform/log_rotation.rs rename to crates/platform/tests/log_rotation.rs index ed3c9f604..f3aa7d750 100644 --- a/tests/platform/log_rotation.rs +++ b/crates/platform/tests/log_rotation.rs @@ -1,6 +1,6 @@ use std::{fs, io::Write}; -use codex_proxy_rs::platform::logging::rotation::{build_file_appender, RotationConfig}; +use codex_proxy_platform::logging::{build_file_appender, RotationConfig}; #[test] fn rolling_appender_should_write_daily_codex_log_file() { diff --git a/tests/platform/storage_schema.rs b/crates/platform/tests/storage_schema.rs similarity index 99% rename from tests/platform/storage_schema.rs rename to crates/platform/tests/storage_schema.rs index 5ad1fd72a..e11962c62 100644 --- a/tests/platform/storage_schema.rs +++ b/crates/platform/tests/storage_schema.rs @@ -1,9 +1,8 @@ use std::str::FromStr; +use codex_proxy_platform::storage::connect_sqlite; use sqlx::{sqlite::SqliteConnectOptions, SqlitePool}; -use codex_proxy_rs::platform::storage::db::connect_sqlite; - #[tokio::test] async fn sqlite_schema_creates_accounts_and_event_tables() { let dir = tempfile::tempdir().unwrap(); diff --git a/crates/runtime/Cargo.toml b/crates/runtime/Cargo.toml new file mode 100644 index 000000000..129389af6 --- /dev/null +++ b/crates/runtime/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "codex-proxy-runtime" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +rust-version.workspace = true + +[dependencies] +base64.workspace = true +bytes.workspace = true +chrono.workspace = true +codex-proxy-adapters = { path = "../adapters" } +codex-proxy-core = { path = "../core" } +codex-proxy-platform = { path = "../platform" } +futures.workspace = true +reqwest.workspace = true +secrecy.workspace = true +serde_json.workspace = true +sqlx.workspace = true +thiserror.workspace = true +tokio.workspace = true +tracing.workspace = true +uuid.workspace = true + +[dev-dependencies] +async-trait.workspace = true +reqwest.workspace = true +tempfile = "3.27.0" +tokio.workspace = true +tokio-tungstenite.workspace = true +wiremock = "0.6.5" + +[lints] +workspace = true diff --git a/crates/runtime/src/bootstrap.rs b/crates/runtime/src/bootstrap.rs new file mode 100644 index 000000000..669ce8743 --- /dev/null +++ b/crates/runtime/src/bootstrap.rs @@ -0,0 +1,12 @@ +//! 启动编排。 + +use codex_proxy_adapters::codex::fingerprint::FingerprintRepository; +use codex_proxy_core::gateway::fingerprint::Fingerprint; + +/// 从指纹仓储加载运行时指纹，必要时回退到默认值。 +pub async fn load_runtime_fingerprint(repository: &FingerprintRepository) -> Fingerprint { + match repository.load_latest_auto_updated().await { + Ok(Some(fingerprint)) => fingerprint, + _ => Fingerprint::default_codex_desktop(), + } +} diff --git a/crates/runtime/src/config.rs b/crates/runtime/src/config.rs new file mode 100644 index 000000000..52d6c9503 --- /dev/null +++ b/crates/runtime/src/config.rs @@ -0,0 +1,27 @@ +//! 配置映射。 + +use codex_proxy_platform::config as platform; + +/// 运行时配置视图。 +#[derive(Debug, Clone)] +pub struct RuntimeConfig { + /// 模型配置。 + pub model: platform::ModelConfig, + /// 管理端配置。 + pub admin: platform::AdminConfig, + /// 认证与刷新配置。 + pub auth: platform::AuthConfig, + /// 配额刷新配置。 + pub quota: platform::QuotaConfig, +} + +impl From for RuntimeConfig { + fn from(config: platform::AppConfig) -> Self { + Self { + model: config.model, + admin: config.admin, + auth: config.auth, + quota: config.quota, + } + } +} diff --git a/crates/runtime/src/lib.rs b/crates/runtime/src/lib.rs new file mode 100644 index 000000000..083649726 --- /dev/null +++ b/crates/runtime/src/lib.rs @@ -0,0 +1,9 @@ +//! 运行时组装层，提供依赖注入、应用状态和后台任务协调。 + +pub mod bootstrap; +pub mod config; +pub mod repositories; +pub mod services; +pub mod state; +pub mod tasks; +pub mod upstream; diff --git a/crates/runtime/src/repositories.rs b/crates/runtime/src/repositories.rs new file mode 100644 index 000000000..5a2164fd1 --- /dev/null +++ b/crates/runtime/src/repositories.rs @@ -0,0 +1,54 @@ +//! 仓储组装。 + +use codex_proxy_adapters::{ + codex::fingerprint::FingerprintRepository, + sqlite::{ + accounts::SqliteAccountStore, admin_sessions::SqliteAdminSessionStore, + client_keys::SqliteClientKeyStore, cookies::SqliteCookieStore, events::SqliteEventLogStore, + models::ModelSnapshotRepository, refresh_leases::SqliteRefreshLeaseStore, + session_affinity::SqliteSessionAffinityStore, + }, +}; +use codex_proxy_platform::{crypto::SecretBox, identity::ApiKeyHasher, storage::SqlitePool}; + +/// 运行时仓储集合。 +#[derive(Clone)] +pub struct Repositories { + /// 账号存储。 + pub accounts: SqliteAccountStore, + /// 模型快照存储。 + pub model_snapshots: ModelSnapshotRepository, + /// 客户端 API Key 存储。 + pub client_keys: SqliteClientKeyStore, + /// 管理员会话存储。 + pub admin_sessions: SqliteAdminSessionStore, + /// 事件日志存储。 + pub event_logs: SqliteEventLogStore, + /// Cookie 存储。 + pub cookies: SqliteCookieStore, + /// 指纹存储。 + pub fingerprints: FingerprintRepository, + /// 会话亲和性存储。 + pub session_affinity: SqliteSessionAffinityStore, + /// 账号刷新租约存储。 + pub refresh_leases: SqliteRefreshLeaseStore, +} + +/// 从 SQLite 连接池构造仓储集合。 +pub fn sqlite_repositories( + pool: SqlitePool, + secret_box: SecretBox, + api_key_hasher: ApiKeyHasher, +) -> Repositories { + Repositories { + accounts: SqliteAccountStore::new(pool.clone(), secret_box.clone()), + model_snapshots: ModelSnapshotRepository::new(pool.clone()), + client_keys: SqliteClientKeyStore::new(pool.clone(), api_key_hasher), + admin_sessions: SqliteAdminSessionStore::new(pool.clone()), + event_logs: SqliteEventLogStore::new(pool.clone()), + cookies: SqliteCookieStore::new(pool.clone(), secret_box), + fingerprints: FingerprintRepository::new(pool.clone()), + session_affinity: SqliteSessionAffinityStore::new(pool.clone()), + refresh_leases: SqliteRefreshLeaseStore::new(pool), + } +} diff --git a/crates/runtime/src/services.rs b/crates/runtime/src/services.rs new file mode 100644 index 000000000..073f148c9 --- /dev/null +++ b/crates/runtime/src/services.rs @@ -0,0 +1,6592 @@ +//! 服务组装。 + +use std::{ + collections::BTreeMap, + path::PathBuf, + pin::Pin, + sync::{Arc, RwLock as StdRwLock}, + task::{Context, Poll}, + time::{Duration as StdDuration, Instant}, +}; + +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use bytes::Bytes; +use chrono::{DateTime, Duration, Utc}; +use codex_proxy_adapters::{ + codex::client::{ + CodexBackendClient, CodexBackendResponse, CodexBackendSseStream, + CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, + CodexCompactResponse, CodexRateLimitHeaderUpdates, CodexRequestContext, + CodexTurnStateUpdate, + }, + codex::fingerprint::FingerprintRepository, + oauth::openai::default_openai_oauth_client, + sqlite::{ + accounts::{ + AccountClaimsUpdate, AccountQuotaSnapshot, AccountUsageListRecord, AccountUsageSummary, + NewAccount, SqliteAccountStore, StoredAccount, StoredAccountMetadata, + }, + admin_sessions::SqliteAdminSessionStore, + client_keys::{CreatedClientApiKey, SqliteClientKeyStore, StoredClientApiKey}, + cookies::SqliteCookieStore, + events::{EventLogFilter, SqliteEventLogStore}, + refresh_leases::SqliteRefreshLeaseStore, + session_affinity::{SqliteSessionAffinityStore, SqliteSessionAffinityStoreError}, + }, +}; +use codex_proxy_core::{ + accounts::{ + cloudflare::CloudflarePathBlockTracker, + model::{Account, AccountStatus}, + pool::{ + AccountAcquireRequest, AccountCapacitySummary, AccountPool, AccountPoolOptions, + AccountWindowUsageDelta, AcquiredAccount, RotationStrategy, + }, + ports::{AccountStore, AccountStoreError}, + }, + admin::{ + auth::AdminAuthService, + client_keys::ClientKeyService, + settings::{ + AdminQuotaWarningThresholds, AdminSettings, AdminSettingsPatch, SettingsService, + SettingsServiceError, + }, + }, + auth::{ + oauth::{ + DeviceCode, OAuthConfig, OAuthError, PkceLogin, PkceSessionStore, RefreshFailure, + TokenPair, + }, + ports::{OAuthClient, TokenRefresher}, + }, + events::{ + model::{EventLevel, EventLog}, + service::EventLogService, + }, + gateway::{ + conversation::{build_conversation_identity, ensure_prompt_cache_key}, + fingerprint::Fingerprint, + ports::CodexModelCatalogClient, + }, + models::{ + model::ModelConfig, + ports::ModelSnapshotStore, + service::{ModelRefreshResult, ModelService, ModelServiceError}, + }, + protocol::{ + codex::{ + events::{extract_sse_usage, parse_rate_limit_headers, rate_limit_quota, TokenUsage}, + responses::{CodexCompactRequest, CodexResponsesRequest}, + sse::{encode_sse_event, parse_sse_events, SseError}, + }, + openai::chat::chat_completion_from_codex_sse, + openai::responses::{ + completed_response_metadata, reconvert_responses_sse_event_tuple_values, + response_from_codex_sse, CollectedResponse, ResponsesSseFailure, + }, + }, + serving::{ + affinity::{ + compute_variant_hash, hash_instructions, prepare_variant_identity, + SessionAffinityEntry, SessionAffinityMap, + }, + fallback::{ + status_code_is_quota_exhausted, status_code_is_rate_limited, + status_code_is_transient_upstream, + }, + implicit_resume::{ + continuation_input_start, implicit_resume_allowed, ImplicitResumeSnapshot, + }, + quota::{quota_from_usage, quota_reached}, + reasoning_replay::ReasoningReplayCache, + recovery::status_code_allows_same_account_retry, + responses::apply_response_model_options, + }, + usage::service::UsageService, +}; +use codex_proxy_platform::{ + config::{AppConfig, ConfigWriteError, QuotaWarningThresholds}, + identity::{hash_admin_password, verify_admin_password}, + json::Page, +}; +use futures::{stream, Stream, StreamExt}; +use secrecy::{ExposeSecret, SecretString}; +use serde_json::{Map, Value}; +use thiserror::Error; +use tokio::{ + sync::{mpsc, oneshot, RwLock}, + time::sleep, +}; + +use crate::{repositories::Repositories, upstream}; + +#[derive(Clone, Copy)] +enum ExhaustedAccountClass { + QuotaExhausted, + RateLimited, + Expired, + CloudflareChallenge, + CloudflarePathBlocked, + ModelUnsupported, +} + +/// 运行时服务集合。 +#[derive(Clone)] +pub struct Services { + /// 模型目录服务。 + pub models: Arc, + /// 管理端模型服务。 + pub admin_models: Arc, + /// 账号存储。 + pub accounts: Arc, + /// 模型快照存储。 + pub model_snapshots: Arc, + /// 客户端 API Key 服务。 + pub client_keys: Arc, + /// 管理端客户端 API Key 服务。 + pub admin_client_keys: Arc, + /// 管理员会话服务。 + pub admin_sessions: Arc, + /// 管理端设置服务。 + pub settings: Arc, + /// 管理端账号服务。 + pub admin_accounts: Arc, + /// 管理端 OAuth 服务。 + pub admin_oauth: Arc, + /// 管理端日志服务。 + pub logs: Arc, + /// 管理端用量服务。 + pub usage: Arc, + /// 运行时账号池服务。 + pub account_pool: Arc, + /// OpenAI Chat Completions 调度服务。 + pub chat: Arc, + /// OpenAI Responses 调度服务。 + pub responses: Arc, + /// 运行时会话亲和性服务。 + pub session_affinity: Arc, + /// Codex 上游客户端。 + pub codex: Arc, + /// 当前运行时指纹。 + pub fingerprint: Fingerprint, + /// 当前运行时 Codex installation id。 + pub installation_id: Option, + /// 后台任务需要的具体存储适配器。 + pub background_tasks: BackgroundTaskStores, +} + +/// Codex upstream connectivity probe result. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RuntimeUpstreamProbe { + /// Probed target name. + pub target: &'static str, + /// Configured backend base URL. + pub backend_base_url: String, + /// Full endpoint URL. + pub endpoint: String, + /// Whether the endpoint responded at transport level. + pub reachable: bool, + /// HTTP status code, when available. + pub status_code: Option, + /// Authentication outcome inferred from status. + pub authorization: &'static str, +} + +/// 后台任务需要的具体存储适配器集合。 +#[derive(Clone)] +pub struct BackgroundTaskStores { + /// 账号存储。 + pub accounts: SqliteAccountStore, + /// 管理员会话存储。 + pub admin_sessions: SqliteAdminSessionStore, + /// Cookie 存储。 + pub cookies: SqliteCookieStore, + /// 指纹存储。 + pub fingerprints: FingerprintRepository, + /// 会话亲和性存储。 + pub session_affinity: SqliteSessionAffinityStore, + /// 账号刷新租约存储。 + pub refresh_leases: SqliteRefreshLeaseStore, +} + +impl Services { + /// 从仓储和上游适配器构造运行时服务集合。 + pub fn new(config: &AppConfig, repositories: Repositories, fingerprint: Fingerprint) -> Self { + Self::with_installation_id(config, repositories, fingerprint, None) + } + + /// 从仓储、上游适配器和 installation id 构造运行时服务集合。 + pub fn with_installation_id( + config: &AppConfig, + repositories: Repositories, + fingerprint: Fingerprint, + installation_id: Option, + ) -> Self { + Self::with_installation_id_and_local_config_path( + config, + repositories, + fingerprint, + installation_id, + "local.yaml", + ) + } + + /// 从仓储、上游适配器、installation id 和本地配置路径构造运行时服务集合。 + pub fn with_installation_id_and_local_config_path( + config: &AppConfig, + repositories: Repositories, + fingerprint: Fingerprint, + installation_id: Option, + local_config_path: impl Into, + ) -> Self { + let token_refresher: Arc = + Arc::new(default_openai_oauth_client(oauth_config(config))); + let oauth_client: Arc = + Arc::new(default_openai_oauth_client(oauth_config(config))); + Self::with_installation_id_local_config_path_and_oauth_clients( + config, + repositories, + fingerprint, + installation_id, + local_config_path, + token_refresher, + oauth_client, + ) + } + + /// 从仓储、上游适配器、installation id、本地配置路径和 token refresher 构造服务集合。 + pub fn with_installation_id_local_config_path_and_token_refresher( + config: &AppConfig, + repositories: Repositories, + fingerprint: Fingerprint, + installation_id: Option, + local_config_path: impl Into, + token_refresher: Arc, + ) -> Self { + let oauth_client: Arc = + Arc::new(default_openai_oauth_client(oauth_config(config))); + Self::with_installation_id_local_config_path_and_oauth_clients( + config, + repositories, + fingerprint, + installation_id, + local_config_path, + token_refresher, + oauth_client, + ) + } + + /// 从仓储、上游适配器、installation id、本地配置路径和 OAuth 端口构造服务集合。 + pub fn with_installation_id_local_config_path_and_oauth_clients( + config: &AppConfig, + repositories: Repositories, + fingerprint: Fingerprint, + installation_id: Option, + local_config_path: impl Into, + token_refresher: Arc, + oauth_client: Arc, + ) -> Self { + let installation_id = installation_id.filter(|id| !id.trim().is_empty()); + let runtime_fingerprint = fingerprint.clone(); + let model_config = ModelConfig { + default_model: config.model.default_model.clone(), + default_reasoning_effort: config.model.default_reasoning_effort.clone(), + service_tier: config.model.service_tier.clone(), + aliases: config.model.aliases.clone(), + }; + let account_store = repositories.accounts; + let cookies = repositories.cookies; + let background_tasks = BackgroundTaskStores { + accounts: account_store.clone(), + admin_sessions: repositories.admin_sessions.clone(), + cookies: cookies.clone(), + fingerprints: repositories.fingerprints.clone(), + session_affinity: repositories.session_affinity.clone(), + refresh_leases: repositories.refresh_leases.clone(), + }; + let accounts = Arc::new(account_store.clone()) as Arc; + let account_pool = Arc::new(RuntimeAccountPoolService::new( + accounts.clone(), + account_pool_options(config), + config.auth.request_interval_ms, + )); + let model_snapshots = Arc::new(repositories.model_snapshots) as Arc; + let model_catalog_client = + upstream::model_catalog_client(config.api.base_url.clone(), fingerprint.clone()) + as Arc; + let models = Arc::new(ModelService::new( + model_config, + Some(model_snapshots.clone()), + Some(model_catalog_client), + Some(Arc::new(tokio::sync::Mutex::new(BTreeMap::new()))), + )); + let client_key_store = repositories.client_keys; + let client_keys = Arc::new(ClientKeyService::new(Arc::new(client_key_store.clone()))); + let admin_client_keys = Arc::new(AdminClientKeyService::new(client_key_store)); + let admin_sessions = Arc::new(AdminSessionService::new( + repositories.admin_sessions, + config.admin.default_username.clone(), + config.admin.session_ttl_minutes, + )); + let settings = Arc::new(RuntimeSettingsService::with_local_config_path( + config.clone(), + local_config_path, + )); + let admin_models = Arc::new(AdminModelService::new( + models.clone(), + accounts.clone(), + installation_id.clone(), + )); + let codex = upstream::codex_backend_client( + config.api.base_url.clone(), + fingerprint, + &config.ws_pool, + ); + let admin_accounts = Arc::new(AdminAccountService::new( + account_store.clone(), + cookies.clone(), + config.quota.warning_thresholds.clone(), + codex.clone(), + account_pool.clone(), + token_refresher, + installation_id.clone(), + )); + let admin_oauth = Arc::new(AdminOAuthService::new(oauth_config(config), oauth_client)); + let logs = Arc::new(AdminLogService::new( + repositories.event_logs, + config.logging.enabled, + config.logging.capacity, + config.logging.capture_body, + )); + let session_affinity = Arc::new(RuntimeSessionAffinityService::new( + repositories.session_affinity, + )); + let usage = Arc::new(AdminUsageService::new(account_store)); + let cloudflare_recovery = + CloudflareRecovery::new(cookies, CloudflarePathBlockTracker::new()); + let chat = Arc::new(ChatDispatchService::new( + account_pool.clone(), + models.clone(), + codex.clone(), + installation_id.clone(), + cloudflare_recovery.clone(), + )); + let responses = Arc::new(ResponseDispatchService::new( + account_pool.clone(), + models.clone(), + codex.clone(), + session_affinity.clone(), + logs.clone(), + installation_id.clone(), + cloudflare_recovery, + )); + + Self { + models, + admin_models, + accounts, + model_snapshots, + client_keys, + admin_client_keys, + admin_sessions, + settings, + admin_accounts, + admin_oauth, + logs, + usage, + account_pool, + chat, + responses, + session_affinity, + codex, + fingerprint: runtime_fingerprint, + installation_id, + background_tasks, + } + } +} + +impl Services { + /// Probe the Codex models endpoint without exposing account secrets. + pub async fn probe_codex_models_endpoint(&self, request_id: &str) -> RuntimeUpstreamProbe { + let config = self.settings.current(); + let probe = self + .codex + .probe_models_endpoint(CodexRequestContext { + access_token: "", + account_id: None, + 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, + installation_id: self.installation_id.as_deref(), + session_id: None, + }) + .await; + + match probe { + Ok(probe) => RuntimeUpstreamProbe { + target: "codexModels", + backend_base_url: config.api.base_url.clone(), + endpoint: probe.endpoint, + reachable: true, + status_code: Some(probe.status.as_u16()), + authorization: if probe.status.as_u16() == 401 { + "rejected" + } else { + "unknown" + }, + }, + Err(_) => RuntimeUpstreamProbe { + target: "codexModels", + backend_base_url: config.api.base_url.clone(), + endpoint: format!( + "{}/codex/models?client_version={}", + config.api.base_url.trim_end_matches('/'), + self.fingerprint.app_version + ), + reachable: false, + status_code: None, + authorization: "unknown", + }, + } + } +} + +fn account_pool_options(config: &AppConfig) -> AccountPoolOptions { + AccountPoolOptions { + max_concurrent_per_account: config.auth.max_concurrent_per_account, + stale_slot_ttl: 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: BTreeMap::new(), + } +} + +/// 运行时账号池服务。 +#[derive(Clone)] +pub struct RuntimeAccountPoolService { + accounts: Arc, + pool: Arc>, + request_interval: StdDuration, +} + +impl RuntimeAccountPoolService { + /// 构造运行时账号池服务。 + pub fn new( + accounts: Arc, + options: AccountPoolOptions, + request_interval_ms: u64, + ) -> Self { + Self { + accounts, + pool: Arc::new(tokio::sync::Mutex::new(AccountPool::with_options(options))), + request_interval: StdDuration::from_millis(request_interval_ms), + } + } + + /// 从账号存储恢复账号池内容。 + pub async fn restore_from_repository(&self) -> Result { + let accounts = self.accounts.list_pool_accounts().await?; + let restored = accounts.len(); + let mut pool = self.pool.lock().await; + pool.clear(); + for account in accounts { + pool.insert(account); + } + Ok(restored) + } + + /// 从账号存储同步单个账号到运行时账号池；账号已不存在时从池中移除。 + pub async fn sync_account_from_repository( + &self, + account_id: &str, + ) -> Result { + let account = self.accounts.get_pool_account(account_id).await?; + let mut pool = self.pool.lock().await; + if let Some(account) = account { + pool.insert(account); + return Ok(true); + } + Ok(pool.remove(account_id)) + } + + /// 从运行时账号池移除账号。 + pub async fn remove_account(&self, account_id: &str) -> bool { + self.pool.lock().await.remove(account_id) + } + + /// 清空运行时账号池。 + pub async fn clear(&self) { + self.pool.lock().await.clear(); + } + + /// 从账号池获取指定模型可用账号。 + pub async fn acquire(&self, model: &str, now: DateTime) -> Option { + self.acquire_with(AccountAcquireRequest::new(model, now)) + .await + } + + /// 使用完整获取请求从账号池获取账号。 + pub async fn acquire_with(&self, request: AccountAcquireRequest) -> Option { + let acquired = self.pool.lock().await.acquire_with(request)?; + if let Err(error) = self.accounts.record_request(&acquired.account.id).await { + tracing::warn!( + account_id = acquired.account.id, + error = %error, + "failed to persist account request usage" + ); + } + Some(acquired) + } + + /// 等待同一账号前一个在途请求满足配置的发送间隔。 + pub async fn wait_for_request_interval(&self, acquired: &AcquiredAccount) { + if self.request_interval.is_zero() { + return; + } + let Some(previous_slot_at) = acquired.previous_slot_at else { + return; + }; + let elapsed = Utc::now() + .signed_duration_since(previous_slot_at) + .to_std() + .unwrap_or_default(); + if elapsed < self.request_interval { + sleep(self.request_interval - elapsed).await; + } + } + + /// 释放账号的一个在途槽位。 + pub async fn release(&self, account_id: &str) { + self.pool.lock().await.release(account_id); + } + + /// Return a snapshot of runtime account-pool capacity. + pub async fn capacity_summary(&self, now: DateTime) -> AccountCapacitySummary { + self.pool.lock().await.capacity_summary(now) + } + + /// Return a snapshot of runtime account-pool capacity using the current time. + pub async fn capacity_summary_now(&self) -> AccountCapacitySummary { + self.capacity_summary(Utc::now()).await + } + + /// 标记账号因配额限流进入冷却。 + pub async fn mark_quota_limited_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> bool { + let persisted = match self + .accounts + .mark_quota_limited_until(account_id, cooldown_until) + .await + { + Ok(persisted) => persisted, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to persist quota cooldown state" + ); + false + } + }; + let in_memory = self + .pool + .lock() + .await + .mark_quota_limited_until(account_id, cooldown_until); + persisted || in_memory + } + + /// 标记账号处于 Cloudflare 冷却期。 + pub async fn set_cloudflare_cooldown_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> bool { + let persisted = match self + .accounts + .set_cloudflare_cooldown_until(account_id, cooldown_until) + .await + { + Ok(persisted) => persisted, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to persist Cloudflare cooldown state" + ); + false + } + }; + let in_memory = self + .pool + .lock() + .await + .set_cloudflare_cooldown_until(account_id, cooldown_until); + persisted || in_memory + } + + /// 更新账号状态。 + pub async fn set_status(&self, account_id: &str, status: AccountStatus) -> bool { + let persisted = match self.accounts.set_status(account_id, status).await { + Ok(persisted) => persisted, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to persist account status" + ); + false + } + }; + let in_memory = self.pool.lock().await.set_status(account_id, status); + persisted || in_memory + } + + /// 清零运行时账号池中的累计和窗口用量。 + pub async fn reset_usage(&self, account_id: &str) -> bool { + self.pool.lock().await.reset_usage(account_id) + } + + /// 记录账号成功响应的 token 用量。 + pub async fn record_token_usage(&self, account_id: &str, usage: TokenUsage) { + self.record_response_usage(account_id, usage, false).await; + } + + /// 记录 Responses 成功响应的 token 与工具用量。 + pub async fn record_response_usage( + &self, + account_id: &str, + usage: TokenUsage, + image_generation_requested: bool, + ) { + let image_request_succeeded = image_generation_requested && usage.image_output_tokens > 0; + let image_request_failed = image_generation_requested && !image_request_succeeded; + let mut persisted_usage = UsageService::account_delta_from_token_usage(usage); + persisted_usage.image_requests = bool_to_u64(image_request_succeeded); + persisted_usage.image_request_failures = bool_to_u64(image_request_failed); + if let Err(error) = self + .accounts + .record_usage_delta(account_id, persisted_usage) + .await + { + tracing::warn!( + account_id, + error = %error, + "failed to persist account token usage" + ); + } + self.pool.lock().await.record_window_token_usage( + account_id, + AccountWindowUsageDelta { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_request_succeeded, + image_request_failed, + }, + ); + } + + /// 记录 Responses 空响应尝试。 + pub async fn record_empty_response_attempt( + &self, + account_id: &str, + image_generation_requested: bool, + ) { + let usage = codex_proxy_core::accounts::usage::AccountUsageDelta { + empty_responses: 1, + image_request_failures: bool_to_u64(image_generation_requested), + ..codex_proxy_core::accounts::usage::AccountUsageDelta::default() + }; + if let Err(error) = self.accounts.record_usage_delta(account_id, usage).await { + tracing::warn!( + account_id, + error = %error, + "failed to persist empty response usage" + ); + } + if image_generation_requested { + self.pool.lock().await.record_window_token_usage( + account_id, + AccountWindowUsageDelta { + image_request_failed: true, + ..AccountWindowUsageDelta::default() + }, + ); + } + } + + /// 将上游成功响应头里的 rate-limit 状态被动写回配额和窗口缓存。 + pub async fn sync_passive_rate_limit_headers( + &self, + account: &Account, + headers: &[(String, String)], + ) { + let Some(rate_limits) = parse_rate_limit_headers(headers) else { + return; + }; + let existing_quota = match self.accounts.get_quota_json(&account.id).await { + Ok(Some(quota_json)) => serde_json::from_str::("a_json).ok(), + Ok(None) => None, + Err(error) => { + tracing::warn!( + account_id = %account.id, + error = %error, + "failed to read existing quota json before passive rate-limit sync" + ); + None + } + }; + let quota = rate_limit_quota( + &rate_limits, + account.plan_type.as_deref(), + existing_quota.as_ref(), + ); + if let Err(error) = self + .accounts + .update_quota_json(&account.id, "a.to_string()) + .await + { + tracing::warn!( + account_id = %account.id, + error = %error, + "failed to persist passive rate-limit quota" + ); + } + + let Some(primary) = rate_limits.primary else { + return; + }; + let Some(reset_at) = primary.reset_at_datetime() else { + return; + }; + let limit_window_seconds = primary.limit_window_seconds(); + if let Err(error) = self + .accounts + .sync_rate_limit_window(&account.id, reset_at, limit_window_seconds) + .await + { + tracing::warn!( + account_id = %account.id, + error = %error, + "failed to persist passive rate-limit window" + ); + } + self.pool + .lock() + .await + .sync_rate_limit_window(&account.id, reset_at, limit_window_seconds); + let limit_reached = quota + .pointer("/rate_limit/limit_reached") + .and_then(Value::as_bool) + .unwrap_or(false); + if limit_reached { + self.mark_quota_limited_until(&account.id, reset_at).await; + } + } +} + +/// 运行时账号池错误。 +#[derive(Debug, Error)] +pub enum RuntimeAccountPoolError { + /// 账号存储访问失败。 + #[error("account store error: {0}")] + Store(#[from] AccountStoreError), +} + +const CLOUDFLARE_CHALLENGE_COOLDOWN_SECONDS: i64 = 10; + +#[derive(Clone)] +pub(crate) struct CloudflareRecovery { + cookies: SqliteCookieStore, + path_block_tracker: CloudflarePathBlockTracker, +} + +impl CloudflareRecovery { + fn new(cookies: SqliteCookieStore, path_block_tracker: CloudflarePathBlockTracker) -> Self { + Self { + cookies, + path_block_tracker, + } + } + + async fn cookie_header_for_request( + &self, + account_id: &str, + request_path: &str, + ) -> Option { + match self + .cookies + .cookie_header_for_request(account_id, "chatgpt.com", request_path) + .await + { + Ok(cookie_header) => cookie_header, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to read account cookies for upstream request" + ); + None + } + } + } + + async fn capture_set_cookie_headers(&self, account_id: &str, headers: &[String]) { + for header in headers { + if let Err(error) = self.cookies.capture_set_cookie(account_id, header).await { + tracing::warn!( + account_id, + error = %error, + "failed to persist upstream set-cookie header" + ); + } + } + } + + async fn apply_challenge(&self, account_pool: &RuntimeAccountPoolService, account_id: &str) { + self.delete_account_cookies(account_id, "Cloudflare challenge") + .await; + let cooldown_until = Utc::now() + Duration::seconds(CLOUDFLARE_CHALLENGE_COOLDOWN_SECONDS); + account_pool + .set_cloudflare_cooldown_until(account_id, cooldown_until) + .await; + } + + async fn apply_path_block(&self, account_pool: &RuntimeAccountPoolService, account_id: &str) { + self.delete_account_cookies(account_id, "Cloudflare path-block") + .await; + let now = Utc::now(); + let count = self + .path_block_tracker + .record_path_block(account_id, now) + .await; + if self + .path_block_tracker + .should_disable(account_id, now) + .await + { + account_pool + .set_status(account_id, AccountStatus::Disabled) + .await; + } + tracing::warn!( + account_id, + path_block_count = count, + "upstream returned Cloudflare path-block" + ); + } + + async fn reset_path_block(&self, account_id: &str) { + self.path_block_tracker.reset(account_id).await; + } + + async fn delete_account_cookies(&self, account_id: &str, reason: &str) { + if let Err(error) = self.cookies.delete_account_cookies(account_id).await { + tracing::warn!( + account_id, + reason, + error = %error, + "failed to delete account cookies after Cloudflare recovery signal" + ); + } + } +} + +/// 默认会话亲和性 TTL 秒数。 +const DEFAULT_SESSION_AFFINITY_TTL_SECS: i64 = 4 * 60 * 60; + +/// 运行时会话亲和性服务。 +#[derive(Clone)] +pub struct RuntimeSessionAffinityService { + store: SqliteSessionAffinityStore, + map: Arc>, + ttl: Duration, +} + +/// 默认 reasoning replay TTL 秒数。 +const DEFAULT_REASONING_REPLAY_TTL_SECS: i64 = 55 * 60; +const IMPLICIT_RESUME_MAX_AGE_SECS: i64 = DEFAULT_REASONING_REPLAY_TTL_SECS; + +impl RuntimeSessionAffinityService { + /// 构造运行时会话亲和性服务。 + pub fn new(store: SqliteSessionAffinityStore) -> Self { + let ttl = Duration::seconds(DEFAULT_SESSION_AFFINITY_TTL_SECS); + Self { + store, + map: Arc::new(tokio::sync::RwLock::new(SessionAffinityMap::new(ttl))), + ttl, + } + } + + /// 从 SQLite 恢复未过期的会话亲和性记录。 + pub async fn restore_from_repository( + &self, + now: DateTime, + ) -> Result { + let records = self.store.list_active(now).await?; + Ok(self.map.write().await.restore(records, now)) + } + + /// 记录并持久化响应 ID 的亲和性条目。 + pub async fn record( + &self, + response_id: String, + entry: SessionAffinityEntry, + ) -> Result<(), RuntimeSessionAffinityError> { + self.store.upsert(&response_id, &entry, self.ttl).await?; + self.map.write().await.record(response_id, entry); + Ok(()) + } + + /// 根据响应 ID 查找账号 ID。 + pub async fn lookup_account(&self, response_id: &str, now: DateTime) -> Option { + self.map.read().await.lookup_account(response_id, now) + } + + /// 根据响应 ID 查找对话 ID。 + pub async fn lookup_conversation_id( + &self, + response_id: &str, + now: DateTime, + ) -> Option { + self.map + .read() + .await + .lookup_conversation_id(response_id, now) + } + + /// 根据响应 ID 查找 turn state。 + pub async fn lookup_turn_state(&self, response_id: &str, now: DateTime) -> Option { + self.map.read().await.lookup_turn_state(response_id, now) + } + + /// 根据响应 ID 查找指令哈希。 + pub async fn lookup_instructions_hash( + &self, + response_id: &str, + now: DateTime, + ) -> Option { + self.map + .read() + .await + .lookup_instructions_hash(response_id, now) + } + + /// 根据响应 ID 查找函数调用 ID 列表。 + pub async fn lookup_function_call_ids( + &self, + response_id: &str, + now: DateTime, + ) -> Vec { + self.map + .read() + .await + .lookup_function_call_ids(response_id, now) + } + + /// 查找指定对话和变体下最新的响应 ID。 + pub async fn lookup_latest_response_by_conversation( + &self, + conversation_id: &str, + max_age: Option, + variant_hash: Option<&str>, + now: DateTime, + ) -> Option { + self.map + .read() + .await + .lookup_latest_response_by_conversation(conversation_id, max_age, variant_hash, now) + } + + /// 删除响应 ID 的内存亲和性映射。 + pub async fn forget(&self, response_id: &str) -> bool { + self.map.write().await.forget(response_id) + } +} + +/// 运行时会话亲和性错误。 +#[derive(Debug, Error)] +pub enum RuntimeSessionAffinityError { + /// 存储访问失败。 + #[error("session affinity store error: {0}")] + Store(#[from] SqliteSessionAffinityStoreError), +} + +pub(crate) fn oauth_config(config: &AppConfig) -> OAuthConfig { + OAuthConfig { + client_id: config.auth.oauth_client_id.clone(), + auth_endpoint: config.auth.oauth_auth_endpoint.clone(), + device_code_endpoint: oauth_device_code_endpoint(&config.auth.oauth_token_endpoint), + token_endpoint: config.auth.oauth_token_endpoint.clone(), + } +} + +fn oauth_device_code_endpoint(token_endpoint: &str) -> String { + token_endpoint + .strip_suffix("/token") + .map(|prefix| format!("{prefix}/device/code")) + .unwrap_or_else(|| "https://auth.openai.com/oauth/device/code".to_string()) +} + +/// 管理端模型服务。 +#[derive(Clone)] +pub struct AdminModelService { + models: Arc, + accounts: Arc, + installation_id: Option, +} + +impl AdminModelService { + /// 构造管理端模型服务。 + pub fn new( + models: Arc, + accounts: Arc, + installation_id: Option, + ) -> Self { + Self { + models, + accounts, + installation_id, + } + } + + /// 用当前活跃账号刷新上游模型快照。 + pub async fn refresh_backend_models( + &self, + request_id: &str, + ) -> Result { + let accounts = self + .accounts + .list_pool_accounts() + .await + .map_err(|_| AdminModelError::ListAccounts)?; + self.models + .refresh_backend_models_with_installation_id( + &accounts, + request_id, + self.installation_id.as_deref(), + ) + .await + .map_err(AdminModelError::from) + } +} + +/// 管理端模型错误。 +#[derive(Debug, Error)] +pub enum AdminModelError { + /// 列出账号失败。 + #[error("failed to list accounts")] + ListAccounts, + /// 没有可用账号。 + #[error("no active accounts available for model refresh")] + NoAccounts, + /// 模型快照存储不可用。 + #[error("model snapshot store is unavailable")] + SnapshotStoreUnavailable, + /// 上游客户端不可用。 + #[error("model upstream client is unavailable")] + UpstreamClientUnavailable, + /// 写入快照失败。 + #[error("failed to store model snapshot")] + StoreSnapshot, + /// 读取快照失败。 + #[error("failed to load model snapshots")] + LoadSnapshots, + /// 所有计划刷新失败。 + #[error("all model refresh plans failed")] + AllPlansFailed(ModelRefreshResult), +} + +impl From for AdminModelError { + fn from(error: ModelServiceError) -> Self { + match error { + ModelServiceError::SnapshotStoreUnavailable => Self::SnapshotStoreUnavailable, + ModelServiceError::UpstreamClientUnavailable => Self::UpstreamClientUnavailable, + ModelServiceError::NoAccounts => Self::NoAccounts, + ModelServiceError::StoreSnapshot => Self::StoreSnapshot, + ModelServiceError::LoadSnapshots => Self::LoadSnapshots, + ModelServiceError::AllPlansFailed(result) => Self::AllPlansFailed(result), + } + } +} + +/// 管理端 OAuth 服务。 +#[derive(Clone)] +pub struct AdminOAuthService { + config: OAuthConfig, + client: Arc, + sessions: Arc>, +} + +impl AdminOAuthService { + /// 构造管理端 OAuth 服务。 + pub fn new(config: OAuthConfig, client: Arc) -> Self { + Self { + config, + client, + sessions: Arc::new(tokio::sync::Mutex::new(PkceSessionStore::default())), + } + } + + /// 开始 PKCE 登录。 + pub async fn start_pkce_login(&self, return_host: &str) -> PkceLogin { + self.sessions + .lock() + .await + .start_login(return_host, &self.config) + } + + /// 请求设备码登录信息。 + pub async fn request_device_code(&self) -> Result { + self.client + .request_device_code() + .await + .map_err(AdminOAuthError::OAuth) + } + + /// 轮询设备码 token。 + pub async fn poll_device_token( + &self, + device_code: &str, + ) -> Result { + match self.client.poll_device_token(device_code).await { + Ok(tokens) => Ok(AdminDevicePoll::Authorized(tokens)), + Err(error) => { + if let Some(code) = error.pending_code() { + Ok(AdminDevicePoll::Pending { code }) + } else { + Err(AdminOAuthError::OAuth(error)) + } + } + } + } + + /// 使用回调 code/state 完成 PKCE token 交换。 + pub async fn exchange_callback( + &self, + code: &str, + state: &str, + ) -> Result { + let session = self + .sessions + .lock() + .await + .try_acquire(state) + .ok_or(AdminOAuthError::InvalidState)?; + + match self + .client + .exchange_code(code, &session.code_verifier, &session.redirect_uri) + .await + { + Ok(tokens) => { + self.sessions.lock().await.complete(state); + Ok(AdminOAuthCallback { + tokens, + return_host: session.return_host, + }) + } + Err(error) => { + self.sessions.lock().await.release(state); + Err(AdminOAuthError::OAuth(error)) + } + } + } +} + +/// 设备码轮询结果。 +#[derive(Debug, Clone)] +pub enum AdminDevicePoll { + /// 授权还未完成。 + Pending { + /// OAuth 标准 pending 错误码。 + code: &'static str, + }, + /// 已换取 token。 + Authorized(TokenPair), +} + +/// PKCE 回调换取的 token 和返回 host。 +#[derive(Debug, Clone)] +pub struct AdminOAuthCallback { + /// OAuth token 对。 + pub tokens: TokenPair, + /// 登录前的管理端 host。 + pub return_host: String, +} + +/// 管理端 OAuth 错误。 +#[derive(Debug, Error)] +pub enum AdminOAuthError { + /// callback URL 或 query 缺少必需字段。 + #[error("invalid OAuth callback")] + InvalidCallback, + /// OAuth state 不存在、过期或正在处理。 + #[error("invalid OAuth state")] + InvalidState, + /// OAuth 上游错误。 + #[error("{0}")] + OAuth(OAuthError), +} + +/// 管理端用量服务。 +#[derive(Clone)] +pub struct AdminUsageService { + store: SqliteAccountStore, +} + +impl AdminUsageService { + /// 构造管理端用量服务。 + pub fn new(store: SqliteAccountStore) -> Self { + Self { store } + } + + /// 分页列出账号用量统计。 + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> Result, AdminUsageError> { + let page = self + .store + .list_usage(cursor, limit) + .await + .map_err(|_| AdminUsageError::List)?; + Ok(Page { + items: page.items.into_iter().map(AdminUsageRecord::from).collect(), + next_cursor: page.next_cursor, + }) + } + + /// 汇总账号用量统计。 + pub async fn summary(&self) -> Result { + self.store + .usage_summary() + .await + .map(AdminUsageSummary::from) + .map_err(|_| AdminUsageError::Summary) + } +} + +/// 管理端用量错误。 +#[derive(Debug, Error)] +pub enum AdminUsageError { + /// 列表失败。 + #[error("failed to list account usage")] + List, + /// 汇总失败。 + #[error("failed to summarize account usage")] + Summary, +} + +/// 管理端用量记录。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminUsageRecord { + /// 账号 ID。 + pub account_id: String, + /// 邮箱。 + pub email: Option, + /// 管理标签。 + pub label: Option, + /// 计划类型。 + pub plan_type: Option, + /// 请求数。 + pub request_count: i64, + /// 空响应数。 + pub empty_response_count: i64, + /// 输入 token 数。 + pub input_tokens: i64, + /// 输出 token 数。 + pub output_tokens: i64, + /// 缓存 token 数。 + pub cached_tokens: i64, + /// 图片输入 token 数。 + pub image_input_tokens: i64, + /// 图片输出 token 数。 + pub image_output_tokens: i64, + /// 图片请求数。 + pub image_request_count: i64, + /// 图片请求失败数。 + pub image_request_failed_count: i64, + /// 最近使用时间。 + pub last_used_at: Option>, +} + +/// 管理端用量汇总。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AdminUsageSummary { + /// 有用量记录的账号数。 + pub account_count: i64, + /// 请求总数。 + pub request_count: i64, + /// 空响应总数。 + pub empty_response_count: i64, + /// 输入 token 总数。 + pub input_tokens: i64, + /// 输出 token 总数。 + pub output_tokens: i64, + /// 缓存 token 总数。 + pub cached_tokens: i64, + /// 图片输入 token 总数。 + pub image_input_tokens: i64, + /// 图片输出 token 总数。 + pub image_output_tokens: i64, + /// 图片请求总数。 + pub image_request_count: i64, + /// 图片请求失败总数。 + pub image_request_failed_count: i64, +} + +/// OpenAI Chat Completions 调度服务。 +#[derive(Clone)] +pub struct ChatDispatchService { + account_pool: Arc, + models: Arc, + codex: Arc, + installation_id: Option, + cloudflare: CloudflareRecovery, +} + +impl ChatDispatchService { + /// 构造 Chat Completions 调度服务。 + pub(crate) fn new( + account_pool: Arc, + models: Arc, + codex: Arc, + installation_id: Option, + cloudflare: CloudflareRecovery, + ) -> Self { + Self { + account_pool, + models, + codex, + installation_id, + cloudflare, + } + } + + /// 调度非流式 Chat Completions 请求到 Codex Responses 上游。 + pub async fn complete( + &self, + request_id: &str, + mut request: CodexResponsesRequest, + requested_model: &str, + ) -> Result { + let catalog = self.models.catalog().await; + let parsed_model = catalog.parse_model_name(requested_model); + let display_model = + codex_proxy_core::models::catalog::ModelCatalog::build_display_model_name( + &parsed_model, + ); + request.model = parsed_model.model_id; + let include_reasoning = request.reasoning.is_some(); + let tuple_schema = request.tuple_schema.clone(); + let mut excluded_account_ids = Vec::new(); + let mut rate_limited_count = 0usize; + let mut last_rate_limit_error = None; + let mut quota_exhausted_count = 0usize; + let mut last_quota_error = None; + let mut expired_count = 0usize; + let mut last_auth_error = None; + let mut cloudflare_challenge_count = 0usize; + let mut last_cloudflare_challenge_error = None; + let mut cloudflare_path_block_count = 0usize; + let mut last_cloudflare_path_block_error = None; + let mut model_unsupported_count = 0usize; + let mut last_model_unsupported_error = None; + let mut model_unsupported_retry_used = false; + let (account_id, response) = loop { + let acquire_request = AccountAcquireRequest::new(&request.model, Utc::now()) + .with_exclude_account_ids(excluded_account_ids.iter().cloned()); + let acquired = match self.account_pool.acquire_with(acquire_request).await { + Some(acquired) => acquired, + None if quota_exhausted_count > 0 => { + return Err(ChatDispatchError::QuotaExhausted { + count: quota_exhausted_count, + upstream_error: last_quota_error.unwrap_or_default(), + }); + } + None if rate_limited_count > 0 => { + return Err(ChatDispatchError::RateLimited { + count: rate_limited_count, + upstream_error: last_rate_limit_error.unwrap_or_default(), + }); + } + None if expired_count > 0 => { + return Err(ChatDispatchError::Expired { + count: expired_count, + upstream_error: last_auth_error.unwrap_or_default(), + }); + } + None if cloudflare_challenge_count > 0 => { + return Err(ChatDispatchError::CloudflareChallenge { + count: cloudflare_challenge_count, + upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), + }); + } + None if cloudflare_path_block_count > 0 => { + return Err(ChatDispatchError::CloudflarePathBlocked { + count: cloudflare_path_block_count, + upstream_error: last_cloudflare_path_block_error.unwrap_or_default(), + }); + } + None if model_unsupported_count > 0 => { + return Err(ChatDispatchError::ModelUnsupported { + count: model_unsupported_count, + upstream_error: last_model_unsupported_error.unwrap_or_default(), + }); + } + None => return Err(ChatDispatchError::NoActiveAccount), + }; + self.account_pool.wait_for_request_interval(&acquired).await; + let account = acquired.account; + let release_account_id = account.id.clone(); + let response_result = create_response_with_account( + &self.codex, + self.installation_id.as_deref(), + &self.cloudflare, + &request, + request_id, + &account, + ) + .await; + self.account_pool.release(&release_account_id).await; + + match response_result { + Ok(response) => break (release_account_id, response), + Err(error) if is_rate_limit_upstream_error(&error) => { + rate_limited_count += 1; + last_rate_limit_error = Some(upstream_error_body(&error)); + let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); + self.account_pool + .mark_quota_limited_until(&release_account_id, cooldown_until) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_quota_exhausted_upstream_error(&error) => { + quota_exhausted_count += 1; + last_quota_error = Some(upstream_error_body(&error)); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_auth_upstream_error(&error) => { + expired_count += 1; + last_auth_error = Some(upstream_error_body(&error)); + let account_status = auth_failure_account_status(&error); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_challenge_upstream_error(&error) => { + cloudflare_challenge_count += 1; + last_cloudflare_challenge_error = + Some(cloudflare_challenge_error_message().to_string()); + self.cloudflare + .apply_challenge(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_path_block_upstream_error(&error) => { + cloudflare_path_block_count += 1; + last_cloudflare_path_block_error = + Some(cloudflare_path_block_error_message().to_string()); + self.cloudflare + .apply_path_block(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_model_unsupported_upstream_error(&error) => { + let upstream_error = upstream_error_body(&error); + if model_unsupported_retry_used { + return Err(ChatDispatchError::ModelUnsupported { + count: model_unsupported_count + 1, + upstream_error, + }); + } + model_unsupported_count += 1; + last_model_unsupported_error = Some(upstream_error); + model_unsupported_retry_used = true; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_banned_upstream_error(&error) => { + self.account_pool + .set_status(&release_account_id, AccountStatus::Banned) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) => return Err(ChatDispatchError::Upstream(error)), + } + }; + let body = chat_completion_from_codex_sse( + &response.body, + &display_model, + include_reasoning, + tuple_schema.as_ref(), + ) + .map_err(ChatDispatchError::InvalidSse)? + .ok_or(ChatDispatchError::EmptyUpstreamResponse)?; + self.cloudflare.reset_path_block(&account_id).await; + if let Some(usage) = response.usage { + self.account_pool + .record_token_usage(&account_id, usage) + .await; + } + Ok(body) + } +} + +/// Chat Completions 调度错误。 +#[derive(Debug, Error)] +pub enum ChatDispatchError { + /// 账号存储失败。 + #[error("failed to list runtime accounts")] + AccountStore, + /// 没有活跃账号。 + #[error("no active account is available")] + NoActiveAccount, + /// 所有账号都因配额耗尽不可用。 + #[error("all accounts exhausted by quota")] + QuotaExhausted { + /// 配额耗尽账号数量。 + count: usize, + /// 最后一个上游错误体。 + upstream_error: String, + }, + /// 所有账号都因限流不可用。 + #[error("all accounts exhausted by rate limit")] + RateLimited { + /// 限流账号数量。 + count: usize, + /// 最后一个上游错误体。 + upstream_error: String, + }, + /// 所有账号都因认证失效不可用。 + #[error("all accounts exhausted by expired auth")] + Expired { + /// 认证失效账号数量。 + count: usize, + /// 最后一个上游错误体。 + upstream_error: String, + }, + /// 所有账号都因 Cloudflare challenge 暂不可用。 + #[error("all accounts exhausted by Cloudflare challenge")] + CloudflareChallenge { + /// Cloudflare challenge 账号数量。 + count: usize, + /// 最后一个上游错误说明。 + upstream_error: String, + }, + /// 所有账号都因 Cloudflare path-block 暂不可用。 + #[error("all accounts exhausted by Cloudflare path-block")] + CloudflarePathBlocked { + /// Cloudflare path-block 账号数量。 + count: usize, + /// 最后一个上游错误说明。 + upstream_error: String, + }, + /// 所有账号都不支持请求模型。 + #[error("all accounts do not support the requested model")] + ModelUnsupported { + /// 不支持模型的账号数量。 + count: usize, + /// 最后一个上游错误体。 + upstream_error: String, + }, + /// 上游请求失败。 + #[error("upstream request failed: {0}")] + Upstream(#[from] CodexClientError), + /// 上游 SSE 无法解析。 + #[error("invalid upstream SSE response: {0}")] + InvalidSse(#[from] SseError), + /// 上游响应没有完成事件。 + #[error("upstream response did not include response.completed")] + EmptyUpstreamResponse, +} + +/// OpenAI Responses 调度服务。 +#[derive(Clone)] +pub struct ResponseDispatchService { + account_pool: Arc, + models: Arc, + codex: Arc, + session_affinity: Arc, + reasoning_replay: Arc>, + logs: Arc, + installation_id: Option, + cloudflare: CloudflareRecovery, +} + +/// Responses live SSE 响应体流。 +pub type ResponseBodyStream = + Pin> + Send + 'static>>; + +/// Responses live SSE 调度结果。 +pub struct ResponseDispatchStream { + /// 可直接转为 HTTP body 的 SSE 字节流。 + pub body: ResponseBodyStream, +} + +struct MpscResponseBodyStream { + receiver: mpsc::Receiver>, + cancel: Option>, +} + +impl Drop for MpscResponseBodyStream { + fn drop(&mut self) { + if let Some(cancel) = self.cancel.take() { + let _ = cancel.send(()); + } + } +} + +impl Stream for MpscResponseBodyStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.receiver.poll_recv(cx) + } +} + +impl ResponseDispatchService { + /// 构造 Responses 调度服务。 + pub(crate) fn new( + account_pool: Arc, + models: Arc, + codex: Arc, + session_affinity: Arc, + logs: Arc, + installation_id: Option, + cloudflare: CloudflareRecovery, + ) -> Self { + Self { + account_pool, + models, + codex, + session_affinity, + reasoning_replay: Arc::new(tokio::sync::Mutex::new(ReasoningReplayCache::new( + Duration::seconds(DEFAULT_REASONING_REPLAY_TTL_SECS), + ))), + logs, + installation_id, + cloudflare, + } + } + + async fn prepare_response_session( + &self, + request: &mut CodexResponsesRequest, + ) -> Option { + prepare_variant_identity(request); + if let Some(previous_response_id) = request.previous_response_id.clone() { + if request.prompt_cache_key.is_none() { + request.prompt_cache_key = self + .session_affinity + .lookup_conversation_id(&previous_response_id, Utc::now()) + .await; + } + if request.turn_state.as_deref().is_none_or(str::is_empty) { + request.turn_state = self + .session_affinity + .lookup_turn_state(&previous_response_id, Utc::now()) + .await; + } + ensure_prompt_cache_key(request); + return None; + } + + ensure_prompt_cache_key(request); + self.apply_implicit_resume(request).await + } + + async fn apply_implicit_resume( + &self, + request: &mut CodexResponsesRequest, + ) -> Option { + let continuation_start = continuation_input_start(&request.input); + if continuation_start == 0 || continuation_start >= request.input.len() { + return None; + } + let conversation_id = request + .prompt_cache_key + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty())? + .to_string(); + let snapshot = ImplicitResumeSnapshot::capture(request); + let variant_hash = compute_variant_hash(request); + let now = Utc::now(); + let previous_response_id = self + .session_affinity + .lookup_latest_response_by_conversation( + &conversation_id, + Some(Duration::seconds(IMPLICIT_RESUME_MAX_AGE_SECS)), + Some(&variant_hash), + now, + ) + .await?; + let current_instructions_hash = hash_instructions(Some(&request.instructions)); + if self + .session_affinity + .lookup_instructions_hash(&previous_response_id, now) + .await + .as_deref() + != Some(current_instructions_hash.as_str()) + { + return None; + } + let stored_function_call_ids = self + .session_affinity + .lookup_function_call_ids(&previous_response_id, now) + .await; + if !implicit_resume_allowed( + &request.input[continuation_start..], + &request.input, + &stored_function_call_ids, + ) { + return None; + } + let account_id = self + .session_affinity + .lookup_account(&previous_response_id, now) + .await?; + let replay_items = self.reasoning_replay.lock().await.lookup( + &previous_response_id, + &account_id, + &conversation_id, + &variant_hash, + now, + ); + let continuation = request.input[continuation_start..].to_vec(); + let mut input = replay_items; + input.extend(continuation); + + request.previous_response_id = Some(previous_response_id.clone()); + request.use_websocket = true; + request.force_http_sse = false; + request.input = input; + if let Some(turn_state) = self + .session_affinity + .lookup_turn_state(&previous_response_id, now) + .await + { + request.turn_state = Some(turn_state); + } + + Some(snapshot) + } + + async fn preferred_account_id_for_request( + &self, + request: &CodexResponsesRequest, + now: DateTime, + ) -> Option { + let previous_response_id = request.previous_response_id.as_deref()?; + self.session_affinity + .lookup_account(previous_response_id, now) + .await + } + + async fn recover_request_history( + &self, + request: &mut CodexResponsesRequest, + implicit_resume: &mut Option, + ) { + if let Some(previous_response_id) = request.previous_response_id.as_deref() { + self.session_affinity.forget(previous_response_id).await; + } + if let Some(snapshot) = implicit_resume.take() { + snapshot.restore(request); + request.previous_response_id = None; + request.turn_state = None; + request.turn_metadata = None; + } else { + strip_request_history(request); + } + } + + async fn evict_reasoning_replay(&self, request: &CodexResponsesRequest, account_id: &str) { + evict_reasoning_replay(&self.reasoning_replay, request, account_id).await; + } + + async fn record_response_affinity( + &self, + request: &CodexResponsesRequest, + account_id: &str, + body: &str, + turn_state: Option, + usage: Option, + ) { + record_response_affinity( + &self.session_affinity, + &self.reasoning_replay, + request, + account_id, + body, + turn_state, + usage, + ) + .await; + } + + /// 调度非流式 Responses 请求到 Codex Responses 上游。 + pub async fn complete( + &self, + request_id: &str, + mut request: CodexResponsesRequest, + requested_model: &str, + ) -> Result { + let started_at = Instant::now(); + let catalog = self.models.catalog().await; + let parsed_model = catalog.parse_model_name(requested_model); + apply_response_model_options(&mut request, &parsed_model, self.models.config()); + let tuple_schema = request.tuple_schema.clone(); + let image_generation_requested = request.expects_image_generation(); + let now = Utc::now(); + let mut implicit_resume = self.prepare_response_session(&mut request).await; + let preferred_account_id = self.preferred_account_id_for_request(&request, now).await; + let mut acquire_request = AccountAcquireRequest::new(&request.model, now); + if let Some(preferred_account_id) = preferred_account_id { + acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); + } + let mut excluded_account_ids = Vec::new(); + let mut rate_limited_count = 0usize; + let mut last_rate_limit_error = None; + let mut quota_exhausted_count = 0usize; + let mut last_quota_error = None; + let mut expired_count = 0usize; + let mut last_auth_error = None; + let mut cloudflare_challenge_count = 0usize; + let mut last_cloudflare_challenge_error = None; + let mut cloudflare_path_block_count = 0usize; + let mut last_cloudflare_path_block_error = None; + let mut model_unsupported_count = 0usize; + let mut last_model_unsupported_error = None; + let mut model_unsupported_retry_used = false; + let mut history_recovery_used = false; + let mut last_exhausted_account_class = None; + let mut empty_response_retries = 0u8; + const MAX_EMPTY_RESPONSE_RETRIES: u8 = 2; + let (account, response, collected_response) = loop { + let mut attempt_acquire_request = acquire_request + .clone() + .with_exclude_account_ids(excluded_account_ids.iter().cloned()); + attempt_acquire_request.now = Utc::now(); + let acquired = match self + .account_pool + .acquire_with(attempt_acquire_request) + .await + { + Some(acquired) => acquired, + None => match last_exhausted_account_class { + Some(ExhaustedAccountClass::QuotaExhausted) => { + return Err(ResponseDispatchError::QuotaExhausted { + count: quota_exhausted_count, + upstream_error: last_quota_error.unwrap_or_default(), + }); + } + Some(ExhaustedAccountClass::RateLimited) => { + return Err(ResponseDispatchError::RateLimited { + count: rate_limited_count, + upstream_error: last_rate_limit_error.unwrap_or_default(), + }); + } + Some(ExhaustedAccountClass::Expired) => { + return Err(ResponseDispatchError::Expired { + count: expired_count, + upstream_error: last_auth_error.unwrap_or_default(), + }); + } + Some(ExhaustedAccountClass::CloudflareChallenge) => { + return Err(ResponseDispatchError::CloudflareChallenge { + count: cloudflare_challenge_count, + upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), + }); + } + Some(ExhaustedAccountClass::CloudflarePathBlocked) => { + return Err(ResponseDispatchError::CloudflarePathBlocked { + count: cloudflare_path_block_count, + upstream_error: last_cloudflare_path_block_error.unwrap_or_default(), + }); + } + Some(ExhaustedAccountClass::ModelUnsupported) => { + return Err(ResponseDispatchError::ModelUnsupported { + count: model_unsupported_count, + upstream_error: last_model_unsupported_error.unwrap_or_default(), + }); + } + None => return Err(ResponseDispatchError::NoActiveAccount), + }, + }; + self.account_pool.wait_for_request_interval(&acquired).await; + let account = acquired.account; + let release_account_id = account.id.clone(); + let response_result = create_response_with_account_retrying_5xx( + &self.codex, + self.installation_id.as_deref(), + &self.cloudflare, + &request, + request_id, + &account, + ) + .await; + self.account_pool.release(&release_account_id).await; + + match response_result { + Ok(response) => { + let collected_response = + response_from_codex_sse(&response.body, tuple_schema.as_ref()) + .map_err(ResponseDispatchError::InvalidSse)?; + if matches!(collected_response, CollectedResponse::Empty) { + self.account_pool + .record_empty_response_attempt( + &release_account_id, + image_generation_requested, + ) + .await; + empty_response_retries += 1; + if empty_response_retries <= MAX_EMPTY_RESPONSE_RETRIES { + continue; + } + } + if let CollectedResponse::Failed(failure) = &collected_response { + if is_history_recovery_sse_failure(failure) && !history_recovery_used { + if sse_failure_invalid_reasoning_replay(failure) { + self.evict_reasoning_replay(&request, &release_account_id) + .await; + } + self.recover_request_history(&mut request, &mut implicit_resume) + .await; + history_recovery_used = true; + continue; + } + if is_model_unsupported_sse_failure(failure) { + let upstream_error = sse_failure_error_body(failure); + if model_unsupported_retry_used { + return Err(ResponseDispatchError::ModelUnsupported { + count: model_unsupported_count + 1, + upstream_error, + }); + } + model_unsupported_count += 1; + last_model_unsupported_error = Some(upstream_error); + last_exhausted_account_class = + Some(ExhaustedAccountClass::ModelUnsupported); + model_unsupported_retry_used = true; + excluded_account_ids.push(release_account_id); + continue; + } + if is_quota_exhausted_sse_failure(failure) { + quota_exhausted_count += 1; + last_quota_error = Some(failure.message.clone()); + last_exhausted_account_class = + Some(ExhaustedAccountClass::QuotaExhausted); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + if is_auth_sse_failure(failure) { + expired_count += 1; + last_auth_error = Some(sse_failure_error_body(failure)); + last_exhausted_account_class = Some(ExhaustedAccountClass::Expired); + let account_status = auth_sse_failure_account_status(failure); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + } + break (account, response, collected_response); + } + Err(error) if is_rate_limit_upstream_error(&error) => { + rate_limited_count += 1; + last_rate_limit_error = Some(upstream_error_body(&error)); + last_exhausted_account_class = Some(ExhaustedAccountClass::RateLimited); + let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); + self.account_pool + .mark_quota_limited_until(&release_account_id, cooldown_until) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) + if is_history_recovery_upstream_error(&error) && !history_recovery_used => + { + if client_error_invalid_reasoning_replay(&error) { + self.evict_reasoning_replay(&request, &release_account_id) + .await; + } + self.recover_request_history(&mut request, &mut implicit_resume) + .await; + history_recovery_used = true; + } + Err(error) if is_quota_exhausted_upstream_error(&error) => { + quota_exhausted_count += 1; + last_quota_error = Some(upstream_error_body(&error)); + last_exhausted_account_class = Some(ExhaustedAccountClass::QuotaExhausted); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_auth_upstream_error(&error) => { + expired_count += 1; + last_auth_error = Some(upstream_error_body(&error)); + last_exhausted_account_class = Some(ExhaustedAccountClass::Expired); + let account_status = auth_failure_account_status(&error); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_challenge_upstream_error(&error) => { + cloudflare_challenge_count += 1; + last_cloudflare_challenge_error = + Some(cloudflare_challenge_error_message().to_string()); + last_exhausted_account_class = Some(ExhaustedAccountClass::CloudflareChallenge); + self.cloudflare + .apply_challenge(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_path_block_upstream_error(&error) => { + cloudflare_path_block_count += 1; + last_cloudflare_path_block_error = + Some(cloudflare_path_block_error_message().to_string()); + last_exhausted_account_class = + Some(ExhaustedAccountClass::CloudflarePathBlocked); + self.cloudflare + .apply_path_block(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_model_unsupported_upstream_error(&error) => { + let upstream_error = upstream_error_body(&error); + if model_unsupported_retry_used { + return Err(ResponseDispatchError::ModelUnsupported { + count: model_unsupported_count + 1, + upstream_error, + }); + } + model_unsupported_count += 1; + last_model_unsupported_error = Some(upstream_error); + last_exhausted_account_class = Some(ExhaustedAccountClass::ModelUnsupported); + model_unsupported_retry_used = true; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_banned_upstream_error(&error) => { + self.account_pool + .set_status(&release_account_id, AccountStatus::Banned) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) => return Err(ResponseDispatchError::Upstream(error)), + } + }; + + match collected_response { + CollectedResponse::Completed(body) => { + self.cloudflare.reset_path_block(&account.id).await; + self.cloudflare + .capture_set_cookie_headers(&account.id, &response.set_cookie_headers) + .await; + self.account_pool + .sync_passive_rate_limit_headers(&account, &response.rate_limit_headers) + .await; + if let Some(usage) = response.usage { + self.account_pool + .record_response_usage(&account.id, usage, image_generation_requested) + .await; + } + self.record_response_affinity( + &request, + &account.id, + &response.body, + response.turn_state.clone(), + response.usage, + ) + .await; + record_response_event(ResponseEventRecord { + logs: &self.logs, + request_id, + account_id: &account.id, + model: requested_model, + started_at, + status_code: 200, + level: EventLevel::Info, + message: "v1 responses completed", + metadata: serde_json::json!({ + "stream": false, + "usage": response.usage, + }), + rate_limit_headers: &response.rate_limit_headers, + }) + .await; + Ok(body) + } + CollectedResponse::Failed(failure) => Err(ResponseDispatchError::Failed(failure)), + CollectedResponse::MissingCompleted => Err(ResponseDispatchError::MissingCompleted), + CollectedResponse::Empty => Err(ResponseDispatchError::EmptyUpstreamResponse), + } + } + + /// 调度 Responses compact 请求到 Codex compact 上游。 + pub async fn compact( + &self, + request_id: &str, + mut request: CodexCompactRequest, + requested_model: &str, + ) -> Result { + let catalog = self.models.catalog().await; + let parsed_model = catalog.parse_model_name(requested_model); + request.model = parsed_model.model_id; + let mut excluded_account_ids = Vec::new(); + let mut rate_limited_count = 0usize; + let mut last_rate_limit_error = None; + let mut quota_exhausted_count = 0usize; + let mut last_quota_error = None; + let mut expired_count = 0usize; + let mut last_auth_error = None; + let mut cloudflare_challenge_count = 0usize; + let mut last_cloudflare_challenge_error = None; + let mut cloudflare_path_block_count = 0usize; + let mut last_cloudflare_path_block_error = None; + let mut model_unsupported_count = 0usize; + let mut last_model_unsupported_error = None; + let mut model_unsupported_retry_used = false; + + loop { + let acquire_request = AccountAcquireRequest::new(&request.model, Utc::now()) + .with_exclude_account_ids(excluded_account_ids.iter().cloned()); + let acquired = match self.account_pool.acquire_with(acquire_request).await { + Some(acquired) => acquired, + None if quota_exhausted_count > 0 => { + return Err(ResponseDispatchError::QuotaExhausted { + count: quota_exhausted_count, + upstream_error: last_quota_error.unwrap_or_default(), + }); + } + None if rate_limited_count > 0 => { + return Err(ResponseDispatchError::RateLimited { + count: rate_limited_count, + upstream_error: last_rate_limit_error.unwrap_or_default(), + }); + } + None if expired_count > 0 => { + return Err(ResponseDispatchError::Expired { + count: expired_count, + upstream_error: last_auth_error.unwrap_or_default(), + }); + } + None if cloudflare_challenge_count > 0 => { + return Err(ResponseDispatchError::CloudflareChallenge { + count: cloudflare_challenge_count, + upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), + }); + } + None if cloudflare_path_block_count > 0 => { + return Err(ResponseDispatchError::CloudflarePathBlocked { + count: cloudflare_path_block_count, + upstream_error: last_cloudflare_path_block_error.unwrap_or_default(), + }); + } + None if model_unsupported_count > 0 => { + return Err(ResponseDispatchError::ModelUnsupported { + count: model_unsupported_count, + upstream_error: last_model_unsupported_error.unwrap_or_default(), + }); + } + None => return Err(ResponseDispatchError::NoActiveAccount), + }; + let account = acquired.account; + let release_account_id = account.id.clone(); + let response_result = create_compact_response_with_account_retrying_5xx( + &self.codex, + self.installation_id.as_deref(), + &self.cloudflare, + &request, + request_id, + &account, + ) + .await; + self.account_pool.release(&release_account_id).await; + + match response_result { + Ok(response) => { + self.cloudflare.reset_path_block(&account.id).await; + return Ok(response.body); + } + Err(error) if is_rate_limit_upstream_error(&error) => { + rate_limited_count += 1; + last_rate_limit_error = Some(upstream_error_body(&error)); + let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); + self.account_pool + .mark_quota_limited_until(&release_account_id, cooldown_until) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_quota_exhausted_upstream_error(&error) => { + quota_exhausted_count += 1; + last_quota_error = Some(upstream_error_body(&error)); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_auth_upstream_error(&error) => { + expired_count += 1; + last_auth_error = Some(upstream_error_body(&error)); + let account_status = auth_failure_account_status(&error); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_challenge_upstream_error(&error) => { + cloudflare_challenge_count += 1; + last_cloudflare_challenge_error = + Some(cloudflare_challenge_error_message().to_string()); + self.cloudflare + .apply_challenge(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_path_block_upstream_error(&error) => { + cloudflare_path_block_count += 1; + last_cloudflare_path_block_error = + Some(cloudflare_path_block_error_message().to_string()); + self.cloudflare + .apply_path_block(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_model_unsupported_upstream_error(&error) => { + let upstream_error = upstream_error_body(&error); + if model_unsupported_retry_used { + return Err(ResponseDispatchError::ModelUnsupported { + count: model_unsupported_count + 1, + upstream_error, + }); + } + model_unsupported_count += 1; + last_model_unsupported_error = Some(upstream_error); + model_unsupported_retry_used = true; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_banned_upstream_error(&error) => { + self.account_pool + .set_status(&release_account_id, AccountStatus::Banned) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) => return Err(ResponseDispatchError::Upstream(error)), + } + } + } + + /// 调度流式 Responses 请求到 Codex Responses 上游。 + pub async fn stream( + &self, + request_id: &str, + mut request: CodexResponsesRequest, + requested_model: &str, + ) -> Result { + let catalog = self.models.catalog().await; + let parsed_model = catalog.parse_model_name(requested_model); + apply_response_model_options(&mut request, &parsed_model, self.models.config()); + request.stream = true; + let started_at = Instant::now(); + let tuple_schema = request.tuple_schema.clone(); + let now = Utc::now(); + let mut implicit_resume = self.prepare_response_session(&mut request).await; + let preferred_account_id = self.preferred_account_id_for_request(&request, now).await; + let mut acquire_request = AccountAcquireRequest::new(&request.model, now); + if let Some(preferred_account_id) = preferred_account_id { + acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); + } + let mut excluded_account_ids = Vec::new(); + let mut rate_limited_count = 0usize; + let mut last_rate_limit_error = None; + let mut quota_exhausted_count = 0usize; + let mut last_quota_error = None; + let mut expired_count = 0usize; + let mut last_auth_error = None; + let mut cloudflare_challenge_count = 0usize; + let mut last_cloudflare_challenge_error = None; + let mut cloudflare_path_block_count = 0usize; + let mut last_cloudflare_path_block_error = None; + let mut model_unsupported_count = 0usize; + let mut last_model_unsupported_error = None; + let mut model_unsupported_retry_used = false; + let mut history_recovery_used = false; + let mut last_exhausted_account_class = None; + loop { + let mut attempt_acquire_request = acquire_request + .clone() + .with_exclude_account_ids(excluded_account_ids.iter().cloned()); + attempt_acquire_request.now = Utc::now(); + let acquired = match self + .account_pool + .acquire_with(attempt_acquire_request) + .await + { + Some(acquired) => acquired, + None => match last_exhausted_account_class { + Some(ExhaustedAccountClass::QuotaExhausted) => { + return Err(ResponseDispatchError::QuotaExhausted { + count: quota_exhausted_count, + upstream_error: last_quota_error.unwrap_or_default(), + }); + } + Some(ExhaustedAccountClass::RateLimited) => { + return Err(ResponseDispatchError::RateLimited { + count: rate_limited_count, + upstream_error: last_rate_limit_error.unwrap_or_default(), + }); + } + Some(ExhaustedAccountClass::Expired) => { + return Err(ResponseDispatchError::Expired { + count: expired_count, + upstream_error: last_auth_error.unwrap_or_default(), + }); + } + Some(ExhaustedAccountClass::CloudflareChallenge) => { + return Err(ResponseDispatchError::CloudflareChallenge { + count: cloudflare_challenge_count, + upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), + }); + } + Some(ExhaustedAccountClass::CloudflarePathBlocked) => { + return Err(ResponseDispatchError::CloudflarePathBlocked { + count: cloudflare_path_block_count, + upstream_error: last_cloudflare_path_block_error.unwrap_or_default(), + }); + } + Some(ExhaustedAccountClass::ModelUnsupported) => { + return Err(ResponseDispatchError::ModelUnsupported { + count: model_unsupported_count, + upstream_error: last_model_unsupported_error.unwrap_or_default(), + }); + } + None => return Err(ResponseDispatchError::NoActiveAccount), + }, + }; + self.account_pool.wait_for_request_interval(&acquired).await; + let account = acquired.account; + let release_account_id = account.id.clone(); + let response_result = create_response_stream_with_account_retrying_5xx( + &self.codex, + self.installation_id.as_deref(), + &self.cloudflare, + &request, + request_id, + &account, + ) + .await; + + match response_result { + Ok(response) => { + let transport = response.transport; + let rate_limit_headers = response.rate_limit_headers; + let rate_limit_header_updates = response.rate_limit_header_updates; + let turn_state_update = response.turn_state_update; + let turn_state = response.turn_state; + let (prefetched, body) = match prefetch_first_sse_chunk(response.body).await { + Ok(prefetched) => prefetched, + Err(ResponseDispatchError::Upstream(error)) + if is_history_recovery_upstream_error(&error) + && !history_recovery_used => + { + self.account_pool.release(&release_account_id).await; + if client_error_invalid_reasoning_replay(&error) { + self.evict_reasoning_replay(&request, &release_account_id) + .await; + } + self.recover_request_history(&mut request, &mut implicit_resume) + .await; + history_recovery_used = true; + continue; + } + Err(error) => { + self.account_pool.release(&release_account_id).await; + return Err(error); + } + }; + if let Some(failure) = first_sse_failure(&prefetched)? { + if is_history_recovery_sse_failure(&failure) && !history_recovery_used { + self.account_pool.release(&release_account_id).await; + if sse_failure_invalid_reasoning_replay(&failure) { + self.evict_reasoning_replay(&request, &release_account_id) + .await; + } + self.recover_request_history(&mut request, &mut implicit_resume) + .await; + history_recovery_used = true; + continue; + } + if is_model_unsupported_sse_failure(&failure) { + let upstream_error = sse_failure_error_body(&failure); + if model_unsupported_retry_used { + self.account_pool.release(&release_account_id).await; + return Err(ResponseDispatchError::ModelUnsupported { + count: model_unsupported_count + 1, + upstream_error, + }); + } + model_unsupported_count += 1; + last_model_unsupported_error = Some(upstream_error); + model_unsupported_retry_used = true; + excluded_account_ids.push(release_account_id); + self.account_pool.release(&account.id).await; + continue; + } + if is_quota_exhausted_sse_failure(&failure) { + quota_exhausted_count += 1; + last_quota_error = Some(failure.message.clone()); + last_exhausted_account_class = + Some(ExhaustedAccountClass::QuotaExhausted); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + self.account_pool.release(&account.id).await; + continue; + } + if is_auth_sse_failure(&failure) { + expired_count += 1; + last_auth_error = Some(sse_failure_error_body(&failure)); + last_exhausted_account_class = Some(ExhaustedAccountClass::Expired); + let account_status = auth_sse_failure_account_status(&failure); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + self.account_pool.release(&account.id).await; + continue; + } + self.account_pool.release(&release_account_id).await; + return Err(ResponseDispatchError::Failed(failure.clone())); + } + return Ok(spawn_live_response_stream( + LiveResponseStreamContext { + account_pool: Arc::clone(&self.account_pool), + session_affinity: Arc::clone(&self.session_affinity), + reasoning_replay: Arc::clone(&self.reasoning_replay), + logs: Arc::clone(&self.logs), + cloudflare: self.cloudflare.clone(), + account_id: account.id, + request_id: request_id.to_string(), + model: requested_model.to_string(), + request, + tuple_schema, + transport, + rate_limit_headers, + rate_limit_header_updates, + turn_state_update, + turn_state, + started_at, + }, + prefetched, + body, + )); + } + Err(error) if is_rate_limit_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + rate_limited_count += 1; + last_rate_limit_error = Some(upstream_error_body(&error)); + last_exhausted_account_class = Some(ExhaustedAccountClass::RateLimited); + let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); + self.account_pool + .mark_quota_limited_until(&release_account_id, cooldown_until) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_quota_exhausted_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + quota_exhausted_count += 1; + last_quota_error = Some(upstream_error_body(&error)); + last_exhausted_account_class = Some(ExhaustedAccountClass::QuotaExhausted); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) + if is_history_recovery_upstream_error(&error) && !history_recovery_used => + { + self.account_pool.release(&release_account_id).await; + if client_error_invalid_reasoning_replay(&error) { + self.evict_reasoning_replay(&request, &release_account_id) + .await; + } + self.recover_request_history(&mut request, &mut implicit_resume) + .await; + history_recovery_used = true; + } + Err(error) if is_auth_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + expired_count += 1; + last_auth_error = Some(upstream_error_body(&error)); + last_exhausted_account_class = Some(ExhaustedAccountClass::Expired); + let account_status = auth_failure_account_status(&error); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_challenge_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + cloudflare_challenge_count += 1; + last_cloudflare_challenge_error = + Some(cloudflare_challenge_error_message().to_string()); + last_exhausted_account_class = Some(ExhaustedAccountClass::CloudflareChallenge); + self.cloudflare + .apply_challenge(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_path_block_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + cloudflare_path_block_count += 1; + last_cloudflare_path_block_error = + Some(cloudflare_path_block_error_message().to_string()); + last_exhausted_account_class = + Some(ExhaustedAccountClass::CloudflarePathBlocked); + self.cloudflare + .apply_path_block(self.account_pool.as_ref(), &release_account_id) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_model_unsupported_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + let upstream_error = upstream_error_body(&error); + if model_unsupported_retry_used { + return Err(ResponseDispatchError::ModelUnsupported { + count: model_unsupported_count + 1, + upstream_error, + }); + } + model_unsupported_count += 1; + last_model_unsupported_error = Some(upstream_error); + last_exhausted_account_class = Some(ExhaustedAccountClass::ModelUnsupported); + model_unsupported_retry_used = true; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_banned_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + self.account_pool + .set_status(&release_account_id, AccountStatus::Banned) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) => { + self.account_pool.release(&release_account_id).await; + return Err(ResponseDispatchError::Upstream(error)); + } + } + } + } +} + +async fn create_response_with_account( + codex: &CodexBackendClient, + installation_id: Option<&str>, + cloudflare: &CloudflareRecovery, + request: &CodexResponsesRequest, + request_id: &str, + account: &Account, +) -> Result { + let cookie_header = cloudflare + .cookie_header_for_request(&account.id, "/codex/responses") + .await; + let identity = build_conversation_identity( + request.prompt_cache_key.as_deref(), + request.codex_window_id.as_deref(), + &account.id, + ); + codex + .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: identity.window_id.as_deref(), + parent_thread_id: request.parent_thread_id.as_deref(), + cookie_header: cookie_header.as_deref(), + installation_id, + session_id: identity.conversation_id.as_deref(), + }, + ) + .await +} + +async fn create_response_stream_with_account( + codex: &CodexBackendClient, + installation_id: Option<&str>, + cloudflare: &CloudflareRecovery, + request: &CodexResponsesRequest, + request_id: &str, + account: &Account, +) -> Result { + let cookie_header = cloudflare + .cookie_header_for_request(&account.id, "/codex/responses") + .await; + let identity = build_conversation_identity( + request.prompt_cache_key.as_deref(), + request.codex_window_id.as_deref(), + &account.id, + ); + codex + .create_response_stream( + 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: identity.window_id.as_deref(), + parent_thread_id: request.parent_thread_id.as_deref(), + cookie_header: cookie_header.as_deref(), + installation_id, + session_id: identity.conversation_id.as_deref(), + }, + ) + .await +} + +async fn create_compact_response_with_account( + codex: &CodexBackendClient, + installation_id: Option<&str>, + cloudflare: &CloudflareRecovery, + request: &CodexCompactRequest, + request_id: &str, + account: &Account, +) -> Result { + let cookie_header = cloudflare + .cookie_header_for_request(&account.id, "/codex/responses/compact") + .await; + codex + .create_compact_response( + request, + 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: cookie_header.as_deref(), + installation_id, + session_id: None, + }, + ) + .await +} + +const MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT: usize = 2; + +async fn create_response_with_account_retrying_5xx( + codex: &CodexBackendClient, + installation_id: Option<&str>, + cloudflare: &CloudflareRecovery, + request: &CodexResponsesRequest, + request_id: &str, + account: &Account, +) -> Result { + let mut retries = 0; + loop { + let result = create_response_with_account( + codex, + installation_id, + cloudflare, + request, + request_id, + account, + ) + .await; + match result { + Err(error) + if is_retryable_upstream_5xx_error(&error) + && retries < MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT => + { + retries += 1; + } + result => return result, + } + } +} + +async fn create_response_stream_with_account_retrying_5xx( + codex: &CodexBackendClient, + installation_id: Option<&str>, + cloudflare: &CloudflareRecovery, + request: &CodexResponsesRequest, + request_id: &str, + account: &Account, +) -> Result { + let mut retries = 0; + loop { + let result = create_response_stream_with_account( + codex, + installation_id, + cloudflare, + request, + request_id, + account, + ) + .await; + match result { + Err(error) + if is_retryable_upstream_5xx_error(&error) + && retries < MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT => + { + retries += 1; + } + result => return result, + } + } +} + +async fn create_compact_response_with_account_retrying_5xx( + codex: &CodexBackendClient, + installation_id: Option<&str>, + cloudflare: &CloudflareRecovery, + request: &CodexCompactRequest, + request_id: &str, + account: &Account, +) -> Result { + let mut retries = 0; + loop { + let result = create_compact_response_with_account( + codex, + installation_id, + cloudflare, + request, + request_id, + account, + ) + .await; + match result { + Err(error) + if is_retryable_upstream_5xx_error(&error) + && retries < MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT => + { + retries += 1; + } + result => return result, + } + } +} + +fn is_rate_limit_upstream_error(error: &CodexClientError) -> bool { + matches!( + error, + CodexClientError::Upstream { status, .. } if status_code_is_rate_limited(status.as_u16()) + ) +} + +fn is_retryable_upstream_5xx_error(error: &CodexClientError) -> bool { + matches!( + error, + CodexClientError::Upstream { status, body, .. } + if status_code_is_transient_upstream(status.as_u16()) + && status_code_allows_same_account_retry(status.as_u16()) + && !is_history_recovery_signal(body) + ) +} + +fn is_quota_exhausted_upstream_error(error: &CodexClientError) -> bool { + matches!( + error, + CodexClientError::Upstream { status, .. } if status_code_is_quota_exhausted(status.as_u16()) + ) +} + +fn is_auth_upstream_error(error: &CodexClientError) -> bool { + matches!( + error, + CodexClientError::Upstream { status, .. } if status.as_u16() == 401 + ) +} + +fn is_cloudflare_challenge_upstream_error(error: &CodexClientError) -> bool { + matches!( + error, + CodexClientError::Upstream { status, body, .. } + if status.as_u16() == 403 && is_cloudflare_challenge_signal(body) + ) +} + +fn is_cloudflare_path_block_upstream_error(error: &CodexClientError) -> bool { + matches!( + error, + CodexClientError::Upstream { status, body, .. } + if status.as_u16() == 404 && body.trim().is_empty() + ) +} + +fn is_model_unsupported_upstream_error(error: &CodexClientError) -> bool { + matches!( + error, + CodexClientError::Upstream { status, body, .. } + if status.is_client_error() + && !matches!(status.as_u16(), 401 | 402 | 403 | 404 | 429) + && is_model_unsupported_signal(body) + ) +} + +fn is_history_recovery_upstream_error(error: &CodexClientError) -> bool { + matches!( + error, + CodexClientError::Upstream { body, .. } if is_history_recovery_signal(body) + ) +} + +fn is_banned_upstream_error(error: &CodexClientError) -> bool { + matches!( + error, + CodexClientError::Upstream { status, body, .. } + if status.as_u16() == 403 && is_banned_auth_signal(body) + ) +} + +fn auth_failure_account_status(error: &CodexClientError) -> AccountStatus { + match error { + CodexClientError::Upstream { body, .. } if is_banned_auth_signal(body) => { + AccountStatus::Banned + } + _ => AccountStatus::Expired, + } +} + +fn upstream_error_body(error: &CodexClientError) -> String { + match error { + CodexClientError::Upstream { body, .. } => body.clone(), + error => error.to_string(), + } +} + +fn sse_failure_error_body(failure: &ResponsesSseFailure) -> String { + match failure.upstream_code.as_deref() { + Some(code) => serde_json::json!({ + "error": { + "code": code, + "message": failure.message.as_str(), + } + }) + .to_string(), + None => failure.message.clone(), + } +} + +fn is_quota_exhausted_sse_failure(failure: &ResponsesSseFailure) -> bool { + failure + .upstream_code + .as_deref() + .is_some_and(|code| matches!(code, "quota_exceeded" | "insufficient_quota")) + || failure.message.to_ascii_lowercase().contains("quota") +} + +fn is_auth_sse_failure(failure: &ResponsesSseFailure) -> bool { + failure.upstream_code.as_deref().is_some_and(|code| { + let code = code.to_ascii_lowercase(); + matches!( + code.as_str(), + "token_invalid" + | "token_expired" + | "token_revoked" + | "account_deactivated" + | "unauthorized" + | "invalid_api_key" + ) + }) || { + let message = failure.message.to_ascii_lowercase(); + message.contains("token revoked") + || message.contains("token invalid") + || message.contains("token expired") + } +} + +fn is_model_unsupported_sse_failure(failure: &ResponsesSseFailure) -> bool { + failure + .upstream_code + .as_deref() + .is_some_and(is_model_unsupported_signal) + || is_model_unsupported_signal(&failure.message) +} + +fn is_history_recovery_sse_failure(failure: &ResponsesSseFailure) -> bool { + failure + .upstream_code + .as_deref() + .is_some_and(is_history_recovery_signal) + || is_history_recovery_signal(&failure.message) +} + +fn sse_failure_invalid_reasoning_replay(failure: &ResponsesSseFailure) -> bool { + failure + .upstream_code + .as_deref() + .is_some_and(is_invalid_encrypted_content_signal) + || is_invalid_encrypted_content_signal(&failure.message) +} + +fn client_error_invalid_reasoning_replay(error: &CodexClientError) -> bool { + matches!( + error, + CodexClientError::Upstream { body, .. } if is_invalid_encrypted_content_signal(body) + ) +} + +fn auth_sse_failure_account_status(failure: &ResponsesSseFailure) -> AccountStatus { + if failure + .upstream_code + .as_deref() + .is_some_and(is_banned_auth_signal) + || is_banned_auth_signal(&failure.message) + { + AccountStatus::Banned + } else { + AccountStatus::Expired + } +} + +fn is_banned_auth_signal(value: &str) -> bool { + let value = value.to_ascii_lowercase(); + value.contains("account_deactivated") + || value.contains("account deactivated") + || value.contains("account has been deactivated") + || value.contains("deactivated") + || value.contains("banned") +} + +fn is_cloudflare_challenge_signal(value: &str) -> bool { + let value = value.to_ascii_lowercase(); + value.contains("cf-mitigated") + || value.contains("cf-chl-bypass") + || value.contains("_cf_chl") + || value.contains("cf_chl") + || value.contains("attention required") + || value.contains("just a moment") +} + +fn is_model_unsupported_signal(value: &str) -> bool { + let value = value.to_ascii_lowercase(); + value.contains("model_not_supported") + || value.contains("model_not_available") + || (value.contains("model") + && (value.contains("not supported") + || value.contains("not available") + || value.contains("not_supported") + || value.contains("not_available"))) +} + +fn is_history_recovery_signal(value: &str) -> bool { + let value = value.to_ascii_lowercase(); + value.contains("previous_response_not_found") + || (value.contains("previous response") && value.contains("not found")) + || value.contains("no tool output found for function call") + || is_invalid_encrypted_content_signal(&value) +} + +fn is_invalid_encrypted_content_signal(value: &str) -> bool { + let value = value.to_ascii_lowercase(); + value.contains("invalid_encrypted_content") + || (value.contains("invalid") && value.contains("encrypted") && value.contains("content")) +} + +fn strip_request_history(request: &mut CodexResponsesRequest) { + request.previous_response_id = None; + request.turn_state = None; + request.turn_metadata = None; + if !request.explicit_prompt_cache_key { + request.prompt_cache_key = None; + } +} + +fn cloudflare_challenge_error_message() -> &'static str { + "Upstream blocked the request (Cloudflare challenge)" +} + +fn cloudflare_path_block_error_message() -> &'static str { + "Upstream blocked the request (Cloudflare path-block)" +} + +fn rate_limit_cooldown_until(error: &CodexClientError, now: DateTime) -> DateTime { + let retry_after_seconds = match error { + CodexClientError::Upstream { + retry_after_seconds, + .. + } => retry_after_seconds.unwrap_or(60), + _ => 60, + }; + now + Duration::seconds(retry_after_seconds.min(i64::MAX as u64) as i64) +} + +const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024; +const DONE_SSE_FRAME_TEXT: &str = "data: [DONE]\n\n"; +const DONE_SSE_FRAME: &[u8] = b"data: [DONE]\n\n"; + +async fn prefetch_first_sse_chunk( + mut body: CodexBackendSseStream, +) -> Result<(Bytes, CodexBackendSseStream), ResponseDispatchError> { + let mut prefetched = Vec::new(); + while !contains_sse_event_separator(&prefetched) { + let Some(next) = body.next().await else { + if prefetched.is_empty() { + return Err(ResponseDispatchError::EmptyUpstreamResponse); + } + return Err(ResponseDispatchError::MissingCompleted); + }; + let chunk = next.map_err(ResponseDispatchError::Upstream)?; + prefetched.extend_from_slice(&chunk); + if prefetched.len() > MAX_STREAM_PREFETCH_BYTES { + return Err(ResponseDispatchError::InvalidSse( + SseError::BufferExceeded { + max_bytes: MAX_STREAM_PREFETCH_BYTES, + }, + )); + } + } + + Ok((Bytes::from(prefetched), body)) +} + +fn contains_sse_event_separator(bytes: &[u8]) -> bool { + bytes.windows(2).any(|window| window == b"\n\n") + || bytes.windows(4).any(|window| window == b"\r\n\r\n") +} + +fn first_sse_failure(prefetched: &[u8]) -> Result, SseError> { + let body = String::from_utf8_lossy(prefetched); + match response_from_codex_sse(&body, None)? { + CollectedResponse::Failed(failure) => Ok(Some(failure)), + CollectedResponse::Completed(_) + | CollectedResponse::MissingCompleted + | CollectedResponse::Empty => Ok(None), + } +} + +struct LiveResponseStreamContext { + account_pool: Arc, + session_affinity: Arc, + reasoning_replay: Arc>, + logs: Arc, + cloudflare: CloudflareRecovery, + account_id: String, + request_id: String, + model: String, + request: CodexResponsesRequest, + tuple_schema: Option, + transport: CodexBackendTransport, + rate_limit_headers: Vec<(String, String)>, + rate_limit_header_updates: Option, + turn_state_update: Option, + turn_state: Option, + started_at: Instant, +} + +fn spawn_live_response_stream( + context: LiveResponseStreamContext, + prefetched: Bytes, + mut body: CodexBackendSseStream, +) -> ResponseDispatchStream { + let (sender, receiver) = mpsc::channel(8); + let (cancel_sender, mut cancel_receiver) = oneshot::channel(); + tokio::spawn(async move { + let mut tuple_transformer = context + .tuple_schema + .clone() + .map(TupleSseEventTransformer::new); + let mut body_bytes = Vec::new(); + if !send_live_response_stream_chunk( + &sender, + &mut body_bytes, + tuple_transformer.as_mut(), + prefetched, + ) + .await + { + context.account_pool.release(&context.account_id).await; + return; + } + + loop { + let next = tokio::select! { + _ = &mut cancel_receiver => { + context.account_pool.release(&context.account_id).await; + return; + } + next = body.next() => next, + }; + let Some(next) = next else { + break; + }; + match next { + Ok(chunk) => { + if !send_live_response_stream_chunk( + &sender, + &mut body_bytes, + tuple_transformer.as_mut(), + chunk, + ) + .await + { + context.account_pool.release(&context.account_id).await; + return; + } + } + Err(error) => { + if !flush_live_response_stream_transformer( + &sender, + &mut body_bytes, + tuple_transformer.as_mut(), + ) + .await + { + context.account_pool.release(&context.account_id).await; + return; + } + let detail = error.to_string(); + let Some(body_text) = + send_live_response_stream_tail(&sender, &mut body_bytes, Some(&detail)) + .await + else { + context.account_pool.release(&context.account_id).await; + return; + }; + finalize_live_response_stream(context, body_text).await; + return; + } + } + } + + if !flush_live_response_stream_transformer( + &sender, + &mut body_bytes, + tuple_transformer.as_mut(), + ) + .await + { + context.account_pool.release(&context.account_id).await; + return; + } + let Some(body_text) = send_live_response_stream_tail(&sender, &mut body_bytes, None).await + else { + context.account_pool.release(&context.account_id).await; + return; + }; + + finalize_live_response_stream(context, body_text).await; + }); + + ResponseDispatchStream { + body: Box::pin(MpscResponseBodyStream { + receiver, + cancel: Some(cancel_sender), + }), + } +} + +async fn send_live_response_stream_chunk( + sender: &mpsc::Sender>, + body_bytes: &mut Vec, + transformer: Option<&mut TupleSseEventTransformer>, + chunk: Bytes, +) -> bool { + let chunks = match transformer { + Some(transformer) => transformer.push(&chunk), + None => vec![chunk], + }; + send_live_response_stream_chunks(sender, body_bytes, chunks).await +} + +async fn flush_live_response_stream_transformer( + sender: &mpsc::Sender>, + body_bytes: &mut Vec, + transformer: Option<&mut TupleSseEventTransformer>, +) -> bool { + let Some(transformer) = transformer else { + return true; + }; + send_live_response_stream_chunks(sender, body_bytes, transformer.finish()).await +} + +async fn send_live_response_stream_chunks( + sender: &mpsc::Sender>, + body_bytes: &mut Vec, + chunks: Vec, +) -> bool { + for chunk in chunks { + body_bytes.extend_from_slice(&chunk); + if sender.send(Ok(chunk)).await.is_err() { + return false; + } + } + true +} + +struct TupleSseEventTransformer { + tuple_schema: Value, + pending: Vec, +} + +impl TupleSseEventTransformer { + fn new(tuple_schema: Value) -> Self { + Self { + tuple_schema, + pending: Vec::new(), + } + } + + fn push(&mut self, chunk: &[u8]) -> Vec { + self.pending.extend_from_slice(chunk); + let mut chunks = Vec::new(); + while let Some(frame_end) = next_sse_frame_end(&self.pending) { + let frame = self.pending.drain(..frame_end).collect::>(); + chunks.push(self.transform_frame(&frame)); + } + chunks + } + + fn finish(&mut self) -> Vec { + if self.pending.is_empty() { + return Vec::new(); + } + let frame = std::mem::take(&mut self.pending); + vec![self.transform_frame(&frame)] + } + + fn transform_frame(&self, frame: &[u8]) -> Bytes { + let frame_text = String::from_utf8_lossy(frame); + let Ok(events) = parse_sse_events(&frame_text) else { + return Bytes::copy_from_slice(frame); + }; + let [event] = events.as_slice() else { + return Bytes::copy_from_slice(frame); + }; + let Ok(data) = serde_json::from_str::(&event.data) else { + return Bytes::copy_from_slice(frame); + }; + let transformed = reconvert_responses_sse_event_tuple_values( + event.event.as_deref(), + data, + &self.tuple_schema, + ); + Bytes::from(encode_sse_event( + event.event.as_deref().unwrap_or_default(), + &transformed.to_string(), + )) + } +} + +fn next_sse_frame_end(bytes: &[u8]) -> Option { + let lf_lf = bytes + .windows(2) + .position(|window| window == b"\n\n") + .map(|position| position + 2); + let crlf_crlf = bytes + .windows(4) + .position(|window| window == b"\r\n\r\n") + .map(|position| position + 4); + match (lf_lf, crlf_crlf) { + (Some(left), Some(right)) => Some(left.min(right)), + (Some(end), None) | (None, Some(end)) => Some(end), + (None, None) => None, + } +} + +async fn send_live_response_stream_tail( + sender: &mpsc::Sender>, + body_bytes: &mut Vec, + failure_detail: Option<&str>, +) -> Option { + let mut body_text = String::from_utf8_lossy(body_bytes).to_string(); + if !sse_body_has_terminal_event(&body_text) { + if let Some(separator) = missing_sse_event_separator(&body_text) { + body_text.push_str(separator); + body_bytes.extend_from_slice(separator.as_bytes()); + if sender + .send(Ok(Bytes::copy_from_slice(separator.as_bytes()))) + .await + .is_err() + { + return None; + } + } + let failure = premature_close_failed_event(failure_detail); + body_text.push_str(&failure); + body_bytes.extend_from_slice(failure.as_bytes()); + if sender.send(Ok(Bytes::from(failure))).await.is_err() { + return None; + } + } + + if !sse_body_has_done(&body_text) { + body_text.push_str(DONE_SSE_FRAME_TEXT); + body_bytes.extend_from_slice(DONE_SSE_FRAME); + if sender + .send(Ok(Bytes::from_static(DONE_SSE_FRAME))) + .await + .is_err() + { + return None; + } + } + + Some(body_text) +} + +fn sse_body_has_terminal_event(body: &str) -> bool { + parse_sse_events(body).is_ok_and(|events| { + events.iter().any(|event| { + matches!( + event.event.as_deref(), + Some("response.completed" | "response.failed" | "error") + ) + }) + }) +} + +fn missing_sse_event_separator(body: &str) -> Option<&'static str> { + if body.is_empty() + || body.ends_with("\n\n") + || body.ends_with("\r\n\r\n") + || body.ends_with("\r\r") + { + None + } else if body.ends_with('\n') || body.ends_with('\r') { + Some("\n") + } else { + Some("\n\n") + } +} + +fn premature_close_failed_event(detail: Option<&str>) -> String { + let message = match detail.filter(|value| !value.trim().is_empty()) { + Some(detail) => format!("Upstream stream closed before response.completed: {detail}"), + None => "Upstream stream closed before response.completed".to_string(), + }; + response_stream_failed_sse_event("server_error", "stream_disconnected", &message) +} + +fn response_stream_failed_sse_event(error_type: &str, code: &str, message: &str) -> String { + let error = serde_json::json!({ + "type": error_type, + "code": code, + "message": message, + }); + let data = serde_json::json!({ + "type": "response.failed", + "response": { + "id": format!("resp_proxy_{}", uuid::Uuid::new_v4().simple()), + "status": "failed", + "error": error, + }, + "error": error, + }); + encode_sse_event("response.failed", &data.to_string()) +} + +fn sse_body_has_done(body: &str) -> bool { + body.trim_end_matches(['\r', '\n']) + .ends_with("data: [DONE]") +} + +async fn finalize_live_response_stream(context: LiveResponseStreamContext, body: String) { + let rate_limit_headers = live_response_rate_limit_headers(&context).await; + let turn_state = live_response_turn_state(&context).await; + let usage = match extract_sse_usage(&body) { + Ok(Some(usage)) => { + context + .account_pool + .record_token_usage(&context.account_id, usage) + .await; + Some(usage) + } + Ok(None) => None, + Err(error) => { + tracing::warn!(account_id = %context.account_id, error = %error, "failed to extract streaming token usage"); + None + } + }; + + match response_from_codex_sse(&body, context.tuple_schema.as_ref()) { + Ok(CollectedResponse::Completed(completed)) => { + context + .cloudflare + .reset_path_block(&context.account_id) + .await; + let response_id = completed.get("id").and_then(Value::as_str); + record_response_affinity( + &context.session_affinity, + &context.reasoning_replay, + &context.request, + &context.account_id, + &body, + turn_state, + usage, + ) + .await; + record_live_response_stream_event( + &context, + 200, + EventLevel::Info, + "v1 responses stream completed", + serde_json::json!({ + "stream": true, + "completed": true, + "responseId": response_id, + "usage": usage, + }), + &rate_limit_headers, + &body, + ) + .await; + } + Ok(CollectedResponse::Failed(failure)) => { + if sse_failure_invalid_reasoning_replay(&failure) { + evict_reasoning_replay( + &context.reasoning_replay, + &context.request, + &context.account_id, + ) + .await; + } + tracing::warn!( + account_id = %context.account_id, + event = %failure.event, + code = ?failure.upstream_code.as_deref(), + "live upstream stream ended with response.failed" + ); + record_live_response_stream_event( + &context, + status_code_for_stream_failure(&failure), + EventLevel::Error, + "v1 responses stream failed", + stream_failure_metadata(&failure, usage), + &rate_limit_headers, + &body, + ) + .await; + } + Ok(CollectedResponse::MissingCompleted | CollectedResponse::Empty) => { + tracing::warn!( + account_id = %context.account_id, + "live upstream stream ended without response.completed" + ); + record_live_response_stream_event( + &context, + 502, + EventLevel::Error, + "v1 responses stream ended without response.completed", + serde_json::json!({ + "stream": true, + "failed": true, + "upstreamCode": "missing_completed", + "usage": usage, + }), + &rate_limit_headers, + &body, + ) + .await; + } + Err(error) => { + tracing::warn!(account_id = %context.account_id, error = %error, "failed to parse completed live stream"); + record_live_response_stream_event( + &context, + 502, + EventLevel::Warn, + "v1 responses stream SSE response invalid", + serde_json::json!({ + "stream": true, + "sseParseError": error.to_string(), + "usage": usage, + }), + &rate_limit_headers, + &body, + ) + .await; + } + } + + context.account_pool.release(&context.account_id).await; +} + +async fn record_response_affinity( + session_affinity: &Arc, + reasoning_replay: &Arc>, + request: &CodexResponsesRequest, + account_id: &str, + body: &str, + turn_state: Option, + usage: Option, +) { + let Some(conversation_id) = request + .prompt_cache_key + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + return; + }; + let metadata = match completed_response_metadata(body) { + Ok(Some(metadata)) => metadata, + Ok(None) => return, + Err(error) => { + tracing::warn!( + error = %error, + "failed to parse completed response metadata for session affinity" + ); + return; + } + }; + + let variant_hash = compute_variant_hash(request); + let entry = SessionAffinityEntry { + account_id: account_id.to_string(), + conversation_id: conversation_id.to_string(), + turn_state: turn_state + .filter(|value| !value.trim().is_empty()) + .or_else(|| request.turn_state.clone()), + instructions_hash: Some(hash_instructions(Some(&request.instructions))), + input_tokens: usage.map(|usage| usage.input_tokens), + function_call_ids: metadata.function_call_ids, + variant_hash: Some(variant_hash.clone()), + created_at: Utc::now(), + }; + if let Err(error) = session_affinity + .record(metadata.response_id.clone(), entry) + .await + { + tracing::warn!( + error = %error, + response_id = %metadata.response_id, + account_id = %account_id, + "failed to record session affinity" + ); + } + + reasoning_replay.lock().await.record( + metadata.response_id, + account_id, + conversation_id, + &variant_hash, + &metadata.replay_items, + Utc::now(), + ); +} + +async fn evict_reasoning_replay( + reasoning_replay: &Arc>, + request: &CodexResponsesRequest, + account_id: &str, +) { + let Some(conversation_id) = request + .prompt_cache_key + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + return; + }; + let variant_hash = compute_variant_hash(request); + let evicted = reasoning_replay.lock().await.evict_by_identity( + account_id, + conversation_id, + &variant_hash, + Utc::now(), + ); + if evicted > 0 { + tracing::info!( + account_id = %account_id, + conversation_id = %conversation_id, + variant_hash = %variant_hash, + evicted, + "evicted reasoning replay after invalid encrypted content" + ); + } +} + +struct ResponseEventRecord<'a> { + logs: &'a AdminLogService, + request_id: &'a str, + account_id: &'a str, + model: &'a str, + started_at: Instant, + status_code: i64, + level: EventLevel, + message: &'a str, + metadata: Value, + rate_limit_headers: &'a [(String, String)], +} + +async fn record_response_event(mut record: ResponseEventRecord<'_>) { + if let Some(object) = record.metadata.as_object_mut() { + if !record.rate_limit_headers.is_empty() { + object + .entry("rateLimitHeaders".to_string()) + .or_insert_with(|| serde_json::json!(record.rate_limit_headers)); + } + } + let mut event = EventLog::new("v1.response", record.level, record.message); + event.request_id = Some(record.request_id.to_string()); + event.account_id = Some(record.account_id.to_string()); + event.route = Some("/v1/responses".to_string()); + event.model = Some(record.model.to_string()); + event.status_code = Some(record.status_code); + event.latency_ms = Some(elapsed_millis_i64(record.started_at)); + event.metadata = record.metadata; + if let Err(error) = record.logs.record(event).await { + tracing::warn!(account_id = record.account_id, error = %error, "failed to record response event"); + } +} + +async fn record_live_response_stream_event( + context: &LiveResponseStreamContext, + status_code: i64, + level: EventLevel, + message: &str, + mut metadata: Value, + rate_limit_headers: &[(String, String)], + body: &str, +) { + ensure_stream_metadata_flag(&mut metadata); + enrich_live_response_stream_metadata(context, rate_limit_headers, &mut metadata, body); + 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(status_code); + event.latency_ms = Some(elapsed_millis_i64(context.started_at)); + event.metadata = metadata; + if let Err(error) = context.logs.record(event).await { + tracing::warn!(account_id = %context.account_id, error = %error, "failed to record live response stream event"); + } +} + +fn ensure_stream_metadata_flag(metadata: &mut Value) { + let Some(object) = metadata.as_object_mut() else { + *metadata = serde_json::json!({ "stream": true }); + return; + }; + object + .entry("stream".to_string()) + .or_insert(Value::Bool(true)); +} + +fn enrich_live_response_stream_metadata( + context: &LiveResponseStreamContext, + rate_limit_headers: &[(String, String)], + metadata: &mut Value, + body: &str, +) { + let Some(object) = metadata.as_object_mut() else { + return; + }; + if context.transport == CodexBackendTransport::WebSocket { + object + .entry("transport".to_string()) + .or_insert_with(|| Value::String("websocket".to_string())); + } + if !rate_limit_headers.is_empty() { + object + .entry("rateLimitHeaders".to_string()) + .or_insert_with(|| serde_json::json!(rate_limit_headers)); + } + object + .entry("requestBody".to_string()) + .or_insert_with(|| serde_json::json!(context.request)); + object + .entry("responseBody".to_string()) + .or_insert_with(|| Value::String(body.to_string())); +} + +async fn live_response_rate_limit_headers( + context: &LiveResponseStreamContext, +) -> Vec<(String, String)> { + let mut headers = context.rate_limit_headers.clone(); + if let Some(updates) = &context.rate_limit_header_updates { + headers.extend(updates.lock().await.iter().cloned()); + } + headers +} + +async fn live_response_turn_state(context: &LiveResponseStreamContext) -> Option { + if let Some(update) = &context.turn_state_update { + return update.lock().await.clone(); + } + context.turn_state.clone() +} + +fn elapsed_millis_i64(started_at: Instant) -> i64 { + started_at.elapsed().as_millis().min(i64::MAX as u128) as i64 +} + +fn bool_to_u64(value: bool) -> u64 { + if value { + 1 + } else { + 0 + } +} + +fn stream_failure_metadata(failure: &ResponsesSseFailure, usage: Option) -> Value { + serde_json::json!({ + "stream": true, + "failed": true, + "failureEvent": failure.event, + "failureMessage": failure.message, + "upstreamCode": failure.upstream_code, + "usage": usage, + }) +} + +fn status_code_for_stream_failure(failure: &ResponsesSseFailure) -> i64 { + let code = failure + .upstream_code + .as_deref() + .unwrap_or("error") + .to_ascii_lowercase(); + if code.contains("model") && (code.contains("not_supported") || code.contains("not_available")) + { + return 400; + } + if code.contains("invalid_request") || code.contains("not_found") { + return 400; + } + if code.contains("rate_limit") || code.contains("usage_limit") { + return 429; + } + if code.contains("unauthorized") + || code.contains("invalid_api_key") + || code == "token_invalid" + || code == "token_expired" + || code == "account_deactivated" + { + return 401; + } + if code.contains("forbidden") || code.contains("banned") { + return 403; + } + if code.contains("payment") || code.contains("quota") { + return 402; + } + 502 +} + +/// Responses 调度错误。 +#[derive(Debug, Error)] +pub enum ResponseDispatchError { + /// 账号存储失败。 + #[error("failed to list runtime accounts")] + AccountStore, + /// 没有活跃账号。 + #[error("no active account is available")] + NoActiveAccount, + /// 所有账号都因配额耗尽不可用。 + #[error("all accounts exhausted by quota")] + QuotaExhausted { + /// 配额耗尽账号数量。 + count: usize, + /// 最后一个上游错误体。 + upstream_error: String, + }, + /// 所有账号都因限流不可用。 + #[error("all accounts exhausted by rate limit")] + RateLimited { + /// 限流账号数量。 + count: usize, + /// 最后一个上游错误体。 + upstream_error: String, + }, + /// 所有账号都因认证失效不可用。 + #[error("all accounts exhausted by expired auth")] + Expired { + /// 认证失效账号数量。 + count: usize, + /// 最后一个上游错误体。 + upstream_error: String, + }, + /// 所有账号都因 Cloudflare challenge 暂不可用。 + #[error("all accounts exhausted by Cloudflare challenge")] + CloudflareChallenge { + /// Cloudflare challenge 账号数量。 + count: usize, + /// 最后一个上游错误说明。 + upstream_error: String, + }, + /// 所有账号都因 Cloudflare path-block 暂不可用。 + #[error("all accounts exhausted by Cloudflare path-block")] + CloudflarePathBlocked { + /// Cloudflare path-block 账号数量。 + count: usize, + /// 最后一个上游错误说明。 + upstream_error: String, + }, + /// 所有候选账号都不支持当前模型。 + #[error("all accounts exhausted by unsupported model")] + ModelUnsupported { + /// 模型不支持账号数量。 + count: usize, + /// 最后一个上游错误体。 + upstream_error: String, + }, + /// 上游请求失败。 + #[error("upstream request failed: {0}")] + Upstream(#[from] CodexClientError), + /// 上游 SSE 无法解析。 + #[error("invalid upstream SSE response: {0}")] + InvalidSse(#[from] SseError), + /// 上游响应没有完成事件。 + #[error("upstream response did not include response.completed")] + MissingCompleted, + /// 上游响应完成但没有可见输出。 + #[error("upstream response did not include visible output")] + EmptyUpstreamResponse, + /// 上游返回失败事件。 + #[error("upstream response failed: {0:?}")] + Failed(ResponsesSseFailure), +} + +/// Responses live SSE body stream error. +#[derive(Debug, Error)] +pub enum ResponseDispatchStreamError { + /// 上游字节流读取失败。 + #[error("upstream stream failed: {0}")] + Upstream(#[from] CodexClientError), +} + +/// 管理端日志服务。 +#[derive(Clone)] +pub struct AdminLogService { + store: SqliteEventLogStore, + settings: Arc>, +} + +#[derive(Debug, Clone, Copy)] +struct AdminLogSettings { + enabled: bool, + capacity: u32, + capture_body: bool, +} + +impl AdminLogService { + /// 构造管理端日志服务。 + pub fn new( + store: SqliteEventLogStore, + enabled: bool, + capacity: u32, + capture_body: bool, + ) -> Self { + Self { + store, + settings: Arc::new(RwLock::new(AdminLogSettings { + enabled, + capacity, + capture_body, + })), + } + } + + /// 分页查询日志。 + pub async fn list( + &self, + cursor: Option, + limit: u32, + filter: AdminLogFilter, + ) -> Result, AdminLogError> { + self.store + .list(filter.into(), cursor, limit) + .await + .map_err(|_| AdminLogError::List) + } + + /// 按 ID 读取日志。 + pub async fn get(&self, id: &str) -> Result, AdminLogError> { + self.store.get(id).await.map_err(|_| AdminLogError::Get) + } + + /// 读取日志状态。 + pub async fn state(&self) -> Result { + let settings = *self.settings.read().await; + Ok(AdminLogState { + enabled: settings.enabled, + capacity: settings.capacity, + capture_body: settings.capture_body, + stored_count: self.store.count().await.map_err(|_| AdminLogError::Count)?, + }) + } + + /// 更新日志状态。 + pub async fn update_state( + &self, + update: AdminLogStateUpdate, + ) -> Result { + if matches!(update.capacity, Some(0)) { + return Err(AdminLogError::InvalidCapacity); + } + + let trim_capacity = { + let mut settings = self.settings.write().await; + if let Some(enabled) = update.enabled { + settings.enabled = enabled; + } + if let Some(capacity) = update.capacity { + settings.capacity = capacity; + } + if let Some(capture_body) = update.capture_body { + settings.capture_body = capture_body; + } + update.capacity + }; + + if let Some(capacity) = trim_capacity { + self.store + .trim_to_capacity(capacity) + .await + .map_err(|_| AdminLogError::Trim)?; + } + + self.state().await + } + + /// 清空日志。 + pub async fn clear(&self) -> Result { + self.store + .clear() + .await + .map(|cleared| AdminClearLogs { cleared }) + .map_err(|_| AdminLogError::Clear) + } + + async fn record(&self, mut event: EventLog) -> Result<(), AdminLogError> { + let settings = *self.settings.read().await; + let policy = EventLogService::new(settings.enabled); + if !policy.should_record(&event) { + return Ok(()); + } + apply_capture_body_policy(&mut event, settings.capture_body); + self.store + .append(&event) + .await + .map_err(|_| AdminLogError::Append)?; + self.store + .trim_to_capacity(settings.capacity) + .await + .map_err(|_| AdminLogError::Trim)?; + Ok(()) + } +} + +fn apply_capture_body_policy(event: &mut EventLog, capture_body: bool) { + if capture_body { + return; + } + let Some(metadata) = event.metadata.as_object_mut() else { + return; + }; + for key in [ + "body", + "rawBody", + "requestBody", + "responseBody", + "upstreamBody", + ] { + metadata.remove(key); + } +} + +/// 日志查询过滤器。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct AdminLogFilter { + /// 事件类别。 + pub kind: Option, + /// 事件等级。 + pub level: Option, + /// 请求 ID。 + pub request_id: Option, + /// 账号 ID。 + pub account_id: Option, + /// 路由。 + pub route: Option, + /// 模型。 + pub model: Option, + /// HTTP 状态码。 + pub status_code: Option, + /// 搜索关键词。 + pub search: Option, +} + +/// 日志状态。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminLogState { + /// 是否启用。 + pub enabled: bool, + /// 内存容量。 + pub capacity: u32, + /// 是否捕获请求体。 + pub capture_body: bool, + /// 已存储数量。 + pub stored_count: u64, +} + +/// 日志状态更新。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct AdminLogStateUpdate { + /// 是否启用。 + pub enabled: Option, + /// 日志容量。 + pub capacity: Option, + /// 是否捕获请求体。 + pub capture_body: Option, +} + +/// 清空日志结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminClearLogs { + /// 清理数量。 + pub cleared: u64, +} + +/// 管理端日志错误。 +#[derive(Debug, Error)] +pub enum AdminLogError { + /// 列表失败。 + #[error("failed to list event logs")] + List, + /// 读取失败。 + #[error("failed to get event log")] + Get, + /// 计数失败。 + #[error("failed to count event logs")] + Count, + /// 清空失败。 + #[error("failed to clear event logs")] + Clear, + /// 写入失败。 + #[error("failed to append event log")] + Append, + /// 裁剪失败。 + #[error("failed to trim event logs")] + Trim, + /// 日志容量非法。 + #[error("log capacity must be greater than zero")] + InvalidCapacity, +} + +/// 管理端账号服务。 +#[derive(Clone)] +pub struct AdminAccountService { + store: SqliteAccountStore, + cookies: SqliteCookieStore, + quota_thresholds: QuotaWarningThresholds, + codex: Arc, + account_pool: Arc, + token_refresher: Arc, + installation_id: Option, +} + +impl AdminAccountService { + /// 构造管理端账号服务。 + pub fn new( + store: SqliteAccountStore, + cookies: SqliteCookieStore, + quota_thresholds: QuotaWarningThresholds, + codex: Arc, + account_pool: Arc, + token_refresher: Arc, + installation_id: Option, + ) -> Self { + Self { + store, + cookies, + quota_thresholds, + codex, + account_pool, + token_refresher, + installation_id, + } + } + + /// 分页列出账号元数据。 + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> Result, AdminAccountError> { + let page = self + .store + .list_metadata(cursor, limit) + .await + .map_err(|_| AdminAccountError::List)?; + Ok(Page { + items: page + .items + .into_iter() + .map(AdminAccountMetadata::from) + .collect(), + next_cursor: page.next_cursor, + }) + } + + /// 返回管理端认证状态摘要。 + pub async fn auth_status(&self) -> Result { + let mut cursor = None; + let mut pool = AdminAuthPoolStatus::default(); + let mut user = None; + loop { + let page = self + .store + .list_metadata(cursor, 200) + .await + .map_err(|_| AdminAccountError::List)?; + for account in page.items { + pool.record(account.status); + if user.is_none() && account.status == AccountStatus::Active { + user = Some(AdminAccountMetadata::from(account)); + } + } + if page.next_cursor.is_none() { + break; + } + cursor = page.next_cursor; + } + + Ok(AdminAuthStatus { + authenticated: pool.total > 0, + user, + pool, + }) + } + + /// 清空管理端 OAuth 登录账号。 + pub async fn logout(&self) -> Result { + let deleted = self + .store + .delete_all() + .await + .map_err(|_| AdminAccountError::Delete)?; + self.account_pool.clear().await; + Ok(AdminAuthLogout { + success: true, + deleted, + }) + } + + /// 导出账号；包含可重新导入的 token，只应暴露给管理端会话。 + pub async fn export( + &self, + ids: Vec, + ) -> Result, AdminAccountError> { + if ids.is_empty() { + let mut accounts = Vec::new(); + let mut cursor = None; + loop { + let page = self + .store + .list(cursor, 200) + .await + .map_err(|_| AdminAccountError::Export)?; + accounts.extend(page.items.into_iter().map(AdminStoredAccount::from)); + if page.next_cursor.is_none() { + return Ok(accounts); + } + cursor = page.next_cursor; + } + } + + let mut accounts = Vec::with_capacity(ids.len()); + for id in ids { + match self.store.get(&id).await { + Ok(Some(account)) => accounts.push(AdminStoredAccount::from(account)), + Ok(None) => {} + Err(_) => return Err(AdminAccountError::Export), + } + } + Ok(accounts) + } + + /// 导入 native 账号导出数据。 + pub async fn import(&self, payload: &Value) -> Result { + let entries = parse_account_import_payload(payload)?; + if entries.is_empty() { + return Err(AdminAccountError::NoImportableAccounts); + } + + let mut imported = 0u32; + let mut skipped = 0u32; + for entry in entries { + match self.import_entry(entry).await? { + ImportedAccountState::Imported(account_id) => { + imported += 1; + self.sync_account_pool(&account_id).await?; + } + ImportedAccountState::Skipped => skipped += 1, + } + } + + Ok(ImportedAccounts { imported, skipped }) + } + + /// 手动创建或更新一个经 JWT claims 校验的账号。 + pub async fn create( + &self, + token: Option, + refresh_token: Option, + ) -> Result { + let provided_refresh_token = normalize_nonempty(refresh_token); + let tokens = + if let Some(access_token) = normalize_nonempty(token.map(normalize_bearer_token)) { + ManualCreateTokens { + access_token, + refresh_token_for_new: provided_refresh_token.clone(), + refresh_token_for_existing: provided_refresh_token, + } + } else if let Some(refresh_token) = provided_refresh_token { + let token_pair = self + .token_refresher + .refresh(&refresh_token) + .await + .map_err(AdminAccountError::RefreshTokenExchange)?; + let access_token = + normalize_nonempty(Some(normalize_bearer_token(token_pair.access_token))) + .ok_or(AdminAccountError::TokenRequired)?; + ManualCreateTokens { + access_token, + refresh_token_for_new: token_pair + .refresh_token + .clone() + .or_else(|| Some(refresh_token.clone())), + refresh_token_for_existing: token_pair.refresh_token, + } + } else { + return Err(AdminAccountError::TokenRequired); + }; + + let claims = manual_account_claims(&tokens.access_token, chrono::Utc::now()) + .map_err(AdminAccountError::InvalidToken)?; + let existing = self + .store + .find_by_chatgpt_identity(&claims.account_id, claims.user_id.as_deref()) + .await + .map_err(|_| AdminAccountError::Inspect)?; + + let account_id = if let Some(existing) = existing { + let refresh_token = tokens.refresh_token_for_existing; + let updated = self + .store + .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(tokens.access_token.into()), + refresh_token: refresh_token.map(|token| SecretString::new(token.into())), + access_token_expires_at: Some(claims.expires_at), + status: AccountStatus::Active, + }, + ) + .await + .map_err(|_| AdminAccountError::UpdateClaims)?; + if !updated { + return Err(AdminAccountError::NotFound); + } + existing.id + } else { + let id = normalized_account_id(None); + let refresh_token = tokens.refresh_token_for_new; + self.store + .insert(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(tokens.access_token.into()), + refresh_token: refresh_token.map(|token| SecretString::new(token.into())), + access_token_expires_at: Some(claims.expires_at), + status: AccountStatus::Active, + }) + .await + .map_err(|_| AdminAccountError::Import)?; + id + }; + + self.sync_account_pool(&account_id).await?; + + self.store + .get(&account_id) + .await + .map_err(|_| AdminAccountError::Inspect)? + .map(AdminAccountMetadata::from) + .ok_or(AdminAccountError::NotFound) + } + + /// 导入 Codex CLI 的 auth.json 内容。 + pub async fn import_codex_cli_auth( + &self, + payload: &Value, + ) -> Result { + let token = first_string(payload, &["access_token", "accessToken", "token"]); + let refresh_token = first_string(payload, &["refresh_token", "refreshToken"]); + if token.is_none() && refresh_token.is_none() { + return Err(AdminAccountError::NoImportableAccounts); + } + self.create(token, refresh_token).await + } + + async fn import_entry( + &self, + entry: AccountImportEntry, + ) -> Result { + let Some(access_token) = normalize_nonempty(entry.token) else { + return Ok(ImportedAccountState::Skipped); + }; + let id = normalized_account_id(entry.id); + + match self.store.get(&id).await { + Ok(Some(_)) => return Ok(ImportedAccountState::Skipped), + Ok(None) => {} + Err(_) => return Err(AdminAccountError::Inspect), + } + + let label = normalize_label(entry.label); + if label + .as_ref() + .is_some_and(|label| label.chars().count() > 64) + { + return Err(AdminAccountError::LabelTooLong); + } + + let access_token_expires_at = entry + .access_token_expires_at + .as_deref() + .map(parse_account_import_datetime) + .transpose()?; + let status = parse_account_import_status(entry.status.as_deref())?; + let account_id = id.clone(); + let account = NewAccount { + id, + email: normalize_nonempty(entry.email), + account_id: normalize_nonempty(entry.account_id), + user_id: normalize_nonempty(entry.user_id), + label, + plan_type: normalize_nonempty(entry.plan_type), + access_token: SecretString::new(access_token.into()), + refresh_token: normalize_nonempty(entry.refresh_token) + .map(|token| SecretString::new(token.into())), + access_token_expires_at, + status, + }; + self.store + .insert(account) + .await + .map_err(|_| AdminAccountError::Import)?; + + Ok(ImportedAccountState::Imported(account_id)) + } + + /// 更新账号标签。 + 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(AdminAccountError::LabelTooLong); + } + let updated = self + .store + .set_label(account_id, label) + .await + .map_err(|_| AdminAccountError::UpdateLabel)?; + if updated { + self.sync_account_pool_best_effort(account_id, "account label update") + .await; + } + Ok(updated) + } + + /// 更新账号状态。 + pub async fn update_status( + &self, + account_id: &str, + status: &str, + ) -> Result, AdminAccountError> { + let status = parse_account_status(status)?; + match self.store.set_status(account_id, status).await { + Ok(true) => { + self.sync_account_pool_best_effort(account_id, "account status update") + .await; + self.evict_account_websocket_pool(account_id).await; + Ok(Some(UpdatedAccountStatus { + id: account_id.to_string(), + status, + })) + } + Ok(false) => Ok(None), + Err(_) => Err(AdminAccountError::UpdateStatus), + } + } + + /// 删除账号。 + pub async fn delete(&self, account_id: &str) -> Result { + let deleted = self + .store + .delete(account_id) + .await + .map_err(|_| AdminAccountError::Delete)?; + if deleted { + self.account_pool.remove_account(account_id).await; + } + Ok(deleted) + } + + /// 批量删除账号。 + pub async fn batch_delete( + &self, + ids: Vec, + ) -> Result { + if ids.is_empty() { + return Err(AdminAccountError::EmptyIds); + } + + let mut deleted = 0u32; + let mut not_found = Vec::new(); + for id in ids { + match self.store.delete(&id).await { + Ok(true) => { + deleted += 1; + self.account_pool.remove_account(&id).await; + } + Ok(false) => not_found.push(id), + Err(_) => return Err(AdminAccountError::Delete), + } + } + + Ok(BatchDeleteAccounts { deleted, not_found }) + } + + /// 批量更新账号状态。 + pub async fn batch_update_status( + &self, + ids: Vec, + status: &str, + ) -> Result { + if ids.is_empty() { + return Err(AdminAccountError::EmptyIds); + } + let status = parse_batch_account_status(status)?; + + let mut updated = 0u32; + let mut not_found = Vec::new(); + for id in ids { + match self.store.set_status(&id, status).await { + Ok(true) => { + updated += 1; + self.sync_account_pool_best_effort(&id, "account batch status update") + .await; + self.evict_account_websocket_pool(&id).await; + } + Ok(false) => not_found.push(id), + Err(_) => return Err(AdminAccountError::UpdateStatus), + } + } + + Ok(BatchUpdateAccountStatus { updated, not_found }) + } + + /// 使用账号 refresh token 刷新 access token。 + pub async fn refresh_account( + &self, + account_id: &str, + ) -> Result { + let account = match self.store.get(account_id).await { + Ok(Some(account)) => account, + Ok(None) => return Err(AdminAccountError::NotFound), + Err(_) => return Err(AdminAccountError::Inspect), + }; + let previous_status = account.status; + let Some(refresh_token) = account.refresh_token.as_ref() else { + return Err(AdminAccountError::TokenRequired); + }; + + match self + .token_refresher + .refresh(refresh_token.expose_secret()) + .await + { + Ok(tokens) => { + let access_token = + normalize_nonempty(Some(normalize_bearer_token(tokens.access_token))) + .ok_or(AdminAccountError::TokenRequired)?; + let claims = manual_account_claims(&access_token, chrono::Utc::now()) + .map_err(AdminAccountError::InvalidToken)?; + let updated = self + .store + .update_from_claims( + account_id, + AccountClaimsUpdate { + email: claims.email, + account_id: Some(claims.account_id), + user_id: claims.user_id, + plan_type: claims.plan_type, + access_token: SecretString::new(access_token.into()), + refresh_token: tokens + .refresh_token + .map(|token| SecretString::new(token.into())), + access_token_expires_at: Some(claims.expires_at), + status: AccountStatus::Active, + }, + ) + .await + .map_err(|_| AdminAccountError::UpdateClaims)?; + if !updated { + return Err(AdminAccountError::NotFound); + } + self.sync_account_pool(account_id).await?; + Ok(AdminAccountRefresh { + id: account_id.to_string(), + previous_status, + outcome: AdminAccountProbeOutcome::Alive, + status: Some(AccountStatus::Active), + error: None, + }) + } + Err(failure) => { + let status = refresh_failure_status(failure); + let updated = self + .store + .set_status(account_id, status) + .await + .map_err(|_| AdminAccountError::UpdateStatus)?; + if !updated { + return Err(AdminAccountError::NotFound); + } + self.sync_account_pool_best_effort(account_id, "account refresh failure") + .await; + Ok(AdminAccountRefresh { + id: account_id.to_string(), + previous_status, + outcome: AdminAccountProbeOutcome::Dead, + status: Some(status), + error: Some(failure.to_string()), + }) + } + } + } + + /// 重置账号本地用量计数。 + pub async fn reset_usage( + &self, + account_id: &str, + ) -> Result { + match self.store.get(account_id).await { + Ok(Some(_)) => {} + Ok(None) => return Err(AdminAccountError::NotFound), + Err(_) => return Err(AdminAccountError::Inspect), + } + + self.store + .reset_usage(account_id) + .await + .map_err(|_| AdminAccountError::ResetUsage)?; + self.sync_account_pool(account_id).await?; + + Ok(AdminAccountResetUsage { + id: account_id.to_string(), + reset: true, + }) + } + + /// 读取账号 Cookie 请求头。 + pub async fn cookies(&self, account_id: &str) -> Result, AdminAccountError> { + self.ensure_cookie_account_exists(account_id).await?; + self.cookies + .cookie_header(account_id, "chatgpt.com") + .await + .map_err(|_| AdminAccountError::LoadCookies) + } + + /// 设置账号 Cookie 请求头。 + pub async fn set_cookies( + &self, + account_id: &str, + cookie_header: &str, + ) -> Result, AdminAccountError> { + self.ensure_cookie_account_exists(account_id).await?; + match self + .cookies + .set_cookie_header(account_id, cookie_header) + .await + { + Ok(0) => Err(AdminAccountError::NoValidCookies), + Ok(_) => self + .cookies + .cookie_header(account_id, "chatgpt.com") + .await + .map_err(|_| AdminAccountError::LoadCookies), + Err(_) => Err(AdminAccountError::StoreCookies), + } + } + + /// 删除账号 Cookie。 + pub async fn delete_cookies(&self, account_id: &str) -> Result<(), AdminAccountError> { + self.ensure_cookie_account_exists(account_id).await?; + self.cookies + .delete_account_cookies(account_id) + .await + .map(|_| ()) + .map_err(|_| AdminAccountError::DeleteCookies) + } + + /// 返回基于缓存配额快照的账号预警。 + pub async fn quota_warnings(&self) -> Result { + let snapshots = self + .store + .list_quota_snapshots() + .await + .map_err(|_| AdminAccountError::QuotaWarnings)?; + Ok(quota_warnings_from_snapshots( + snapshots, + &self.quota_thresholds, + )) + } + + /// 拉取并持久化单个账号的 Codex usage 配额快照。 + pub async fn account_quota( + &self, + account_id: &str, + request_id: &str, + ) -> Result { + let account = match self.store.get(account_id).await { + Ok(Some(account)) => account, + Ok(None) => return Err(AdminAccountError::NotFound), + Err(_) => return Err(AdminAccountError::Inspect), + }; + if account.status != AccountStatus::Active { + return Err(AdminAccountError::Inactive(account.status)); + } + + let raw = self + .codex + .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: None, + installation_id: self.installation_id.as_deref(), + session_id: None, + }) + .await + .map_err(|error| AdminAccountError::FetchQuota(error.to_string()))?; + let quota = quota_from_usage(&raw); + let updated = self + .store + .update_quota_json(&account.id, "a.to_string()) + .await + .map_err(|_| AdminAccountError::StoreQuota)?; + if !updated { + return Err(AdminAccountError::NotFound); + } + + Ok(AdminAccountQuota { quota, raw }) + } + + /// 对账号执行 refresh-token 健康探测。 + pub async fn health_check_accounts( + &self, + ids: Option>, + concurrency: usize, + stagger_ms: u64, + _request_id: &str, + ) -> Result, AdminAccountError> { + let accounts = self.health_check_candidates(ids).await?; + let concurrency = concurrency.max(1); + let results = stream::iter(accounts.into_iter().enumerate()) + .map(|(index, account)| { + let service = self.clone(); + async move { + if stagger_ms > 0 && index > 0 { + let multiplier = index.min(concurrency); + tokio::time::sleep(std::time::Duration::from_millis( + stagger_ms.saturating_mul(multiplier as u64), + )) + .await; + } + service.probe_account_refresh(account).await + } + }) + .buffer_unordered(concurrency) + .collect::>() + .await; + Ok(results) + } + + async fn health_check_candidates( + &self, + ids: Option>, + ) -> Result, AdminAccountError> { + if let Some(ids) = ids { + let mut accounts = Vec::with_capacity(ids.len()); + for id in ids { + match self.store.get(&id).await { + Ok(Some(account)) => accounts.push(account), + Ok(None) => {} + Err(_) => return Err(AdminAccountError::HealthCheck), + } + } + return Ok(accounts); + } + + let mut accounts = Vec::new(); + let mut cursor = None; + loop { + let page = self + .store + .list(cursor, 200) + .await + .map_err(|_| AdminAccountError::HealthCheck)?; + accounts.extend(page.items); + if page.next_cursor.is_none() { + return Ok(accounts); + } + cursor = page.next_cursor; + } + } + + async fn probe_account_refresh(&self, account: StoredAccount) -> AdminAccountProbeResult { + let started_at = Instant::now(); + let previous_status = account.status; + if account.status == AccountStatus::Disabled { + return skipped_admin_account_probe_result(account, "manually disabled"); + } + let Some(refresh_token) = account.refresh_token.as_ref() else { + return skipped_admin_account_probe_result(account, "no refresh token"); + }; + + match self + .token_refresher + .refresh(refresh_token.expose_secret()) + .await + { + Ok(tokens) => { + let Some(access_token) = + normalize_nonempty(Some(normalize_bearer_token(tokens.access_token))) + else { + return dead_admin_account_probe_result( + account, + previous_status, + "token or refreshToken is required".to_string(), + started_at, + ); + }; + let claims = match manual_account_claims(&access_token, Utc::now()) { + Ok(claims) => claims, + Err(error) => { + return dead_admin_account_probe_result( + account, + previous_status, + error.to_string(), + started_at, + ); + } + }; + + match self + .store + .update_from_claims( + &account.id, + AccountClaimsUpdate { + email: claims.email, + account_id: Some(claims.account_id), + user_id: claims.user_id, + plan_type: claims.plan_type, + access_token: SecretString::new(access_token.into()), + refresh_token: tokens + .refresh_token + .map(|token| SecretString::new(token.into())), + access_token_expires_at: Some(claims.expires_at), + status: AccountStatus::Active, + }, + ) + .await + { + Ok(true) => { + self.sync_account_pool_best_effort(&account.id, "account health refresh") + .await; + AdminAccountProbeResult { + id: account.id, + email: account.email, + previous_status, + outcome: AdminAccountProbeOutcome::Alive, + status: Some(AccountStatus::Active), + error: None, + duration_ms: Some(started_at.elapsed().as_millis()), + } + } + Ok(false) => dead_admin_account_probe_result( + account, + previous_status, + AdminAccountError::NotFound.to_string(), + started_at, + ), + Err(_) => dead_admin_account_probe_result( + account, + previous_status, + AdminAccountError::UpdateClaims.to_string(), + started_at, + ), + } + } + Err(failure) => { + let status = health_check_failure_status(failure); + if let Some(status) = status { + match self.store.set_status(&account.id, status).await { + Ok(true) => { + self.sync_account_pool_best_effort( + &account.id, + "account health refresh failure", + ) + .await; + } + Ok(false) => { + return dead_admin_account_probe_result( + account, + previous_status, + AdminAccountError::NotFound.to_string(), + started_at, + ); + } + Err(_) => { + return dead_admin_account_probe_result( + account, + previous_status, + AdminAccountError::UpdateStatus.to_string(), + started_at, + ); + } + } + } + AdminAccountProbeResult { + id: account.id, + email: account.email, + previous_status, + outcome: AdminAccountProbeOutcome::Dead, + status, + error: Some(failure.to_string()), + duration_ms: Some(started_at.elapsed().as_millis()), + } + } + } + } + + async fn ensure_cookie_account_exists( + &self, + account_id: &str, + ) -> Result<(), AdminAccountError> { + match self.cookies.account_exists(account_id).await { + Ok(true) => Ok(()), + Ok(false) => Err(AdminAccountError::NotFound), + Err(_) => Err(AdminAccountError::Inspect), + } + } + + async fn sync_account_pool(&self, account_id: &str) -> Result<(), AdminAccountError> { + self.account_pool + .sync_account_from_repository(account_id) + .await + .map(|_| ()) + .map_err(|_| AdminAccountError::SyncAccountPool) + } + + async fn sync_account_pool_best_effort(&self, account_id: &str, operation: &str) { + if let Err(error) = self + .account_pool + .sync_account_from_repository(account_id) + .await + { + tracing::warn!( + account_id, + operation, + error = %error, + "failed to sync runtime account pool after admin account update" + ); + } + } + + async fn evict_account_websocket_pool(&self, account_id: &str) { + self.codex.evict_websocket_account(account_id).await; + match self.store.get(account_id).await { + Ok(Some(account)) => { + if let Some(upstream_account_id) = account + .account_id + .as_deref() + .filter(|value| *value != account_id) + { + self.codex + .evict_websocket_account(upstream_account_id) + .await; + } + } + Ok(None) => {} + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to inspect account while evicting websocket pool" + ); + } + } + } +} + +fn skipped_admin_account_probe_result( + account: StoredAccount, + error: &str, +) -> AdminAccountProbeResult { + AdminAccountProbeResult { + id: account.id, + email: account.email, + previous_status: account.status, + outcome: AdminAccountProbeOutcome::Skipped, + status: Some(account.status), + error: Some(error.to_string()), + duration_ms: None, + } +} + +fn dead_admin_account_probe_result( + account: StoredAccount, + previous_status: AccountStatus, + error: String, + started_at: Instant, +) -> AdminAccountProbeResult { + AdminAccountProbeResult { + id: account.id, + email: account.email, + previous_status, + outcome: AdminAccountProbeOutcome::Dead, + status: None, + error: Some(error), + duration_ms: Some(started_at.elapsed().as_millis()), + } +} + +fn health_check_failure_status(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, + } +} + +/// 管理端账号错误。 +#[derive(Debug, Error)] +pub enum AdminAccountError { + /// 列表失败。 + #[error("failed to list accounts")] + List, + /// 导出失败。 + #[error("failed to export accounts")] + Export, + /// 导入失败。 + #[error("failed to import accounts")] + Import, + /// 检查账号失败。 + #[error("failed to inspect account")] + Inspect, + /// 更新标签失败。 + #[error("failed to update account label")] + UpdateLabel, + /// 更新状态失败。 + #[error("failed to update account status")] + UpdateStatus, + /// 删除失败。 + #[error("failed to delete account")] + Delete, + /// 重置用量失败。 + #[error("failed to reset account usage")] + ResetUsage, + /// 账号不存在。 + #[error("account not found")] + NotFound, + /// 读取 Cookie 失败。 + #[error("failed to load account cookies")] + LoadCookies, + /// 写入 Cookie 失败。 + #[error("failed to store account cookies")] + StoreCookies, + /// 删除 Cookie 失败。 + #[error("failed to delete account cookies")] + DeleteCookies, + /// 根据 JWT claims 更新账号失败。 + #[error("failed to update account claims")] + UpdateClaims, + /// 读取配额预警失败。 + #[error("failed to load account quota warnings")] + QuotaWarnings, + /// 写入配额快照失败。 + #[error("failed to store account quota")] + StoreQuota, + /// 拉取配额失败。 + #[error("failed to fetch account quota: {0}")] + FetchQuota(String), + /// 健康检查失败。 + #[error("failed to health-check accounts")] + HealthCheck, + /// 账号非 active，不能执行需要上游访问的操作。 + #[error("account is {0:?}, cannot query quota")] + Inactive(AccountStatus), + /// token 为空。 + #[error("token or refreshToken is required")] + TokenRequired, + /// token 非法。 + #[error("{0}")] + InvalidToken(&'static str), + /// refresh token 换取 access token 失败。 + #[error("failed to exchange refreshToken: {0}")] + RefreshTokenExchange(RefreshFailure), + /// 同步运行时账号池失败。 + #[error("failed to sync runtime account pool")] + SyncAccountPool, + /// 没有有效 Cookie。 + #[error("No valid cookies found")] + NoValidCookies, + /// 标签过长。 + #[error("account label must be 64 characters or fewer")] + LabelTooLong, + /// 状态值无效。 + #[error("unsupported account status: {0}")] + InvalidStatus(String), + /// ID 列表为空。 + #[error("account ids are required")] + EmptyIds, + /// 没有可导入账号。 + #[error("No importable accounts found")] + NoImportableAccounts, + /// access token 过期时间非法。 + #[error("invalid accessTokenExpiresAt")] + InvalidAccessTokenExpiresAt, +} + +/// 管理端账号元数据。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminAccountMetadata { + /// 账号 ID。 + pub id: String, + /// 邮箱。 + pub email: Option, + /// ChatGPT 账号 ID。 + pub account_id: Option, + /// ChatGPT 用户 ID。 + pub user_id: Option, + /// 管理标签。 + pub label: Option, + /// 计划类型。 + pub plan_type: Option, + /// access token 过期时间。 + pub access_token_expires_at: Option>, + /// 账号状态。 + pub status: codex_proxy_core::accounts::model::AccountStatus, + /// 创建时间。 + pub added_at: chrono::DateTime, + /// 更新时间。 + pub updated_at: chrono::DateTime, +} + +/// 管理端账号配额拉取结果。 +#[derive(Debug, Clone, PartialEq)] +pub struct AdminAccountQuota { + /// 归一化后的配额快照。 + pub quota: Value, + /// Codex usage 原始响应。 + pub raw: Value, +} + +/// 管理端可导出的完整账号数据。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminStoredAccount { + /// 账号 ID。 + pub id: String, + /// 邮箱。 + pub email: Option, + /// ChatGPT 账号 ID。 + pub account_id: Option, + /// ChatGPT 用户 ID。 + pub user_id: Option, + /// 管理标签。 + pub label: Option, + /// 计划类型。 + pub plan_type: Option, + /// access token 明文。 + pub access_token: String, + /// refresh token 明文。 + pub refresh_token: Option, + /// access token 过期时间。 + pub access_token_expires_at: Option>, + /// 账号状态。 + pub status: codex_proxy_core::accounts::model::AccountStatus, + /// 创建时间。 + pub added_at: chrono::DateTime, + /// 更新时间。 + pub updated_at: chrono::DateTime, +} + +/// 管理端账号配额预警集合。 +#[derive(Debug, Clone, PartialEq)] +pub struct AdminAccountQuotaWarnings { + /// 预警列表。 + pub warnings: Vec, + /// 产生预警的快照中最新的拉取时间。 + pub updated_at: Option>, +} + +/// 管理端账号配额预警。 +#[derive(Debug, Clone, PartialEq)] +pub struct AdminAccountQuotaWarning { + /// 账号 ID。 + pub account_id: String, + /// 邮箱。 + pub email: Option, + /// 配额窗口。 + pub window: AdminQuotaWarningWindow, + /// 预警级别。 + pub level: AdminQuotaWarningLevel, + /// 已使用百分比。 + pub used_percent: f64, + /// 重置时间戳。 + pub reset_at: Option, +} + +/// 配额预警窗口。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdminQuotaWarningWindow { + /// 主窗口。 + Primary, + /// 次窗口。 + Secondary, +} + +impl AdminQuotaWarningWindow { + /// 返回 API 字符串值。 + pub fn as_str(self) -> &'static str { + match self { + Self::Primary => "primary", + Self::Secondary => "secondary", + } + } +} + +/// 配额预警级别。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdminQuotaWarningLevel { + /// 普通预警。 + Warning, + /// 临界预警。 + Critical, +} + +impl AdminQuotaWarningLevel { + /// 返回 API 字符串值。 + pub fn as_str(self) -> &'static str { + match self { + Self::Warning => "warning", + Self::Critical => "critical", + } + } +} + +/// 账号健康探测结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminAccountProbeResult { + /// 账号 ID。 + pub id: String, + /// 邮箱。 + pub email: Option, + /// 探测前状态。 + pub previous_status: AccountStatus, + /// 探测结果。 + pub outcome: AdminAccountProbeOutcome, + /// 探测后的状态。 + pub status: Option, + /// 错误信息。 + pub error: Option, + /// 耗时毫秒。 + pub duration_ms: Option, +} + +/// 账号健康探测结果类型。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdminAccountProbeOutcome { + /// 上游 usage 请求成功。 + Alive, + /// 上游 usage 请求失败。 + Dead, + /// 未执行上游探测。 + Skipped, +} + +impl AdminAccountProbeOutcome { + /// 返回 API 字符串值。 + pub fn as_str(self) -> &'static str { + match self { + Self::Alive => "alive", + Self::Dead => "dead", + Self::Skipped => "skipped", + } + } +} + +/// 账号导入结果。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ImportedAccounts { + /// 成功导入数量。 + pub imported: u32, + /// 跳过数量。 + pub skipped: u32, +} + +/// 管理端认证状态摘要。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminAuthStatus { + /// 是否存在已导入账号。 + pub authenticated: bool, + /// 当前可展示的 active 账号元数据。 + pub user: Option, + /// 账号池状态计数。 + pub pool: AdminAuthPoolStatus, +} + +/// 管理端认证状态中的账号池计数。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct AdminAuthPoolStatus { + /// 账号总数。 + pub total: u32, + /// active 账号数。 + pub active: u32, + /// expired 账号数。 + pub expired: u32, + /// quota_exhausted 账号数。 + pub quota_exhausted: u32, + /// refreshing 账号数。 + pub refreshing: u32, + /// disabled 账号数。 + pub disabled: u32, + /// banned 账号数。 + pub banned: u32, +} + +impl AdminAuthPoolStatus { + fn record(&mut self, status: AccountStatus) { + self.total += 1; + match status { + AccountStatus::Active => self.active += 1, + AccountStatus::Expired => self.expired += 1, + AccountStatus::QuotaExhausted => self.quota_exhausted += 1, + AccountStatus::Refreshing => self.refreshing += 1, + AccountStatus::Disabled => self.disabled += 1, + AccountStatus::Banned => self.banned += 1, + } + } +} + +/// 管理端登出结果。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AdminAuthLogout { + /// 是否成功。 + pub success: bool, + /// 删除账号数。 + pub deleted: u64, +} + +#[derive(Debug, Clone)] +struct ManualCreateTokens { + access_token: String, + refresh_token_for_new: Option, + refresh_token_for_existing: Option, +} + +#[derive(Debug, Clone)] +struct AccountImportEntry { + id: Option, + email: Option, + account_id: Option, + user_id: Option, + label: Option, + plan_type: Option, + token: Option, + refresh_token: Option, + access_token_expires_at: Option, + status: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum ImportedAccountState { + Imported(String), + Skipped, +} + +/// 账号状态更新结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UpdatedAccountStatus { + /// 账号 ID。 + pub id: String, + /// 新状态。 + pub status: codex_proxy_core::accounts::model::AccountStatus, +} + +/// 批量删除账号结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BatchDeleteAccounts { + /// 成功删除数量。 + pub deleted: u32, + /// 未找到的账号 ID。 + pub not_found: Vec, +} + +/// 批量状态更新结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BatchUpdateAccountStatus { + /// 成功更新数量。 + pub updated: u32, + /// 未找到的账号 ID。 + pub not_found: Vec, +} + +/// 管理端手动刷新账号结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminAccountRefresh { + /// 账号 ID。 + pub id: String, + /// 刷新前状态。 + pub previous_status: codex_proxy_core::accounts::model::AccountStatus, + /// 刷新结果。 + pub outcome: AdminAccountProbeOutcome, + /// 刷新后状态。 + pub status: Option, + /// 错误信息。 + pub error: Option, +} + +/// 管理端重置用量结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminAccountResetUsage { + /// 账号 ID。 + pub id: String, + /// 是否已处理。 + pub reset: bool, +} + +/// 运行时设置服务。 +#[derive(Clone)] +pub struct RuntimeSettingsService { + current: Arc>>, + local_config_path: Arc, +} + +impl RuntimeSettingsService { + /// 构造运行时设置服务。 + pub fn new(config: AppConfig) -> Self { + Self::with_local_config_path(config, "local.yaml") + } + + /// 构造带本地配置覆盖路径的运行时设置服务。 + pub fn with_local_config_path( + config: AppConfig, + local_config_path: impl Into, + ) -> Self { + Self { + current: Arc::new(StdRwLock::new(Arc::new(config))), + local_config_path: Arc::new(local_config_path.into()), + } + } + + /// 返回当前配置快照。 + pub fn current(&self) -> Arc { + self.current + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + + /// Return the local settings overlay path configured for this runtime. + pub fn local_config_path(&self) -> Arc { + self.local_config_path.clone() + } + + /// 更新当前设置并写入本地配置覆盖文件。 + pub async fn update( + &self, + patch: AdminSettingsPatch, + ) -> Result, RuntimeSettingsError> { + let mut next = (*self.current()).clone(); + let mut settings = admin_settings_from_config(&next); + SettingsService::apply_patch(&mut settings, patch)?; + apply_admin_settings_to_config(&mut next, settings); + next.write_settings_overlay(self.local_config_path.as_ref()) + .await?; + let next = Arc::new(next); + *self + .current + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = next.clone(); + Ok(next) + } +} + +/// 运行时设置错误。 +#[derive(Debug, Error)] +pub enum RuntimeSettingsError { + /// 设置补丁验证失败。 + #[error(transparent)] + InvalidField(#[from] SettingsServiceError), + /// 本地配置覆盖写入失败。 + #[error(transparent)] + Persist(#[from] ConfigWriteError), +} + +fn admin_settings_from_config(config: &AppConfig) -> AdminSettings { + AdminSettings { + 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: AdminQuotaWarningThresholds { + primary: config.quota.warning_thresholds.primary.clone(), + secondary: config.quota.warning_thresholds.secondary.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, + } +} + +fn apply_admin_settings_to_config(config: &mut AppConfig, settings: AdminSettings) { + config.model.default_model = settings.default_model; + config.model.default_reasoning_effort = settings.default_reasoning_effort; + config.model.service_tier = settings.service_tier; + config.model.aliases = settings.model_aliases; + config.auth.refresh_enabled = settings.refresh_enabled; + config.auth.refresh_margin_seconds = settings.refresh_margin_seconds; + config.auth.refresh_concurrency = settings.refresh_concurrency; + config.auth.max_concurrent_per_account = settings.max_concurrent_per_account; + config.auth.request_interval_ms = settings.request_interval_ms; + config.auth.rotation_strategy = settings.rotation_strategy; + config.auth.tier_priority = settings.tier_priority; + config.quota.refresh_interval_minutes = settings.quota_refresh_interval_minutes; + config.quota.warning_thresholds = QuotaWarningThresholds { + primary: settings.quota_warning_thresholds.primary, + secondary: settings.quota_warning_thresholds.secondary, + }; + config.quota.skip_exhausted = settings.quota_skip_exhausted; + config.logging.enabled = settings.logs_enabled; + config.logging.capacity = settings.logs_capacity; + config.logging.capture_body = settings.logs_capture_body; + config.usage_stats.history_retention_days = settings.usage_history_retention_days; +} + +/// 管理员会话服务。 +#[derive(Clone)] +pub struct AdminSessionService { + store: SqliteAdminSessionStore, + auth: AdminAuthService, + default_username: String, + session_ttl_minutes: u64, +} + +impl AdminSessionService { + /// 构造管理员会话服务。 + pub fn new( + store: SqliteAdminSessionStore, + default_username: String, + session_ttl_minutes: u64, + ) -> Self { + Self { + store, + auth: AdminAuthService::new(default_username.clone()), + default_username, + session_ttl_minutes, + } + } + + /// 校验管理员会话是否存在且未过期。 + pub async fn validate(&self, session_id: Option<&str>) -> Result { + let Some(session_id) = session_id else { + return Ok(false); + }; + self.store + .validate_session(session_id) + .await + .map_err(|_| AdminSessionError::Validate) + } + + /// 如果还没有管理员用户，则根据配置密码创建默认管理员。 + pub async fn ensure_default_admin(&self, password: &str) -> Result { + let password_hash = + hash_admin_password(password).map_err(|_| AdminSessionError::HashPassword)?; + self.store + .ensure_default_admin(&password_hash) + .await + .map_err(|_| AdminSessionError::CreateAdmin) + } + + /// 使用管理员用户名和密码创建会话。 + pub async fn login( + &self, + username: Option<&str>, + password: &str, + ) -> Result, AdminSessionError> { + let username = username.unwrap_or(&self.default_username); + if !self.auth.username_matches(username) { + return Ok(None); + } + + let Some(admin) = self + .store + .load_first_admin() + .await + .map_err(|_| AdminSessionError::LoadAdmin)? + else { + return Ok(None); + }; + let password_matches = verify_admin_password(password, &admin.password_hash) + .map_err(|_| AdminSessionError::VerifyPassword)?; + if !password_matches { + return Ok(None); + } + + let session_id = format!("sess_{}", uuid::Uuid::new_v4().simple()); + let ttl_minutes = self.session_ttl_minutes.min(i64::MAX as u64) as i64; + let expires_at = Utc::now() + Duration::minutes(ttl_minutes); + self.store + .create_session(&session_id, &admin.id, expires_at) + .await + .map_err(|_| AdminSessionError::CreateSession)?; + + Ok(Some(AdminLoginSession { + session_id, + expires_at, + })) + } +} + +/// 管理员登录成功后的会话。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminLoginSession { + /// 会话 ID。 + pub session_id: String, + /// 过期时间。 + pub expires_at: DateTime, +} + +/// 管理员会话错误。 +#[derive(Debug, Error)] +pub enum AdminSessionError { + /// 会话校验失败。 + #[error("failed to validate admin session")] + Validate, + /// 密码哈希失败。 + #[error("failed to hash admin password")] + HashPassword, + /// 创建管理员失败。 + #[error("failed to create default admin user")] + CreateAdmin, + /// 读取管理员失败。 + #[error("failed to load admin user")] + LoadAdmin, + /// 密码校验失败。 + #[error("failed to verify admin password")] + VerifyPassword, + /// 创建会话失败。 + #[error("failed to create admin session")] + CreateSession, +} + +/// 管理端客户端 API Key 服务。 +#[derive(Clone)] +pub struct AdminClientKeyService { + store: SqliteClientKeyStore, +} + +impl AdminClientKeyService { + /// 构造管理端客户端 API Key 服务。 + pub fn new(store: SqliteClientKeyStore) -> Self { + Self { store } + } + + /// 创建新的客户端 API Key。 + pub async fn create( + &self, + name: &str, + ) -> Result { + let name = name.trim(); + if name.is_empty() { + return Err(AdminClientKeyError::EmptyName); + } + self.store + .create(name) + .await + .map(AdminCreatedClientApiKey::from) + .map_err(|_| AdminClientKeyError::Create) + } + + /// 分页列出客户端 API Key。 + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> Result, AdminClientKeyError> { + let page = self + .store + .list(cursor, limit) + .await + .map_err(|_| AdminClientKeyError::List)?; + Ok(Page { + items: page + .items + .into_iter() + .map(AdminStoredClientApiKey::from) + .collect(), + next_cursor: page.next_cursor, + }) + } + + /// 更新客户端 API Key 标签。 + pub async fn update_label( + &self, + key_id: &str, + label: Option, + ) -> Result, AdminClientKeyError> { + if label + .as_ref() + .is_some_and(|label| label.chars().count() > 64) + { + return Err(AdminClientKeyError::LabelTooLong); + } + self.store + .set_label(key_id, label) + .await + .map(|key| key.map(AdminStoredClientApiKey::from)) + .map_err(|_| AdminClientKeyError::UpdateLabel) + } + + /// 更新客户端 API Key 启用状态。 + pub async fn update_status( + &self, + key_id: &str, + status: &str, + ) -> Result, AdminClientKeyError> { + let enabled = parse_client_key_status(status)?; + match self.store.set_enabled(key_id, enabled).await { + Ok(true) => Ok(Some(UpdatedClientApiKeyStatus { + id: key_id.to_string(), + enabled, + })), + Ok(false) => Ok(None), + Err(_) => Err(AdminClientKeyError::UpdateStatus), + } + } + + /// 删除客户端 API Key。 + pub async fn delete(&self, key_id: &str) -> Result { + self.store + .delete(key_id) + .await + .map_err(|_| AdminClientKeyError::Delete) + } + + /// 批量删除客户端 API Key。 + pub async fn batch_delete( + &self, + ids: Vec, + ) -> Result { + if ids.is_empty() { + return Err(AdminClientKeyError::EmptyIds); + } + + let mut deleted = 0u32; + let mut not_found = Vec::new(); + for id in ids { + match self.store.delete(&id).await { + Ok(true) => deleted += 1, + Ok(false) => not_found.push(id), + Err(_) => return Err(AdminClientKeyError::Delete), + } + } + + Ok(BatchDeleteClientApiKeys { deleted, not_found }) + } + + /// 导出客户端 API Key 元数据，不包含明文和哈希材料。 + pub async fn export( + &self, + ids: Vec, + ) -> Result, AdminClientKeyError> { + if ids.is_empty() { + let mut all_keys = Vec::new(); + let mut cursor = None; + loop { + let page = self + .store + .list(cursor, 200) + .await + .map_err(|_| AdminClientKeyError::Export)?; + all_keys.extend(page.items.into_iter().map(AdminStoredClientApiKey::from)); + if page.next_cursor.is_none() { + return Ok(all_keys); + } + cursor = page.next_cursor; + } + } + + let mut keys = Vec::with_capacity(ids.len()); + for id in ids { + match self.store.get(&id).await { + Ok(Some(key)) => keys.push(AdminStoredClientApiKey::from(key)), + Ok(None) => {} + Err(_) => return Err(AdminClientKeyError::Export), + } + } + Ok(keys) + } + + /// 导入导出的客户端 API Key 元数据，并轮换为新的本地明文。 + pub async fn import( + &self, + payload: &Value, + ) -> Result { + let entries = parse_client_key_import_payload(payload); + if entries.is_empty() { + return Err(AdminClientKeyError::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(AdminClientKeyError::LabelTooLong); + } + + let mut created = self.create(name).await?; + if entry.label.is_some() || !entry.enabled { + self.store + .set_label(&created.id, entry.label) + .await + .map_err(|_| AdminClientKeyError::Import)?; + if !entry.enabled { + self.store + .set_enabled(&created.id, false) + .await + .map_err(|_| AdminClientKeyError::Import)?; + } + let Some(stored) = self + .store + .get(&created.id) + .await + .map_err(|_| AdminClientKeyError::Import)? + else { + return Err(AdminClientKeyError::Import); + }; + created.label = stored.label; + created.enabled = stored.enabled; + } + imported += 1; + keys.push(ImportedClientApiKey { + source_id: entry.source_id, + source_prefix: entry.source_prefix, + key: created, + }); + } + + Ok(ImportedClientApiKeys { + imported, + skipped, + keys, + }) + } +} + +/// 管理端可见的客户端 API Key 元数据。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminStoredClientApiKey { + /// API Key 记录 ID。 + pub id: String, + /// API Key 名称。 + pub name: String, + /// 管理员可见标签。 + pub label: Option, + /// 明文 API Key 的短前缀。 + pub prefix: String, + /// 是否允许用于 `/v1` 认证。 + pub enabled: bool, + /// 创建时间。 + pub created_at: String, + /// 最近一次成功使用时间。 + pub last_used_at: Option, +} + +/// 新建客户端 API Key 后的一次性结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminCreatedClientApiKey { + /// API Key 记录 ID。 + pub id: String, + /// API Key 名称。 + pub name: String, + /// 管理员可见标签。 + pub label: Option, + /// 明文 API Key 的短前缀。 + pub prefix: String, + /// 是否允许用于 `/v1` 认证。 + pub enabled: bool, + /// 创建时间。 + pub created_at: String, + /// 最近一次成功使用时间。 + pub last_used_at: Option, + /// 仅返回一次的明文 API Key。 + pub plaintext: String, +} + +/// 客户端 API Key 状态更新结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UpdatedClientApiKeyStatus { + /// API Key 记录 ID。 + pub id: String, + /// 是否启用。 + pub enabled: bool, +} + +/// 批量删除客户端 API Key 的结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BatchDeleteClientApiKeys { + /// 成功删除数量。 + pub deleted: u32, + /// 未找到的 ID。 + pub not_found: Vec, +} + +/// 导入后的客户端 API Key。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ImportedClientApiKey { + /// 来源 ID。 + pub source_id: Option, + /// 来源短前缀。 + pub source_prefix: Option, + /// 新建的本地 API Key。 + pub key: AdminCreatedClientApiKey, +} + +/// 客户端 API Key 导入结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ImportedClientApiKeys { + /// 成功导入数量。 + pub imported: u32, + /// 跳过数量。 + pub skipped: u32, + /// 新建 API Key 列表。 + pub keys: Vec, +} + +/// 管理端客户端 API Key 错误。 +#[derive(Debug, Error)] +pub enum AdminClientKeyError { + /// 列表失败。 + #[error("failed to list client API keys")] + List, + /// 导出失败。 + #[error("failed to export client API keys")] + Export, + /// 导入失败。 + #[error("failed to import client API keys")] + Import, + /// 创建失败。 + #[error("failed to create client API key")] + Create, + /// 删除失败。 + #[error("failed to delete client API key")] + Delete, + /// 更新标签失败。 + #[error("failed to update client API key label")] + UpdateLabel, + /// 更新状态失败。 + #[error("failed to update client API key status")] + UpdateStatus, + /// 状态值无效。 + #[error("unsupported client API key status: {0}")] + InvalidStatus(String), + /// 名称为空。 + #[error("client API key name is required")] + EmptyName, + /// ID 列表为空。 + #[error("client API key ids are required")] + EmptyIds, + /// 标签过长。 + #[error("client API key label must be 64 characters or fewer")] + LabelTooLong, + /// 没有可导入的 API Key。 + #[error("no importable client API keys found")] + NoImportableKeys, +} + +#[derive(Debug, Clone)] +struct ClientApiKeyImportEntry { + source_id: Option, + source_prefix: Option, + name: String, + label: Option, + enabled: bool, +} + +fn parse_client_key_status(status: &str) -> Result { + match status.trim().to_ascii_lowercase().as_str() { + "active" => Ok(true), + "disabled" => Ok(false), + other => Err(AdminClientKeyError::InvalidStatus(other.to_string())), + } +} + +fn parse_account_status( + status: &str, +) -> Result { + match status.trim().to_ascii_lowercase().as_str() { + "active" => Ok(codex_proxy_core::accounts::model::AccountStatus::Active), + "disabled" => Ok(codex_proxy_core::accounts::model::AccountStatus::Disabled), + "expired" => Ok(codex_proxy_core::accounts::model::AccountStatus::Expired), + "quota_exhausted" => Ok(codex_proxy_core::accounts::model::AccountStatus::QuotaExhausted), + "refreshing" => Ok(codex_proxy_core::accounts::model::AccountStatus::Refreshing), + "banned" => Ok(codex_proxy_core::accounts::model::AccountStatus::Banned), + other => Err(AdminAccountError::InvalidStatus(other.to_string())), + } +} + +fn parse_batch_account_status( + status: &str, +) -> Result { + match status.trim().to_ascii_lowercase().as_str() { + "active" => Ok(codex_proxy_core::accounts::model::AccountStatus::Active), + "disabled" => Ok(codex_proxy_core::accounts::model::AccountStatus::Disabled), + other => Err(AdminAccountError::InvalidStatus(other.to_string())), + } +} + +fn refresh_failure_status( + failure: RefreshFailure, +) -> codex_proxy_core::accounts::model::AccountStatus { + match failure { + RefreshFailure::InvalidGrant => codex_proxy_core::accounts::model::AccountStatus::Expired, + RefreshFailure::QuotaExhausted => { + codex_proxy_core::accounts::model::AccountStatus::QuotaExhausted + } + RefreshFailure::Banned => codex_proxy_core::accounts::model::AccountStatus::Banned, + RefreshFailure::Disabled => codex_proxy_core::accounts::model::AccountStatus::Disabled, + RefreshFailure::Transport => codex_proxy_core::accounts::model::AccountStatus::Active, + } +} + +const ACCOUNT_IMPORT_ENVELOPE_KEYS: &[&str] = + &["code", "message", "data", "requestId", "request_id"]; +const ACCOUNT_IMPORT_CONTAINER_KEYS: &[&str] = &["sourceFormat", "source_format", "accounts"]; +const ACCOUNT_IMPORT_ACCOUNT_KEYS: &[&str] = &[ + "id", + "email", + "accountId", + "account_id", + "userId", + "user_id", + "label", + "planType", + "plan_type", + "token", + "accessToken", + "access_token", + "refreshToken", + "refresh_token", + "accessTokenExpiresAt", + "access_token_expires_at", + "status", + "addedAt", + "added_at", + "updatedAt", + "updated_at", +]; + +fn parse_account_import_payload( + payload: &Value, +) -> Result, AdminAccountError> { + let payload = payload + .get("data") + .filter(|data| data.get("accounts").is_some()) + .map(|data| { + ensure_account_import_keys(payload, ACCOUNT_IMPORT_ENVELOPE_KEYS)?; + Ok(data) + }) + .transpose()? + .unwrap_or(payload); + + if let Some(accounts) = payload.get("accounts") { + ensure_account_import_keys(payload, ACCOUNT_IMPORT_CONTAINER_KEYS)?; + ensure_native_account_import_source(payload)?; + let accounts = accounts + .as_array() + .ok_or(AdminAccountError::NoImportableAccounts)?; + return parse_account_import_entries(accounts); + } + if let Some(accounts) = payload.as_array() { + return parse_account_import_entries(accounts); + } + + Ok(account_import_entry_from_value(payload)? + .into_iter() + .collect()) +} + +fn parse_account_import_entries( + accounts: &[Value], +) -> Result, AdminAccountError> { + let mut entries = Vec::new(); + for account in accounts { + if let Some(entry) = account_import_entry_from_value(account)? { + entries.push(entry); + } + } + Ok(entries) +} + +fn account_import_entry_from_value( + value: &Value, +) -> Result, AdminAccountError> { + let Some(account) = value.as_object() else { + return Ok(None); + }; + if account + .keys() + .any(|key| !ACCOUNT_IMPORT_ACCOUNT_KEYS.contains(&key.as_str())) + { + return Err(AdminAccountError::NoImportableAccounts); + } + + let token = first_string(value, &["token", "accessToken", "access_token"]); + let refresh_token = first_string(value, &["refreshToken", "refresh_token"]); + if token.is_none() && refresh_token.is_none() { + return Ok(None); + } + + Ok(Some(AccountImportEntry { + id: first_string(value, &["id"]), + email: first_string(value, &["email"]), + account_id: first_string(value, &["accountId", "account_id"]), + user_id: first_string(value, &["userId", "user_id"]), + label: first_string(value, &["label"]), + plan_type: first_string(value, &["planType", "plan_type"]), + token, + refresh_token, + access_token_expires_at: first_string( + value, + &["accessTokenExpiresAt", "access_token_expires_at"], + ), + status: first_string(value, &["status"]), + })) +} + +fn ensure_account_import_keys( + value: &Value, + allowed_keys: &[&str], +) -> Result<(), AdminAccountError> { + let Some(object) = value.as_object() else { + return Err(AdminAccountError::NoImportableAccounts); + }; + if object + .keys() + .all(|key| allowed_keys.contains(&key.as_str())) + { + Ok(()) + } else { + Err(AdminAccountError::NoImportableAccounts) + } +} + +fn ensure_native_account_import_source(value: &Value) -> Result<(), AdminAccountError> { + if let Some(source_format) = first_string(value, &["sourceFormat", "source_format"]) { + if !source_format.eq_ignore_ascii_case("native") { + return Err(AdminAccountError::NoImportableAccounts); + } + } + Ok(()) +} + +fn parse_account_import_status( + status: Option<&str>, +) -> Result { + parse_account_status(status.unwrap_or("active")) +} + +fn parse_account_import_datetime( + value: &str, +) -> Result, AdminAccountError> { + chrono::DateTime::parse_from_rfc3339(value) + .map(|value| value.with_timezone(&chrono::Utc)) + .map_err(|_| AdminAccountError::InvalidAccessTokenExpiresAt) +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct ManualAccountClaims { + account_id: String, + user_id: Option, + email: Option, + plan_type: Option, + expires_at: chrono::DateTime, +} + +fn manual_account_claims( + token: &str, + now: chrono::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 = + chrono::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; + } + 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) +} + +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() +} + +fn normalized_account_id(id: Option) -> String { + normalize_nonempty(id).unwrap_or_else(|| format!("acct_{}", uuid::Uuid::new_v4().simple())) +} + +fn normalize_label(value: Option) -> Option { + normalize_nonempty(value) +} + +fn normalize_nonempty(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +fn parse_client_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_key_import_entry_from_value) + .collect(); + } + if let Some(keys) = payload.get("keys").and_then(Value::as_array) { + return keys + .iter() + .filter_map(client_key_import_entry_from_value) + .collect(); + } + if let Some(keys) = payload.as_array() { + return keys + .iter() + .filter_map(client_key_import_entry_from_value) + .collect(); + } + + client_key_import_entry_from_value(payload) + .into_iter() + .collect() +} + +fn client_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_key_import_enabled(value), + }) +} + +fn client_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 first_string(value: &Value, keys: &[&str]) -> Option { + keys.iter() + .find_map(|key| value.get(key).and_then(Value::as_str)) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) +} + +fn quota_warnings_from_snapshots( + snapshots: Vec, + thresholds: &QuotaWarningThresholds, +) -> AdminAccountQuotaWarnings { + let primary_thresholds = sorted_thresholds(&thresholds.primary); + let secondary_thresholds = sorted_thresholds(&thresholds.secondary); + let mut warnings = Vec::new(); + let mut updated_at = None; + + for snapshot in snapshots { + let Ok(quota) = serde_json::from_str::(&snapshot.quota_json) else { + continue; + }; + let before_len = warnings.len(); + if let Some(warning) = warning_from_quota_window( + &snapshot.account_id, + snapshot.email.as_deref(), + "a, + "rate_limit", + AdminQuotaWarningWindow::Primary, + &primary_thresholds, + ) { + warnings.push(warning); + } + if let Some(warning) = warning_from_quota_window( + &snapshot.account_id, + snapshot.email.as_deref(), + "a, + "secondary_rate_limit", + AdminQuotaWarningWindow::Secondary, + &secondary_thresholds, + ) { + warnings.push(warning); + } + if warnings.len() > before_len { + updated_at = max_optional_datetime(updated_at, snapshot.quota_fetched_at); + } + } + + AdminAccountQuotaWarnings { + warnings, + updated_at, + } +} + +fn warning_from_quota_window( + account_id: &str, + email: Option<&str>, + quota: &Value, + field: &str, + window: AdminQuotaWarningWindow, + thresholds: &[u8], +) -> Option { + let quota_window = quota.get(field).filter(|value| !value.is_null())?; + let used_percent = quota_window + .get("used_percent") + .and_then(Value::as_f64) + .filter(|value| value.is_finite())?; + let level = warning_level(used_percent, thresholds)?; + + Some(AdminAccountQuotaWarning { + account_id: account_id.to_string(), + email: email.map(ToString::to_string), + window, + level, + used_percent, + reset_at: quota_window.get("reset_at").and_then(Value::as_i64), + }) +} + +fn warning_level(used_percent: f64, thresholds: &[u8]) -> Option { + let matched_index = thresholds + .iter() + .rposition(|threshold| quota_reached(used_percent, f64::from(*threshold)))?; + if matched_index + 1 == thresholds.len() { + Some(AdminQuotaWarningLevel::Critical) + } else { + Some(AdminQuotaWarningLevel::Warning) + } +} + +fn sorted_thresholds(thresholds: &[u8]) -> Vec { + let mut thresholds = thresholds.to_vec(); + thresholds.sort_unstable(); + thresholds.dedup(); + thresholds +} + +fn max_optional_datetime( + current: Option>, + candidate: Option>, +) -> Option> { + match (current, candidate) { + (Some(current), Some(candidate)) => Some(current.max(candidate)), + (Some(current), None) => Some(current), + (None, Some(candidate)) => Some(candidate), + (None, None) => None, + } +} + +impl From for AdminStoredAccount { + fn from(account: StoredAccount) -> 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, + access_token: account.access_token.expose_secret().to_string(), + refresh_token: account + .refresh_token + .map(|token| token.expose_secret().to_string()), + access_token_expires_at: account.access_token_expires_at, + status: account.status, + added_at: account.added_at, + updated_at: account.updated_at, + } + } +} + +impl From for AdminAccountMetadata { + fn from(account: StoredAccount) -> 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, + access_token_expires_at: account.access_token_expires_at, + status: account.status, + added_at: account.added_at, + updated_at: account.updated_at, + } + } +} + +impl From for AdminStoredClientApiKey { + 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 AdminCreatedClientApiKey { + fn from(key: CreatedClientApiKey) -> 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: key.plaintext, + } + } +} + +impl From for AdminAccountMetadata { + 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, + access_token_expires_at: account.access_token_expires_at, + status: account.status, + added_at: account.added_at, + updated_at: account.updated_at, + } + } +} + +impl From for AdminUsageRecord { + 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, + empty_response_count: usage.empty_response_count, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_request_count: usage.image_request_count, + image_request_failed_count: usage.image_request_failed_count, + last_used_at: usage.last_used_at, + } + } +} + +impl From for AdminUsageSummary { + fn from(summary: AccountUsageSummary) -> Self { + Self { + account_count: summary.account_count, + request_count: summary.request_count, + empty_response_count: summary.empty_response_count, + input_tokens: summary.input_tokens, + output_tokens: summary.output_tokens, + cached_tokens: summary.cached_tokens, + image_input_tokens: summary.image_input_tokens, + image_output_tokens: summary.image_output_tokens, + image_request_count: summary.image_request_count, + image_request_failed_count: summary.image_request_failed_count, + } + } +} + +impl From for EventLogFilter { + fn from(filter: AdminLogFilter) -> Self { + Self { + kind: filter.kind, + level: filter.level, + request_id: filter.request_id, + account_id: filter.account_id, + route: filter.route, + model: filter.model, + status_code: filter.status_code, + search: filter.search, + } + } +} diff --git a/crates/runtime/src/state.rs b/crates/runtime/src/state.rs new file mode 100644 index 000000000..9c8e1dee1 --- /dev/null +++ b/crates/runtime/src/state.rs @@ -0,0 +1,316 @@ +//! 运行时状态。 + +use std::{path::PathBuf, sync::Arc}; + +use chrono::{DateTime, Utc}; +use codex_proxy_core::{ + auth::ports::{OAuthClient, TokenRefresher}, + gateway::fingerprint::Fingerprint, +}; +use codex_proxy_platform::{ + config::AppConfig, crypto::SecretBox, identity::ApiKeyHasher, storage::SqlitePool, +}; + +use crate::{ + config::RuntimeConfig, + repositories::sqlite_repositories, + services::{RuntimeAccountPoolError, RuntimeSessionAffinityError, Services}, +}; + +/// 应用状态。 +#[derive(Clone)] +pub struct AppState { + /// 运行时配置。 + pub config: RuntimeConfig, + /// 运行时服务。 + pub services: Services, +} + +struct AppStateServiceOptions { + fingerprint: Fingerprint, + installation_id: Option, + local_config_path: PathBuf, + token_refresher: Option>, + oauth_client: Option>, +} + +impl AppState { + /// 从连接池和密钥箱构造状态。 + pub fn with_pool_and_secret_box( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + ) -> Self { + Self::with_pool_secret_api_key_hasher_and_fingerprint( + config, + pool, + secret_box, + ApiKeyHasher::new([0u8; 32]), + Fingerprint::default_codex_desktop(), + ) + } + + /// 从连接池、密钥箱和 API key hasher 构造状态。 + pub fn with_pool_secret_and_api_key_hasher( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + hasher: ApiKeyHasher, + ) -> Self { + Self::with_pool_secret_api_key_hasher_and_fingerprint( + config, + pool, + secret_box, + hasher, + Fingerprint::default_codex_desktop(), + ) + } + + /// 从连接池、密钥箱、API key hasher 和 installation id 构造状态。 + pub fn with_pool_secret_api_key_hasher_and_installation_id( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + hasher: ApiKeyHasher, + installation_id: String, + ) -> Self { + Self::with_pool_secret_api_key_hasher_fingerprint_and_installation_id( + config, + pool, + secret_box, + hasher, + Fingerprint::default_codex_desktop(), + Some(installation_id), + ) + } + + /// 从连接池、密钥箱、API key hasher 和指纹构造状态。 + pub fn with_pool_secret_api_key_hasher_and_fingerprint( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + hasher: ApiKeyHasher, + fingerprint: Fingerprint, + ) -> Self { + Self::with_pool_secret_api_key_hasher_fingerprint_and_installation_id( + config, + pool, + secret_box, + hasher, + fingerprint, + None, + ) + } + + /// 从连接池、密钥箱、API key hasher、指纹和 installation id 构造状态。 + pub fn with_pool_secret_api_key_hasher_fingerprint_and_installation_id( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + hasher: ApiKeyHasher, + fingerprint: Fingerprint, + installation_id: Option, + ) -> Self { + Self::with_pool_secret_api_key_hasher_fingerprint_installation_id_and_local_config_path( + config, + pool, + secret_box, + hasher, + fingerprint, + installation_id, + "local.yaml", + ) + } + + /// 从连接池、密钥箱、API key hasher 和本地配置路径构造状态。 + pub fn with_pool_secret_api_key_hasher_and_local_config_path( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + hasher: ApiKeyHasher, + local_config_path: impl Into, + ) -> Self { + Self::with_pool_secret_api_key_hasher_fingerprint_installation_id_and_local_config_path( + config, + pool, + secret_box, + hasher, + Fingerprint::default_codex_desktop(), + None, + local_config_path, + ) + } + + /// 从连接池、密钥箱、API key hasher、指纹、installation id 和本地配置路径构造状态。 + pub fn with_pool_secret_api_key_hasher_fingerprint_installation_id_and_local_config_path( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + hasher: ApiKeyHasher, + fingerprint: Fingerprint, + installation_id: Option, + local_config_path: impl Into, + ) -> Self { + Self::with_pool_secret_api_key_hasher_and_options( + config, + pool, + secret_box, + hasher, + AppStateServiceOptions { + fingerprint, + installation_id, + local_config_path: local_config_path.into(), + token_refresher: None, + oauth_client: None, + }, + ) + } + + /// 从连接池、密钥箱、API key hasher 和 token refresher 构造状态。 + pub fn with_pool_secret_api_key_hasher_and_token_refresher( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + hasher: ApiKeyHasher, + token_refresher: C, + ) -> Self + where + C: TokenRefresher, + { + Self::with_pool_secret_api_key_hasher_and_options( + config, + pool, + secret_box, + hasher, + AppStateServiceOptions { + fingerprint: Fingerprint::default_codex_desktop(), + installation_id: None, + local_config_path: PathBuf::from("local.yaml"), + token_refresher: Some(Arc::new(token_refresher)), + oauth_client: None, + }, + ) + } + + /// 从连接池、密钥箱、API key hasher 和 OAuth client 构造状态。 + pub fn with_pool_secret_api_key_hasher_and_oauth_client( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + hasher: ApiKeyHasher, + oauth_client: C, + ) -> Self + where + C: OAuthClient + TokenRefresher + Clone, + { + let token_refresher: Arc = Arc::new(oauth_client.clone()); + let oauth_client: Arc = Arc::new(oauth_client); + Self::with_pool_secret_api_key_hasher_and_options( + config, + pool, + secret_box, + hasher, + AppStateServiceOptions { + fingerprint: Fingerprint::default_codex_desktop(), + installation_id: None, + local_config_path: PathBuf::from("local.yaml"), + token_refresher: Some(token_refresher), + oauth_client: Some(oauth_client), + }, + ) + } + + fn with_pool_secret_api_key_hasher_and_options( + config: AppConfig, + pool: SqlitePool, + secret_box: SecretBox, + hasher: ApiKeyHasher, + options: AppStateServiceOptions, + ) -> Self { + let AppStateServiceOptions { + fingerprint, + installation_id, + local_config_path, + token_refresher, + oauth_client, + } = options; + let runtime_config: RuntimeConfig = config.clone().into(); + let repositories = sqlite_repositories(pool, secret_box, hasher); + let services = match (token_refresher, oauth_client) { + (Some(token_refresher), Some(oauth_client)) => { + Services::with_installation_id_local_config_path_and_oauth_clients( + &config, + repositories, + fingerprint, + installation_id, + local_config_path, + token_refresher, + oauth_client, + ) + } + (Some(token_refresher), None) => { + Services::with_installation_id_local_config_path_and_token_refresher( + &config, + repositories, + fingerprint, + installation_id, + local_config_path, + token_refresher, + ) + } + (None, Some(oauth_client)) => { + let token_refresher: Arc = Arc::new( + codex_proxy_adapters::oauth::openai::default_openai_oauth_client( + crate::services::oauth_config(&config), + ), + ); + Services::with_installation_id_local_config_path_and_oauth_clients( + &config, + repositories, + fingerprint, + installation_id, + local_config_path, + token_refresher, + oauth_client, + ) + } + (None, None) => Services::with_installation_id_and_local_config_path( + &config, + repositories, + fingerprint, + installation_id, + local_config_path, + ), + }; + Self { + config: runtime_config, + services, + } + } + + /// 从持久化存储恢复未过期的会话亲和性映射。 + pub async fn restore_session_affinity_from_repository( + &self, + now: DateTime, + ) -> Result { + self.services + .session_affinity + .restore_from_repository(now) + .await + } + + /// 使用当前时间从持久化存储恢复未过期的会话亲和性映射。 + pub async fn restore_session_affinity_from_repository_now( + &self, + ) -> Result { + self.restore_session_affinity_from_repository(Utc::now()) + .await + } + + /// 从持久化存储恢复运行时账号池。 + pub async fn restore_account_pool_from_repository( + &self, + ) -> Result { + self.services.account_pool.restore_from_repository().await + } +} diff --git a/src/codex/tasks/cookie_cleanup.rs b/crates/runtime/src/tasks/cookie_cleanup.rs similarity index 51% rename from src/codex/tasks/cookie_cleanup.rs rename to crates/runtime/src/tasks/cookie_cleanup.rs index f7b55753b..467fa0d71 100644 --- a/src/codex/tasks/cookie_cleanup.rs +++ b/crates/runtime/src/tasks/cookie_cleanup.rs @@ -1,56 +1,64 @@ +//! Cookie 清理任务。 + use std::time::Duration; -use chrono::Utc; +use chrono::{DateTime, Utc}; +use codex_proxy_adapters::sqlite::cookies::{SqliteCookieStore, SqliteCookieStoreResult}; use tokio::time::interval; use tracing::{debug, info, warn}; -use crate::{ - codex::accounts::cookies::repository::{CookieRepository, CookieResult}, - runtime::tasks::types::SchedulerHandle, -}; +use super::coordinator::SchedulerHandle; const DEFAULT_INTERVAL_SECS: u64 = 5 * 60; -pub struct CookieCleanupScheduler { - repository: CookieRepository, +/// Cookie 清理任务。 +pub struct CookieCleanupTask { + store: SqliteCookieStore, interval_secs: u64, } -impl CookieCleanupScheduler { - pub fn new(repository: CookieRepository) -> Self { +impl CookieCleanupTask { + /// 构造默认 Cookie 清理任务。 + pub fn new(store: SqliteCookieStore) -> Self { Self { - repository, + store, interval_secs: DEFAULT_INTERVAL_SECS, } } + /// 使用自定义间隔构造 Cookie 清理任务。 + pub fn with_interval_secs(store: SqliteCookieStore, interval_secs: u64) -> Self { + Self { + store, + interval_secs, + } + } + + /// 启动后台清理任务。 pub fn start(self) -> SchedulerHandle { let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); tokio::spawn(async move { let mut ticker = interval(Duration::from_secs(self.interval_secs)); - info!( - interval_secs = self.interval_secs, - "cookie 清理调度器已启动" - ); + info!(interval_secs = self.interval_secs, "Cookie 清理任务已启动"); loop { tokio::select! { _ = ticker.tick() => { match self.cleanup_once().await { Ok(deleted) if deleted > 0 => { - info!(deleted, "已清理过期 cookie"); + info!(deleted, "已清理过期 Cookie"); } Ok(_) => { - debug!("没有需要清理的过期 cookie"); + debug!("没有需要清理的过期 Cookie"); } Err(error) => { - warn!(error = %error, "清理过期 cookie 失败"); + warn!(error = %error, "清理过期 Cookie 失败"); } } } _ = shutdown_rx.recv() => { - info!("cookie 清理调度器已关闭"); + info!("Cookie 清理任务已关闭"); break; } } @@ -60,7 +68,13 @@ impl CookieCleanupScheduler { SchedulerHandle::new(shutdown_tx) } - pub async fn cleanup_once(&self) -> CookieResult { - self.repository.cleanup_expired(Utc::now()).await + /// 执行一次清理。 + pub async fn cleanup_once(&self) -> SqliteCookieStoreResult { + self.cleanup_once_at(Utc::now()).await + } + + /// 在指定时间点执行一次清理。 + pub async fn cleanup_once_at(&self, now: DateTime) -> SqliteCookieStoreResult { + self.store.cleanup_expired(now).await } } diff --git a/crates/runtime/src/tasks/coordinator.rs b/crates/runtime/src/tasks/coordinator.rs new file mode 100644 index 000000000..8d0e620f8 --- /dev/null +++ b/crates/runtime/src/tasks/coordinator.rs @@ -0,0 +1,177 @@ +//! 后台任务协调器。 + +use std::path::PathBuf; + +use codex_proxy_adapters::oauth::openai::default_openai_oauth_client; +use codex_proxy_core::auth::oauth::{OAuthConfig, RefreshPolicy}; +use tokio::task::JoinHandle; + +use crate::{config::RuntimeConfig, state::AppState}; + +use super::{ + cookie_cleanup::CookieCleanupTask, + fingerprint_update::{ + start_fingerprint_update_task, CODEX_DESKTOP_APPCAST_URL, + DEFAULT_EXTRACTED_FINGERPRINT_PATH, + }, + model_refresh::ModelRefreshTask, + quota_refresh::QuotaRefreshTask, + session_affinity_cleanup::SessionAffinityCleanupTask, + session_cleanup::SessionCleanupTask, + token_refresh::TokenRefreshTask, +}; + +/// 后台任务协调器。 +#[derive(Default)] +pub struct BackgroundTaskCoordinator { + handles: Vec<(&'static str, SchedulerHandle)>, +} + +impl BackgroundTaskCoordinator { + /// 注册一个后台任务句柄。 + pub fn push(&mut self, name: &'static str, handle: SchedulerHandle) { + self.handles.push((name, handle)); + } + + /// 返回已注册任务名称。 + pub fn task_names(&self) -> Vec<&'static str> { + self.handles.iter().map(|(name, _)| *name).collect() + } + + /// 关闭所有后台任务。 + pub async fn shutdown(self) { + tracing::info!("正在关闭后台任务"); + for (name, handle) in self.handles { + handle.shutdown().await; + tracing::info!(task = name, "后台任务已停止"); + } + tracing::info!("所有后台任务已停止"); + } +} + +/// 启动后台任务。 +pub async fn start_background_tasks(state: &AppState) -> BackgroundTaskCoordinator { + let mut coordinator = BackgroundTaskCoordinator::default(); + let stores = &state.services.background_tasks; + + coordinator.push( + "cookie_cleanup", + CookieCleanupTask::new(stores.cookies.clone()).start(), + ); + coordinator.push( + "session_cleanup", + SessionCleanupTask::new( + stores.admin_sessions.clone(), + state.config.admin.session_cleanup_interval_secs, + ) + .start(), + ); + coordinator.push( + "session_affinity_cleanup", + SessionAffinityCleanupTask::new( + stores.session_affinity.clone(), + state.config.admin.session_cleanup_interval_secs, + ) + .start(), + ); + coordinator.push( + "model_refresh", + ModelRefreshTask::new( + state.services.models.clone(), + state.services.accounts.clone(), + ) + .with_installation_id(state.services.installation_id.clone()) + .start(), + ); + if state.config.auth.refresh_enabled { + coordinator.push( + "token_refresh", + TokenRefreshTask::new( + stores.accounts.clone(), + RefreshPolicy { + refresh_margin_seconds: state.config.auth.refresh_margin_seconds, + refresh_concurrency: state.config.auth.refresh_concurrency, + }, + default_openai_oauth_client(oauth_config(&state.config)), + ) + .with_refresh_lease_store(stores.refresh_leases.clone()) + .start(), + ); + } + coordinator.push( + "quota_refresh", + QuotaRefreshTask::with_intervals( + stores.accounts.clone(), + state.services.codex.clone(), + state + .config + .quota + .refresh_interval_minutes + .saturating_mul(60) + .max(1), + 30 * 60, + ) + .with_installation_id(state.services.installation_id.clone()) + .start(), + ); + coordinator.push( + "fingerprint_update", + start_fingerprint_update_task( + Some(stores.fingerprints.clone()), + CODEX_DESKTOP_APPCAST_URL.to_string(), + PathBuf::from(DEFAULT_EXTRACTED_FINGERPRINT_PATH), + state.services.fingerprint.app_version.clone(), + state.services.fingerprint.build_number.clone(), + ), + ); + + coordinator +} + +fn oauth_config(config: &RuntimeConfig) -> OAuthConfig { + OAuthConfig { + client_id: config.auth.oauth_client_id.clone(), + auth_endpoint: config.auth.oauth_auth_endpoint.clone(), + device_code_endpoint: oauth_device_code_endpoint(&config.auth.oauth_token_endpoint), + token_endpoint: config.auth.oauth_token_endpoint.clone(), + } +} + +fn oauth_device_code_endpoint(token_endpoint: &str) -> String { + token_endpoint + .strip_suffix("/token") + .map(|prefix| format!("{prefix}/device/code")) + .unwrap_or_else(|| "https://auth.openai.com/oauth/device/code".to_string()) +} + +/// 后台任务关闭句柄。 +pub enum SchedulerHandle { + /// 通过 channel 发送关闭信号。 + Channel(tokio::sync::mpsc::Sender<()>), + /// 直接持有 tokio 任务句柄。 + JoinHandle(JoinHandle<()>), +} + +impl SchedulerHandle { + /// 使用关闭 channel 构造句柄。 + pub fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>) -> Self { + Self::Channel(shutdown_tx) + } + + /// 使用 `JoinHandle` 构造句柄。 + pub fn from_join_handle(handle: JoinHandle<()>) -> Self { + Self::JoinHandle(handle) + } + + /// 关闭任务。 + pub async fn shutdown(self) { + match self { + Self::Channel(shutdown_tx) => { + let _ = shutdown_tx.send(()).await; + } + Self::JoinHandle(handle) => { + handle.abort(); + } + } + } +} diff --git a/crates/runtime/src/tasks/fingerprint_update.rs b/crates/runtime/src/tasks/fingerprint_update.rs new file mode 100644 index 000000000..c955b4d80 --- /dev/null +++ b/crates/runtime/src/tasks/fingerprint_update.rs @@ -0,0 +1,35 @@ +//! 指纹更新任务接线。 + +use std::path::PathBuf; + +use codex_proxy_adapters::codex::{ + default_http_client, + fingerprint::{FingerprintRepository, UpdateChecker}, +}; + +use super::coordinator::SchedulerHandle; + +/// Codex Desktop 官方 appcast 地址。 +pub const CODEX_DESKTOP_APPCAST_URL: &str = + "https://persistent.oaistatic.com/codex-app-prod/appcast.xml"; +/// 本地可选完整指纹提取文件路径。 +pub const DEFAULT_EXTRACTED_FINGERPRINT_PATH: &str = "data/extracted-fingerprint.json"; + +/// 启动指纹自动更新后台任务。 +pub fn start_fingerprint_update_task( + repository: Option, + appcast_url: String, + extracted_fingerprint_path: PathBuf, + current_version: String, + current_build: String, +) -> SchedulerHandle { + let update_checker = UpdateChecker::with_client( + repository, + default_http_client(), + appcast_url, + extracted_fingerprint_path, + current_version, + current_build, + ); + SchedulerHandle::from_join_handle(update_checker.start_background_checker()) +} diff --git a/crates/runtime/src/tasks/mod.rs b/crates/runtime/src/tasks/mod.rs new file mode 100644 index 000000000..80fb24831 --- /dev/null +++ b/crates/runtime/src/tasks/mod.rs @@ -0,0 +1,10 @@ +//! 后台任务入口。 + +pub mod cookie_cleanup; +pub mod coordinator; +pub mod fingerprint_update; +pub mod model_refresh; +pub mod quota_refresh; +pub mod session_affinity_cleanup; +pub mod session_cleanup; +pub mod token_refresh; diff --git a/src/codex/tasks/model_refresh.rs b/crates/runtime/src/tasks/model_refresh.rs similarity index 52% rename from src/codex/tasks/model_refresh.rs rename to crates/runtime/src/tasks/model_refresh.rs index dfc92dfa8..6c8e38ba6 100644 --- a/src/codex/tasks/model_refresh.rs +++ b/crates/runtime/src/tasks/model_refresh.rs @@ -1,66 +1,74 @@ -use std::sync::Arc; -use std::time::Duration; +//! 模型刷新任务。 + +use std::{sync::Arc, time::Duration}; + +use codex_proxy_core::{ + accounts::{model::Account, ports::AccountStore}, + models::service::{ModelService, ModelServiceError}, +}; use tokio::time::{interval, sleep}; use tracing::{info, warn}; -use crate::{codex::models::service::ModelService, runtime::tasks::types::SchedulerHandle}; +use super::coordinator::SchedulerHandle; -/// 模型刷新调度器 - 定期从 Codex 后端刷新模型列表 -/// -/// 功能： -/// - 启动后1秒进行首次刷新 -/// - 首次刷新失败时每10秒重试，最多12次 -/// - 首次成功后，每1小时刷新一次 -/// - 为每个计划类型刷新一次模型列表 -pub struct ModelRefresher { +/// 模型刷新任务接线器。 +pub struct ModelRefreshTask { model_service: Arc, + account_store: Arc, refresh_interval_secs: u64, initial_delay_ms: u64, retry_delay_ms: u64, max_retries: u32, + installation_id: Option, } -const DEFAULT_REFRESH_INTERVAL_SECS: u64 = 3600; // 1小时 -const INITIAL_DELAY_MS: u64 = 1000; // 1秒 -const RETRY_DELAY_MS: u64 = 10_000; // 10秒 +const DEFAULT_REFRESH_INTERVAL_SECS: u64 = 3600; +const INITIAL_DELAY_MS: u64 = 1000; +const RETRY_DELAY_MS: u64 = 10_000; const MAX_RETRIES: u32 = 12; -impl ModelRefresher { - pub fn new(model_service: Arc) -> Self { +impl ModelRefreshTask { + /// 构造默认任务。 + pub fn new(model_service: Arc, account_store: Arc) -> Self { Self { model_service, + account_store, refresh_interval_secs: DEFAULT_REFRESH_INTERVAL_SECS, initial_delay_ms: INITIAL_DELAY_MS, retry_delay_ms: RETRY_DELAY_MS, max_retries: MAX_RETRIES, + installation_id: None, } } - /// 启动模型刷新调度器 + /// 设置 Codex installation id。 + pub fn with_installation_id(mut self, installation_id: Option) -> Self { + self.installation_id = installation_id.filter(|id| !id.trim().is_empty()); + self + } + + /// 启动后台刷新任务。 pub fn start(self) -> SchedulerHandle { let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); tokio::spawn(async move { - info!("模型刷新器已启动"); + info!("模型刷新任务已启动"); - // 初始延迟后开始首次刷新 sleep(Duration::from_millis(self.initial_delay_ms)).await; - // 尝试首次刷新，失败时重试 let mut attempt = 0; let mut has_fetched_once = false; while attempt < self.max_retries { tokio::select! { _ = shutdown_rx.recv() => { - info!("模型刷新器在首次拉取前关闭"); + info!("模型刷新任务在首次拉取前关闭"); return; } _ = async {} => {} } - let request_id = uuid::Uuid::new_v4().to_string(); - match self.model_service.refresh_backend_models(&request_id).await { + match self.refresh_once().await { Ok(result) => { info!( refreshed_plans = result.refreshed_plans, @@ -70,16 +78,9 @@ impl ModelRefresher { has_fetched_once = true; break; } - Err(e) => { + Err(error) => { attempt += 1; - - // 如果是 NoAccounts 错误，跳过重试直接进入定期刷新 - let is_no_accounts = matches!( - e, - crate::codex::models::service::ModelServiceError::NoAccounts - ); - - if is_no_accounts { + if matches!(error, ModelServiceError::NoAccounts) { warn!("首次模型刷新跳过：没有可用账户，将进入周期重试"); break; } @@ -87,14 +88,14 @@ impl ModelRefresher { warn!( attempt, max_retries = self.max_retries, - error = ?e, + error = ?error, "首次模型刷新失败" ); if attempt < self.max_retries { tokio::select! { _ = shutdown_rx.recv() => { - info!("模型刷新器在重试等待期间关闭"); + info!("模型刷新任务在重试等待期间关闭"); return; } _ = sleep(Duration::from_millis(self.retry_delay_ms)) => {} @@ -105,20 +106,21 @@ impl ModelRefresher { } if !has_fetched_once { - warn!("模型刷新器首次尝试全部失败，切换到周期刷新"); + warn!("模型刷新任务首次尝试全部失败，切换到周期刷新"); } - // 定期刷新 let mut ticker = interval(Duration::from_secs(self.refresh_interval_secs)); - ticker.tick().await; // 跳过第一次立即触发 + ticker.tick().await; loop { tokio::select! { _ = ticker.tick() => { - self.tick().await; + if let Err(error) = self.refresh_once().await { + warn!(error = ?error, "刷新模型列表失败"); + } } _ = shutdown_rx.recv() => { - info!("模型刷新器正在关闭"); + info!("模型刷新任务正在关闭"); break; } } @@ -128,20 +130,27 @@ impl ModelRefresher { SchedulerHandle::new(shutdown_tx) } - async fn tick(&self) { + async fn refresh_once( + &self, + ) -> Result { let request_id = uuid::Uuid::new_v4().to_string(); - match self.model_service.refresh_backend_models(&request_id).await { - Ok(result) => { - info!( - refreshed_plans = result.refreshed_plans, - model_count = result.model_count, - failed_plans = result.failed_plans, - "模型列表已刷新" - ); - } - Err(e) => { - warn!(error = ?e, "刷新模型列表失败"); - } - } + let accounts = self.load_accounts().await?; + self.model_service + .refresh_backend_models_with_installation_id( + &accounts, + &request_id, + self.installation_id.as_deref(), + ) + .await + } + + async fn load_accounts(&self) -> Result, ModelServiceError> { + self.account_store + .list_pool_accounts() + .await + .map_err(|error| { + tracing::warn!(error = %error, "加载账号列表失败"); + ModelServiceError::NoAccounts + }) } } diff --git a/crates/runtime/src/tasks/quota_refresh.rs b/crates/runtime/src/tasks/quota_refresh.rs new file mode 100644 index 000000000..854fddb3a --- /dev/null +++ b/crates/runtime/src/tasks/quota_refresh.rs @@ -0,0 +1,252 @@ +//! 配额刷新任务。 + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use codex_proxy_adapters::{ + codex::client::{CodexBackendClient, CodexClientError, CodexRequestContext}, + sqlite::accounts::{SqliteAccountStore, SqliteAccountStoreError}, +}; +use codex_proxy_core::{ + accounts::{ + model::AccountStatus, + ports::{AccountStore, AccountStoreError}, + }, + serving::quota::quota_from_usage, +}; +use thiserror::Error; +use tokio::time::{interval, sleep, Instant}; +use tracing::{debug, info, warn}; + +use super::coordinator::SchedulerHandle; + +const DEFAULT_INTERVAL_SECS: u64 = 15 * 60; +const MIN_REFRESH_INTERVAL_SECS: u64 = 30 * 60; +const DEFAULT_REQUEST_SPACING_SECS: u64 = 3; + +/// 主动配额刷新任务。 +pub struct QuotaRefreshTask { + store: SqliteAccountStore, + codex: Arc, + interval_secs: u64, + min_refresh_interval_secs: u64, + request_spacing: Duration, + installation_id: Option, +} + +impl QuotaRefreshTask { + /// 构造默认配额刷新任务。 + pub fn new(store: SqliteAccountStore, codex: Arc) -> Self { + Self { + store, + codex, + interval_secs: DEFAULT_INTERVAL_SECS, + min_refresh_interval_secs: MIN_REFRESH_INTERVAL_SECS, + request_spacing: Self::default_request_spacing(), + installation_id: None, + } + } + + /// 返回旧调度器用于错开连续 quota 请求的默认间隔。 + pub fn default_request_spacing() -> Duration { + Duration::from_secs(DEFAULT_REQUEST_SPACING_SECS) + } + + /// 使用自定义间隔构造配额刷新任务。 + pub fn with_intervals( + store: SqliteAccountStore, + codex: Arc, + interval_secs: u64, + min_refresh_interval_secs: u64, + ) -> Self { + Self { + store, + codex, + interval_secs, + min_refresh_interval_secs, + request_spacing: Self::default_request_spacing(), + installation_id: None, + } + } + + /// 设置连续 quota 请求之间的错峰间隔。 + pub fn with_request_spacing(mut self, request_spacing: Duration) -> Self { + self.request_spacing = request_spacing; + self + } + + /// 设置 Codex installation id。 + pub fn with_installation_id(mut self, installation_id: Option) -> Self { + self.installation_id = installation_id.filter(|id| !id.trim().is_empty()); + self + } + + /// 启动后台刷新任务。 + pub fn start(self) -> SchedulerHandle { + let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); + + tokio::spawn(async move { + let mut ticker = interval(Duration::from_secs(self.interval_secs)); + let mut last_refreshed = HashMap::new(); + info!(interval_secs = self.interval_secs, "quota 刷新任务已启动"); + + loop { + tokio::select! { + _ = ticker.tick() => { + match self.refresh_locked_accounts(&mut last_refreshed).await { + Ok(summary) if summary.refreshed > 0 => { + info!( + refreshed = summary.refreshed, + failed = summary.failed, + "quota 刷新任务完成" + ); + } + Ok(summary) => { + debug!( + candidates = summary.candidates, + skipped_recent = summary.skipped_recent, + failed = summary.failed, + "没有需要刷新的 quota 锁定账号" + ); + } + Err(error) => { + warn!(error = %error, "quota 刷新任务失败"); + } + } + } + _ = shutdown_rx.recv() => { + info!("quota 刷新任务已关闭"); + break; + } + } + } + }); + + SchedulerHandle::new(shutdown_tx) + } + + /// 执行一次配额锁定账号刷新。 + pub async fn refresh_locked_accounts_once( + &self, + ) -> QuotaRefreshTaskResult { + let mut last_refreshed = HashMap::new(); + self.refresh_locked_accounts(&mut last_refreshed).await + } + + async fn refresh_locked_accounts( + &self, + last_refreshed: &mut HashMap, + ) -> QuotaRefreshTaskResult { + let accounts = self + .store + .list_pool_accounts() + .await + .map_err(QuotaRefreshTaskError::AccountStore)?; + let mut summary = QuotaRefreshSummary { + scanned: accounts.len(), + ..QuotaRefreshSummary::default() + }; + let min_interval = Duration::from_secs(self.min_refresh_interval_secs); + let now = Instant::now(); + + let mut candidates = accounts + .into_iter() + .filter(|account| { + account.status == AccountStatus::Active && account.quota_limit_reached + }) + .peekable(); + + while let Some(account) = candidates.next() { + summary.candidates += 1; + if last_refreshed + .get(&account.id) + .is_some_and(|last| now.duration_since(*last) < min_interval) + { + summary.skipped_recent += 1; + continue; + } + last_refreshed.insert(account.id.clone(), now); + + let request_id = uuid::Uuid::new_v4().to_string(); + match self + .codex + .fetch_usage(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, + installation_id: self.installation_id.as_deref(), + session_id: None, + }) + .await + { + Ok(raw) => { + let quota = quota_from_usage(&raw); + if self + .store + .update_quota_json(&account.id, "a.to_string()) + .await? + { + summary.refreshed += 1; + } else { + summary.missing += 1; + } + } + Err(error) => { + warn!( + account_id = %account.id, + error = %error, + "quota usage 拉取失败" + ); + summary.failed += 1; + } + } + + if self.request_spacing > Duration::ZERO && candidates.peek().is_some() { + sleep(self.request_spacing).await; + } + } + + Ok(summary) + } +} + +/// 单次配额刷新摘要。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct QuotaRefreshSummary { + /// 扫描账号数。 + pub scanned: usize, + /// 配额锁定候选账号数。 + pub candidates: usize, + /// 成功刷新账号数。 + pub refreshed: usize, + /// 最近刷新过而跳过的账号数。 + pub skipped_recent: usize, + /// 刷新期间账号消失数量。 + pub missing: usize, + /// 上游失败账号数。 + pub failed: usize, +} + +/// 配额刷新任务错误。 +#[derive(Debug, Error)] +pub enum QuotaRefreshTaskError { + /// 账号读取失败。 + #[error("failed to list accounts for quota refresh: {0}")] + AccountStore(#[from] AccountStoreError), + /// 账号写入失败。 + #[error("failed to persist quota refresh result: {0}")] + Store(#[from] SqliteAccountStoreError), + /// 上游 usage 请求失败。 + #[error("failed to fetch quota usage: {0}")] + Codex(#[from] CodexClientError), +} + +/// 配额刷新任务结果。 +pub type QuotaRefreshTaskResult = Result; diff --git a/crates/runtime/src/tasks/session_affinity_cleanup.rs b/crates/runtime/src/tasks/session_affinity_cleanup.rs new file mode 100644 index 000000000..5a7e59bf7 --- /dev/null +++ b/crates/runtime/src/tasks/session_affinity_cleanup.rs @@ -0,0 +1,78 @@ +//! 会话亲和性清理任务。 + +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use codex_proxy_adapters::sqlite::session_affinity::{ + SqliteSessionAffinityStore, SqliteSessionAffinityStoreResult, +}; +use tokio::time::interval; +use tracing::{debug, info, warn}; + +use super::coordinator::SchedulerHandle; + +/// 会话亲和性清理任务。 +pub struct SessionAffinityCleanupTask { + store: SqliteSessionAffinityStore, + interval_secs: u64, +} + +impl SessionAffinityCleanupTask { + /// 构造会话亲和性清理任务。 + pub fn new(store: SqliteSessionAffinityStore, interval_secs: u64) -> Self { + Self { + store, + interval_secs, + } + } + + /// 启动后台清理任务。 + pub fn start(self) -> SchedulerHandle { + let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); + + tokio::spawn(async move { + let mut ticker = interval(Duration::from_secs(self.interval_secs)); + info!( + interval_secs = self.interval_secs, + "会话亲和性清理任务已启动" + ); + + loop { + tokio::select! { + _ = ticker.tick() => { + match self.cleanup_once().await { + Ok(deleted) if deleted > 0 => { + info!(deleted, "已清理过期会话亲和性"); + } + Ok(_) => { + debug!("没有需要清理的过期会话亲和性"); + } + Err(error) => { + warn!(error = %error, "清理过期会话亲和性失败"); + } + } + } + _ = shutdown_rx.recv() => { + info!("会话亲和性清理任务已关闭"); + break; + } + } + } + }); + + SchedulerHandle::new(shutdown_tx) + } + + /// 执行一次清理。 + pub async fn cleanup_once(&self) -> SqliteSessionAffinityStoreResult { + self.cleanup_once_at(Utc::now()).await + } + + /// 在指定时间点执行一次清理。 + pub async fn cleanup_once_at( + &self, + now: DateTime, + ) -> SqliteSessionAffinityStoreResult { + self.store.delete_expired(now).await + } +} diff --git a/crates/runtime/src/tasks/session_cleanup.rs b/crates/runtime/src/tasks/session_cleanup.rs new file mode 100644 index 000000000..b8d39c7c1 --- /dev/null +++ b/crates/runtime/src/tasks/session_cleanup.rs @@ -0,0 +1,75 @@ +//! 会话清理任务。 + +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use codex_proxy_adapters::sqlite::admin_sessions::{ + SqliteAdminSessionStore, SqliteAdminSessionStoreResult, +}; +use tokio::time::interval; +use tracing::{debug, info, warn}; + +use super::coordinator::SchedulerHandle; + +/// 管理员会话清理任务。 +pub struct SessionCleanupTask { + store: SqliteAdminSessionStore, + interval_secs: u64, +} + +impl SessionCleanupTask { + /// 构造管理员会话清理任务。 + pub fn new(store: SqliteAdminSessionStore, interval_secs: u64) -> Self { + Self { + store, + interval_secs, + } + } + + /// 启动后台清理任务。 + pub fn start(self) -> SchedulerHandle { + let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); + + tokio::spawn(async move { + let mut ticker = interval(Duration::from_secs(self.interval_secs)); + info!( + interval_secs = self.interval_secs, + "管理员会话清理任务已启动" + ); + + loop { + tokio::select! { + _ = ticker.tick() => { + match self.cleanup_once().await { + Ok(deleted) if deleted > 0 => { + info!(deleted, "已清理过期管理员会话"); + } + Ok(_) => { + debug!("没有需要清理的过期管理员会话"); + } + Err(error) => { + warn!(error = %error, "清理过期管理员会话失败"); + } + } + } + _ = shutdown_rx.recv() => { + info!("管理员会话清理任务已关闭"); + break; + } + } + } + }); + + SchedulerHandle::new(shutdown_tx) + } + + /// 执行一次清理。 + pub async fn cleanup_once(&self) -> SqliteAdminSessionStoreResult { + self.cleanup_once_at(Utc::now()).await + } + + /// 在指定时间点执行一次清理。 + pub async fn cleanup_once_at(&self, now: DateTime) -> SqliteAdminSessionStoreResult { + self.store.cleanup_expired_sessions(now).await + } +} diff --git a/crates/runtime/src/tasks/token_refresh.rs b/crates/runtime/src/tasks/token_refresh.rs new file mode 100644 index 000000000..b7bee95ca --- /dev/null +++ b/crates/runtime/src/tasks/token_refresh.rs @@ -0,0 +1,834 @@ +//! 令牌刷新任务。 + +use std::{ + collections::{HashMap, HashSet}, + future::Future, + pin::Pin, + sync::Arc, + time::Duration, +}; + +use chrono::{DateTime, Utc}; +use codex_proxy_adapters::sqlite::{ + accounts::{AccountClaimsUpdate, SqliteAccountStore, SqliteAccountStoreError, StoredAccount}, + refresh_leases::{SqliteRefreshLeaseStore, SqliteRefreshLeaseStoreError}, +}; +use codex_proxy_core::{ + accounts::{ + jwt::jwt_expiration, + model::{Account, AccountStatus}, + ports::{AccountStore, AccountStoreError}, + }, + auth::{ + oauth::{RefreshError, RefreshPolicy, RefreshScheduler, RefreshTrigger}, + ports::TokenRefresher, + }, +}; +use secrecy::{ExposeSecret, SecretString}; +use thiserror::Error; +use tokio::{ + sync::Mutex, + task::JoinHandle, + time::{interval, sleep}, +}; +use tracing::{debug, info, warn}; +use uuid::Uuid; + +use super::coordinator::SchedulerHandle; + +const DEFAULT_INTERVAL_SECS: u64 = 60; +const REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60; +const MAX_REFRESH_ATTEMPTS: usize = 5; +const REFRESH_RETRY_BASE_DELAY_MILLIS: u64 = 5_000; +const REFRESH_RETRY_MAX_DELAY_MILLIS: u64 = 300_000; +const PERMANENT_FAILURE_CONFIRMATION_THRESHOLD: usize = 2; +const RECOVERY_DELAY_SECONDS: i64 = 10 * 60; + +/// 令牌刷新任务。 +pub struct TokenRefreshTask +where + C: TokenRefresher, +{ + store: SqliteAccountStore, + scheduler: Arc>, + refresh_margin_seconds: u64, + interval_secs: u64, + refresh_leases: Option, + lease_owner: String, + retry_delays: Vec, + in_flight: Arc>>, + recovery_not_before: Arc>>>, + timers: Arc>>>, +} + +impl Clone for TokenRefreshTask +where + C: TokenRefresher, +{ + fn clone(&self) -> Self { + Self { + store: self.store.clone(), + scheduler: self.scheduler.clone(), + refresh_margin_seconds: self.refresh_margin_seconds, + interval_secs: self.interval_secs, + refresh_leases: self.refresh_leases.clone(), + lease_owner: self.lease_owner.clone(), + retry_delays: self.retry_delays.clone(), + in_flight: self.in_flight.clone(), + recovery_not_before: self.recovery_not_before.clone(), + timers: self.timers.clone(), + } + } +} + +impl TokenRefreshTask +where + C: TokenRefresher, +{ + /// 构造默认令牌刷新任务。 + pub fn new(store: SqliteAccountStore, policy: RefreshPolicy, client: C) -> Self { + Self { + store, + scheduler: Arc::new(RefreshScheduler::new(policy, client)), + refresh_margin_seconds: policy.refresh_margin_seconds, + interval_secs: DEFAULT_INTERVAL_SECS, + refresh_leases: None, + lease_owner: refresh_lease_owner(), + retry_delays: default_refresh_retry_delays(), + in_flight: Arc::new(Mutex::new(HashSet::new())), + recovery_not_before: Arc::new(Mutex::new(HashMap::new())), + timers: Arc::new(Mutex::new(HashMap::new())), + } + } + + /// 使用自定义间隔构造令牌刷新任务。 + pub fn with_interval_secs( + store: SqliteAccountStore, + policy: RefreshPolicy, + client: C, + interval_secs: u64, + ) -> Self { + Self { + store, + scheduler: Arc::new(RefreshScheduler::new(policy, client)), + refresh_margin_seconds: policy.refresh_margin_seconds, + interval_secs, + refresh_leases: None, + lease_owner: refresh_lease_owner(), + retry_delays: default_refresh_retry_delays(), + in_flight: Arc::new(Mutex::new(HashSet::new())), + recovery_not_before: Arc::new(Mutex::new(HashMap::new())), + timers: Arc::new(Mutex::new(HashMap::new())), + } + } + + /// 使用刷新租约存储保护账号刷新。 + pub fn with_refresh_lease_store(mut self, refresh_leases: SqliteRefreshLeaseStore) -> Self { + self.refresh_leases = Some(refresh_leases); + self + } + + /// 使用自定义刷新重试延迟。 + pub fn with_retry_delays(mut self, retry_delays: I) -> Self + where + I: IntoIterator, + { + self.retry_delays = retry_delays.into_iter().collect(); + self + } + + /// 启动后台刷新任务。 + pub fn start(self) -> SchedulerHandle { + let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); + + tokio::spawn(async move { + let mut ticker = interval(Duration::from_secs(self.interval_secs)); + info!(interval_secs = self.interval_secs, "token 刷新任务已启动"); + + loop { + tokio::select! { + _ = ticker.tick() => { + match self.schedule_account_timers_once().await { + Ok(summary) if summary.changed() > 0 => { + info!( + scheduled = summary.scheduled, + immediate = summary.immediate, + recovery_scheduled = summary.recovery_scheduled, + replaced = summary.replaced, + "token 刷新定时器已调度" + ); + } + Ok(summary) => { + debug!( + scanned = summary.scanned, + skipped = summary.skipped, + "没有需要调度刷新的 token" + ); + } + Err(error) => { + warn!(error = %error, "token 刷新定时器调度失败"); + } + } + } + _ = shutdown_rx.recv() => { + self.clear_scheduled_timers().await; + info!("token 刷新任务已关闭"); + break; + } + } + } + }); + + SchedulerHandle::new(shutdown_tx) + } + + /// 执行一次到期刷新扫描。 + pub async fn refresh_due_accounts_once(&self) -> TokenRefreshTaskResult { + self.refresh_due_accounts_once_at(Utc::now()).await + } + + /// 执行一次账号级刷新定时器调度。 + pub async fn schedule_account_timers_once(&self) -> TokenRefreshTaskResult { + self.schedule_account_timers_once_at(Utc::now()).await + } + + /// 在指定时间点执行一次账号级刷新定时器调度。 + pub async fn schedule_account_timers_once_at( + &self, + now: DateTime, + ) -> TokenRefreshTaskResult { + let accounts = self + .store + .list_pool_accounts() + .await + .map_err(TokenRefreshTaskError::AccountStore)?; + let mut summary = TokenTimerSummary { + scanned: accounts.len(), + ..TokenTimerSummary::default() + }; + let mut expired_recovery_index = 0; + + for account in accounts { + if account.refresh_token.is_none() + || matches!( + account.status, + AccountStatus::Disabled | AccountStatus::Banned + ) + { + summary.skipped += 1; + continue; + } + + let timer = match account.status { + AccountStatus::Refreshing => Some((Duration::ZERO, RefreshTrigger::Unauthorized)), + AccountStatus::Expired => { + let delay = Duration::from_secs(30 + expired_recovery_index * 2); + expired_recovery_index += 1; + Some((delay, RefreshTrigger::Unauthorized)) + } + _ => self + .timer_delay_for_account(&account, now) + .map(|delay| (delay, RefreshTrigger::BeforeExpiry)), + }; + + let Some((delay, trigger)) = timer else { + summary.skipped += 1; + continue; + }; + + if delay.is_zero() && account.status != AccountStatus::Expired { + if self.clear_scheduled_timer(&account.id).await { + summary.replaced += 1; + } + let outcome = self + .refresh_scheduled_account(&account.id, trigger, now) + .await?; + if matches!(outcome, TokenRefreshOutcome::Skipped) { + summary.skipped += 1; + } else { + summary.immediate += 1; + } + continue; + } + + let scheduled_at = scheduled_at_from_delay(now, delay); + if self + .schedule_account_timer(account.id.clone(), delay, trigger, scheduled_at) + .await + { + summary.replaced += 1; + } + + if account.status == AccountStatus::Expired { + summary.recovery_scheduled += 1; + } else if delay.is_zero() { + summary.immediate += 1; + } else { + summary.scheduled += 1; + } + } + + Ok(summary) + } + + /// 返回当前账号级刷新定时器数量。 + pub async fn scheduled_timer_count(&self) -> usize { + self.timers.lock().await.len() + } + + /// 在指定时间点执行一次到期刷新扫描。 + pub async fn refresh_due_accounts_once_at( + &self, + now: DateTime, + ) -> TokenRefreshTaskResult { + let accounts = self + .store + .list_pool_accounts() + .await + .map_err(TokenRefreshTaskError::AccountStore)?; + let mut summary = TokenRefreshSummary { + scanned: accounts.len(), + ..TokenRefreshSummary::default() + }; + + for account in accounts { + if self.is_recovery_delayed(&account.id, now).await { + summary.skipped += 1; + continue; + } + + let Some(trigger) = self.refresh_trigger_for_account(&account, now) else { + summary.skipped += 1; + continue; + }; + + if !self.try_mark_in_flight(&account.id).await { + summary.skipped += 1; + continue; + } + + let lease_acquired = match self.try_acquire_refresh_lease(&account.id, now).await { + Ok(acquired) => acquired, + Err(error) => { + self.release_in_flight(&account.id).await; + return Err(error); + } + }; + if !lease_acquired { + self.release_in_flight(&account.id).await; + summary.skipped += 1; + continue; + } + + let outcome = self + .refresh_account_with_status_transitions(&account, trigger, now) + .await; + + self.release_in_flight(&account.id).await; + self.release_refresh_lease(&account.id).await?; + + let outcome = outcome?; + self.schedule_next_refresh_after_outcome(&outcome, Utc::now()) + .await; + + match outcome { + TokenRefreshOutcome::Refreshed(_) => summary.refreshed += 1, + TokenRefreshOutcome::StatusUpdated => summary.status_updated += 1, + TokenRefreshOutcome::Skipped => summary.skipped += 1, + TokenRefreshOutcome::Failed => summary.failed += 1, + }; + } + + Ok(summary) + } + + fn timer_delay_for_account(&self, account: &Account, now: DateTime) -> Option { + let expires_at = account + .access_token_expires_at + .or_else(|| jwt_expiration(&account.access_token))?; + let margin_seconds = self.refresh_margin_seconds.min(i64::MAX as u64) as i64; + let refresh_at = expires_at - chrono::Duration::seconds(margin_seconds); + + if refresh_at <= now { + return Some(Duration::ZERO); + } + + (refresh_at - now).to_std().ok() + } + + fn schedule_account_timer( + &self, + account_id: String, + delay: Duration, + trigger: RefreshTrigger, + scheduled_at: DateTime, + ) -> Pin + Send + '_>> { + Box::pin(async move { + let replaced = self.clear_scheduled_timer(&account_id).await; + let task = self.clone(); + let timer_account_id = account_id.clone(); + let handle = tokio::spawn(async move { + if !delay.is_zero() { + sleep(delay).await; + } + task.timers.lock().await.remove(&timer_account_id); + if let Err(error) = task + .refresh_scheduled_account(&timer_account_id, trigger, scheduled_at) + .await + { + warn!( + account_id = %timer_account_id, + error = %error, + "scheduled token refresh failed" + ); + } + }); + + self.timers.lock().await.insert(account_id, handle); + replaced + }) + } + + async fn clear_scheduled_timer(&self, account_id: &str) -> bool { + if let Some(handle) = self.timers.lock().await.remove(account_id) { + handle.abort(); + true + } else { + false + } + } + + async fn clear_scheduled_timers(&self) { + let mut timers = self.timers.lock().await; + for (_, handle) in timers.drain() { + handle.abort(); + } + } + + async fn refresh_scheduled_account( + &self, + account_id: &str, + trigger: RefreshTrigger, + now: DateTime, + ) -> TokenRefreshTaskResult { + if self.is_recovery_delayed(account_id, now).await { + return Ok(TokenRefreshOutcome::Skipped); + } + + let Some(account) = self.store.get(account_id).await? else { + return Err(TokenRefreshTaskError::AccountNotFound( + account_id.to_string(), + )); + }; + let account = stored_account_to_refresh_account(account); + + if !self.try_mark_in_flight(&account.id).await { + return Ok(TokenRefreshOutcome::Skipped); + } + + let lease_acquired = match self.try_acquire_refresh_lease(&account.id, now).await { + Ok(acquired) => acquired, + Err(error) => { + self.release_in_flight(&account.id).await; + return Err(error); + } + }; + if !lease_acquired { + self.release_in_flight(&account.id).await; + return Ok(TokenRefreshOutcome::Skipped); + } + + let outcome = self + .refresh_account_with_status_transitions(&account, trigger, now) + .await; + + self.release_in_flight(&account.id).await; + self.release_refresh_lease(&account.id).await?; + + let outcome = outcome?; + self.schedule_next_refresh_after_outcome(&outcome, Utc::now()) + .await; + Ok(outcome) + } + + async fn schedule_next_refresh_after_outcome( + &self, + outcome: &TokenRefreshOutcome, + now: DateTime, + ) { + let TokenRefreshOutcome::Refreshed(account) = outcome else { + return; + }; + + let Some(delay) = self.timer_delay_for_account(account, now) else { + return; + }; + + if delay.is_zero() { + debug!( + account_id = %account.id, + "refreshed token is already within refresh margin; skipping next timer" + ); + return; + } + + let scheduled_at = scheduled_at_from_delay(now, delay); + self.schedule_account_timer( + account.id.clone(), + delay, + RefreshTrigger::BeforeExpiry, + scheduled_at, + ) + .await; + } + + fn refresh_trigger_for_account( + &self, + account: &Account, + now: DateTime, + ) -> Option { + if account.status == AccountStatus::Refreshing { + return Some(RefreshTrigger::Unauthorized); + } + + self.scheduler + .should_refresh_account_at(account, RefreshTrigger::BeforeExpiry, now) + .then_some(RefreshTrigger::BeforeExpiry) + } + + async fn refresh_account_with_status_transitions( + &self, + account: &Account, + trigger: RefreshTrigger, + now: DateTime, + ) -> TokenRefreshTaskResult { + let attempt = self.prepare_refresh_attempt(account).await?; + + if attempt.refresh_token.is_none() { + persist_status(&self.store, &account.id, AccountStatus::Expired).await?; + self.clear_recovery_delay(&account.id).await; + return Ok(TokenRefreshOutcome::StatusUpdated); + } + + let max_attempts = self.retry_delays.len() + 1; + let mut permanent_failures = 0; + + for attempt_index in 0..max_attempts { + match self + .scheduler + .refresh_account_at(&attempt, trigger, now) + .await + { + Ok(mut updated) if account_requires_token_write(account, &updated) => { + updated.access_token_expires_at = + jwt_expiration(&updated.access_token).or(updated.access_token_expires_at); + persist_token_update(&self.store, &updated).await?; + self.clear_recovery_delay(&account.id).await; + return Ok(TokenRefreshOutcome::Refreshed(Box::new(updated))); + } + Ok(updated) if account.status != updated.status => { + if is_permanent_refresh_status(updated.status) { + permanent_failures += 1; + if permanent_failures < PERMANENT_FAILURE_CONFIRMATION_THRESHOLD + && self.sleep_before_retry(attempt_index).await + { + continue; + } + } + + persist_status(&self.store, &account.id, updated.status).await?; + self.clear_recovery_delay(&account.id).await; + return Ok(TokenRefreshOutcome::StatusUpdated); + } + Ok(_) => return Ok(TokenRefreshOutcome::Skipped), + Err(RefreshError::Transport) if self.sleep_before_retry(attempt_index).await => { + continue; + } + Err(RefreshError::Transport) => { + persist_status(&self.store, &account.id, AccountStatus::Active).await?; + self.schedule_recovery(&account.id, now).await; + return Ok(TokenRefreshOutcome::Failed); + } + Err(error) => return Err(TokenRefreshTaskError::Refresh(error)), + } + } + + Ok(TokenRefreshOutcome::Skipped) + } + + async fn prepare_refresh_attempt(&self, account: &Account) -> TokenRefreshTaskResult { + if account.status != AccountStatus::Refreshing { + persist_status(&self.store, &account.id, AccountStatus::Refreshing).await?; + } + + let mut attempt = account.clone(); + attempt.status = AccountStatus::Active; + Ok(attempt) + } + + async fn sleep_before_retry(&self, attempt_index: usize) -> bool { + let Some(delay) = self.retry_delays.get(attempt_index).copied() else { + return false; + }; + + if !delay.is_zero() { + sleep(delay).await; + } + true + } + + async fn try_mark_in_flight(&self, account_id: &str) -> bool { + self.in_flight.lock().await.insert(account_id.to_string()) + } + + async fn release_in_flight(&self, account_id: &str) { + self.in_flight.lock().await.remove(account_id); + } + + async fn is_recovery_delayed(&self, account_id: &str, now: DateTime) -> bool { + let mut recovery_not_before = self.recovery_not_before.lock().await; + let Some(not_before) = recovery_not_before.get(account_id).copied() else { + return false; + }; + + if now < not_before { + return true; + } + + recovery_not_before.remove(account_id); + false + } + + async fn schedule_recovery(&self, account_id: &str, now: DateTime) { + self.recovery_not_before.lock().await.insert( + account_id.to_string(), + now + chrono::Duration::seconds(RECOVERY_DELAY_SECONDS), + ); + } + + async fn clear_recovery_delay(&self, account_id: &str) { + self.recovery_not_before.lock().await.remove(account_id); + } + + async fn try_acquire_refresh_lease( + &self, + account_id: &str, + now: DateTime, + ) -> TokenRefreshTaskResult { + let Some(refresh_leases) = self.refresh_leases.as_ref() else { + return Ok(true); + }; + refresh_leases + .try_acquire( + account_id, + &self.lease_owner, + now + chrono::Duration::seconds(REFRESH_LEASE_TTL_SECONDS), + now, + ) + .await + .map_err(TokenRefreshTaskError::Lease) + } + + async fn release_refresh_lease(&self, account_id: &str) -> TokenRefreshTaskResult<()> { + let Some(refresh_leases) = self.refresh_leases.as_ref() else { + return Ok(()); + }; + refresh_leases + .release(account_id, &self.lease_owner) + .await + .map(|_| ()) + .map_err(TokenRefreshTaskError::Lease) + } +} + +enum TokenRefreshOutcome { + Refreshed(Box), + StatusUpdated, + Skipped, + Failed, +} + +/// 单次刷新扫描摘要。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct TokenRefreshSummary { + /// 扫描账号数。 + pub scanned: usize, + /// 成功刷新 token 的账号数。 + pub refreshed: usize, + /// 仅更新状态的账号数。 + pub status_updated: usize, + /// 无需刷新的账号数。 + pub skipped: usize, + /// 传输失败账号数。 + pub failed: usize, +} + +/// 账号级刷新定时器调度摘要。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct TokenTimerSummary { + /// 扫描账号数。 + pub scanned: usize, + /// 已调度的未来刷新定时器数。 + pub scheduled: usize, + /// 已调度的立即刷新定时器数。 + pub immediate: usize, + /// 已调度的过期账号恢复定时器数。 + pub recovery_scheduled: usize, + /// 无需调度的账号数。 + pub skipped: usize, + /// 被替换的既有定时器数。 + pub replaced: usize, +} + +impl TokenTimerSummary { + fn changed(&self) -> usize { + self.scheduled + self.immediate + self.recovery_scheduled + } +} + +/// 令牌刷新任务错误。 +#[derive(Debug, Error)] +pub enum TokenRefreshTaskError { + /// 账号存储读取失败。 + #[error("failed to list accounts for token refresh: {0}")] + AccountStore(#[from] AccountStoreError), + /// 账号存储写入失败。 + #[error("failed to persist token refresh result: {0}")] + Store(#[from] SqliteAccountStoreError), + /// 刷新租约存储失败。 + #[error("failed to coordinate token refresh lease: {0}")] + Lease(#[from] SqliteRefreshLeaseStoreError), + /// core 刷新调度失败。 + #[error("token refresh scheduler failed: {0}")] + Refresh(#[from] RefreshError), + /// 刷新期间账号不存在。 + #[error("account disappeared during token refresh: {0}")] + AccountNotFound(String), +} + +/// 令牌刷新任务结果。 +pub type TokenRefreshTaskResult = Result; + +async fn persist_token_update( + store: &SqliteAccountStore, + account: &Account, +) -> TokenRefreshTaskResult<()> { + let updated = store + .update_from_claims( + &account.id, + AccountClaimsUpdate { + email: account.email.clone(), + account_id: account.account_id.clone(), + user_id: account.user_id.clone(), + plan_type: account.plan_type.clone(), + access_token: SecretString::new(account.access_token.clone().into()), + refresh_token: account + .refresh_token + .clone() + .map(|token| SecretString::new(token.into())), + access_token_expires_at: jwt_expiration(&account.access_token) + .or(account.access_token_expires_at), + status: account.status, + }, + ) + .await?; + if updated { + Ok(()) + } else { + Err(TokenRefreshTaskError::AccountNotFound(account.id.clone())) + } +} + +async fn persist_status( + store: &SqliteAccountStore, + account_id: &str, + status: AccountStatus, +) -> TokenRefreshTaskResult<()> { + let updated = store.set_status(account_id, status).await?; + if updated { + Ok(()) + } else { + Err(TokenRefreshTaskError::AccountNotFound( + account_id.to_string(), + )) + } +} + +fn account_requires_token_write(previous: &Account, updated: &Account) -> bool { + previous.access_token != updated.access_token + || previous.refresh_token != updated.refresh_token + || (updated.status == AccountStatus::Active && previous.status != updated.status) +} + +fn is_permanent_refresh_status(status: AccountStatus) -> bool { + matches!( + status, + AccountStatus::Expired + | AccountStatus::QuotaExhausted + | AccountStatus::Disabled + | AccountStatus::Banned + ) +} + +fn refresh_lease_owner() -> String { + format!("runtime-token-refresh:{}", Uuid::new_v4().simple()) +} + +fn default_refresh_retry_delays() -> Vec { + (0..MAX_REFRESH_ATTEMPTS.saturating_sub(1)) + .map(|attempt_index| { + let multiplier = 3_u64.saturating_pow(attempt_index as u32); + let millis = REFRESH_RETRY_BASE_DELAY_MILLIS + .saturating_mul(multiplier) + .min(REFRESH_RETRY_MAX_DELAY_MILLIS); + Duration::from_millis(millis) + }) + .collect() +} + +fn scheduled_at_from_delay(now: DateTime, delay: Duration) -> DateTime { + match chrono::Duration::from_std(delay) { + Ok(delay) => now + delay, + Err(_) => now, + } +} + +fn stored_account_to_refresh_account(stored: StoredAccount) -> Account { + Account { + id: stored.id, + email: stored.email, + account_id: stored.account_id, + user_id: stored.user_id, + label: stored.label, + plan_type: stored.plan_type, + access_token: stored.access_token.expose_secret().to_string(), + refresh_token: stored + .refresh_token + .map(|token| token.expose_secret().to_string()), + access_token_expires_at: stored.access_token_expires_at, + status: stored.status, + quota_limit_reached: false, + quota_verify_required: false, + quota_cooldown_until: None, + cloudflare_cooldown_until: None, + request_count: 0, + empty_response_count: 0, + image_input_tokens: 0, + image_output_tokens: 0, + image_request_count: 0, + image_request_failed_count: 0, + window_request_count: 0, + window_input_tokens: 0, + window_output_tokens: 0, + window_cached_tokens: 0, + window_image_input_tokens: 0, + window_image_output_tokens: 0, + window_image_request_count: 0, + window_image_request_failed_count: 0, + window_started_at: None, + window_reset_at: None, + limit_window_seconds: None, + added_at: stored.added_at.to_rfc3339(), + last_used_at: None, + } +} diff --git a/crates/runtime/src/upstream.rs b/crates/runtime/src/upstream.rs new file mode 100644 index 000000000..7e353496a --- /dev/null +++ b/crates/runtime/src/upstream.rs @@ -0,0 +1,50 @@ +//! 上游适配器组装。 + +use std::{sync::Arc, time::Duration}; + +use codex_proxy_adapters::codex::{ + client::CodexBackendClient, + default_http_client, + websocket::pool::{CodexWebSocketPool, CodexWebSocketPoolConfig as AdapterWebSocketPoolConfig}, +}; +use codex_proxy_core::gateway::{fingerprint::Fingerprint, ports::CodexModelCatalogClient}; +use codex_proxy_platform::config::WebSocketPoolConfig; + +/// 构造 Codex 模型目录客户端。 +pub fn model_catalog_client( + base_url: String, + fingerprint: Fingerprint, +) -> Arc { + Arc::new(CodexBackendClient::new( + default_http_client(), + base_url, + fingerprint, + )) +} + +/// 构造 Codex 后端请求客户端。 +pub fn codex_backend_client( + base_url: String, + fingerprint: Fingerprint, + ws_pool: &WebSocketPoolConfig, +) -> Arc { + let client = CodexBackendClient::new(default_http_client(), base_url, fingerprint); + if ws_pool.enabled { + Arc::new( + client.with_websocket_pool(Arc::new(CodexWebSocketPool::with_config( + adapter_websocket_pool_config(ws_pool), + ))), + ) + } else { + Arc::new(client) + } +} + +fn adapter_websocket_pool_config(config: &WebSocketPoolConfig) -> AdapterWebSocketPoolConfig { + AdapterWebSocketPoolConfig { + enabled: config.enabled, + max_age: Duration::from_millis(config.max_age_ms), + max_per_account: config.max_per_account, + ..AdapterWebSocketPoolConfig::default() + } +} diff --git a/crates/runtime/tests/account_pool_restore.rs b/crates/runtime/tests/account_pool_restore.rs new file mode 100644 index 000000000..495cd6b05 --- /dev/null +++ b/crates/runtime/tests/account_pool_restore.rs @@ -0,0 +1,126 @@ +use std::collections::BTreeMap; + +use chrono::{TimeZone, Utc}; +use codex_proxy_platform::{ + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, + UsageStatsConfig, WebSocketPoolConfig, + }, + crypto::SecretBox, + storage::connect_sqlite, +}; +use codex_proxy_runtime::state::AppState; +use secrecy::SecretString; + +#[tokio::test] +async fn app_state_should_restore_active_accounts_into_runtime_pool() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("runtime-account-pool.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.expect("sqlite pool"); + let secret_box = SecretBox::new([34u8; 32]); + insert_account(&pool, &secret_box, "acct_pool", "active").await; + + let state = AppState::with_pool_and_secret_box(test_config(url), pool, secret_box); + let restored = state + .restore_account_pool_from_repository() + .await + .expect("account pool should restore"); + let acquired = state + .services + .account_pool + .acquire( + "gpt-5.5", + Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(), + ) + .await + .expect("active account should be acquired"); + + assert_eq!(restored, 1); + assert_eq!(acquired.account.id, "acct_pool"); +} + +async fn insert_account(pool: &sqlx::SqlitePool, secret_box: &SecretBox, id: &str, status: &str) { + let access_token = secret_box + .encrypt(&SecretString::new(format!("access-{id}").into())) + .expect("access token should encrypt"); + sqlx::query( + "insert into accounts (id, email, account_id, user_id, access_token_cipher, access_token_expires_at, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(id) + .bind(format!("{id}@example.com")) + .bind(format!("chatgpt-{id}")) + .bind(format!("user-{id}")) + .bind(access_token) + .bind("2026-06-18T13:00:00Z") + .bind(status) + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .expect("account should be inserted"); +} + +fn test_config(database_url: String) -> AppConfig { + AppConfig { + server: ServerConfig { + host: "127.0.0.1".to_string(), + port: 0, + }, + api: ApiConfig { + base_url: "https://example.invalid".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(), + oauth_client_id: "app_id".to_string(), + oauth_auth_endpoint: "https://auth.invalid".to_string(), + oauth_token_endpoint: "https://token.invalid".to_string(), + }, + 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, + }, + ws_pool: WebSocketPoolConfig::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + }, + } +} diff --git a/crates/runtime/tests/fingerprint.rs b/crates/runtime/tests/fingerprint.rs new file mode 100644 index 000000000..09ad089f2 --- /dev/null +++ b/crates/runtime/tests/fingerprint.rs @@ -0,0 +1,198 @@ +use codex_proxy_adapters::codex::fingerprint::FingerprintRepository; +use codex_proxy_core::models::model::{BackendModelEntry, ModelPlanSnapshot}; +use codex_proxy_platform::storage::connect_sqlite; +use std::collections::BTreeMap; + +use codex_proxy_platform::config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, +}; +use codex_proxy_platform::crypto::SecretBox; + +#[tokio::test] +async fn runtime_should_load_latest_auto_updated_fingerprint_when_present() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("fingerprints.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let repo = FingerprintRepository::new(pool); + repo.upsert_auto_update("26.900.1", "7001", Some("147")) + .await + .expect("insert fingerprint"); + + let stored = codex_proxy_runtime::bootstrap::load_runtime_fingerprint(&repo).await; + + assert_eq!(stored.app_version, "26.900.1"); +} + +#[tokio::test] +async fn runtime_should_build_state_with_model_service() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("runtime-state.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let state = codex_proxy_runtime::state::AppState::with_pool_and_secret_box( + AppConfig { + server: ServerConfig { + host: "127.0.0.1".to_string(), + port: 0, + }, + api: ApiConfig { + base_url: "https://example.invalid".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(), + oauth_client_id: "app_id".to_string(), + oauth_auth_endpoint: "https://auth.invalid".to_string(), + oauth_token_endpoint: "https://token.invalid".to_string(), + }, + 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: format!("sqlite://{}", db.display()), + }, + 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, + }, + ws_pool: Default::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + }, + }, + pool, + SecretBox::new([42u8; 32]), + ); + + let catalog = state.services.models.catalog().await; + assert_eq!(catalog.models().len(), 1); +} + +#[tokio::test] +async fn runtime_state_should_expose_backend_model_snapshot_through_model_service() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("runtime-models.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let config = AppConfig { + server: ServerConfig { + host: "127.0.0.1".to_string(), + port: 0, + }, + api: ApiConfig { + base_url: "https://example.invalid".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(), + oauth_client_id: "app_id".to_string(), + oauth_auth_endpoint: "https://auth.invalid".to_string(), + oauth_token_endpoint: "https://token.invalid".to_string(), + }, + 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: format!("sqlite://{}", db.display()), + }, + 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, + }, + ws_pool: Default::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + }, + }; + let snapshot_store = + codex_proxy_adapters::sqlite::models::ModelSnapshotRepository::new(pool.clone()); + snapshot_store + .replace_plan_snapshot(&ModelPlanSnapshot::from_backend_entries( + "plus", + vec![BackendModelEntry { + id: Some("gpt-6".to_string()), + name: Some("GPT-6".to_string()), + ..BackendModelEntry::default() + }], + )) + .await + .expect("replace snapshot"); + + let state = codex_proxy_runtime::state::AppState::with_pool_and_secret_box( + config, + pool, + SecretBox::new([43u8; 32]), + ); + + let catalog = state.services.models.catalog().await; + assert_eq!(catalog.models()[0].id, "gpt-6"); +} diff --git a/crates/runtime/tests/quota_refresh.rs b/crates/runtime/tests/quota_refresh.rs new file mode 100644 index 000000000..9de90b49f --- /dev/null +++ b/crates/runtime/tests/quota_refresh.rs @@ -0,0 +1,218 @@ +use std::{sync::Arc, time::Duration as StdDuration}; + +use chrono::{Duration, Utc}; +use codex_proxy_adapters::{ + codex::client::CodexBackendClient, + sqlite::accounts::{NewAccount, SqliteAccountStore}, +}; +use codex_proxy_core::{accounts::model::AccountStatus, gateway::fingerprint::Fingerprint}; +use codex_proxy_platform::{crypto::SecretBox, storage::connect_sqlite}; +use secrecy::SecretString; +use serde_json::{json, Value}; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +#[test] +fn quota_refresh_task_should_keep_old_default_request_spacing() { + assert_eq!( + codex_proxy_runtime::tasks::quota_refresh::QuotaRefreshTask::default_request_spacing(), + StdDuration::from_secs(3) + ); +} + +#[tokio::test] +async fn quota_refresh_task_should_fetch_usage_for_quota_locked_accounts_and_store_quota() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/codex/usage")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "plan_type": "plus", + "rate_limit": { + "allowed": true, + "limit_reached": false, + "primary_window": { + "used_percent": 72.4, + "reset_at": 1_800_000_000, + "limit_window_seconds": 18_000 + }, + "secondary_window": { + "used_percent": 88.2, + "reset_at": 1_800_000_100, + "limit_window_seconds": 3_600 + } + } + }))) + .mount(&server) + .await; + + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("quota-refresh.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool.clone(), SecretBox::new([12u8; 32])); + store + .insert(NewAccount { + id: "acct-quota".to_string(), + email: Some("user@example.com".to_string()), + account_id: Some("chatgpt-account".to_string()), + user_id: None, + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new("access-token".to_string().into()), + refresh_token: None, + access_token_expires_at: Some(Utc::now() + Duration::hours(1)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + sqlx::query("update accounts set quota_limit_reached = 1 where id = ?") + .bind("acct-quota") + .execute(&pool) + .await + .expect("quota lock should be set"); + let codex = CodexBackendClient::new( + reqwest::Client::new(), + server.uri(), + Fingerprint::default_for_tests(), + ); + let task = codex_proxy_runtime::tasks::quota_refresh::QuotaRefreshTask::new( + store.clone(), + Arc::new(codex), + ); + + let summary = task + .refresh_locked_accounts_once() + .await + .expect("quota refresh should succeed"); + let stored = store + .get_quota_json("acct-quota") + .await + .expect("quota json should load") + .expect("quota json should be present"); + let quota: Value = serde_json::from_str(&stored).expect("quota json should parse"); + + assert_eq!( + ( + summary.refreshed, + quota + .pointer("/rate_limit/limit_reached") + .and_then(Value::as_bool), + quota + .pointer("/rate_limit/remaining_percent") + .and_then(Value::as_i64), + quota + .pointer("/secondary_rate_limit/limit_reached") + .and_then(Value::as_bool), + ), + (1, Some(false), Some(28), Some(false)) + ); +} + +#[tokio::test] +async fn quota_refresh_task_should_stagger_multiple_locked_account_requests() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/codex/usage")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "plan_type": "plus", + "rate_limit": { + "allowed": true, + "limit_reached": false, + "primary_window": { + "used_percent": 42.0, + "reset_at": 1_800_000_000, + "limit_window_seconds": 18_000 + } + } + }))) + .mount(&server) + .await; + + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("quota-refresh-spacing.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool.clone(), SecretBox::new([13u8; 32])); + insert_quota_locked_account(&store, &pool, "acct-quota-1", "access-token-1").await; + insert_quota_locked_account(&store, &pool, "acct-quota-2", "access-token-2").await; + let codex = CodexBackendClient::new( + reqwest::Client::new(), + server.uri(), + Fingerprint::default_for_tests(), + ); + let task = + codex_proxy_runtime::tasks::quota_refresh::QuotaRefreshTask::new(store, Arc::new(codex)) + .with_request_spacing(StdDuration::from_millis(200)); + + let refresh = tokio::spawn(async move { task.refresh_locked_accounts_once().await }); + wait_for_usage_requests(&server, 1).await; + tokio::time::sleep(StdDuration::from_millis(50)).await; + + assert_eq!(usage_request_count(&server).await, 1); + + let summary = refresh + .await + .expect("quota refresh task should join") + .expect("quota refresh should succeed"); + + assert_eq!( + (summary.refreshed, usage_request_count(&server).await), + (2, 2) + ); +} + +async fn insert_quota_locked_account( + store: &SqliteAccountStore, + pool: &sqlx::SqlitePool, + id: &str, + access_token: &str, +) { + store + .insert(NewAccount { + id: id.to_string(), + email: Some(format!("{id}@example.com")), + account_id: Some(format!("chatgpt-{id}")), + user_id: None, + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(access_token.to_string().into()), + refresh_token: None, + access_token_expires_at: Some(Utc::now() + Duration::hours(1)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + sqlx::query("update accounts set quota_limit_reached = 1 where id = ?") + .bind(id) + .execute(pool) + .await + .expect("quota lock should be set"); +} + +async fn wait_for_usage_requests(server: &MockServer, expected_count: usize) { + let deadline = tokio::time::Instant::now() + StdDuration::from_secs(2); + loop { + if usage_request_count(server).await >= expected_count { + return; + } + assert!( + tokio::time::Instant::now() < deadline, + "quota refresh task did not request usage before timeout" + ); + tokio::time::sleep(StdDuration::from_millis(25)).await; + } +} + +async fn usage_request_count(server: &MockServer) -> usize { + server + .received_requests() + .await + .expect("received requests should load") + .iter() + .filter(|request| request.url.path() == "/codex/usage") + .count() +} diff --git a/crates/runtime/tests/session_affinity.rs b/crates/runtime/tests/session_affinity.rs new file mode 100644 index 000000000..8cd9e82a3 --- /dev/null +++ b/crates/runtime/tests/session_affinity.rs @@ -0,0 +1,121 @@ +use chrono::{Duration, TimeZone, Utc}; +use codex_proxy_adapters::sqlite::session_affinity::SqliteSessionAffinityStore; +use codex_proxy_core::serving::affinity::SessionAffinityEntry; +use codex_proxy_platform::{ + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, + UsageStatsConfig, WebSocketPoolConfig, + }, + crypto::SecretBox, + storage::connect_sqlite, +}; +use codex_proxy_runtime::state::AppState; + +#[tokio::test] +async fn app_state_should_restore_session_affinity_from_sqlite() { + let dir = tempfile::tempdir().expect("tempdir"); + let db = dir.path().join("runtime-session-affinity.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.expect("sqlite pool"); + let store = SqliteSessionAffinityStore::new(pool.clone()); + let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(); + + store + .upsert( + "resp_restore", + &SessionAffinityEntry { + account_id: "acct_restore".to_string(), + conversation_id: "conv_restore".to_string(), + turn_state: Some("turn_restore".to_string()), + instructions_hash: Some("hash_restore".to_string()), + input_tokens: Some(7), + function_call_ids: vec!["call_restore".to_string()], + variant_hash: Some("variant_restore".to_string()), + created_at: now, + }, + Duration::hours(4), + ) + .await + .expect("affinity should be stored"); + + let state = + AppState::with_pool_and_secret_box(test_config(url), pool, SecretBox::new([33u8; 32])); + let restored = state + .restore_session_affinity_from_repository(now + Duration::minutes(1)) + .await + .expect("session affinity should restore"); + + assert_eq!(restored, 1); + assert_eq!( + state + .services + .session_affinity + .lookup_account("resp_restore", now + Duration::minutes(1)) + .await, + Some("acct_restore".to_string()) + ); +} + +fn test_config(database_url: String) -> AppConfig { + AppConfig { + server: ServerConfig { + host: "127.0.0.1".to_string(), + port: 0, + }, + api: ApiConfig { + base_url: "https://example.invalid".to_string(), + }, + model: ModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: None, + service_tier: None, + aliases: Default::default(), + }, + 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(), + oauth_client_id: "app_id".to_string(), + oauth_auth_endpoint: "https://auth.invalid".to_string(), + oauth_token_endpoint: "https://token.invalid".to_string(), + }, + 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, + }, + ws_pool: WebSocketPoolConfig::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + }, + } +} diff --git a/crates/runtime/tests/tasks.rs b/crates/runtime/tests/tasks.rs new file mode 100644 index 000000000..c09f7e797 --- /dev/null +++ b/crates/runtime/tests/tasks.rs @@ -0,0 +1,604 @@ +use std::{sync::Arc, time::Duration as StdDuration}; + +use async_trait::async_trait; +use chrono::{Duration, Utc}; +use codex_proxy_adapters::{ + codex::fingerprint::FingerprintRepository, + sqlite::{ + admin_sessions::SqliteAdminSessionStore, cookies::SqliteCookieStore, + session_affinity::SqliteSessionAffinityStore, + }, +}; +use codex_proxy_core::{ + accounts::{ + model::{Account, AccountStatus}, + ports::{AccountStore, AccountStoreResult}, + usage::AccountUsageDelta, + }, + models::{model::ModelConfig, service::ModelService}, + serving::affinity::SessionAffinityEntry, +}; +use codex_proxy_platform::config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, + ModelConfig as PlatformModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, + ServerConfig, TlsConfig, UsageStatsConfig, +}; +use codex_proxy_platform::{crypto::SecretBox, storage::connect_sqlite}; +use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, +}; + +#[tokio::test] +async fn fingerprint_update_task_should_start_background_checker() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("fingerprints.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let repo = FingerprintRepository::new(pool); + + let handle = codex_proxy_runtime::tasks::fingerprint_update::start_fingerprint_update_task( + Some(repo), + "https://example.invalid/appcast.xml".to_string(), + dir.path().join("extracted-fingerprint.json"), + "26.800.1".to_string(), + "6001".to_string(), + ); + + handle.shutdown().await; +} + +#[tokio::test] +async fn fingerprint_update_task_should_apply_initial_appcast_update_to_repository() { + let server = MockServer::start().await; + mount_appcast( + &server, + r#" + + + + + + + + "#, + ) + .await; + + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("fingerprints-initial-update.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let repo = FingerprintRepository::new(pool); + let extracted_path = dir.path().join("extracted-fingerprint.json"); + std::fs::write( + &extracted_path, + r#"{"app_version":"26.900.1","build_number":"7001","chromium_version":"147"}"#, + ) + .expect("extracted fingerprint should be written"); + + let handle = codex_proxy_runtime::tasks::fingerprint_update::start_fingerprint_update_task( + Some(repo.clone()), + format!("{}/appcast.xml", server.uri()), + extracted_path, + "26.800.1".to_string(), + "6001".to_string(), + ); + + let stored = wait_for_auto_updated_fingerprint(&repo).await; + handle.shutdown().await; + + assert_eq!( + ( + stored.app_version.as_str(), + stored.build_number.as_str(), + stored.chromium_version.as_str() + ), + ("26.900.1", "7001", "147") + ); +} + +#[tokio::test] +async fn fingerprint_update_task_should_not_persist_when_appcast_matches_current_fingerprint() { + let server = MockServer::start().await; + mount_appcast( + &server, + r#" + + + + + + + + "#, + ) + .await; + + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("fingerprints-no-update.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let repo = FingerprintRepository::new(pool); + let handle = codex_proxy_runtime::tasks::fingerprint_update::start_fingerprint_update_task( + Some(repo.clone()), + format!("{}/appcast.xml", server.uri()), + dir.path().join("missing-extracted-fingerprint.json"), + "26.900.1".to_string(), + "7001".to_string(), + ); + + wait_for_appcast_requests(&server, 1).await; + handle.shutdown().await; + let stored = repo + .load_latest_auto_updated() + .await + .expect("latest auto update should load"); + + assert!(stored.is_none()); +} + +#[tokio::test] +async fn fingerprint_update_task_should_check_appcast_without_repository() { + let server = MockServer::start().await; + mount_appcast( + &server, + r#" + + + + + + + + "#, + ) + .await; + + let dir = tempfile::tempdir().expect("temp dir"); + let handle = codex_proxy_runtime::tasks::fingerprint_update::start_fingerprint_update_task( + None, + format!("{}/appcast.xml", server.uri()), + dir.path().join("missing-extracted-fingerprint.json"), + "26.800.1".to_string(), + "6001".to_string(), + ); + + wait_for_appcast_requests(&server, 1).await; + handle.shutdown().await; +} + +#[tokio::test] +async fn model_refresh_task_should_start_and_shutdown() { + let model_service = Arc::new(ModelService::new( + ModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: None, + service_tier: None, + aliases: Default::default(), + }, + None, + None, + None, + )); + let account_store = Arc::new(FakeAccountStore); + + let handle = codex_proxy_runtime::tasks::model_refresh::ModelRefreshTask::new( + model_service, + account_store, + ) + .start(); + + handle.shutdown().await; +} + +#[tokio::test] +async fn cookie_cleanup_task_should_delete_only_expired_cookie_rows() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("cookies.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let secret_box = SecretBox::new([7u8; 32]); + let store = SqliteCookieStore::new(pool.clone(), secret_box); + insert_account(&pool, "acct-cookie").await; + let now = Utc::now(); + insert_cookie( + &pool, + "expired-cookie", + "acct-cookie", + "old", + now - Duration::minutes(1), + ) + .await; + insert_cookie( + &pool, + "active-cookie", + "acct-cookie", + "fresh", + now + Duration::minutes(10), + ) + .await; + + let deleted = codex_proxy_runtime::tasks::cookie_cleanup::CookieCleanupTask::new(store) + .cleanup_once_at(now) + .await + .expect("cookie cleanup should succeed"); + + let remaining = cookie_ids(&pool).await; + + assert_eq!((deleted, remaining), (1, vec!["active-cookie".to_string()])); +} + +#[tokio::test] +async fn session_cleanup_task_should_delete_only_expired_sessions() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("sessions.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAdminSessionStore::new(pool.clone()); + let now = Utc::now(); + insert_admin_user(&pool, "admin").await; + insert_admin_session( + &pool, + "expired-session", + "admin", + now - Duration::minutes(1), + ) + .await; + insert_admin_session( + &pool, + "active-session", + "admin", + now + Duration::minutes(10), + ) + .await; + + let deleted = codex_proxy_runtime::tasks::session_cleanup::SessionCleanupTask::new(store, 3600) + .cleanup_once_at(now) + .await + .expect("session cleanup should succeed"); + + let remaining = admin_session_ids(&pool).await; + + assert_eq!( + (deleted, remaining), + (1, vec!["active-session".to_string()]) + ); +} + +#[tokio::test] +async fn session_affinity_cleanup_task_should_delete_only_expired_affinities() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("session-affinity-cleanup.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteSessionAffinityStore::new(pool.clone()); + let now = Utc::now(); + store + .upsert( + "expired-response", + &session_affinity_entry("expired-account", now - Duration::hours(2)), + Duration::hours(1), + ) + .await + .expect("expired affinity should be inserted"); + store + .upsert( + "active-response", + &session_affinity_entry("active-account", now), + Duration::hours(1), + ) + .await + .expect("active affinity should be inserted"); + + let deleted = + codex_proxy_runtime::tasks::session_affinity_cleanup::SessionAffinityCleanupTask::new( + store, 3600, + ) + .cleanup_once_at(now) + .await + .expect("session affinity cleanup should succeed"); + + let remaining = session_affinity_response_ids(&pool).await; + + assert_eq!( + (deleted, remaining), + (1, vec!["active-response".to_string()]) + ); +} + +#[tokio::test] +async fn start_background_tasks_should_register_migrated_runtime_tasks() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("background-tasks.sqlite"); + let database_url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&database_url).await.expect("sqlite pool"); + let state = codex_proxy_runtime::state::AppState::with_pool_and_secret_box( + test_config(database_url), + pool, + SecretBox::new([9u8; 32]), + ); + + let coordinator = codex_proxy_runtime::tasks::coordinator::start_background_tasks(&state).await; + let task_names = coordinator.task_names(); + + assert_eq!( + task_names, + [ + "cookie_cleanup", + "session_cleanup", + "session_affinity_cleanup", + "model_refresh", + "token_refresh", + "quota_refresh", + "fingerprint_update" + ] + ); + + coordinator.shutdown().await; +} + +struct FakeAccountStore; + +#[async_trait] +impl AccountStore for FakeAccountStore { + async fn list_pool_accounts(&self) -> AccountStoreResult> { + Ok(vec![Account::test("acct-1", AccountStatus::Active)]) + } + + async fn mark_quota_limited_until( + &self, + _account_id: &str, + _cooldown_until: chrono::DateTime, + ) -> AccountStoreResult { + Ok(true) + } + + async fn set_cloudflare_cooldown_until( + &self, + _account_id: &str, + _cooldown_until: chrono::DateTime, + ) -> AccountStoreResult { + Ok(true) + } + + async fn set_status( + &self, + _account_id: &str, + _status: AccountStatus, + ) -> AccountStoreResult { + Ok(true) + } + + async fn record_usage_delta( + &self, + _account_id: &str, + _usage: AccountUsageDelta, + ) -> AccountStoreResult<()> { + Ok(()) + } +} + +async fn mount_appcast(server: &MockServer, body: &str) { + Mock::given(method("GET")) + .and(path("/appcast.xml")) + .respond_with(ResponseTemplate::new(200).set_body_raw(body, "application/xml")) + .mount(server) + .await; +} + +async fn wait_for_auto_updated_fingerprint( + repo: &FingerprintRepository, +) -> codex_proxy_core::gateway::fingerprint::Fingerprint { + let deadline = tokio::time::Instant::now() + StdDuration::from_secs(2); + loop { + if let Some(stored) = repo + .load_latest_auto_updated() + .await + .expect("latest auto update should load") + { + return stored; + } + + assert!( + tokio::time::Instant::now() < deadline, + "fingerprint update task did not persist auto update before timeout" + ); + tokio::time::sleep(StdDuration::from_millis(25)).await; + } +} + +async fn wait_for_appcast_requests(server: &MockServer, expected_count: usize) { + let deadline = tokio::time::Instant::now() + StdDuration::from_secs(2); + loop { + let request_count = server + .received_requests() + .await + .expect("received requests should load") + .iter() + .filter(|request| request.url.path() == "/appcast.xml") + .count(); + if request_count >= expected_count { + return; + } + + assert!( + tokio::time::Instant::now() < deadline, + "fingerprint update task did not request appcast before timeout" + ); + tokio::time::sleep(StdDuration::from_millis(25)).await; + } +} + +async fn insert_account(pool: &sqlx::SqlitePool, account_id: &str) { + let now = Utc::now().to_rfc3339(); + sqlx::query( + "insert into accounts (id, access_token_cipher, status, added_at, updated_at) values (?, ?, 'active', ?, ?)", + ) + .bind(account_id) + .bind("cipher") + .bind(&now) + .bind(&now) + .execute(pool) + .await + .expect("account should be inserted"); +} + +async fn insert_cookie( + pool: &sqlx::SqlitePool, + id: &str, + account_id: &str, + name: &str, + expires_at: chrono::DateTime, +) { + sqlx::query( + "insert into account_cookies (id, account_id, domain, name, value_cipher, path, expires_at, updated_at) values (?, ?, 'chatgpt.com', ?, 'cipher', '/', ?, ?)", + ) + .bind(id) + .bind(account_id) + .bind(name) + .bind(expires_at.to_rfc3339()) + .bind(Utc::now().to_rfc3339()) + .execute(pool) + .await + .expect("cookie should be inserted"); +} + +async fn insert_admin_user(pool: &sqlx::SqlitePool, user_id: &str) { + let now = Utc::now().to_rfc3339(); + sqlx::query( + "insert into admin_users (id, password_hash, created_at, updated_at) values (?, 'hash', ?, ?)", + ) + .bind(user_id) + .bind(&now) + .bind(&now) + .execute(pool) + .await + .expect("admin user should be inserted"); +} + +async fn insert_admin_session( + pool: &sqlx::SqlitePool, + session_id: &str, + user_id: &str, + expires_at: chrono::DateTime, +) { + 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(Utc::now().to_rfc3339()) + .execute(pool) + .await + .expect("admin session should be inserted"); +} + +async fn cookie_ids(pool: &sqlx::SqlitePool) -> Vec { + sqlx::query_scalar::<_, String>("select id from account_cookies order by id") + .fetch_all(pool) + .await + .expect("cookie ids should load") +} + +async fn admin_session_ids(pool: &sqlx::SqlitePool) -> Vec { + sqlx::query_scalar::<_, String>("select id from admin_sessions order by id") + .fetch_all(pool) + .await + .expect("session ids should load") +} + +async fn session_affinity_response_ids(pool: &sqlx::SqlitePool) -> Vec { + sqlx::query_scalar::<_, String>( + "select response_id from session_affinities order by response_id", + ) + .fetch_all(pool) + .await + .expect("session affinity response ids should load") +} + +fn session_affinity_entry( + account_id: &str, + created_at: chrono::DateTime, +) -> SessionAffinityEntry { + SessionAffinityEntry { + account_id: account_id.to_string(), + conversation_id: format!("conversation-{account_id}"), + turn_state: Some(format!("turn-{account_id}")), + instructions_hash: None, + input_tokens: Some(1), + function_call_ids: Vec::new(), + variant_hash: None, + created_at, + } +} + +fn test_config(database_url: String) -> AppConfig { + AppConfig { + server: ServerConfig { + host: "127.0.0.1".to_string(), + port: 0, + }, + api: ApiConfig { + base_url: "https://example.invalid".to_string(), + }, + model: PlatformModelConfig { + default_model: "gpt-5.5".to_string(), + default_reasoning_effort: None, + service_tier: None, + aliases: Default::default(), + }, + 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(), + oauth_client_id: "app_id".to_string(), + oauth_auth_endpoint: "https://auth.invalid".to_string(), + oauth_token_endpoint: "https://token.invalid".to_string(), + }, + 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, + }, + ws_pool: Default::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + }, + } +} diff --git a/crates/runtime/tests/token_refresh.rs b/crates/runtime/tests/token_refresh.rs new file mode 100644 index 000000000..1e3e4260a --- /dev/null +++ b/crates/runtime/tests/token_refresh.rs @@ -0,0 +1,983 @@ +use std::{ + collections::VecDeque, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, + time::Duration as StdDuration, +}; + +use async_trait::async_trait; +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use chrono::{Duration, TimeZone, Utc}; +use codex_proxy_adapters::sqlite::{ + accounts::{NewAccount, SqliteAccountStore}, + refresh_leases::SqliteRefreshLeaseStore, +}; +use codex_proxy_core::{ + accounts::model::AccountStatus, + auth::{ + oauth::{RefreshFailure, RefreshPolicy, TokenPair}, + ports::TokenRefresher, + }, +}; +use codex_proxy_platform::{crypto::SecretBox, storage::connect_sqlite}; +use secrecy::{ExposeSecret, SecretString}; +use serde_json::json; +use tokio::{ + sync::{Mutex, Notify}, + time::{sleep, timeout}, +}; + +#[tokio::test] +async fn token_refresh_task_should_persist_refreshed_access_token_and_keep_refresh_token() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool, SecretBox::new([11u8; 32])); + let now = Utc.with_ymd_and_hms(2026, 6, 18, 8, 0, 0).unwrap(); + let new_expires_at = now + Duration::hours(1); + let new_access_token = test_jwt(new_expires_at.timestamp()); + store + .insert(NewAccount { + id: "acct-refresh".to_string(), + email: Some("user@example.com".to_string()), + account_id: Some("chatgpt-account".to_string()), + user_id: Some("user-id".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new( + test_jwt((now + Duration::seconds(30)).timestamp()).into(), + ), + refresh_token: Some(SecretString::new("refresh-old".to_string().into())), + access_token_expires_at: Some(now + Duration::seconds(30)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let refresher = StaticTokenRefresher { + response: Arc::new(Mutex::new(Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: None, + }))), + }; + let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher, + ); + + let summary = task + .refresh_due_accounts_once_at(now) + .await + .expect("refresh should succeed"); + let stored = store + .get("acct-refresh") + .await + .expect("account should load") + .expect("account should exist"); + + assert_eq!( + ( + summary.refreshed, + stored.access_token.expose_secret(), + stored + .refresh_token + .as_ref() + .map(ExposeSecret::expose_secret), + stored.access_token_expires_at, + stored.status, + ), + ( + 1, + new_access_token.as_str(), + Some("refresh-old"), + Some(new_expires_at), + AccountStatus::Active, + ) + ); +} + +#[tokio::test] +async fn token_refresh_task_should_skip_account_when_refresh_lease_is_held() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-lease.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool.clone(), SecretBox::new([12u8; 32])); + let leases = SqliteRefreshLeaseStore::new(pool.clone()); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 8, 0, 0).unwrap(); + let old_access_token = test_jwt((now + Duration::seconds(30)).timestamp()); + store + .insert(NewAccount { + id: "acct-lease-held".to_string(), + email: Some("lease@example.com".to_string()), + account_id: Some("chatgpt-lease".to_string()), + user_id: Some("user-lease".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(old_access_token.clone().into()), + refresh_token: Some(SecretString::new("refresh-held".to_string().into())), + access_token_expires_at: Some(now + Duration::seconds(30)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + assert!(leases + .try_acquire( + "acct-lease-held", + "external-owner", + now + Duration::minutes(5), + now, + ) + .await + .expect("external owner should acquire lease")); + let refresher = CountingTokenRefresher::default(); + let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher.clone(), + ) + .with_refresh_lease_store(leases); + + let summary = task + .refresh_due_accounts_once_at(now) + .await + .expect("refresh scan should skip lease-held account"); + let stored = store + .get("acct-lease-held") + .await + .expect("account should load") + .expect("account should exist"); + + assert_eq!(summary.skipped, 1); + assert_eq!(summary.refreshed, 0); + assert_eq!(refresher.calls.load(Ordering::SeqCst), 0); + assert_eq!( + stored.access_token.expose_secret(), + old_access_token.as_str() + ); + assert_eq!(stored.status, AccountStatus::Active); +} + +#[tokio::test] +async fn token_refresh_task_should_mark_due_account_refreshing_before_refresher_call() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-status.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool, SecretBox::new([13u8; 32])); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 9, 0, 0).unwrap(); + let new_expires_at = now + Duration::hours(1); + let new_access_token = test_jwt(new_expires_at.timestamp()); + store + .insert(NewAccount { + id: "acct-refreshing-marker".to_string(), + email: Some("marker@example.com".to_string()), + account_id: Some("chatgpt-marker".to_string()), + user_id: Some("user-marker".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new( + test_jwt((now + Duration::seconds(30)).timestamp()).into(), + ), + refresh_token: Some(SecretString::new("refresh-marker".to_string().into())), + access_token_expires_at: Some(now + Duration::seconds(30)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let observed_statuses = Arc::new(Mutex::new(Vec::new())); + let refresher = StoreObservingTokenRefresher { + store: store.clone(), + account_id: "acct-refreshing-marker".to_string(), + observed_statuses: observed_statuses.clone(), + response: Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: None, + }), + }; + let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher, + ); + + let summary = task + .refresh_due_accounts_once_at(now) + .await + .expect("refresh should succeed"); + let stored = store + .get("acct-refreshing-marker") + .await + .expect("account should load") + .expect("account should exist"); + let observed_statuses = observed_statuses.lock().await.clone(); + + assert_eq!(summary.refreshed, 1); + assert_eq!(observed_statuses, [AccountStatus::Refreshing]); + assert_eq!(stored.status, AccountStatus::Active); + assert_eq!( + stored.access_token.expose_secret(), + new_access_token.as_str() + ); +} + +#[tokio::test] +async fn token_refresh_task_should_recover_refreshing_account_after_restart() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-recovery.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool, SecretBox::new([14u8; 32])); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 10, 0, 0).unwrap(); + let old_access_token = test_jwt((now + Duration::hours(1)).timestamp()); + let new_access_token = test_jwt((now + Duration::hours(2)).timestamp()); + store + .insert(NewAccount { + id: "acct-refreshing-recovery".to_string(), + email: Some("recovery@example.com".to_string()), + account_id: Some("chatgpt-recovery".to_string()), + user_id: Some("user-recovery".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(old_access_token.into()), + refresh_token: Some(SecretString::new("refresh-recovery".to_string().into())), + access_token_expires_at: Some(now + Duration::hours(1)), + status: AccountStatus::Refreshing, + }) + .await + .expect("account should be inserted"); + let refresher = StaticTokenRefresher { + response: Arc::new(Mutex::new(Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: None, + }))), + }; + let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher, + ); + + let summary = task + .refresh_due_accounts_once_at(now) + .await + .expect("refreshing account should recover"); + let stored = store + .get("acct-refreshing-recovery") + .await + .expect("account should load") + .expect("account should exist"); + + assert_eq!(summary.refreshed, 1); + assert_eq!(stored.status, AccountStatus::Active); + assert_eq!( + stored.access_token.expose_secret(), + new_access_token.as_str() + ); +} + +#[tokio::test] +async fn token_refresh_task_should_restore_active_status_after_transport_failure() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-transport-failure.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool, SecretBox::new([15u8; 32])); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 11, 0, 0).unwrap(); + let old_access_token = test_jwt((now + Duration::seconds(30)).timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-transport".to_string(), + email: Some("transport@example.com".to_string()), + account_id: Some("chatgpt-transport".to_string()), + user_id: Some("user-transport".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(old_access_token.clone().into()), + refresh_token: Some(SecretString::new("refresh-transport".to_string().into())), + access_token_expires_at: Some(now + Duration::seconds(30)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let observed_statuses = Arc::new(Mutex::new(Vec::new())); + let refresher = StoreObservingTokenRefresher { + store: store.clone(), + account_id: "acct-refresh-transport".to_string(), + observed_statuses: observed_statuses.clone(), + response: Err(RefreshFailure::Transport), + }; + let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher, + ) + .with_retry_delays(vec![StdDuration::ZERO; 4]); + + let summary = task + .refresh_due_accounts_once_at(now) + .await + .expect("transport failure should be summarized"); + let stored = store + .get("acct-refresh-transport") + .await + .expect("account should load") + .expect("account should exist"); + let observed_statuses = observed_statuses.lock().await.clone(); + + assert_eq!(summary.failed, 1); + assert_eq!(observed_statuses, [AccountStatus::Refreshing; 5]); + assert_eq!(stored.status, AccountStatus::Active); + assert_eq!( + stored.access_token.expose_secret(), + old_access_token.as_str() + ); +} + +#[tokio::test] +async fn token_refresh_task_should_skip_duplicate_in_flight_refresh() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-in-flight.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool, SecretBox::new([18u8; 32])); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 14, 0, 0).unwrap(); + let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-in-flight".to_string(), + email: Some("in-flight@example.com".to_string()), + account_id: Some("chatgpt-in-flight".to_string()), + user_id: Some("user-in-flight".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new( + test_jwt((now + Duration::seconds(30)).timestamp()).into(), + ), + refresh_token: Some(SecretString::new("refresh-in-flight".to_string().into())), + access_token_expires_at: Some(now + Duration::seconds(30)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let refresher = BlockingTokenRefresher { + calls: Arc::new(AtomicUsize::new(0)), + started: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + response: Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: None, + }), + }; + let task = Arc::new( + codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher.clone(), + ), + ); + + let first = { + let task = task.clone(); + tokio::spawn(async move { task.refresh_due_accounts_once_at(now).await }) + }; + refresher.started.notified().await; + + let second = { + let task = task.clone(); + tokio::spawn(async move { task.refresh_due_accounts_once_at(now).await }) + }; + let second = match timeout(StdDuration::from_millis(100), second).await { + Ok(result) => result.expect("second scan should join"), + Err(_) => { + refresher.release.notify_waiters(); + let _ = first.await; + panic!("second scan should skip instead of waiting on duplicate refresh"); + } + } + .expect("second scan should succeed"); + + refresher.release.notify_waiters(); + let first = first + .await + .expect("first scan should join") + .expect("first scan should succeed"); + let stored = store + .get("acct-refresh-in-flight") + .await + .expect("account should load") + .expect("account should exist"); + + assert_eq!(first.refreshed, 1); + assert_eq!(second.skipped, 1); + assert_eq!(refresher.calls.load(Ordering::SeqCst), 1); + assert_eq!(stored.status, AccountStatus::Active); + assert_eq!( + stored.access_token.expose_secret(), + new_access_token.as_str() + ); +} + +#[tokio::test] +async fn token_refresh_task_should_persist_success_after_transient_transport_retry() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-transient-success.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool, SecretBox::new([16u8; 32])); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 12, 0, 0).unwrap(); + let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-transient-success".to_string(), + email: Some("transient@example.com".to_string()), + account_id: Some("chatgpt-transient".to_string()), + user_id: Some("user-transient".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new( + test_jwt((now + Duration::seconds(30)).timestamp()).into(), + ), + refresh_token: Some(SecretString::new("refresh-transient".to_string().into())), + access_token_expires_at: Some(now + Duration::seconds(30)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let observed_statuses = Arc::new(Mutex::new(Vec::new())); + let refresher = SequenceTokenRefresher { + store: store.clone(), + account_id: "acct-refresh-transient-success".to_string(), + observed_statuses: observed_statuses.clone(), + responses: Arc::new(Mutex::new(VecDeque::from(vec![ + Err(RefreshFailure::Transport), + Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: None, + }), + ]))), + }; + let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher, + ) + .with_retry_delays(vec![StdDuration::ZERO; 4]); + + let summary = task + .refresh_due_accounts_once_at(now) + .await + .expect("retry should eventually refresh"); + let stored = store + .get("acct-refresh-transient-success") + .await + .expect("account should load") + .expect("account should exist"); + let observed_statuses = observed_statuses.lock().await.clone(); + + assert_eq!(summary.refreshed, 1); + assert_eq!(summary.failed, 0); + assert_eq!(observed_statuses, [AccountStatus::Refreshing; 2]); + assert_eq!(stored.status, AccountStatus::Active); + assert_eq!( + stored.access_token.expose_secret(), + new_access_token.as_str() + ); +} + +#[tokio::test] +async fn token_refresh_task_should_delay_recovery_after_retry_exhaustion() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-delayed-recovery.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool, SecretBox::new([19u8; 32])); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 15, 0, 0).unwrap(); + let old_access_token = test_jwt((now + Duration::seconds(30)).timestamp()); + let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-delayed-recovery".to_string(), + email: Some("delayed-recovery@example.com".to_string()), + account_id: Some("chatgpt-delayed-recovery".to_string()), + user_id: Some("user-delayed-recovery".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(old_access_token.clone().into()), + refresh_token: Some(SecretString::new( + "refresh-delayed-recovery".to_string().into(), + )), + access_token_expires_at: Some(now + Duration::seconds(30)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let mut responses = Vec::new(); + responses.extend((0..5).map(|_| Err(RefreshFailure::Transport))); + responses.push(Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: None, + })); + let observed_statuses = Arc::new(Mutex::new(Vec::new())); + let refresher = SequenceTokenRefresher { + store: store.clone(), + account_id: "acct-refresh-delayed-recovery".to_string(), + observed_statuses: observed_statuses.clone(), + responses: Arc::new(Mutex::new(VecDeque::from(responses))), + }; + let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher, + ) + .with_retry_delays(vec![StdDuration::ZERO; 4]); + + let failed = task + .refresh_due_accounts_once_at(now) + .await + .expect("first scan should summarize retry exhaustion"); + let delayed = task + .refresh_due_accounts_once_at(now + Duration::minutes(5)) + .await + .expect("recovery window should skip refresh"); + let refreshed = task + .refresh_due_accounts_once_at(now + Duration::minutes(10) + Duration::seconds(1)) + .await + .expect("recovery window should allow refresh"); + let stored = store + .get("acct-refresh-delayed-recovery") + .await + .expect("account should load") + .expect("account should exist"); + let observed_statuses = observed_statuses.lock().await.clone(); + + assert_eq!(failed.failed, 1); + assert_eq!(delayed.skipped, 1); + assert_eq!(refreshed.refreshed, 1); + assert_eq!(observed_statuses, [AccountStatus::Refreshing; 6]); + assert_eq!(stored.status, AccountStatus::Active); + assert_eq!( + stored.access_token.expose_secret(), + new_access_token.as_str() + ); +} + +#[tokio::test] +async fn token_refresh_task_should_schedule_future_per_account_timer_without_refreshing_early() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-future-timer.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool, SecretBox::new([20u8; 32])); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 16, 0, 0).unwrap(); + let old_access_token = test_jwt((now + Duration::minutes(10)).timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-future-timer".to_string(), + email: Some("future-timer@example.com".to_string()), + account_id: Some("chatgpt-future-timer".to_string()), + user_id: Some("user-future-timer".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(old_access_token.clone().into()), + refresh_token: Some(SecretString::new("refresh-future-timer".to_string().into())), + access_token_expires_at: Some(now + Duration::minutes(10)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let refresher = CountingTokenRefresher::default(); + let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store, + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher.clone(), + ); + + let summary = task + .schedule_account_timers_once_at(now) + .await + .expect("timer scheduling should succeed"); + tokio::task::yield_now().await; + + assert_eq!(summary.scheduled, 1); + assert_eq!(task.scheduled_timer_count().await, 1); + assert_eq!(refresher.calls.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn token_refresh_task_should_fire_per_account_timer_at_refresh_time() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-per-account-timer.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool, SecretBox::new([21u8; 32])); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 17, 0, 0).unwrap(); + let old_access_token = test_jwt((now + Duration::seconds(1)).timestamp()); + let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-timer".to_string(), + email: Some("timer@example.com".to_string()), + account_id: Some("chatgpt-timer".to_string()), + user_id: Some("user-timer".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(old_access_token.clone().into()), + refresh_token: Some(SecretString::new("refresh-timer".to_string().into())), + access_token_expires_at: Some(now + Duration::seconds(1)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let refresher = NotifyingTokenRefresher { + calls: Arc::new(AtomicUsize::new(0)), + started: Arc::new(Notify::new()), + response: Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: None, + }), + }; + let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 1, + refresh_concurrency: 1, + }, + refresher.clone(), + ) + .with_retry_delays(vec![StdDuration::ZERO; 4]); + + let started = refresher.started.notified(); + tokio::pin!(started); + let summary = task + .schedule_account_timers_once_at(now) + .await + .expect("timer scheduling should succeed"); + timeout(StdDuration::from_secs(2), &mut started) + .await + .expect("scheduled timer should trigger refresh"); + tokio::task::yield_now().await; + let stored = store + .get("acct-refresh-timer") + .await + .expect("account should load") + .expect("account should exist"); + + assert_eq!(summary.immediate, 1); + assert_eq!(refresher.calls.load(Ordering::SeqCst), 1); + assert_eq!(task.scheduled_timer_count().await, 1); + assert_eq!(stored.status, AccountStatus::Active); + assert_eq!( + stored.access_token.expose_secret(), + new_access_token.as_str() + ); +} + +#[tokio::test] +async fn token_refresh_task_should_reschedule_next_timer_after_scheduled_refresh() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-next-timer.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool, SecretBox::new([22u8; 32])); + let now = Utc::now(); + let old_expires_at = now + Duration::milliseconds(50); + let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-next-timer".to_string(), + email: Some("next-timer@example.com".to_string()), + account_id: Some("chatgpt-next-timer".to_string()), + user_id: Some("user-next-timer".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(test_jwt(old_expires_at.timestamp()).into()), + refresh_token: Some(SecretString::new("refresh-next-timer".to_string().into())), + access_token_expires_at: Some(old_expires_at), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let refresher = NotifyingTokenRefresher { + calls: Arc::new(AtomicUsize::new(0)), + started: Arc::new(Notify::new()), + response: Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: None, + }), + }; + let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 0, + refresh_concurrency: 1, + }, + refresher.clone(), + ) + .with_retry_delays(vec![StdDuration::ZERO; 4]); + + let started = refresher.started.notified(); + tokio::pin!(started); + let summary = task + .schedule_account_timers_once_at(now) + .await + .expect("timer scheduling should succeed"); + timeout(StdDuration::from_secs(2), &mut started) + .await + .expect("scheduled timer should trigger refresh"); + + let mut stored = store + .get("acct-refresh-next-timer") + .await + .expect("account should load") + .expect("account should exist"); + for _ in 0..50 { + if stored.access_token.expose_secret() == new_access_token.as_str() { + break; + } + sleep(StdDuration::from_millis(20)).await; + stored = store + .get("acct-refresh-next-timer") + .await + .expect("account should load") + .expect("account should exist"); + } + + assert_eq!(summary.scheduled, 1); + assert_eq!(refresher.calls.load(Ordering::SeqCst), 1); + assert_eq!( + stored.access_token.expose_secret(), + new_access_token.as_str() + ); + assert_eq!(task.scheduled_timer_count().await, 1); +} + +#[tokio::test] +async fn token_refresh_task_should_confirm_invalid_grant_before_expiring_account() { + let dir = tempfile::tempdir().expect("temp dir"); + let db = dir.path().join("token-refresh-permanent-failure.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .expect("sqlite pool"); + let store = SqliteAccountStore::new(pool, SecretBox::new([17u8; 32])); + let now = Utc.with_ymd_and_hms(2026, 6, 19, 13, 0, 0).unwrap(); + let old_access_token = test_jwt((now + Duration::seconds(30)).timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-invalid-grant".to_string(), + email: Some("invalid-grant@example.com".to_string()), + account_id: Some("chatgpt-invalid-grant".to_string()), + user_id: Some("user-invalid-grant".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new(old_access_token.clone().into()), + refresh_token: Some(SecretString::new( + "refresh-invalid-grant".to_string().into(), + )), + access_token_expires_at: Some(now + Duration::seconds(30)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let observed_statuses = Arc::new(Mutex::new(Vec::new())); + let refresher = SequenceTokenRefresher { + store: store.clone(), + account_id: "acct-refresh-invalid-grant".to_string(), + observed_statuses: observed_statuses.clone(), + responses: Arc::new(Mutex::new(VecDeque::from(vec![ + Err(RefreshFailure::InvalidGrant), + Err(RefreshFailure::InvalidGrant), + ]))), + }; + let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher, + ) + .with_retry_delays(vec![StdDuration::ZERO; 4]); + + let summary = task + .refresh_due_accounts_once_at(now) + .await + .expect("confirmed permanent failure should update status"); + let stored = store + .get("acct-refresh-invalid-grant") + .await + .expect("account should load") + .expect("account should exist"); + let observed_statuses = observed_statuses.lock().await.clone(); + + assert_eq!(summary.status_updated, 1); + assert_eq!(observed_statuses, [AccountStatus::Refreshing; 2]); + assert_eq!(stored.status, AccountStatus::Expired); + assert_eq!( + stored.access_token.expose_secret(), + old_access_token.as_str() + ); +} + +#[derive(Clone)] +struct StaticTokenRefresher { + response: Arc>>, +} + +#[async_trait] +impl TokenRefresher for StaticTokenRefresher { + async fn refresh(&self, _refresh_token: &str) -> Result { + self.response.lock().await.clone() + } +} + +#[derive(Clone, Default)] +struct CountingTokenRefresher { + calls: Arc, +} + +#[async_trait] +impl TokenRefresher for CountingTokenRefresher { + async fn refresh(&self, _refresh_token: &str) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(TokenPair { + access_token: "new-access".to_string(), + refresh_token: None, + }) + } +} + +#[derive(Clone)] +struct BlockingTokenRefresher { + calls: Arc, + started: Arc, + release: Arc, + response: Result, +} + +#[async_trait] +impl TokenRefresher for BlockingTokenRefresher { + async fn refresh(&self, _refresh_token: &str) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + self.started.notify_waiters(); + self.release.notified().await; + self.response.clone() + } +} + +#[derive(Clone)] +struct NotifyingTokenRefresher { + calls: Arc, + started: Arc, + response: Result, +} + +#[async_trait] +impl TokenRefresher for NotifyingTokenRefresher { + async fn refresh(&self, _refresh_token: &str) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + self.started.notify_waiters(); + self.response.clone() + } +} + +#[derive(Clone)] +struct StoreObservingTokenRefresher { + store: SqliteAccountStore, + account_id: String, + observed_statuses: Arc>>, + response: Result, +} + +#[async_trait] +impl TokenRefresher for StoreObservingTokenRefresher { + async fn refresh(&self, _refresh_token: &str) -> Result { + let status = self + .store + .get(&self.account_id) + .await + .expect("account should load") + .expect("account should exist") + .status; + self.observed_statuses.lock().await.push(status); + self.response.clone() + } +} + +#[derive(Clone)] +struct SequenceTokenRefresher { + store: SqliteAccountStore, + account_id: String, + observed_statuses: Arc>>, + responses: Arc>>>, +} + +#[async_trait] +impl TokenRefresher for SequenceTokenRefresher { + async fn refresh(&self, _refresh_token: &str) -> Result { + let status = self + .store + .get(&self.account_id) + .await + .expect("account should load") + .expect("account should exist") + .status; + self.observed_statuses.lock().await.push(status); + self.responses + .lock() + .await + .pop_front() + .expect("refresh response should be configured") + } +} + +fn test_jwt(exp: i64) -> String { + let header = json!({"alg": "none", "typ": "JWT"}); + let payload = json!({ "exp": exp }); + format!("{}.{}.", jwt_part(&header), jwt_part(&payload)) +} + +fn jwt_part(value: &serde_json::Value) -> String { + URL_SAFE_NO_PAD.encode(serde_json::to_vec(value).expect("json should encode")) +} diff --git a/crates/runtime/tests/upstream.rs b/crates/runtime/tests/upstream.rs new file mode 100644 index 000000000..f2d844169 --- /dev/null +++ b/crates/runtime/tests/upstream.rs @@ -0,0 +1,132 @@ +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, +}; + +use codex_proxy_adapters::codex::client; +use codex_proxy_core::{ + gateway::fingerprint::Fingerprint, protocol::codex::responses::CodexResponsesRequest, +}; +use codex_proxy_platform::config::WebSocketPoolConfig; +use futures::{SinkExt, StreamExt}; +use serde_json::json; +use tokio::net::TcpListener; +use tokio_tungstenite::{accept_async, tungstenite::Message}; + +#[tokio::test] +async fn codex_backend_client_should_apply_configured_websocket_pool() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut first_websocket = accept_async(first_stream).await.unwrap(); + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_websocket + .send(Message::Text( + response_completed_websocket_message("resp_runtime_pool_first").into(), + )) + .await + .unwrap(); + + tokio::select! { + second_message = first_websocket.next() => { + let _second_message = second_message.unwrap().unwrap(); + first_websocket + .send(Message::Text( + response_completed_websocket_message("resp_runtime_pool_second").into(), + )) + .await + .unwrap(); + first_websocket.close(None).await.unwrap(); + } + accepted = listener.accept() => { + let (second_stream, _) = accepted.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut second_websocket = accept_async(second_stream).await.unwrap(); + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + response_completed_websocket_message("resp_runtime_pool_second").into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + first_websocket.close(None).await.unwrap(); + } + } + }); + let ws_pool = WebSocketPoolConfig { + enabled: true, + max_age_ms: 60_000, + max_per_account: 8, + }; + let backend = codex_proxy_runtime::upstream::codex_backend_client( + format!("http://{addr}"), + Fingerprint::default_for_tests(), + &ws_pool, + ); + let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "be brief", Vec::new()); + request.previous_response_id = Some("resp_runtime_pool_previous".to_string()); + request.prompt_cache_key = Some("conversation-runtime-pool".to_string()); + + let first = backend + .create_response( + &request, + request_context("req_runtime_pool_first", Some("chatgpt-account")), + ) + .await + .expect("first runtime websocket response should succeed"); + let second = backend + .create_response( + &request, + request_context("req_runtime_pool_second", Some("chatgpt-account")), + ) + .await + .expect("second runtime websocket response should reuse configured pool"); + server.await.unwrap(); + + assert!(first.body.contains("resp_runtime_pool_first")); + assert!(second.body.contains("resp_runtime_pool_second")); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 1); +} + +fn request_context<'a>( + request_id: &'a str, + account_id: Option<&'a str>, +) -> client::CodexRequestContext<'a> { + client::CodexRequestContext { + access_token: "access-token", + account_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, + installation_id: None, + session_id: None, + } +} + +fn response_completed_websocket_message(response_id: &str) -> String { + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "object": "response", + "status": "completed", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string() +} diff --git a/crates/server/Cargo.toml b/crates/server/Cargo.toml new file mode 100644 index 000000000..df31490aa --- /dev/null +++ b/crates/server/Cargo.toml @@ -0,0 +1,38 @@ +[package] +name = "codex-proxy-server" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +rust-version.workspace = true + +[dependencies] +axum.workspace = true +codex-proxy-assets = { path = "../assets" } +codex-proxy-core = { path = "../core" } +codex-proxy-platform = { path = "../platform" } +codex-proxy-runtime = { path = "../runtime" } +serde.workspace = true +serde_json.workspace = true +tokio.workspace = true +tower-http.workspace = true +tracing.workspace = true +uuid.workspace = true + +[dev-dependencies] +async-trait.workspace = true +base64.workspace = true +chrono.workspace = true +codex-proxy-adapters = { path = "../adapters" } +futures.workspace = true +secrecy.workspace = true +serde_yml.workspace = true +sqlx.workspace = true +tempfile = "3.27.0" +tokio-tungstenite.workspace = true +tower.workspace = true +tracing-subscriber.workspace = true +wiremock = "0.6.5" + +[lints] +workspace = true diff --git a/src/admin/api/accounts/cookies.rs b/crates/server/src/admin_api/accounts/cookies.rs similarity index 55% rename from src/admin/api/accounts/cookies.rs rename to crates/server/src/admin_api/accounts/cookies.rs index 4ad59882c..f98b3c614 100644 --- a/src/admin/api/accounts/cookies.rs +++ b/crates/server/src/admin_api/accounts/cookies.rs @@ -1,3 +1,5 @@ +//! 账号 Cookie 处理器。 + use axum::{ extract::{Path, State}, http::{HeaderMap, StatusCode}, @@ -7,29 +9,41 @@ use axum::{ use serde::{Deserialize, Serialize}; use serde_json::Value; -use crate::{platform::http::request_id::RequestId, runtime::state::AppState}; +use codex_proxy_runtime::state::AppState; -use super::super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; -use super::account_service_error; +use crate::{ + admin_api::{ + accounts::{account_error, account_not_found}, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, +}; +/// 设置账号 Cookie 请求。 #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] pub struct SetAccountCookiesRequest { + /// Cookie 请求头字符串，或 name/value 对象。 pub cookies: Value, } -#[derive(Debug, Serialize)] +/// 账号 Cookie 响应。 +#[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct AccountCookiesData { + /// 标准 HTTP Cookie 请求头。 pub cookies: Option, } -#[derive(Debug, Serialize)] +/// 删除账号 Cookie 响应。 +#[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct DeleteAccountCookiesData { + /// 是否完成删除。 pub deleted: bool, } +/// `GET /api/admin/accounts/{account_id}/cookies` pub async fn get_account_cookies( State(state): State, Extension(request_id): Extension, @@ -39,19 +53,19 @@ pub async fn get_account_cookies( 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), - )) + match state.services.admin_accounts.cookies(&account_id).await { + Ok(cookies) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountCookiesData { cookies }, request_id), + )), + Err(codex_proxy_runtime::services::AdminAccountError::NotFound) => { + Err(account_not_found(request_id)) + } + Err(error) => Err(account_error(error, request_id)), + } } +/// `POST /api/admin/accounts/{account_id}/cookies` pub async fn set_account_cookies( State(state): State, Extension(request_id): Extension, @@ -60,32 +74,29 @@ pub async fn set_account_cookies( 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, - )); - } - }; + let cookie_header = admin_cookie_header(&payload.cookies).map_err(|message| { + AdminError::new(StatusCode::BAD_REQUEST, 40001, message, request_id.clone()) + })?; require_admin_session(&state, &headers, &request_id).await?; - let cookies = state + match state .services - .accounts + .admin_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), - )) + { + Ok(cookies) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountCookiesData { cookies }, request_id), + )), + Err(codex_proxy_runtime::services::AdminAccountError::NotFound) => { + Err(account_not_found(request_id)) + } + Err(error) => Err(account_error(error, request_id)), + } } +/// `DELETE /api/admin/accounts/{account_id}/cookies` pub async fn delete_account_cookies( State(state): State, Extension(request_id): Extension, @@ -95,17 +106,21 @@ pub async fn delete_account_cookies( let request_id = request_id.as_str().to_string(); require_admin_session(&state, &headers, &request_id).await?; - state + match state .services - .accounts + .admin_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), - )) + { + Ok(()) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(DeleteAccountCookiesData { deleted: true }, request_id), + )), + Err(codex_proxy_runtime::services::AdminAccountError::NotFound) => { + Err(account_not_found(request_id)) + } + Err(error) => Err(account_error(error, request_id)), + } } fn admin_cookie_header(value: &Value) -> Result { @@ -116,12 +131,14 @@ fn admin_cookie_header(value: &Value) -> Result { } 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)| { diff --git a/crates/server/src/admin_api/accounts/create.rs b/crates/server/src/admin_api/accounts/create.rs new file mode 100644 index 000000000..2ff44eb0a --- /dev/null +++ b/crates/server/src/admin_api/accounts/create.rs @@ -0,0 +1,52 @@ +//! 账号创建处理器。 + +use axum::{ + extract::State, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, Json, +}; +use codex_proxy_runtime::state::AppState; +use serde::Deserialize; + +use crate::{ + admin_api::{ + accounts::{account_error, AdminAccountData}, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +/// 手动创建账号请求。 +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CreateAccountRequest { + /// access token，允许带 Bearer 前缀。 + pub token: Option, + /// refresh token。 + pub refresh_token: Option, +} + +/// `POST /api/admin/accounts` +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?; + + match state + .services + .admin_accounts + .create(payload.token, payload.refresh_token) + .await + { + Ok(account) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminAccountData::from(account), request_id), + )), + Err(error) => Err(account_error(error, request_id)), + } +} diff --git a/crates/server/src/admin_api/accounts/export.rs b/crates/server/src/admin_api/accounts/export.rs new file mode 100644 index 000000000..cbc30702a --- /dev/null +++ b/crates/server/src/admin_api/accounts/export.rs @@ -0,0 +1,119 @@ +//! 账号导出处理器。 + +use axum::{ + extract::{Query, State}, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, +}; +use codex_proxy_runtime::{services::AdminStoredAccount, state::AppState}; +use serde::{Deserialize, Serialize}; + +use crate::{ + admin_api::{ + accounts::{account_error, account_status_value}, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +/// 账号导出查询。 +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountExportQuery { + /// 逗号分隔的账号 ID 列表；空表示导出全部。 + pub ids: Option, + /// 导出格式；当前仅支持 `native`。 + pub format: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountExportData { + source_format: &'static str, + accounts: Vec, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountExportEntry { + id: String, + email: Option, + account_id: Option, + user_id: Option, + label: Option, + plan_type: Option, + token: String, + refresh_token: Option, + status: String, + access_token_expires_at: Option, + added_at: String, + updated_at: String, +} + +/// `GET /api/admin/accounts/export` +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(); + require_admin_session(&state, &headers, &request_id).await?; + validate_account_export_format(query.format.as_deref()).map_err(|message| { + AdminError::new(StatusCode::BAD_REQUEST, 40001, message, request_id.clone()) + })?; + let ids = account_export_ids(query.ids.as_deref()); + + match state.services.admin_accounts.export(ids).await { + Ok(accounts) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AccountExportData { + source_format: "native", + accounts: accounts.into_iter().map(AccountExportEntry::from).collect(), + }, + request_id, + ), + )), + Err(error) => Err(account_error(error, request_id)), + } +} + +fn validate_account_export_format(value: Option<&str>) -> Result<(), &'static str> { + match value.unwrap_or("native").trim() { + "" | "native" => Ok(()), + _ => Err("Unsupported account export format"), + } +} + +fn account_export_ids(value: Option<&str>) -> Vec { + value + .into_iter() + .flat_map(|ids| ids.split(',')) + .map(str::trim) + .filter(|id| !id.is_empty()) + .map(ToString::to_string) + .collect() +} + +impl From for AccountExportEntry { + fn from(account: AdminStoredAccount) -> 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, + token: account.access_token, + refresh_token: account.refresh_token, + 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(), + } + } +} diff --git a/src/admin/api/accounts/health.rs b/crates/server/src/admin_api/accounts/health.rs similarity index 50% rename from src/admin/api/accounts/health.rs rename to crates/server/src/admin_api/accounts/health.rs index 11be995fc..a189288f8 100644 --- a/src/admin/api/accounts/health.rs +++ b/crates/server/src/admin_api/accounts/health.rs @@ -1,3 +1,5 @@ +//! 账号健康检查处理器。 + use axum::{ body::Bytes, extract::State, @@ -5,53 +7,77 @@ use axum::{ response::IntoResponse, Extension, }; +use codex_proxy_runtime::{ + services::{AdminAccountProbeOutcome, AdminAccountProbeResult}, + state::AppState, +}; use serde::{Deserialize, Serialize}; use crate::{ - codex::accounts::service::{AccountProbeResult, HealthCheckError}, - platform::http::request_id::RequestId, - runtime::state::AppState, + admin_api::{ + accounts::{account_error, account_status_value}, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, }; -use super::super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; -use super::account_status_value; - +/// 健康检查请求。 #[derive(Debug, Default, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] pub struct HealthCheckRequest { + /// 指定账号 ID；为空时检查全部账号。 pub ids: Option>, + /// 每个探测启动间隔毫秒。 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)] +/// 健康检查汇总。 +#[derive(Debug, Clone, Copy, 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 { + /// 账号 ID。 pub id: String, + /// 邮箱。 pub email: Option, + /// 探测前状态。 pub previous_status: String, + /// 探测结果。 pub result: String, + /// 探测后状态。 pub status: Option, + /// 错误信息。 pub error: Option, + /// 耗时毫秒。 pub duration_ms: Option, } +/// `POST /api/admin/accounts/health-check` pub async fn health_check_accounts( State(state): State, Extension(request_id): Extension, @@ -60,6 +86,67 @@ pub async fn health_check_accounts( ) -> Result { let request_id = request_id.as_str().to_string(); let payload = parse_health_check_request(&body, &request_id)?; + validate_health_check_request(&payload, &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); + match state + .services + .admin_accounts + .health_check_accounts(payload.ids, concurrency, stagger_ms, &request_id) + .await + { + Ok(results) => { + let results = results + .into_iter() + .map(AccountProbeData::from) + .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), + )) + } + Err(error) => Err(account_error(error, request_id)), + } +} + +fn parse_health_check_request( + body: &Bytes, + request_id: &str, +) -> Result { + if body.is_empty() { + return Ok(HealthCheckRequest::default()); + } + serde_json::from_slice(body).map_err(|_| { + AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + "Invalid health check request", + request_id, + ) + }) +} + +fn validate_health_check_request( + payload: &HealthCheckRequest, + request_id: &str, +) -> Result<(), AdminError> { if payload.ids.as_ref().is_some_and(Vec::is_empty) { return Err(AdminError::new( StatusCode::BAD_REQUEST, @@ -90,86 +177,26 @@ pub async fn health_check_accounts( 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) - .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), - )) + Ok(()) } -fn parse_health_check_request( - body: &Bytes, - request_id: &str, -) -> Result { - if body.is_empty() { - return Ok(HealthCheckRequest::default()); - } - serde_json::from_slice(body).map_err(|_| { - AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - "Invalid health check request", - request_id, - ) - }) -} - -pub(super) 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, - ), +impl From for AccountProbeData { + fn from(result: AdminAccountProbeResult) -> Self { + Self { + id: result.id, + email: result.email, + previous_status: account_status_value(result.previous_status).to_string(), + result: probe_outcome_value(result.outcome).to_string(), + status: result + .status + .map(account_status_value) + .map(ToString::to_string), + error: result.error, + duration_ms: result.duration_ms, + } } } -pub(crate) 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, - } +fn probe_outcome_value(outcome: AdminAccountProbeOutcome) -> &'static str { + outcome.as_str() } diff --git a/crates/server/src/admin_api/accounts/import.rs b/crates/server/src/admin_api/accounts/import.rs new file mode 100644 index 000000000..06c895018 --- /dev/null +++ b/crates/server/src/admin_api/accounts/import.rs @@ -0,0 +1,52 @@ +//! 账号导入处理器。 + +use axum::{ + extract::State, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, Json, +}; +use codex_proxy_runtime::state::AppState; +use serde::Serialize; +use serde_json::Value; + +use crate::{ + admin_api::{ + accounts::account_error, require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountImportData { + imported: u32, + skipped: u32, + source_format: &'static str, +} + +/// `POST /api/admin/accounts/import` +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?; + + match state.services.admin_accounts.import(&payload).await { + Ok(imported) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AccountImportData { + imported: imported.imported, + skipped: imported.skipped, + source_format: "native", + }, + request_id, + ), + )), + Err(error) => Err(account_error(error, request_id)), + } +} diff --git a/crates/server/src/admin_api/accounts/import_cli.rs b/crates/server/src/admin_api/accounts/import_cli.rs new file mode 100644 index 000000000..44ca01830 --- /dev/null +++ b/crates/server/src/admin_api/accounts/import_cli.rs @@ -0,0 +1,75 @@ +//! Codex CLI 账号导入处理器。 + +use std::path::{Path, PathBuf}; + +use axum::{ + extract::State, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, Json, +}; +use codex_proxy_runtime::{services::AdminAccountError, state::AppState}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::{ + admin_api::{ + accounts::account_error, require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ImportCliRequest { + codex_home: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountImportCliData { + imported: u32, + skipped: u32, + source_format: &'static str, +} + +/// `POST /api/admin/accounts/import-cli` +pub async fn import_cli_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 auth_path = PathBuf::from(payload.codex_home).join("auth.json"); + let auth_payload = + read_auth_json(&auth_path).map_err(|error| account_error(error, request_id.clone()))?; + + match state + .services + .admin_accounts + .import_codex_cli_auth(&auth_payload) + .await + { + Ok(_) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AccountImportCliData { + imported: 1, + skipped: 0, + source_format: "codex_cli", + }, + request_id, + ), + )), + Err(error) => Err(account_error(error, request_id)), + } +} + +fn read_auth_json(path: &Path) -> Result { + let content = + std::fs::read_to_string(path).map_err(|_| AdminAccountError::NoImportableAccounts)?; + serde_json::from_str(&content).map_err(|_| AdminAccountError::NoImportableAccounts) +} diff --git a/crates/server/src/admin_api/accounts/lifecycle.rs b/crates/server/src/admin_api/accounts/lifecycle.rs new file mode 100644 index 000000000..5b5241e2d --- /dev/null +++ b/crates/server/src/admin_api/accounts/lifecycle.rs @@ -0,0 +1,336 @@ +//! 账号生命周期处理器。 + +use axum::{ + extract::{Path, State}, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, Json, +}; +use serde::{Deserialize, Serialize}; + +use codex_proxy_runtime::{ + services::{AdminAccountProbeOutcome, AdminAccountRefresh}, + state::AppState, +}; + +use crate::{ + admin_api::{ + accounts::{account_error, account_not_found, account_status_value}, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +/// 更新账号标签请求。 +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateAccountLabelRequest { + /// 新标签；`null` 表示清空。 + pub label: Option, +} + +/// 更新账号标签响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateAccountLabelData { + /// 账号 ID。 + pub id: String, + /// 新标签。 + pub label: Option, +} + +/// 更新账号状态请求。 +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateAccountStatusRequest { + /// 新状态。 + pub status: String, +} + +/// 更新账号状态响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct UpdateAccountStatusData { + /// 账号 ID。 + pub id: String, + /// 新状态。 + pub status: String, +} + +/// 批量删除账号请求。 +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchDeleteAccountsRequest { + /// 账号 ID 列表。 + pub ids: Vec, +} + +/// 批量删除账号响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchDeleteAccountsData { + /// 成功删除数量。 + pub deleted: u32, + /// 未找到的账号 ID。 + pub not_found: Vec, +} + +/// 批量更新账号状态请求。 +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchUpdateAccountStatusRequest { + /// 账号 ID 列表。 + pub ids: Vec, + /// 新状态。 + pub status: String, +} + +/// 批量更新账号状态响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct BatchUpdateAccountStatusData { + /// 成功更新数量。 + pub updated: u32, + /// 未找到的账号 ID。 + pub not_found: Vec, +} + +/// 删除账号响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DeleteAccountData { + /// 是否已删除。 + pub deleted: bool, +} + +/// 刷新账号响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct RefreshAccountData { + /// 账号 ID。 + pub id: String, + /// 刷新结果。 + pub result: String, + /// 刷新前状态。 + pub previous_status: String, + /// 刷新后状态。 + pub status: Option, + /// 错误信息。 + pub error: Option, +} + +/// 重置账号用量响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ResetAccountUsageData { + /// 账号 ID。 + pub id: String, + /// 是否已处理。 + pub reset: bool, +} + +/// `PATCH /api/admin/accounts/{account_id}/label` +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; + match state + .services + .admin_accounts + .update_label(&account_id, label.clone()) + .await + { + Ok(true) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + UpdateAccountLabelData { + id: account_id, + label, + }, + request_id, + ), + )), + Ok(false) => Err(account_not_found(request_id)), + Err(error) => Err(account_error(error, request_id)), + } +} + +/// `PATCH /api/admin/accounts/{account_id}/status` +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?; + match state + .services + .admin_accounts + .update_status(&account_id, &payload.status) + .await + { + Ok(Some(updated)) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + UpdateAccountStatusData { + id: updated.id, + status: account_status_value(updated.status).to_string(), + }, + request_id, + ), + )), + Ok(None) => Err(account_not_found(request_id)), + Err(error) => Err(account_error(error, request_id)), + } +} + +/// `DELETE /api/admin/accounts/{account_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?; + match state.services.admin_accounts.delete(&account_id).await { + Ok(true) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(DeleteAccountData { deleted: true }, request_id), + )), + Ok(false) => Err(account_not_found(request_id)), + Err(error) => Err(account_error(error, request_id)), + } +} + +/// `POST /api/admin/accounts/batch-delete` +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(); + require_admin_session(&state, &headers, &request_id).await?; + match state + .services + .admin_accounts + .batch_delete(payload.ids) + .await + { + Ok(deleted) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + BatchDeleteAccountsData { + deleted: deleted.deleted, + not_found: deleted.not_found, + }, + request_id, + ), + )), + Err(error) => Err(account_error(error, request_id)), + } +} + +/// `POST /api/admin/accounts/batch-status` +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(); + require_admin_session(&state, &headers, &request_id).await?; + match state + .services + .admin_accounts + .batch_update_status(payload.ids, &payload.status) + .await + { + Ok(updated) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + BatchUpdateAccountStatusData { + updated: updated.updated, + not_found: updated.not_found, + }, + request_id, + ), + )), + Err(error) => Err(account_error(error, request_id)), + } +} + +/// `POST /api/admin/accounts/{account_id}/refresh` +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?; + match state + .services + .admin_accounts + .refresh_account(&account_id) + .await + { + Ok(refreshed) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(RefreshAccountData::from(refreshed), request_id), + )), + Err(error) => Err(account_error(error, request_id)), + } +} + +/// `POST /api/admin/accounts/{account_id}/reset-usage` +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?; + match state.services.admin_accounts.reset_usage(&account_id).await { + Ok(reset) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + ResetAccountUsageData { + id: reset.id, + reset: reset.reset, + }, + request_id, + ), + )), + Err(error) => Err(account_error(error, request_id)), + } +} + +impl From for RefreshAccountData { + fn from(refreshed: AdminAccountRefresh) -> Self { + Self { + id: refreshed.id, + result: refresh_outcome_value(refreshed.outcome).to_string(), + previous_status: account_status_value(refreshed.previous_status).to_string(), + status: refreshed + .status + .map(account_status_value) + .map(ToString::to_string), + error: refreshed.error, + } + } +} + +fn refresh_outcome_value(outcome: AdminAccountProbeOutcome) -> &'static str { + outcome.as_str() +} diff --git a/crates/server/src/admin_api/accounts/list.rs b/crates/server/src/admin_api/accounts/list.rs new file mode 100644 index 000000000..7ab005a34 --- /dev/null +++ b/crates/server/src/admin_api/accounts/list.rs @@ -0,0 +1,60 @@ +//! 账号列表处理器。 + +use axum::{ + extract::{Query, State}, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, +}; +use codex_proxy_platform::json::{clamp_limit, Page}; +use codex_proxy_runtime::state::AppState; +use serde::Deserialize; + +use crate::{ + admin_api::{ + accounts::{account_error, AdminAccountData}, + require_admin_session, AdminError, AdminPageEnvelope, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +/// 账号列表查询。 +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountsQuery { + /// 分页游标。 + pub cursor: Option, + /// 分页大小。 + pub limit: Option, +} + +/// 分页列出账号元数据。 +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)); + match state + .services + .admin_accounts + .list(query.cursor, limit) + .await + { + Ok(page) => { + let page = Page { + items: page.items.into_iter().map(AdminAccountData::from).collect(), + next_cursor: page.next_cursor, + }; + Ok(AdminResponse::new( + StatusCode::OK, + AdminPageEnvelope::ok(page, limit, request_id), + )) + } + Err(error) => Err(account_error(error, request_id)), + } +} diff --git a/crates/server/src/admin_api/accounts/mod.rs b/crates/server/src/admin_api/accounts/mod.rs new file mode 100644 index 000000000..1bc08bacc --- /dev/null +++ b/crates/server/src/admin_api/accounts/mod.rs @@ -0,0 +1,152 @@ +//! 管理端账号处理器。 + +use axum::http::StatusCode; +use codex_proxy_core::accounts::model::AccountStatus; +use codex_proxy_runtime::services::{AdminAccountError, AdminAccountMetadata}; +use serde::Serialize; + +use crate::admin_api::AdminError; + +pub mod cookies; +pub mod create; +pub mod export; +pub mod health; +pub mod import; +pub mod import_cli; +pub mod lifecycle; +pub mod list; +pub mod oauth; +pub mod quota; + +pub use cookies::{delete_account_cookies, get_account_cookies, set_account_cookies}; +pub use create::create_account; +pub use export::{export_accounts, AccountExportQuery}; +pub use health::health_check_accounts; +pub use import::import_accounts; +pub use import_cli::import_cli_account; +pub use lifecycle::{ + batch_delete_accounts, batch_update_account_status, delete_account, refresh_account, + reset_account_usage, update_account_label, update_account_status, +}; +pub use list::{accounts, AccountsQuery}; +pub use oauth::{ + auth_callback, auth_code_relay, auth_device_login, auth_device_poll, auth_login_start, + auth_logout, auth_status, +}; +pub use quota::{account_quota, quota_warnings}; + +/// 管理端账号元数据响应。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAccountData { + /// 账号 ID。 + pub id: String, + /// 邮箱。 + pub email: Option, + /// ChatGPT 账号 ID。 + pub account_id: Option, + /// ChatGPT 用户 ID。 + pub user_id: Option, + /// 管理标签。 + pub label: Option, + /// 计划类型。 + pub plan_type: Option, + /// 账号状态。 + pub status: String, + /// access token 过期时间。 + pub access_token_expires_at: Option, + /// 创建时间。 + pub added_at: String, + /// 更新时间。 + pub updated_at: String, +} + +impl From for AdminAccountData { + fn from(account: AdminAccountMetadata) -> 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(), + } + } +} + +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", + } +} + +pub(super) fn account_error(error: AdminAccountError, request_id: String) -> AdminError { + match error { + AdminAccountError::InvalidStatus(_) + | AdminAccountError::LabelTooLong + | AdminAccountError::EmptyIds + | AdminAccountError::NoImportableAccounts + | AdminAccountError::InvalidAccessTokenExpiresAt + | AdminAccountError::TokenRequired + | AdminAccountError::InvalidToken(_) + | AdminAccountError::RefreshTokenExchange(_) + | AdminAccountError::NoValidCookies => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + error.to_string(), + request_id, + ), + AdminAccountError::NotFound => account_not_found(request_id), + AdminAccountError::List + | AdminAccountError::Export + | AdminAccountError::Import + | AdminAccountError::Inspect + | AdminAccountError::UpdateLabel + | AdminAccountError::UpdateStatus + | AdminAccountError::Delete + | AdminAccountError::LoadCookies + | AdminAccountError::StoreCookies + | AdminAccountError::DeleteCookies + | AdminAccountError::UpdateClaims + | AdminAccountError::ResetUsage + | AdminAccountError::SyncAccountPool + | AdminAccountError::QuotaWarnings + | AdminAccountError::StoreQuota + | AdminAccountError::FetchQuota(_) + | AdminAccountError::HealthCheck => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + error.to_string(), + request_id, + ), + AdminAccountError::Inactive(status) => AdminError::new( + StatusCode::CONFLICT, + 40901, + format!( + "Account is {}, cannot query quota", + account_status_value(status) + ), + request_id, + ), + } +} + +pub(super) fn account_not_found(request_id: String) -> AdminError { + AdminError::new( + StatusCode::NOT_FOUND, + 40401, + "Account not found", + request_id, + ) +} diff --git a/crates/server/src/admin_api/accounts/oauth.rs b/crates/server/src/admin_api/accounts/oauth.rs new file mode 100644 index 000000000..e19392ebb --- /dev/null +++ b/crates/server/src/admin_api/accounts/oauth.rs @@ -0,0 +1,378 @@ +//! 账号 OAuth 处理器。 + +use axum::{ + extract::{Query, State}, + http::{header::HOST, HeaderMap, StatusCode}, + response::{IntoResponse, Redirect}, + Extension, Json, +}; +use codex_proxy_runtime::{ + services::{AdminDevicePoll, AdminOAuthError}, + state::AppState, +}; +use serde::{Deserialize, Serialize}; + +use crate::{ + admin_api::{ + accounts::{account_error, AdminAccountData}, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +/// OAuth PKCE 登录开始响应。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthLoginStartData { + /// 授权 URL。 + pub auth_url: String, + /// OAuth state。 + pub state: String, +} + +/// 管理端认证状态响应。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthStatusData { + authenticated: bool, + user: Option, + pool: AdminAuthPoolData, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthPoolData { + total: u32, + active: u32, + expired: u32, + quota_exhausted: u32, + refreshing: u32, + disabled: u32, + banned: u32, +} + +/// 管理端登出响应。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthLogoutData { + success: bool, + deleted: u64, +} + +/// 设备码登录响应。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthDeviceLoginData { + device_code: String, + user_code: String, + verification_uri: String, + verification_uri_complete: String, + expires_in: u64, + interval: u64, +} + +/// 设备码轮询响应。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthDevicePollData { + success: bool, + pending: bool, + code: Option<&'static str>, +} + +/// OAuth code relay 请求。 +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthCodeRelayRequest { + callback_url: String, +} + +/// OAuth 登录完成响应。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminAuthSuccessData { + success: bool, +} + +/// OAuth callback query。 +#[derive(Debug, Deserialize)] +pub struct AdminAuthCallbackQuery { + code: String, + state: String, +} + +/// `GET /api/admin/auth/status` +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_accounts + .auth_status() + .await + .map_err(|error| account_error(error, request_id.clone()))?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminAuthStatusData::from(status), request_id), + )) +} + +/// `POST /api/admin/auth/logout` +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 result = state + .services + .admin_accounts + .logout() + .await + .map_err(|error| account_error(error, request_id.clone()))?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AdminAuthLogoutData { + success: result.success, + deleted: result.deleted, + }, + request_id, + ), + )) +} + +/// `POST /api/admin/auth/login-start` +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 login = state + .services + .admin_oauth + .start_pkce_login(&request_host(&headers)) + .await; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AdminAuthLoginStartData { + auth_url: login.auth_url, + state: login.state, + }, + request_id, + ), + )) +} + +/// `POST /api/admin/auth/device-login` +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_oauth + .request_device_code() + .await + .map_err(|error| oauth_error(error, request_id.clone()))?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AdminAuthDeviceLoginData { + device_code: device.device_code, + user_code: device.user_code, + verification_uri: device.verification_uri, + verification_uri_complete: device.verification_uri_complete, + expires_in: device.expires_in, + interval: device.interval, + }, + request_id, + ), + )) +} + +/// `GET /api/admin/auth/device-poll/{device_code}` +pub async fn auth_device_poll( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + axum::extract::Path(device_code): axum::extract::Path, +) -> Result { + let request_id = request_id.as_str().to_string(); + require_admin_session(&state, &headers, &request_id).await?; + + match state + .services + .admin_oauth + .poll_device_token(&device_code) + .await + .map_err(|error| oauth_error(error, request_id.clone()))? + { + AdminDevicePoll::Pending { code } => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AdminAuthDevicePollData { + success: false, + pending: true, + code: Some(code), + }, + request_id, + ), + )), + AdminDevicePoll::Authorized(tokens) => { + state + .services + .admin_accounts + .create(Some(tokens.access_token), tokens.refresh_token) + .await + .map_err(|error| account_error(error, request_id.clone()))?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AdminAuthDevicePollData { + success: true, + pending: false, + code: None, + }, + request_id, + ), + )) + } + } +} + +/// `POST /api/admin/auth/code-relay` +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 (code, oauth_state) = callback_params(&payload.callback_url) + .ok_or_else(|| oauth_error(AdminOAuthError::InvalidCallback, request_id.clone()))?; + let callback = state + .services + .admin_oauth + .exchange_callback(&code, &oauth_state) + .await + .map_err(|error| oauth_error(error, request_id.clone()))?; + state + .services + .admin_accounts + .create( + Some(callback.tokens.access_token), + callback.tokens.refresh_token, + ) + .await + .map_err(|error| account_error(error, request_id.clone()))?; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminAuthSuccessData { success: true }, request_id), + )) +} + +/// `GET /auth/callback` +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?; + let callback = state + .services + .admin_oauth + .exchange_callback(&query.code, &query.state) + .await + .map_err(|error| oauth_error(error, request_id.clone()))?; + state + .services + .admin_accounts + .create( + Some(callback.tokens.access_token), + callback.tokens.refresh_token, + ) + .await + .map_err(|error| account_error(error, request_id))?; + + Ok(Redirect::to(&format!("http://{}/", callback.return_host))) +} + +fn request_host(headers: &HeaderMap) -> String { + headers + .get(HOST) + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) + .unwrap_or_else(|| "localhost".to_string()) +} + +fn callback_params(callback_url: &str) -> Option<(String, String)> { + let (_base, query) = callback_url.split_once('?')?; + let mut code = None; + let mut state = None; + for pair in query.split('&') { + let (key, value) = pair.split_once('=').unwrap_or((pair, "")); + match key { + "code" if !value.is_empty() => code = Some(value.to_string()), + "state" if !value.is_empty() => state = Some(value.to_string()), + _ => {} + } + } + Some((code?, state?)) +} + +fn oauth_error(error: AdminOAuthError, request_id: String) -> AdminError { + match error { + AdminOAuthError::InvalidCallback | AdminOAuthError::InvalidState => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + error.to_string(), + request_id, + ), + AdminOAuthError::OAuth(_) => AdminError::new( + StatusCode::BAD_GATEWAY, + 50201, + error.to_string(), + request_id, + ), + } +} + +impl From for AdminAuthStatusData { + fn from(status: codex_proxy_runtime::services::AdminAuthStatus) -> Self { + Self { + authenticated: status.authenticated, + user: status.user.map(AdminAccountData::from), + pool: AdminAuthPoolData { + total: status.pool.total, + active: status.pool.active, + expired: status.pool.expired, + quota_exhausted: status.pool.quota_exhausted, + refreshing: status.pool.refreshing, + disabled: status.pool.disabled, + banned: status.pool.banned, + }, + } + } +} diff --git a/crates/server/src/admin_api/accounts/quota.rs b/crates/server/src/admin_api/accounts/quota.rs new file mode 100644 index 000000000..14bdbf920 --- /dev/null +++ b/crates/server/src/admin_api/accounts/quota.rs @@ -0,0 +1,201 @@ +//! 账号配额处理器。 + +use axum::{ + extract::{Path, State}, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + Extension, +}; +use codex_proxy_runtime::{ + services::{ + AdminAccountError, AdminAccountQuota, AdminAccountQuotaWarning, AdminAccountQuotaWarnings, + AdminQuotaWarningLevel, AdminQuotaWarningWindow, + }, + state::AppState, +}; +use serde::Serialize; +use serde_json::{json, Value}; + +use crate::{ + admin_api::{ + accounts::{account_error, account_not_found, account_status_value}, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +/// 账号配额响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountQuotaData { + /// 归一化后的配额快照。 + pub quota: Value, + /// Codex usage 原始响应。 + pub raw: Value, +} + +/// 账号配额 HTTP 错误。 +pub 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) -> 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(), + } + } +} + +/// 账号配额预警响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountQuotaWarningsData { + /// 预警列表。 + pub warnings: Vec, + /// 产生预警的快照中最新的拉取时间。 + pub updated_at: Option, +} + +/// 单条账号配额预警响应。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountQuotaWarningData { + /// 账号 ID。 + pub account_id: String, + /// 邮箱。 + pub email: Option, + /// 配额窗口。 + pub window: String, + /// 预警级别。 + pub level: String, + /// 已使用百分比。 + pub used_percent: f64, + /// 重置时间戳。 + pub reset_at: Option, +} + +/// `GET /api/admin/accounts/{account_id}/quota` +pub 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?; + + match state + .services + .admin_accounts + .account_quota(&account_id, &request_id) + .await + { + Ok(quota) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountQuotaData::from(quota), request_id), + )), + Err(error) => Err(account_quota_error(error, request_id)), + } +} + +/// `GET /api/admin/accounts/quota-warnings` +pub async fn quota_warnings( + 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.admin_accounts.quota_warnings().await { + Ok(warnings) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountQuotaWarningsData::from(warnings), request_id), + )), + Err(error) => Err(account_error(error, request_id)), + } +} + +impl From for AccountQuotaData { + fn from(quota: AdminAccountQuota) -> Self { + Self { + quota: quota.quota, + raw: quota.raw, + } + } +} + +impl From for AccountQuotaWarningsData { + fn from(warnings: AdminAccountQuotaWarnings) -> Self { + Self { + warnings: warnings.warnings.into_iter().map(Into::into).collect(), + updated_at: warnings.updated_at.map(|value| value.to_rfc3339()), + } + } +} + +impl From for AccountQuotaWarningData { + fn from(warning: AdminAccountQuotaWarning) -> Self { + Self { + account_id: warning.account_id, + email: warning.email, + window: quota_window_value(warning.window).to_string(), + level: quota_level_value(warning.level).to_string(), + used_percent: warning.used_percent, + reset_at: warning.reset_at, + } + } +} + +fn quota_window_value(window: AdminQuotaWarningWindow) -> &'static str { + window.as_str() +} + +fn quota_level_value(level: AdminQuotaWarningLevel) -> &'static str { + level.as_str() +} + +fn account_quota_error(error: AdminAccountError, request_id: String) -> AccountQuotaHttpError { + match error { + AdminAccountError::NotFound => account_not_found(request_id).into(), + AdminAccountError::Inactive(status) => AdminError::new( + StatusCode::CONFLICT, + 40901, + format!( + "Account is {}, cannot query quota", + account_status_value(status) + ), + request_id, + ) + .into(), + AdminAccountError::FetchQuota(error) => { + AccountQuotaHttpError::UpstreamFetch { error, request_id } + } + AdminAccountError::StoreQuota => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to store account quota", + request_id, + ) + .into(), + error => account_error(error, request_id).into(), + } +} diff --git a/crates/server/src/admin_api/client_keys/create.rs b/crates/server/src/admin_api/client_keys/create.rs new file mode 100644 index 000000000..6c4ae1eda --- /dev/null +++ b/crates/server/src/admin_api/client_keys/create.rs @@ -0,0 +1,32 @@ +//! 客户端 key 创建处理器。 + +use axum::{ + extract::State, http::HeaderMap, http::StatusCode, response::IntoResponse, Extension, Json, +}; +use codex_proxy_runtime::state::AppState; + +use crate::{ + admin_api::{ + client_keys::{client_key_error, CreateApiKeyRequest, CreatedClientApiKeyData}, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +/// 创建客户端 API Key。 +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.admin_client_keys.create(&payload.name).await { + Ok(created) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(CreatedClientApiKeyData::from(created), request_id), + )), + Err(error) => Err(client_key_error(error, request_id)), + } +} diff --git a/crates/server/src/admin_api/client_keys/export.rs b/crates/server/src/admin_api/client_keys/export.rs new file mode 100644 index 000000000..33bb8e74e --- /dev/null +++ b/crates/server/src/admin_api/client_keys/export.rs @@ -0,0 +1,56 @@ +//! 客户端 key 导出处理器。 + +use axum::{ + extract::{Query, State}, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, +}; +use codex_proxy_runtime::state::AppState; + +use crate::{ + admin_api::{ + client_keys::{ + client_key_error, ApiKeyExportQuery, ClientApiKeyExportData, ClientApiKeyExportEntry, + }, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +/// 导出客户端 API Key 元数据。 +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 = query + .ids + .unwrap_or_default() + .split(',') + .map(str::trim) + .filter(|id| !id.is_empty()) + .map(ToString::to_string) + .collect::>(); + + match state.services.admin_client_keys.export(ids).await { + Ok(keys) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + ClientApiKeyExportData { + source_format: "rustLocalClientApiKeys", + rotation_required: true, + api_keys: keys + .into_iter() + .map(ClientApiKeyExportEntry::from) + .collect(), + }, + request_id, + ), + )), + Err(error) => Err(client_key_error(error, request_id)), + } +} diff --git a/crates/server/src/admin_api/client_keys/import.rs b/crates/server/src/admin_api/client_keys/import.rs new file mode 100644 index 000000000..d335b5fe3 --- /dev/null +++ b/crates/server/src/admin_api/client_keys/import.rs @@ -0,0 +1,45 @@ +//! 客户端 key 导入处理器。 + +use axum::{ + extract::State, http::HeaderMap, http::StatusCode, response::IntoResponse, Extension, Json, +}; +use codex_proxy_runtime::state::AppState; +use serde_json::Value; + +use crate::{ + admin_api::{ + client_keys::{client_key_error, imported_client_key_data, ClientApiKeyImportData}, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +/// 导入客户端 API Key 元数据并轮换新密钥。 +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?; + match state.services.admin_client_keys.import(&payload).await { + Ok(imported) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + ClientApiKeyImportData { + imported: imported.imported, + skipped: imported.skipped, + rotated: true, + keys: imported + .keys + .into_iter() + .map(imported_client_key_data) + .collect(), + }, + request_id, + ), + )), + Err(error) => Err(client_key_error(error, request_id)), + } +} diff --git a/src/admin/api/client_keys/lifecycle.rs b/crates/server/src/admin_api/client_keys/lifecycle.rs similarity index 58% rename from src/admin/api/client_keys/lifecycle.rs rename to crates/server/src/admin_api/client_keys/lifecycle.rs index d4f8cfa09..de4ea7caf 100644 --- a/src/admin/api/client_keys/lifecycle.rs +++ b/crates/server/src/admin_api/client_keys/lifecycle.rs @@ -1,19 +1,27 @@ +//! 客户端 key 生命周期处理器。 + use axum::{ extract::{Path, State}, http::{HeaderMap, StatusCode}, response::IntoResponse, Extension, Json, }; +use codex_proxy_runtime::state::AppState; -use crate::{platform::http::request_id::RequestId, runtime::state::AppState}; - -use super::{ - api_key_not_found, api_key_service_error, BatchDeleteClientApiKeysData, - BatchDeleteClientApiKeysRequest, ClientApiKeyData, DeleteClientApiKeyData, - UpdateClientApiKeyLabelRequest, UpdateClientApiKeyStatusData, UpdateClientApiKeyStatusRequest, +use crate::{ + admin_api::{ + client_keys::{ + client_key_error, client_key_not_found, BatchDeleteClientApiKeysData, + BatchDeleteClientApiKeysRequest, ClientApiKeyData, DeleteClientApiKeyData, + UpdateClientApiKeyLabelRequest, UpdateClientApiKeyStatusData, + UpdateClientApiKeyStatusRequest, + }, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, }; -use crate::admin::api::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; +/// 批量删除客户端 API Key。 pub async fn batch_delete_api_keys( State(state): State, Extension(request_id): Extension, @@ -22,23 +30,27 @@ pub async fn batch_delete_api_keys( ) -> 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, - ), - )) + match state + .services + .admin_client_keys + .batch_delete(payload.ids) + .await + { + Ok(deleted) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + BatchDeleteClientApiKeysData { + deleted: deleted.deleted, + not_found: deleted.not_found, + }, + request_id, + ), + )), + Err(error) => Err(client_key_error(error, request_id)), + } } +/// 更新客户端 API Key 标签。 pub async fn update_api_key_label( State(state): State, Extension(request_id): Extension, @@ -48,11 +60,9 @@ pub async fn update_api_key_label( ) -> 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 + .admin_client_keys .update_label(&key_id, payload.label) .await { @@ -60,11 +70,12 @@ pub async fn update_api_key_label( 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)), + Ok(None) => Err(client_key_not_found(request_id)), + Err(error) => Err(client_key_error(error, request_id)), } } +/// 更新客户端 API Key 启用状态。 pub async fn update_api_key_status( State(state): State, Extension(request_id): Extension, @@ -74,11 +85,10 @@ pub async fn update_api_key_status( ) -> 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) + .admin_client_keys + .update_status(&key_id, &payload.status) .await { Ok(Some(updated)) => Ok(AdminResponse::new( @@ -91,11 +101,12 @@ pub async fn update_api_key_status( request_id, ), )), - Ok(None) => Err(api_key_not_found(request_id)), - Err(error) => Err(api_key_service_error(error, request_id)), + Ok(None) => Err(client_key_not_found(request_id)), + Err(error) => Err(client_key_error(error, request_id)), } } +/// 删除客户端 API Key。 pub async fn delete_api_key( State(state): State, Extension(request_id): Extension, @@ -104,13 +115,12 @@ pub async fn delete_api_key( ) -> 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 { + match state.services.admin_client_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)), + Ok(false) => Err(client_key_not_found(request_id)), + Err(error) => Err(client_key_error(error, request_id)), } } diff --git a/src/admin/api/client_keys/list.rs b/crates/server/src/admin_api/client_keys/list.rs similarity index 52% rename from src/admin/api/client_keys/list.rs rename to crates/server/src/admin_api/client_keys/list.rs index de34d86b6..e557230b5 100644 --- a/src/admin/api/client_keys/list.rs +++ b/crates/server/src/admin_api/client_keys/list.rs @@ -1,19 +1,23 @@ +//! 客户端 key 列表处理器。 + use axum::{ extract::{Query, State}, http::{HeaderMap, StatusCode}, response::IntoResponse, Extension, }; +use codex_proxy_platform::json::{clamp_limit, Page}; +use codex_proxy_runtime::state::AppState; use crate::{ - platform::http::request_id::RequestId, - runtime::state::AppState, - utils::pagination::{clamp_limit, Page}, + admin_api::{ + client_keys::{client_key_error, ApiKeysQuery, ClientApiKeyData}, + require_admin_session, AdminError, AdminPageEnvelope, AdminResponse, + }, + middleware::request_id::RequestId, }; -use super::{api_key_service_error, ApiKeysQuery, ClientApiKeyData}; -use crate::admin::api::{require_admin_session, AdminError, AdminPageEnvelope, AdminResponse}; - +/// 分页列出客户端 API Key。 pub async fn api_keys( State(state): State, Extension(request_id): Extension, @@ -23,18 +27,22 @@ pub async fn api_keys( 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 { + match state + .services + .admin_client_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, + items: page.items.into_iter().map(ClientApiKeyData::from).collect(), + next_cursor: page.next_cursor, }; Ok(AdminResponse::new( StatusCode::OK, AdminPageEnvelope::ok(page, limit, request_id), )) } - Err(error) => Err(api_key_service_error(error, request_id)), + Err(error) => Err(client_key_error(error, request_id)), } } diff --git a/src/admin/api/client_keys/mod.rs b/crates/server/src/admin_api/client_keys/mod.rs similarity index 53% rename from src/admin/api/client_keys/mod.rs rename to crates/server/src/admin_api/client_keys/mod.rs index 012424ecc..36020f724 100644 --- a/src/admin/api/client_keys/mod.rs +++ b/crates/server/src/admin_api/client_keys/mod.rs @@ -1,12 +1,12 @@ +//! 客户端 key 处理器。 + use axum::http::StatusCode; +use codex_proxy_runtime::services::{ + AdminClientKeyError, AdminCreatedClientApiKey, AdminStoredClientApiKey, ImportedClientApiKey, +}; use serde::{Deserialize, Serialize}; -use crate::{ - admin::client_keys::service::{ApiKeyServiceError, ImportedClientApiKey}, - platform::identity::client_key_repository::StoredClientApiKey, -}; - -use super::AdminError; +use crate::admin_api::AdminError; pub mod create; pub mod export; @@ -22,134 +22,200 @@ pub use lifecycle::{ }; pub use list::api_keys; +/// API Key 列表查询。 #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] pub struct ApiKeysQuery { + /// 分页游标。 pub cursor: Option, + /// 分页大小。 pub limit: Option, } +/// API Key 导出查询。 #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] pub struct ApiKeyExportQuery { + /// 逗号分隔的 ID 列表。 pub ids: Option, } +/// 创建 API Key 请求。 #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] pub struct CreateApiKeyRequest { + /// API Key 名称。 pub name: String, } +/// 客户端 API Key 元数据。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] pub struct ClientApiKeyData { + /// API Key 记录 ID。 pub id: String, + /// 名称。 pub name: String, + /// 标签。 pub label: Option, + /// 前缀。 pub prefix: String, + /// 是否启用。 pub enabled: bool, + /// 创建时间。 pub created_at: String, + /// 最近使用时间。 pub last_used_at: Option, } +/// 创建 API Key 响应数据。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] pub struct CreatedClientApiKeyData { + /// API Key 记录 ID。 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 { + /// `active` 或 `disabled`。 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 { + /// API Key 记录 ID。 pub id: String, + /// 是否启用。 pub enabled: bool, } +/// 批量删除请求。 #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] pub struct BatchDeleteClientApiKeysRequest { + /// 待删除 ID。 pub ids: Vec, } +/// 批量删除响应。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] pub struct BatchDeleteClientApiKeysData { + /// 成功删除数量。 pub deleted: u32, + /// 未找到 ID。 pub not_found: Vec, } +/// 删除响应。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] pub struct DeleteClientApiKeyData { + /// 是否删除。 pub deleted: bool, } +/// API Key 导出响应。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] pub struct ClientApiKeyExportData { + /// 来源格式。 pub source_format: &'static str, + /// 是否需要轮换。 pub rotation_required: bool, + /// 导出的 API Key 元数据。 pub api_keys: Vec, } +/// API Key 导出条目。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] pub struct ClientApiKeyExportEntry { + /// API Key 记录 ID。 pub id: String, + /// 名称。 pub name: String, + /// 标签。 pub label: Option, + /// 前缀。 pub prefix: String, + /// 是否启用。 pub enabled: bool, + /// 创建时间。 pub created_at: String, + /// 最近使用时间。 pub last_used_at: Option, } +/// API Key 导入响应。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] pub struct ClientApiKeyImportData { + /// 导入数量。 pub imported: u32, + /// 跳过数量。 pub skipped: u32, + /// 是否轮换。 pub rotated: bool, + /// 新建 key 列表。 pub keys: Vec, } +/// API Key 导入条目。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] pub struct ImportedClientApiKeyData { + /// 来源 ID。 pub source_id: Option, + /// 来源前缀。 pub source_prefix: Option, + /// API Key 记录 ID。 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 { +impl From for ClientApiKeyData { + fn from(key: AdminStoredClientApiKey) -> Self { Self { id: key.id, name: key.name, @@ -162,8 +228,8 @@ impl From for ClientApiKeyData { } } -impl From for ClientApiKeyExportEntry { - fn from(key: StoredClientApiKey) -> Self { +impl From for ClientApiKeyExportEntry { + fn from(key: AdminStoredClientApiKey) -> Self { Self { id: key.id, name: key.name, @@ -176,8 +242,8 @@ impl From for ClientApiKeyExportEntry { } } -impl CreatedClientApiKeyData { - pub(super) fn new(key: StoredClientApiKey, plaintext: String) -> Self { +impl From for CreatedClientApiKeyData { + fn from(key: AdminCreatedClientApiKey) -> Self { Self { id: key.id, name: key.name, @@ -186,131 +252,64 @@ impl CreatedClientApiKeyData { enabled: key.enabled, created_at: key.created_at, last_used_at: key.last_used_at, - plaintext, + plaintext: key.plaintext, } } } impl ImportedClientApiKeyData { - pub(super) fn new(imported: ImportedClientApiKey) -> Self { - let ImportedClientApiKey { - key, - plaintext, - source_id, - source_prefix, - } = imported; + fn new(imported: ImportedClientApiKey) -> 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, + source_id: imported.source_id, + source_prefix: imported.source_prefix, + id: imported.key.id, + name: imported.key.name, + label: imported.key.label, + prefix: imported.key.prefix, + enabled: imported.key.enabled, + created_at: imported.key.created_at, + last_used_at: imported.key.last_used_at, + plaintext: imported.key.plaintext, } } } -pub(super) fn api_key_service_error(error: ApiKeyServiceError, request_id: String) -> AdminError { +fn client_key_error(error: AdminClientKeyError, 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( + AdminClientKeyError::InvalidStatus(_) + | AdminClientKeyError::EmptyName + | AdminClientKeyError::EmptyIds + | AdminClientKeyError::LabelTooLong + | AdminClientKeyError::NoImportableKeys => AdminError::new( StatusCode::BAD_REQUEST, 40001, - "API key name is required", + error.to_string(), 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", + AdminClientKeyError::List + | AdminClientKeyError::Export + | AdminClientKeyError::Import + | AdminClientKeyError::Create + | AdminClientKeyError::Delete + | AdminClientKeyError::UpdateLabel + | AdminClientKeyError::UpdateStatus => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + error.to_string(), request_id, ), } } -pub(super) fn api_key_not_found(request_id: String) -> AdminError { +fn client_key_not_found(request_id: String) -> AdminError { AdminError::new( StatusCode::NOT_FOUND, 40401, - "API key not found", + "Client API key not found", request_id, ) } + +fn imported_client_key_data(imported: ImportedClientApiKey) -> ImportedClientApiKeyData { + ImportedClientApiKeyData::new(imported) +} diff --git a/src/admin/api/diagnostics.rs b/crates/server/src/admin_api/diagnostics.rs similarity index 67% rename from src/admin/api/diagnostics.rs rename to crates/server/src/admin_api/diagnostics.rs index 89343bd63..f3ab76ccb 100644 --- a/src/admin/api/diagnostics.rs +++ b/crates/server/src/admin_api/diagnostics.rs @@ -1,17 +1,20 @@ +//! 管理端诊断处理器。 + use axum::{ extract::State, http::{HeaderMap, StatusCode}, response::IntoResponse, Extension, }; +use codex_proxy_runtime::state::AppState; use crate::{ - codex::serving::http::diagnostics::diagnostics_data, platform::http::request_id::RequestId, - runtime::state::AppState, + admin_api::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}, + middleware::request_id::RequestId, + openai_api::diagnostics::diagnostics_data, }; -use super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; - +/// `GET /api/admin/diagnostics` pub async fn diagnostics( State(state): State, Extension(request_id): Extension, diff --git a/src/admin/api/logs/detail.rs b/crates/server/src/admin_api/logs/detail.rs similarity index 72% rename from src/admin/api/logs/detail.rs rename to crates/server/src/admin_api/logs/detail.rs index b41ba4a2c..cd9b9bc83 100644 --- a/src/admin/api/logs/detail.rs +++ b/crates/server/src/admin_api/logs/detail.rs @@ -1,15 +1,19 @@ +//! 日志详情处理器。 + use axum::{ extract::{Path, State}, http::{HeaderMap, StatusCode}, response::IntoResponse, Extension, }; +use codex_proxy_runtime::state::AppState; -use crate::{platform::http::request_id::RequestId, runtime::state::AppState}; - -use super::log_service_error; -use crate::admin::api::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; +use crate::{ + admin_api::{logs::log_error, require_admin_session, AdminEnvelope, AdminError, AdminResponse}, + middleware::request_id::RequestId, +}; +/// 查询单条事件日志。 pub async fn log_detail( State(state): State, Extension(request_id): Extension, @@ -18,7 +22,6 @@ pub async fn log_detail( ) -> Result { let request_id = request_id.as_str().to_string(); require_admin_session(&state, &headers, &request_id).await?; - match state.services.logs.get(&log_id).await { Ok(Some(log)) => Ok(AdminResponse::new( StatusCode::OK, @@ -30,6 +33,6 @@ pub async fn log_detail( "Log event not found", request_id, )), - Err(error) => Err(log_service_error(error, request_id)), + Err(error) => Err(log_error(error, request_id)), } } diff --git a/crates/server/src/admin_api/logs/mod.rs b/crates/server/src/admin_api/logs/mod.rs new file mode 100644 index 000000000..779c8b96b --- /dev/null +++ b/crates/server/src/admin_api/logs/mod.rs @@ -0,0 +1,90 @@ +//! 管理端日志处理器。 + +use axum::http::StatusCode; +use codex_proxy_core::events::model::EventLevel; +use codex_proxy_runtime::services::{AdminLogError, AdminLogState}; +use serde::Serialize; + +use crate::admin_api::AdminError; + +pub mod detail; +pub mod query; +pub mod state; + +pub use detail::log_detail; +pub use query::logs; +pub use state::{clear_logs, logs_state, update_logs_state}; + +/// 管理端日志状态响应。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct LogStateData { + /// 是否启用。 + pub enabled: bool, + /// 内存容量。 + pub capacity: u32, + /// 是否捕获请求体。 + pub capture_body: bool, + /// 已存储数量。 + pub stored_count: u64, +} + +/// 清空日志响应。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ClearLogsData { + /// 清理数量。 + pub cleared: u64, +} + +impl From for LogStateData { + fn from(state: AdminLogState) -> Self { + Self { + enabled: state.enabled, + capacity: state.capacity, + capture_body: state.capture_body, + stored_count: state.stored_count, + } + } +} + +fn log_error(error: AdminLogError, request_id: String) -> AdminError { + match error { + AdminLogError::List + | AdminLogError::Get + | AdminLogError::Count + | AdminLogError::Clear + | AdminLogError::Append + | AdminLogError::Trim => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + error.to_string(), + request_id, + ), + AdminLogError::InvalidCapacity => AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + error.to_string(), + request_id, + ), + } +} + +fn level_from_query(value: Option) -> Result, String> { + let Some(value) = non_empty(value) else { + return Ok(None); + }; + match value.as_str() { + "debug" => Ok(Some(EventLevel::Debug)), + "info" => Ok(Some(EventLevel::Info)), + "warn" => Ok(Some(EventLevel::Warn)), + "error" => Ok(Some(EventLevel::Error)), + other => Err(format!("Unsupported log level: {other}")), + } +} + +fn non_empty(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} diff --git a/crates/server/src/admin_api/logs/query.rs b/crates/server/src/admin_api/logs/query.rs new file mode 100644 index 000000000..613ffa9ae --- /dev/null +++ b/crates/server/src/admin_api/logs/query.rs @@ -0,0 +1,82 @@ +//! 日志查询处理器。 + +use axum::{ + extract::{Query, State}, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, +}; +use codex_proxy_platform::json::clamp_limit; +use codex_proxy_runtime::{services::AdminLogFilter, state::AppState}; +use serde::Deserialize; + +use crate::{ + admin_api::{ + logs::{level_from_query, log_error, non_empty}, + require_admin_session, AdminError, AdminPageEnvelope, AdminResponse, + }, + middleware::request_id::RequestId, +}; + +/// 日志查询参数。 +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct LogsQuery { + /// 分页游标。 + pub cursor: Option, + /// 分页大小。 + pub limit: Option, + /// 事件类别。 + pub kind: Option, + /// 事件等级。 + pub level: Option, + /// 请求 ID。 + pub request_id: Option, + /// 账号 ID。 + pub account_id: Option, + /// 路由。 + pub route: Option, + /// 模型。 + pub model: Option, + /// HTTP 状态码。 + pub status_code: Option, + /// 搜索关键词。 + pub search: 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)); + let cursor = query.cursor.clone(); + let filter = filter_from_query(query, &request_id)?; + + match state.services.logs.list(cursor, limit, filter).await { + Ok(page) => Ok(AdminResponse::new( + StatusCode::OK, + AdminPageEnvelope::ok(page, limit, request_id), + )), + Err(error) => Err(log_error(error, request_id)), + } +} + +fn filter_from_query(query: LogsQuery, request_id: &str) -> Result { + Ok(AdminLogFilter { + kind: non_empty(query.kind), + level: level_from_query(query.level).map_err(|message| { + AdminError::new(StatusCode::BAD_REQUEST, 40001, message, request_id) + })?, + request_id: non_empty(query.request_id), + account_id: non_empty(query.account_id), + route: non_empty(query.route), + model: non_empty(query.model), + status_code: query.status_code, + search: non_empty(query.search), + }) +} diff --git a/src/admin/api/logs/state.rs b/crates/server/src/admin_api/logs/state.rs similarity index 60% rename from src/admin/api/logs/state.rs rename to crates/server/src/admin_api/logs/state.rs index 0d7d13105..5541a5d2d 100644 --- a/src/admin/api/logs/state.rs +++ b/crates/server/src/admin_api/logs/state.rs @@ -1,43 +1,35 @@ +//! 日志状态处理器。 + use axum::{ extract::State, http::{HeaderMap, StatusCode}, response::IntoResponse, Extension, Json, }; -use serde::{Deserialize, Serialize}; +use codex_proxy_runtime::state::AppState; +use serde::Deserialize; use crate::{ - codex::events::service::{ClearLogs, LogState, LogStateUpdate}, - platform::http::request_id::RequestId, - runtime::state::AppState, + admin_api::{ + logs::{log_error, ClearLogsData, LogStateData}, + require_admin_session, AdminEnvelope, AdminError, AdminResponse, + }, + middleware::request_id::RequestId, }; -use super::log_service_error; -use crate::admin::api::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct LogStateData { - pub enabled: bool, - pub capacity: u32, - pub capture_body: bool, - pub stored_count: u64, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct ClearLogsData { - pub cleared: u64, -} - +/// 更新日志状态请求。 #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] pub struct UpdateLogStateRequest { + /// 是否启用。 pub enabled: Option, + /// 日志容量。 pub capacity: Option, + /// 是否捕获请求体。 pub capture_body: Option, } +/// 读取日志状态。 pub async fn logs_state( State(state): State, Extension(request_id): Extension, @@ -45,16 +37,16 @@ pub async fn logs_state( ) -> Result { let request_id = request_id.as_str().to_string(); require_admin_session(&state, &headers, &request_id).await?; - match state.services.logs.state().await { Ok(state) => Ok(AdminResponse::new( StatusCode::OK, AdminEnvelope::ok(LogStateData::from(state), request_id), )), - Err(error) => Err(log_service_error(error, request_id)), + Err(error) => Err(log_error(error, request_id)), } } +/// 更新日志状态。 pub async fn update_logs_state( State(state): State, Extension(request_id): Extension, @@ -63,16 +55,16 @@ pub async fn update_logs_state( ) -> Result { let request_id = request_id.as_str().to_string(); require_admin_session(&state, &headers, &request_id).await?; - match state.services.logs.update_state(payload.into()).await { Ok(state) => Ok(AdminResponse::new( StatusCode::OK, AdminEnvelope::ok(LogStateData::from(state), request_id), )), - Err(error) => Err(log_service_error(error, request_id)), + Err(error) => Err(log_error(error, request_id)), } } +/// 清空事件日志。 pub async fn clear_logs( State(state): State, Extension(request_id): Extension, @@ -80,36 +72,21 @@ pub async fn clear_logs( ) -> Result { let request_id = request_id.as_str().to_string(); require_admin_session(&state, &headers, &request_id).await?; - match state.services.logs.clear().await { Ok(cleared) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(ClearLogsData::from(cleared), request_id), + AdminEnvelope::ok( + ClearLogsData { + cleared: cleared.cleared, + }, + request_id, + ), )), - Err(error) => Err(log_service_error(error, request_id)), + Err(error) => Err(log_error(error, request_id)), } } -impl From for LogStateData { - fn from(state: LogState) -> Self { - Self { - enabled: state.enabled, - capacity: state.capacity, - capture_body: state.capture_body, - stored_count: state.stored_count, - } - } -} - -impl From for ClearLogsData { - fn from(cleared: ClearLogs) -> Self { - Self { - cleared: cleared.cleared, - } - } -} - -impl From for LogStateUpdate { +impl From for codex_proxy_runtime::services::AdminLogStateUpdate { fn from(request: UpdateLogStateRequest) -> Self { Self { enabled: request.enabled, diff --git a/crates/server/src/admin_api/mod.rs b/crates/server/src/admin_api/mod.rs new file mode 100644 index 000000000..27c231391 --- /dev/null +++ b/crates/server/src/admin_api/mod.rs @@ -0,0 +1,26 @@ +//! 管理端 HTTP API。 + +/// 账号处理器。 +pub mod accounts; +/// 客户端 key 处理器。 +pub mod client_keys; +/// 诊断处理器。 +pub mod diagnostics; +/// 日志处理器。 +pub mod logs; +/// 模型处理器。 +pub mod models; +/// 响应封装。 +pub mod response; +/// 管理端路由。 +pub mod router; +/// 会话处理器。 +pub mod session; +/// 设置处理器。 +pub mod settings; +/// 用量处理器。 +pub mod usage; + +pub use response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse, PageMeta}; +pub use router::router; +pub use session::require_admin_session; diff --git a/crates/server/src/admin_api/models.rs b/crates/server/src/admin_api/models.rs new file mode 100644 index 000000000..cd96de67b --- /dev/null +++ b/crates/server/src/admin_api/models.rs @@ -0,0 +1,94 @@ +//! 管理端模型处理器。 + +use axum::{ + extract::State, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, +}; +use codex_proxy_core::models::service::ModelRefreshResult; +use codex_proxy_runtime::{services::AdminModelError, state::AppState}; +use serde::Serialize; + +use crate::{ + admin_api::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}, + middleware::request_id::RequestId, +}; + +/// 模型刷新响应。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct RefreshModelsData { + /// 成功刷新计划数。 + pub refreshed_plans: usize, + /// 写入模型数。 + pub model_count: usize, + /// 刷新失败计划数。 + pub failed_plans: usize, +} + +/// `POST /api/admin/refresh-models` +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 + .admin_models + .refresh_backend_models(&request_id) + .await + { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(RefreshModelsData::from(result), request_id), + )), + Err(error) => model_error(error, request_id), + } +} + +fn model_error( + error: AdminModelError, + request_id: String, +) -> Result>, AdminError> { + match error { + AdminModelError::NoAccounts => Err(AdminError::new( + StatusCode::BAD_REQUEST, + 40001, + error.to_string(), + request_id, + )), + AdminModelError::AllPlansFailed(result) => Ok(AdminResponse::new( + StatusCode::BAD_GATEWAY, + AdminEnvelope::new( + 50201, + "Failed to refresh backend models", + RefreshModelsData::from(result), + request_id, + ), + )), + AdminModelError::ListAccounts + | AdminModelError::SnapshotStoreUnavailable + | AdminModelError::UpstreamClientUnavailable + | AdminModelError::StoreSnapshot + | AdminModelError::LoadSnapshots => Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + error.to_string(), + request_id, + )), + } +} + +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, + } + } +} diff --git a/src/admin/api/response.rs b/crates/server/src/admin_api/response.rs similarity index 77% rename from src/admin/api/response.rs rename to crates/server/src/admin_api/response.rs index 4da2ad6e6..cbf0ff536 100644 --- a/src/admin/api/response.rs +++ b/crates/server/src/admin_api/response.rs @@ -1,29 +1,35 @@ +//! 管理端响应封装。 + use axum::{ http::StatusCode, response::{IntoResponse, Response}, Json, }; +use codex_proxy_platform::json::Page; 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, + /// 请求 ID。 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(), @@ -32,29 +38,40 @@ impl AdminEnvelope { } } + /// 构造成功响应信封。 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, + /// 请求 ID。 pub request_id: String, } impl AdminPageEnvelope { + /// 构造分页响应信封。 pub fn new( code: u32, message: impl Into, @@ -72,32 +89,43 @@ impl AdminPageEnvelope { } } + /// 构造成功分页响应信封。 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 { + /// HTTP 状态码。 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 { + /// HTTP 状态码。 pub status: StatusCode, + /// 业务状态码。 pub code: u32, + /// 错误消息。 pub message: String, + /// 请求 ID。 pub request_id: String, } impl AdminError { + /// 构造管理端错误。 pub fn new( status: StatusCode, code: u32, diff --git a/src/admin/api/router.rs b/crates/server/src/admin_api/router.rs similarity index 74% rename from src/admin/api/router.rs rename to crates/server/src/admin_api/router.rs index 0086bf0bb..8505e30e0 100644 --- a/src/admin/api/router.rs +++ b/crates/server/src/admin_api/router.rs @@ -1,61 +1,58 @@ +//! 管理端路由。 + use axum::{ routing::{delete, get, patch, post}, Router, }; -use crate::runtime::state::AppState; +use codex_proxy_runtime::state::AppState; -use super::{ - accounts::{ - account_quota, accounts, auth_callback, auth_code_relay, auth_device_login, - auth_device_poll, auth_login_start, 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, - quota_warnings, refresh_account, reset_account_usage, set_account_cookies, - update_account_label, update_account_status, - }, - client_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, - }, - diagnostics::diagnostics, - logs::{clear_logs, log_detail, logs, logs_state, update_logs_state}, - models::refresh_models, - session::{auth_logout, auth_status, login}, - settings::{settings, update_settings}, - usage::{usage_stats, usage_stats_summary}, +use crate::admin_api::accounts::{ + account_quota, accounts, auth_callback, auth_code_relay, auth_device_login, auth_device_poll, + auth_login_start, auth_logout, auth_status, 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_account, quota_warnings, refresh_account, + reset_account_usage, set_account_cookies, update_account_label, update_account_status, }; +use crate::admin_api::client_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, +}; +use crate::admin_api::diagnostics::diagnostics; +use crate::admin_api::logs::{clear_logs, log_detail, logs, logs_state, update_logs_state}; +use crate::admin_api::models::refresh_models; +use crate::admin_api::session::login; +use crate::admin_api::settings::{settings, update_settings}; +use crate::admin_api::usage::{usage_stats, usage_stats_summary}; +/// 构造管理端路由。 pub fn router() -> Router { Router::new() .route("/api/admin/login", post(login)) + .route("/api/admin/settings", get(settings).patch(update_settings)) + .route("/api/admin/diagnostics", get(diagnostics)) .route("/api/admin/auth/status", get(auth_status)) .route("/api/admin/auth/logout", post(auth_logout)) .route("/api/admin/auth/login-start", post(auth_login_start)) - .route("/api/admin/auth/code-relay", post(auth_code_relay)) - .route("/auth/callback", get(auth_callback)) .route("/api/admin/auth/device-login", post(auth_device_login)) .route( "/api/admin/auth/device-poll/{device_code}", get(auth_device_poll), ) - .route("/api/admin/diagnostics", get(diagnostics)) - .route("/api/admin/logs", get(logs).delete(clear_logs)) - .route( - "/api/admin/logs/state", - get(logs_state).patch(update_logs_state), - ) - .route("/api/admin/logs/{log_id}", get(log_detail)) - .route("/api/admin/settings", get(settings).patch(update_settings)) + .route("/api/admin/auth/code-relay", post(auth_code_relay)) + .route("/auth/callback", get(auth_callback)) .route("/api/admin/refresh-models", post(refresh_models)) .route("/api/admin/usage-stats", get(usage_stats)) .route("/api/admin/usage-stats/summary", get(usage_stats_summary)) .route("/api/admin/accounts", get(accounts).post(create_account)) + .route("/api/admin/accounts/export", get(export_accounts)) + .route("/api/admin/accounts/import", post(import_accounts)) + .route("/api/admin/accounts/import-cli", post(import_cli_account)) + .route("/api/admin/accounts/quota-warnings", get(quota_warnings)) .route( "/api/admin/accounts/health-check", post(health_check_accounts), ) - .route("/api/admin/accounts/quota-warnings", get(quota_warnings)) .route( "/api/admin/accounts/batch-delete", post(batch_delete_accounts), @@ -64,7 +61,7 @@ pub fn router() -> Router { "/api/admin/accounts/batch-status", post(batch_update_account_status), ) - .route("/api/admin/accounts/export", get(export_accounts)) + .route("/api/admin/accounts/{account_id}", delete(delete_account)) .route( "/api/admin/accounts/{account_id}/refresh", post(refresh_account), @@ -73,14 +70,13 @@ pub fn router() -> Router { "/api/admin/accounts/{account_id}/reset-usage", post(reset_account_usage), ) - .route("/api/admin/accounts/{account_id}/quota", get(account_quota)) .route( "/api/admin/accounts/{account_id}/cookies", get(get_account_cookies) .post(set_account_cookies) .delete(delete_account_cookies), ) - .route("/api/admin/accounts/{account_id}", delete(delete_account)) + .route("/api/admin/accounts/{account_id}/quota", get(account_quota)) .route( "/api/admin/accounts/{account_id}/label", patch(update_account_label), @@ -89,8 +85,12 @@ pub fn router() -> Router { "/api/admin/accounts/{account_id}/status", patch(update_account_status), ) - .route("/api/admin/accounts/import", post(import_accounts)) - .route("/api/admin/accounts/import-cli", post(import_cli_auth)) + .route("/api/admin/logs", get(logs).delete(clear_logs)) + .route( + "/api/admin/logs/state", + get(logs_state).patch(update_logs_state), + ) + .route("/api/admin/logs/{log_id}", get(log_detail)) .route("/api/admin/api-keys", get(api_keys).post(create_api_key)) .route("/api/admin/api-keys/export", get(export_api_keys)) .route("/api/admin/api-keys/import", post(import_api_keys)) diff --git a/crates/server/src/admin_api/session.rs b/crates/server/src/admin_api/session.rs new file mode 100644 index 000000000..c67250f28 --- /dev/null +++ b/crates/server/src/admin_api/session.rs @@ -0,0 +1,133 @@ +//! 管理端会话处理器。 + +use axum::{ + extract::State, + http::{header::SET_COOKIE, HeaderMap, HeaderValue, StatusCode}, + response::{IntoResponse, Response}, + Extension, Json, +}; +use codex_proxy_runtime::state::AppState; +use serde::{Deserialize, Serialize}; + +use crate::{ + admin_api::response::{AdminEnvelope, AdminError, AdminResponse}, + middleware::request_id::RequestId, +}; + +const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session"; + +/// 管理员登录请求。 +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminLoginRequest { + /// 管理员用户名；缺省时使用配置中的默认管理员。 + pub username: Option, + /// 管理员密码。 + pub password: String, +} + +/// 管理员登录响应。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminLoginData { + /// 会话过期时间。 + pub expires_at: String, +} + +/// 会话登录是否成功。 +pub fn session_login_allowed() -> bool { + true +} + +/// `POST /api/admin/login` +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 session = state + .services + .admin_sessions + .login(payload.username.as_deref(), payload.password.as_str()) + .await + .map_err(|_| { + AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to create admin session", + request_id.clone(), + ) + })? + .ok_or_else(|| { + AdminError::new( + StatusCode::UNAUTHORIZED, + 40102, + "Invalid admin credentials", + request_id.clone(), + ) + })?; + + let mut response = AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok( + AdminLoginData { + expires_at: session.expires_at.to_rfc3339(), + }, + request_id.clone(), + ), + ) + .into_response(); + let cookie = format!( + "{ADMIN_SESSION_COOKIE}={}; Path=/; HttpOnly; SameSite=Lax", + session.session_id + ); + response.headers_mut().insert( + SET_COOKIE, + HeaderValue::from_str(&cookie).map_err(|_| { + AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to create admin session cookie", + request_id.clone(), + ) + })?, + ); + Ok(response) +} + +/// 要求请求携带有效管理员会话。 +pub async fn require_admin_session( + state: &AppState, + headers: &HeaderMap, + request_id: &str, +) -> Result<(), AdminError> { + match state + .services + .admin_sessions + .validate(admin_session_cookie(headers).as_deref()) + .await + { + Ok(true) => Ok(()), + Ok(false) => Err(AdminError::new( + StatusCode::UNAUTHORIZED, + 40101, + "Admin session required", + request_id, + )), + Err(_) => Err(AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to validate admin session", + request_id, + )), + } +} + +fn admin_session_cookie(headers: &HeaderMap) -> Option { + let cookie = headers.get("cookie")?.to_str().ok()?; + cookie.split(';').find_map(|part| { + let (name, value) = part.trim().split_once('=')?; + (name == ADMIN_SESSION_COOKIE).then(|| value.to_string()) + }) +} diff --git a/src/admin/api/settings.rs b/crates/server/src/admin_api/settings.rs similarity index 64% rename from src/admin/api/settings.rs rename to crates/server/src/admin_api/settings.rs index 487700cf1..f0c159379 100644 --- a/src/admin/api/settings.rs +++ b/crates/server/src/admin_api/settings.rs @@ -1,3 +1,7 @@ +//! 管理端设置处理器。 + +use std::collections::BTreeMap; + use axum::{ body::Bytes, extract::State, @@ -5,18 +9,17 @@ use axum::{ response::IntoResponse, Extension, }; -use serde::{Deserialize, Serialize}; +use codex_proxy_core::admin::settings::{AdminSettingsPatch, SettingsServiceError}; +use codex_proxy_platform::config::{AppConfig, QuotaWarningThresholds}; +use codex_proxy_runtime::{services::RuntimeSettingsError, state::AppState}; +use serde::Serialize; use serde_json::Value; use crate::{ - admin::settings::{SettingsServiceError, SettingsUpdate}, - config::{AppConfig, QuotaWarningThresholds}, - platform::http::request_id::RequestId, - runtime::state::AppState, + admin_api::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}, + middleware::request_id::RequestId, }; -use super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; - const ALLOWED_SETTINGS_KEYS: [&str; 18] = [ "defaultModel", "defaultReasoningEffort", @@ -38,26 +41,45 @@ const ALLOWED_SETTINGS_KEYS: [&str; 18] = [ "usageHistoryRetentionDays", ]; +/// 管理端设置响应。 #[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 model_aliases: 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, } @@ -86,54 +108,7 @@ impl AdminSettingsData { } } -#[derive(Debug, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -struct AdminSettingsPatch { - default_model: Option, - default_reasoning_effort: Option, - service_tier: Option, - model_aliases: Option>, - refresh_enabled: Option, - refresh_margin_seconds: Option, - refresh_concurrency: Option, - max_concurrent_per_account: Option, - request_interval_ms: Option, - rotation_strategy: Option, - tier_priority: Option>, - quota_refresh_interval_minutes: Option, - quota_warning_thresholds: Option, - quota_skip_exhausted: Option, - logs_enabled: Option, - logs_capacity: Option, - logs_capture_body: Option, - usage_history_retention_days: Option, -} - -impl From for SettingsUpdate { - fn from(patch: AdminSettingsPatch) -> Self { - Self { - default_model: patch.default_model, - default_reasoning_effort: patch.default_reasoning_effort, - service_tier: patch.service_tier, - model_aliases: patch.model_aliases, - refresh_enabled: patch.refresh_enabled, - refresh_margin_seconds: patch.refresh_margin_seconds, - refresh_concurrency: patch.refresh_concurrency, - max_concurrent_per_account: patch.max_concurrent_per_account, - request_interval_ms: patch.request_interval_ms, - rotation_strategy: patch.rotation_strategy, - tier_priority: patch.tier_priority, - quota_refresh_interval_minutes: patch.quota_refresh_interval_minutes, - quota_warning_thresholds: patch.quota_warning_thresholds, - quota_skip_exhausted: patch.quota_skip_exhausted, - logs_enabled: patch.logs_enabled, - logs_capacity: patch.logs_capacity, - logs_capture_body: patch.logs_capture_body, - usage_history_retention_days: patch.usage_history_retention_days, - } - } -} - +/// 读取管理端设置。 pub async fn settings( State(state): State, Extension(request_id): Extension, @@ -141,7 +116,7 @@ pub async fn settings( ) -> Result { let request_id = request_id.as_str().to_string(); require_admin_session(&state, &headers, &request_id).await?; - let config = state.services.settings.current().await; + let config = state.services.settings.current(); Ok(AdminResponse::new( StatusCode::OK, @@ -149,6 +124,7 @@ pub async fn settings( )) } +/// 更新管理端设置。 pub async fn update_settings( State(state): State, Extension(request_id): Extension, @@ -162,9 +138,9 @@ pub async fn update_settings( let config = state .services .settings - .update(patch.into()) + .update(patch) .await - .map_err(|err| settings_service_error(err, &request_id))?; + .map_err(|error| settings_service_error(error, &request_id))?; Ok(AdminResponse::new( StatusCode::OK, @@ -212,17 +188,18 @@ fn parse_settings_patch(body: &[u8], request_id: &str) -> Result AdminError { +fn settings_service_error(error: RuntimeSettingsError, request_id: &str) -> AdminError { match error { - SettingsServiceError::InvalidField { field, message } => AdminError::new( + RuntimeSettingsError::InvalidField(SettingsServiceError::InvalidField { + field, + message, + }) => AdminError::new( StatusCode::BAD_REQUEST, 40001, format!("Invalid setting {field}: {message}"), request_id, ), - SettingsServiceError::Serialize(_) - | SettingsServiceError::CreateDirectory(_) - | SettingsServiceError::Write(_) => AdminError::new( + RuntimeSettingsError::Persist(_) => AdminError::new( StatusCode::INTERNAL_SERVER_ERROR, 50000, "Failed to persist settings", diff --git a/src/admin/api/usage.rs b/crates/server/src/admin_api/usage.rs similarity index 64% rename from src/admin/api/usage.rs rename to crates/server/src/admin_api/usage.rs index 780b3a754..227138f7a 100644 --- a/src/admin/api/usage.rs +++ b/crates/server/src/admin_api/usage.rs @@ -1,51 +1,154 @@ +//! 管理端用量处理器。 + use axum::{ extract::{Query, State}, http::{HeaderMap, StatusCode}, response::IntoResponse, Extension, }; +use codex_proxy_platform::json::{clamp_limit, Page}; +use codex_proxy_runtime::{ + services::{AdminUsageError, AdminUsageRecord, AdminUsageSummary}, + state::AppState, +}; use serde::{Deserialize, Serialize}; use crate::{ - codex::accounts::repository::{ - AccountRepositoryError, AccountUsageListRecord, AccountUsageSummary, + admin_api::{ + require_admin_session, AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse, }, - codex::usage::service::UsageServiceError, - platform::http::request_id::RequestId, - runtime::state::AppState, - utils::pagination::{clamp_limit, Page}, + middleware::request_id::RequestId, }; -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 { + /// 账号 ID。 pub account_id: String, + /// 邮箱。 pub email: Option, + /// 管理标签。 pub label: Option, + /// 计划类型。 pub plan_type: Option, + /// 请求数。 pub request_count: i64, + /// 空响应数。 pub empty_response_count: i64, + /// 输入 token 数。 pub input_tokens: i64, + /// 输出 token 数。 pub output_tokens: i64, + /// 缓存 token 数。 pub cached_tokens: i64, + /// 图片输入 token 数。 pub image_input_tokens: i64, + /// 图片输出 token 数。 pub image_output_tokens: i64, + /// 图片请求数。 pub image_request_count: i64, + /// 图片请求失败数。 pub image_request_failed_count: i64, + /// 最近使用时间。 pub last_used_at: Option, } -impl From for AdminUsageStatsData { - fn from(usage: AccountUsageListRecord) -> Self { +/// 管理端账号用量汇总响应。 +#[derive(Debug, Clone, Copy, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AdminUsageStatsSummaryData { + /// 有用量记录的账号数。 + pub account_count: i64, + /// 请求总数。 + pub request_count: i64, + /// 空响应总数。 + pub empty_response_count: i64, + /// 输入 token 总数。 + pub input_tokens: i64, + /// 输出 token 总数。 + pub output_tokens: i64, + /// 缓存 token 总数。 + pub cached_tokens: i64, + /// 图片输入 token 总数。 + pub image_input_tokens: i64, + /// 图片输出 token 总数。 + pub image_output_tokens: i64, + /// 图片请求总数。 + pub image_request_count: i64, + /// 图片请求失败总数。 + pub image_request_failed_count: i64, +} + +/// `GET /api/admin/usage-stats` +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 = Page { + items: page + .items + .into_iter() + .map(AdminUsageStatsData::from) + .collect(), + next_cursor: page.next_cursor, + }; + Ok(AdminResponse::new( + StatusCode::OK, + AdminPageEnvelope::ok(page, limit, request_id), + )) + } + Err(error) => Err(usage_error(error, request_id)), + } +} + +/// `GET /api/admin/usage-stats/summary` +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_error(error, request_id)), + } +} + +fn usage_error(error: AdminUsageError, request_id: String) -> AdminError { + AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + error.to_string(), + request_id, + ) +} + +impl From for AdminUsageStatsData { + fn from(usage: AdminUsageRecord) -> Self { Self { account_id: usage.account_id, email: usage.email, @@ -65,23 +168,8 @@ impl From for AdminUsageStatsData { } } -#[derive(Debug, Clone, Copy, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminUsageStatsSummaryData { - pub account_count: i64, - pub request_count: i64, - pub empty_response_count: i64, - pub input_tokens: i64, - pub output_tokens: i64, - pub cached_tokens: i64, - pub image_input_tokens: i64, - pub image_output_tokens: i64, - pub image_request_count: i64, - pub image_request_failed_count: i64, -} - -impl From for AdminUsageStatsSummaryData { - fn from(summary: AccountUsageSummary) -> Self { +impl From for AdminUsageStatsSummaryData { + fn from(summary: AdminUsageSummary) -> Self { Self { account_count: summary.account_count, request_count: summary.request_count, @@ -96,83 +184,3 @@ impl From for AdminUsageStatsSummaryData { } } } - -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/crates/server/src/error/admin.rs b/crates/server/src/error/admin.rs new file mode 100644 index 000000000..d9db6fc75 --- /dev/null +++ b/crates/server/src/error/admin.rs @@ -0,0 +1,6 @@ +//! 管理端错误映射。 + +/// 将领域错误格式化为 HTTP 错误消息。 +pub fn format_admin_error(message: &str) -> String { + message.to_string() +} diff --git a/crates/server/src/error/mod.rs b/crates/server/src/error/mod.rs new file mode 100644 index 000000000..d85dce59a --- /dev/null +++ b/crates/server/src/error/mod.rs @@ -0,0 +1,6 @@ +//! HTTP 错误映射。 + +/// 管理端错误映射。 +pub mod admin; +/// OpenAI 错误映射。 +pub mod openai; diff --git a/crates/server/src/error/openai.rs b/crates/server/src/error/openai.rs new file mode 100644 index 000000000..9082eb01b --- /dev/null +++ b/crates/server/src/error/openai.rs @@ -0,0 +1,6 @@ +//! OpenAI 错误映射。 + +/// 将领域错误格式化为 OpenAI 错误消息。 +pub fn format_openai_error(message: &str) -> String { + message.to_string() +} diff --git a/crates/server/src/lib.rs b/crates/server/src/lib.rs new file mode 100644 index 000000000..e8c7b7a73 --- /dev/null +++ b/crates/server/src/lib.rs @@ -0,0 +1,7 @@ +//! HTTP 边界层，提供 Axum 路由、处理器、中间件和响应映射。 + +pub mod admin_api; +pub mod error; +pub mod middleware; +pub mod openai_api; +pub mod router; diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs new file mode 100644 index 000000000..bd0ef79be --- /dev/null +++ b/crates/server/src/main.rs @@ -0,0 +1,86 @@ +use std::{error::Error, path::Path}; + +use axum::Router; +use codex_proxy_platform::{ + config::AppConfig, + crypto::SecretBox, + identity::ApiKeyHasher, + logging::{init_tracing, LogGuard, RotationConfig}, + storage::{connect_sqlite, ensure_data_dir, load_or_create_installation_id}, +}; +use codex_proxy_runtime::{ + state::AppState, + tasks::coordinator::{start_background_tasks, BackgroundTaskCoordinator}, +}; +use codex_proxy_server::router; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let config = AppConfig::load()?; + let host = config.server.host.clone(); + let port = config.server.port; + let _log_guard = init_logging(&config)?; + let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?; + let (app, task_coordinator) = build_application(config).await?; + + let serve_result = axum::serve(listener, app) + .with_graceful_shutdown(shutdown_signal()) + .await; + + task_coordinator.shutdown().await; + serve_result?; + + Ok(()) +} + +async fn build_application( + config: AppConfig, +) -> Result<(Router, BackgroundTaskCoordinator), Box> { + let pool = connect_sqlite(&config.database.url).await?; + let secret_box = SecretBox::load_or_create(Path::new(&config.security.master_key_file))?; + let api_key_hasher = + ApiKeyHasher::load_or_create(Path::new(&config.security.api_key_pepper_file))?; + let default_admin_password = config.admin.default_password.clone(); + let data_dir = ensure_data_dir()?; + let installation_id = load_or_create_installation_id(Some(&data_dir))?; + let state = AppState::with_pool_secret_api_key_hasher_and_installation_id( + config, + pool, + secret_box, + api_key_hasher, + installation_id, + ); + let created_default_admin = state + .services + .admin_sessions + .ensure_default_admin(&default_admin_password) + .await?; + tracing::info!( + created = created_default_admin, + "默认管理员账号已完成初始化检查" + ); + let restored_accounts = state.restore_account_pool_from_repository().await?; + tracing::info!(count = restored_accounts, "运行时账号池已从 SQLite 恢复"); + let restored_session_affinities = state.restore_session_affinity_from_repository_now().await?; + tracing::info!( + count = restored_session_affinities, + "会话亲和性映射已从 SQLite 恢复" + ); + let task_coordinator = start_background_tasks(&state).await; + Ok((router::router().with_state(state), task_coordinator)) +} + +fn init_logging(config: &AppConfig) -> Result, Box> { + if !config.logging.enabled { + return Ok(None); + } + + let rotation = RotationConfig::new(&config.logging.directory, config.logging.retention_days); + Ok(Some(init_tracing(rotation)?)) +} + +async fn shutdown_signal() { + if let Err(error) = tokio::signal::ctrl_c().await { + eprintln!("failed to listen for shutdown signal: {error}"); + } +} diff --git a/crates/server/src/middleware/auth.rs b/crates/server/src/middleware/auth.rs new file mode 100644 index 000000000..e850df7b2 --- /dev/null +++ b/crates/server/src/middleware/auth.rs @@ -0,0 +1,5 @@ +//! 认证中间件。 + +/// 认证结果。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Authenticated; diff --git a/crates/server/src/middleware/cors.rs b/crates/server/src/middleware/cors.rs new file mode 100644 index 000000000..d2c736dec --- /dev/null +++ b/crates/server/src/middleware/cors.rs @@ -0,0 +1,4 @@ +//! CORS 中间件。 + +/// 默认是否允许浏览器跨域。 +pub const DEFAULT_CORS_ENABLED: bool = true; diff --git a/crates/server/src/middleware/mod.rs b/crates/server/src/middleware/mod.rs new file mode 100644 index 000000000..09422083c --- /dev/null +++ b/crates/server/src/middleware/mod.rs @@ -0,0 +1,10 @@ +//! HTTP 中间件。 + +/// 认证中间件。 +pub mod auth; +/// CORS 中间件。 +pub mod cors; +/// 请求 ID 中间件。 +pub mod request_id; +/// tracing 中间件。 +pub mod trace; diff --git a/src/platform/http/request_id.rs b/crates/server/src/middleware/request_id.rs similarity index 82% rename from src/platform/http/request_id.rs rename to crates/server/src/middleware/request_id.rs index 17cbb8ce1..6b3f16562 100644 --- a/src/platform/http/request_id.rs +++ b/crates/server/src/middleware/request_id.rs @@ -1,16 +1,22 @@ +//! 请求 ID 中间件。 + use axum::{extract::Request, http::HeaderValue, middleware::Next, response::Response}; use uuid::Uuid; +/// 请求 ID 头。 pub const REQUEST_ID_HEADER: &str = "x-request-id"; +/// 请求 ID。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct RequestId(String); impl RequestId { + /// 生成新的请求 ID。 pub fn generate() -> Self { Self(format!("req_{}", Uuid::new_v4())) } + /// 从 HTTP 头解析请求 ID。 pub fn from_header(value: &HeaderValue) -> Option { let value = value.to_str().ok()?.trim(); if value.is_empty() { @@ -19,11 +25,13 @@ impl RequestId { Some(Self(value.to_string())) } + /// 返回字符串形式。 pub fn as_str(&self) -> &str { &self.0 } } +/// 为请求附加请求 ID，并在响应头中回写。 pub async fn attach_request_id(mut request: Request, next: Next) -> Response { let request_id = request .headers() @@ -31,7 +39,6 @@ pub async fn attach_request_id(mut request: Request, next: Next) -> Response { .and_then(RequestId::from_header) .unwrap_or_else(RequestId::generate); - // 同一个 requestId 同时进入日志、响应头和 admin body，方便跨层排查。 request.extensions_mut().insert(request_id.clone()); let mut response = next.run(request).await; if let Ok(value) = HeaderValue::from_str(request_id.as_str()) { diff --git a/crates/server/src/middleware/trace.rs b/crates/server/src/middleware/trace.rs new file mode 100644 index 000000000..e7d952861 --- /dev/null +++ b/crates/server/src/middleware/trace.rs @@ -0,0 +1,81 @@ +//! tracing 中间件。 + +use std::time::Duration; + +use axum::{ + body::Body, + http::{Request, Response}, +}; +use tower_http::{ + classify::{ServerErrorsAsFailures, ServerErrorsFailureClass, SharedClassifier}, + trace::TraceLayer, +}; +use tracing::{info, warn, Span}; + +use crate::middleware::request_id::RequestId; + +/// 默认 tracing 事件名。 +pub const TRACE_LAYER_NAME: &str = "http"; + +/// HTTP tracing layer type used by the server router. +pub type HttpTraceLayer = TraceLayer< + SharedClassifier, + fn(&Request) -> Span, + fn(&Request, &Span), + fn(&Response, Duration, &Span), + tower_http::trace::DefaultOnBodyChunk, + tower_http::trace::DefaultOnEos, + fn(ServerErrorsFailureClass, Duration, &Span), +>; + +/// 构造 HTTP tracing layer。 +pub fn http_trace_layer() -> HttpTraceLayer { + TraceLayer::new_for_http() + .make_span_with(make_http_span as fn(&Request) -> Span) + .on_request(on_http_request as fn(&Request, &Span)) + .on_response(on_http_response as fn(&Response, Duration, &Span)) + .on_failure(on_http_failure as fn(ServerErrorsFailureClass, Duration, &Span)) +} + +fn make_http_span(request: &Request) -> Span { + let request_id = request_id(request); + tracing::info_span!( + "http", + request_id = request_id.as_deref(), + method = %request.method(), + uri = %request.uri() + ) +} + +fn on_http_request(request: &Request, _span: &Span) { + let request_id = request_id(request); + info!( + request_id = request_id.as_deref(), + method = %request.method(), + uri = %request.uri(), + "received HTTP request" + ); +} + +fn on_http_response(response: &Response, latency: Duration, _span: &Span) { + info!( + status = response.status().as_u16(), + latency_ms = latency.as_millis(), + "completed HTTP request" + ); +} + +fn on_http_failure(failure_class: ServerErrorsFailureClass, latency: Duration, _span: &Span) { + warn!( + failure = %failure_class, + latency_ms = latency.as_millis(), + "failed HTTP request" + ); +} + +fn request_id(request: &Request) -> Option { + request + .extensions() + .get::() + .map(|request_id| request_id.as_str().to_string()) +} diff --git a/crates/server/src/openai_api/auth.rs b/crates/server/src/openai_api/auth.rs new file mode 100644 index 000000000..5793be399 --- /dev/null +++ b/crates/server/src/openai_api/auth.rs @@ -0,0 +1,29 @@ +//! OpenAI API 认证。 + +use axum::http::HeaderMap; + +use codex_proxy_runtime::state::AppState; + +/// 从请求头提取 Bearer 客户端 API key。 +pub fn bearer_client_api_key(headers: &HeaderMap) -> Option { + let raw = headers + .get(axum::http::header::AUTHORIZATION)? + .to_str() + .ok()?; + let token = raw.strip_prefix("Bearer ")?.trim(); + (!token.is_empty()).then(|| token.to_string()) +} + +/// 校验客户端 API key。 +pub async fn authorize_client_api_key(state: &AppState, headers: &HeaderMap) -> bool { + let Some(api_key) = bearer_client_api_key(headers) else { + return false; + }; + + state + .services + .client_keys + .verify(&api_key) + .await + .unwrap_or(false) +} diff --git a/crates/server/src/openai_api/chat.rs b/crates/server/src/openai_api/chat.rs new file mode 100644 index 000000000..80325ad90 --- /dev/null +++ b/crates/server/src/openai_api/chat.rs @@ -0,0 +1,197 @@ +//! OpenAI 聊天处理器。 + +use axum::{ + body::Bytes, + extract::State, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + Extension, Json, +}; +use codex_proxy_core::protocol::openai::chat::{translate_chat_to_codex, ChatCompletionRequest}; +use codex_proxy_runtime::{services::ChatDispatchError, state::AppState}; +use serde_json::{json, Value}; + +use crate::middleware::request_id::RequestId; + +use super::{auth::authorize_client_api_key, models::model_catalog_for_state}; + +/// `POST /v1/chat/completions` +pub async fn chat_completions( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + body: Bytes, +) -> Response { + 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 openai_error_response( + StatusCode::BAD_REQUEST, + "Invalid chat completion request", + "invalid_request_error", + "invalid_request", + ) + .into_response(); + }; + let model = chat_request.model.clone(); + let catalog = model_catalog_for_state(&state).await; + if !catalog.is_recognized_model_name(&model) { + return model_not_found_response().into_response(); + } + let codex_request = match translate_chat_to_codex(chat_request) { + Ok(request) => request, + Err(_) => { + return openai_error_response( + StatusCode::BAD_REQUEST, + "Invalid chat completion request", + "invalid_request_error", + "invalid_request", + ) + .into_response(); + } + }; + + match state + .services + .chat + .complete(request_id.as_str(), codex_request, &model) + .await + { + Ok(body) => (StatusCode::OK, Json(body)).into_response(), + Err(ChatDispatchError::NoActiveAccount | ChatDispatchError::AccountStore) => { + openai_error_response( + StatusCode::SERVICE_UNAVAILABLE, + "No active upstream account is available", + "server_error", + "upstream_unavailable", + ) + .into_response() + } + Err(ChatDispatchError::Upstream(_)) => openai_error_response( + StatusCode::BAD_GATEWAY, + "Upstream Codex request failed", + "server_error", + "upstream_error", + ) + .into_response(), + Err(ChatDispatchError::QuotaExhausted { + count, + upstream_error, + }) => openai_error_response( + StatusCode::PAYMENT_REQUIRED, + &format!( + "All accounts exhausted ({count} quota-exhausted). Codex upstream error: {upstream_error}" + ), + "server_error", + "upstream_error", + ) + .into_response(), + Err(ChatDispatchError::RateLimited { + count, + upstream_error, + }) => openai_error_response( + StatusCode::TOO_MANY_REQUESTS, + &format!( + "All accounts exhausted ({count} rate-limited). Codex upstream error: {upstream_error}" + ), + "server_error", + "upstream_error", + ) + .into_response(), + Err(ChatDispatchError::Expired { + count, + upstream_error, + }) => openai_error_response( + StatusCode::UNAUTHORIZED, + &format!( + "All accounts exhausted ({count} expired). Codex upstream error: {upstream_error}" + ), + "server_error", + "upstream_error", + ) + .into_response(), + Err(ChatDispatchError::CloudflareChallenge { + count, + upstream_error, + }) => openai_error_response( + StatusCode::BAD_GATEWAY, + &format!( + "All accounts exhausted ({count} cloudflare-challenge). Codex upstream error: {upstream_error}" + ), + "server_error", + "upstream_error", + ) + .into_response(), + Err(ChatDispatchError::CloudflarePathBlocked { + count, + upstream_error, + }) => openai_error_response( + StatusCode::BAD_GATEWAY, + &format!( + "All accounts exhausted ({count} cloudflare-path-block). Codex upstream error: {upstream_error}" + ), + "server_error", + "upstream_error", + ) + .into_response(), + Err(ChatDispatchError::ModelUnsupported { + count, + upstream_error, + }) => openai_error_response( + StatusCode::BAD_REQUEST, + &format!( + "All accounts exhausted ({count} model-unsupported). Codex upstream error: {upstream_error}" + ), + "invalid_request_error", + "upstream_error", + ) + .into_response(), + Err(ChatDispatchError::InvalidSse(_) | ChatDispatchError::EmptyUpstreamResponse) => { + openai_error_response( + StatusCode::BAD_GATEWAY, + "Invalid upstream Codex response", + "server_error", + "invalid_upstream_response", + ) + .into_response() + } + } +} + +fn missing_client_api_key_response() -> (StatusCode, Json) { + openai_error_response( + StatusCode::UNAUTHORIZED, + "Missing client API key", + "invalid_request_error", + "invalid_api_key", + ) +} + +fn model_not_found_response() -> (StatusCode, Json) { + openai_error_response( + StatusCode::NOT_FOUND, + "Model not found", + "invalid_request_error", + "model_not_found", + ) +} + +fn openai_error_response( + status: StatusCode, + message: &str, + error_type: &str, + code: &str, +) -> (StatusCode, Json) { + ( + status, + Json(json!({ + "error": { + "message": message, + "type": error_type, + "code": code + } + })), + ) +} diff --git a/crates/server/src/openai_api/diagnostics.rs b/crates/server/src/openai_api/diagnostics.rs new file mode 100644 index 000000000..5fc3a657f --- /dev/null +++ b/crates/server/src/openai_api/diagnostics.rs @@ -0,0 +1,359 @@ +//! OpenAI 诊断处理器。 + +use axum::{ + extract::State, + http::{HeaderMap, StatusCode}, + response::IntoResponse, + Extension, Json, +}; +use codex_proxy_core::{ + accounts::{ + model::{Account, AccountStatus}, + pool::AccountCapacitySummary, + }, + gateway::fingerprint::Fingerprint, +}; +use codex_proxy_platform::config::AppConfig; +use codex_proxy_runtime::state::AppState; +use serde::Serialize; +use serde_json::json; + +use crate::middleware::request_id::RequestId; + +/// 诊断数据。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DiagnosticsData { + /// 诊断状态。 + pub status: &'static str, + /// 运行时包信息。 + pub runtime: RuntimeDiagnostics, + /// Runtime paths. + pub paths: PathDiagnostics, + /// 上游传输配置。 + pub transport: TransportDiagnostics, + /// 账号状态。 + pub accounts: AccountDiagnostics, + /// 主要运行设置。 + pub settings: SettingsDiagnostics, +} + +/// Runtime path diagnostics. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct PathDiagnostics { + /// Primary config file. + pub config: &'static str, + /// Local override config file. + pub local_config: String, + /// Configured SQLite database URL. + pub database_url: String, +} + +/// 运行时包信息。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct RuntimeDiagnostics { + /// 包名。 + pub package_name: &'static str, + /// 包版本。 + pub package_version: &'static str, +} + +/// 上游传输配置。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct TransportDiagnostics { + /// Codex 后端基础 URL。 + pub backend_base_url: String, + /// TLS 配置。 + pub tls: TlsDiagnostics, + /// Runtime request fingerprint. + pub fingerprint: FingerprintDiagnostics, +} + +/// TLS 配置。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct TlsDiagnostics { + /// 是否强制 HTTP/1.1。 + pub force_http11: bool, +} + +/// 账号诊断数据。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountDiagnostics { + /// 账号仓储是否可用。 + pub repository_available: bool, + /// 账号池摘要。 + pub pool: AccountPoolDiagnostics, + /// Account-pool capacity. + pub capacity: AccountCapacityDiagnostics, +} + +/// 账号池摘要。 +#[derive(Debug, Default, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountPoolDiagnostics { + /// 总数。 + pub total: usize, + /// 活跃数。 + pub active: usize, + /// 过期数。 + pub expired: usize, + /// 配额耗尽数。 + pub quota_exhausted: usize, + /// 刷新中数。 + pub refreshing: usize, + /// 禁用数。 + pub disabled: usize, + /// 封禁数。 + pub banned: usize, +} + +/// Account-pool capacity diagnostics. +#[derive(Debug, Default, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountCapacityDiagnostics { + /// Maximum concurrent slots per account. + pub max_concurrent_per_account: usize, + /// Total slots available across active accounts. + pub total_slots: usize, + /// Currently occupied slots. + pub used_slots: usize, + /// Currently available slots. + pub available_slots: usize, +} + +/// Runtime fingerprint summary. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct FingerprintDiagnostics { + /// Fingerprint source label. + pub source: &'static str, + /// Client originator. + pub originator: String, + /// App version. + pub app_version: String, + /// Build number. + pub build_number: String, + /// Platform name. + pub platform: String, + /// Architecture name. + pub arch: String, + /// Chromium major version. + pub chromium_version: String, + /// Expanded user-agent. + pub user_agent: String, +} + +/// Upstream probe diagnostics. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct UpstreamProbeDiagnostics { + /// Probe target name. + pub target: &'static str, + /// Configured backend base URL. + pub backend_base_url: String, + /// Full endpoint URL. + pub endpoint: String, + /// Whether upstream answered at transport level. + pub reachable: bool, + /// HTTP status code, when available. + pub status_code: Option, + /// Authentication outcome inferred from status. + pub authorization: &'static str, +} + +/// 主要运行设置。 +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct SettingsDiagnostics { + /// 默认模型。 + pub default_model: String, + /// 是否启用刷新。 + pub refresh_enabled: bool, + /// 账号轮换策略。 + pub rotation_strategy: String, + /// 是否跳过配额耗尽账号。 + pub quota_skip_exhausted: bool, + /// 是否启用日志。 + pub logs_enabled: bool, +} + +/// `GET /debug/diagnostics` +pub async fn diagnostics(State(state): State, headers: HeaderMap) -> impl IntoResponse { + if !is_local_debug_request(&headers) { + return ( + StatusCode::FORBIDDEN, + Json(json!({ "error": "debug endpoint is local-only" })), + ) + .into_response(); + } + + (StatusCode::OK, Json(diagnostics_data(&state).await)).into_response() +} + +/// `GET /debug/fingerprint` +pub async fn fingerprint(State(state): State, headers: HeaderMap) -> impl IntoResponse { + if !is_local_debug_request(&headers) { + return ( + StatusCode::FORBIDDEN, + Json(json!({ "error": "debug endpoint is local-only" })), + ) + .into_response(); + } + + ( + StatusCode::OK, + Json(fingerprint_diagnostics(&state.services.fingerprint)), + ) + .into_response() +} + +/// `GET /debug/upstream` +pub async fn upstream( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, +) -> impl IntoResponse { + if !is_local_debug_request(&headers) { + return ( + StatusCode::FORBIDDEN, + Json(json!({ "error": "debug endpoint is local-only" })), + ) + .into_response(); + } + + let probe = state + .services + .probe_codex_models_endpoint(request_id.as_str()) + .await; + + ( + StatusCode::OK, + Json(UpstreamProbeDiagnostics { + target: probe.target, + backend_base_url: probe.backend_base_url, + endpoint: probe.endpoint, + reachable: probe.reachable, + status_code: probe.status_code, + authorization: probe.authorization, + }), + ) + .into_response() +} + +/// 构造诊断数据。 +pub async fn diagnostics_data(state: &AppState) -> DiagnosticsData { + let config = state.services.settings.current(); + let local_config_path = state.services.settings.local_config_path(); + let accounts = state + .services + .accounts + .list_pool_accounts() + .await + .unwrap_or_default(); + let capacity = state.services.account_pool.capacity_summary_now().await; + DiagnosticsData { + status: "ok", + runtime: RuntimeDiagnostics { + package_name: env!("CARGO_PKG_NAME"), + package_version: env!("CARGO_PKG_VERSION"), + }, + paths: PathDiagnostics { + config: "config.yaml", + local_config: local_config_path.to_string_lossy().into_owned(), + database_url: config.database.url.clone(), + }, + transport: transport_diagnostics(&config, &state.services.fingerprint), + accounts: AccountDiagnostics { + repository_available: true, + pool: account_pool_diagnostics(&accounts), + capacity: AccountCapacityDiagnostics::from(capacity), + }, + settings: SettingsDiagnostics::from(config.as_ref()), + } +} + +fn transport_diagnostics(config: &AppConfig, fingerprint: &Fingerprint) -> TransportDiagnostics { + TransportDiagnostics { + backend_base_url: config.api.base_url.clone(), + tls: TlsDiagnostics { + force_http11: config.tls.force_http11, + }, + fingerprint: fingerprint_diagnostics(fingerprint), + } +} + +fn account_pool_diagnostics(accounts: &[Account]) -> AccountPoolDiagnostics { + let mut summary = AccountPoolDiagnostics { + total: accounts.len(), + ..AccountPoolDiagnostics::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 +} + +fn fingerprint_diagnostics(fingerprint: &Fingerprint) -> FingerprintDiagnostics { + FingerprintDiagnostics { + source: "runtime", + originator: fingerprint.originator.clone(), + app_version: fingerprint.app_version.clone(), + build_number: fingerprint.build_number.clone(), + platform: fingerprint.platform.clone(), + arch: fingerprint.arch.clone(), + chromium_version: fingerprint.chromium_version.clone(), + user_agent: fingerprint.user_agent(), + } +} + +fn is_local_debug_request(headers: &HeaderMap) -> bool { + forwarded_header_is_local(headers, "x-forwarded-for") + && forwarded_header_is_local(headers, "x-real-ip") +} + +impl From for AccountCapacityDiagnostics { + fn from(summary: AccountCapacitySummary) -> Self { + Self { + max_concurrent_per_account: summary.max_concurrent_per_account, + total_slots: summary.total_slots, + used_slots: summary.used_slots, + available_slots: summary.available_slots, + } + } +} + +fn forwarded_header_is_local(headers: &HeaderMap, name: &str) -> bool { + let Some(value) = headers.get(name).and_then(|value| value.to_str().ok()) else { + return true; + }; + value.split(',').next().is_some_and(is_local_host) +} + +fn is_local_host(host: &str) -> bool { + let host = host.trim().trim_start_matches('[').trim_end_matches(']'); + host == "localhost" || host == "::1" || host.starts_with("127.") +} + +impl From<&AppConfig> for SettingsDiagnostics { + fn from(config: &AppConfig) -> Self { + Self { + default_model: config.model.default_model.clone(), + refresh_enabled: config.auth.refresh_enabled, + rotation_strategy: config.auth.rotation_strategy.clone(), + quota_skip_exhausted: config.quota.skip_exhausted, + logs_enabled: config.logging.enabled, + } + } +} diff --git a/crates/server/src/openai_api/error.rs b/crates/server/src/openai_api/error.rs new file mode 100644 index 000000000..ed38ac892 --- /dev/null +++ b/crates/server/src/openai_api/error.rs @@ -0,0 +1,6 @@ +//! OpenAI 错误映射。 + +/// 格式化错误字符串。 +pub fn format_error(message: &str) -> String { + message.to_string() +} diff --git a/src/codex/serving/http/mod.rs b/crates/server/src/openai_api/mod.rs similarity index 63% rename from src/codex/serving/http/mod.rs rename to crates/server/src/openai_api/mod.rs index b3716412b..049f1f067 100644 --- a/src/codex/serving/http/mod.rs +++ b/crates/server/src/openai_api/mod.rs @@ -1,9 +1,10 @@ +//! OpenAI 兼容 HTTP API。 + pub mod auth; pub mod chat; pub mod diagnostics; -pub mod errors; +pub mod error; pub mod models; pub mod responses; pub mod router; - -pub use self::router::router; +pub mod sse; diff --git a/src/codex/serving/http/models.rs b/crates/server/src/openai_api/models.rs similarity index 73% rename from src/codex/serving/http/models.rs rename to crates/server/src/openai_api/models.rs index dfd0e6702..8ecff7c78 100644 --- a/src/codex/serving/http/models.rs +++ b/crates/server/src/openai_api/models.rs @@ -1,3 +1,5 @@ +//! OpenAI 模型处理器。 + use axum::{ extract::{Path, State}, http::{HeaderMap, StatusCode}, @@ -6,15 +8,14 @@ use axum::{ }; use serde_json::{json, Value}; -use crate::{codex::models::catalog::ModelCatalog, runtime::state::AppState}; +use codex_proxy_core::models::catalog::ModelCatalog; +use codex_proxy_runtime::state::AppState; -use super::{ - auth::authorize_client_api_key, - errors::{missing_client_api_key_response, model_not_found_response}, -}; +use super::auth::authorize_client_api_key; const MODEL_CREATED_TIMESTAMP: i64 = 1_700_000_000; +/// `GET /v1/models` 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(); @@ -35,6 +36,7 @@ pub async fn models(State(state): State, headers: HeaderMap) -> impl I ) } +/// `GET /v1/models/catalog` 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(); @@ -44,6 +46,7 @@ pub async fn model_catalog(State(state): State, headers: HeaderMap) -> (StatusCode::OK, Json(json!(catalog.models()))) } +/// `GET /v1/models/{model_id}` pub async fn model_detail( State(state): State, headers: HeaderMap, @@ -60,6 +63,7 @@ pub async fn model_detail( (StatusCode::OK, Json(openai_model_json(&model_id))) } +/// `GET /v1/models/{model_id}/info` pub async fn model_info( State(state): State, headers: HeaderMap, @@ -76,6 +80,7 @@ pub async fn model_info( (StatusCode::OK, Json(json!(info))) } +/// `GET /debug/models` 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(); @@ -97,3 +102,29 @@ fn openai_model_json(id: &str) -> Value { "owned_by": "openai" }) } + +fn missing_client_api_key_response() -> (StatusCode, Json) { + ( + StatusCode::UNAUTHORIZED, + Json(json!({ + "error": { + "message": "Missing client API key", + "type": "invalid_request_error", + "code": "invalid_api_key" + } + })), + ) +} + +fn model_not_found_response() -> (StatusCode, Json) { + ( + StatusCode::NOT_FOUND, + Json(json!({ + "error": { + "message": "Model not found", + "type": "invalid_request_error", + "code": "model_not_found" + } + })), + ) +} diff --git a/crates/server/src/openai_api/responses.rs b/crates/server/src/openai_api/responses.rs new file mode 100644 index 000000000..48f0f5bf1 --- /dev/null +++ b/crates/server/src/openai_api/responses.rs @@ -0,0 +1,648 @@ +//! OpenAI Responses 处理器。 + +use axum::{ + body::{Body, Bytes}, + extract::State, + http::{ + header::{CACHE_CONTROL, CONTENT_TYPE}, + HeaderMap, StatusCode, + }, + response::{IntoResponse, Response}, + Extension, Json, +}; +use codex_proxy_core::protocol::openai::responses::{ + translate_response_to_codex, translate_response_to_compact, OpenAiResponsesRequest, +}; +use codex_proxy_runtime::{ + services::{ResponseDispatchError, ResponseDispatchStream}, + state::AppState, +}; +use serde_json::{json, Value}; + +use crate::middleware::request_id::RequestId; + +use super::{ + auth::authorize_client_api_key, + models::model_catalog_for_state, + sse::{done_sse_frame, openai_sse_frame}, +}; + +const OPENAI_SUBAGENT_METADATA_KEY: &str = "x-openai-subagent"; +const OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent"; + +/// `POST /v1/responses` +pub async fn responses( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + body: Bytes, +) -> Response { + handle_responses(state, request_id, headers, body, None).await +} + +/// `POST /v1/responses/review` +pub async fn review_responses( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + body: Bytes, +) -> Response { + handle_responses(state, request_id, headers, body, Some("review")).await +} + +/// `POST /v1/responses/compact` +pub async fn compact_responses( + State(state): State, + Extension(request_id): Extension, + headers: HeaderMap, + body: Bytes, +) -> Response { + handle_compact_responses(state, request_id, headers, body).await +} + +async fn handle_responses( + state: AppState, + request_id: RequestId, + headers: HeaderMap, + body: Bytes, + forced_subagent: Option<&str>, +) -> Response { + if !authorize_client_api_key(&state, &headers).await { + return missing_client_api_key_response().into_response(); + } + + let Ok(mut openai_request) = serde_json::from_slice::(&body) else { + return openai_error_response( + StatusCode::BAD_REQUEST, + "Invalid responses request", + "invalid_request_error", + "invalid_request", + ) + .into_response(); + }; + apply_responses_header_context(&mut openai_request, &headers); + if let Some(subagent) = forced_subagent { + force_openai_subagent(&mut openai_request, subagent); + } else if let Some(subagent) = openai_subagent_from_headers(&headers) { + force_openai_subagent(&mut openai_request, &subagent); + } + let model = openai_request.model.trim().to_string(); + if model.is_empty() { + return openai_error_response( + StatusCode::BAD_REQUEST, + "Invalid responses request", + "invalid_request_error", + "invalid_request", + ) + .into_response(); + }; + let catalog = model_catalog_for_state(&state).await; + if !catalog.is_recognized_model_name(&model) { + return model_not_found_response().into_response(); + } + let stream = openai_request.stream; + let codex_request = translate_response_to_codex(openai_request); + + if stream { + return match state + .services + .responses + .stream(request_id.as_str(), codex_request, &model) + .await + { + Ok(stream) => live_event_stream_response(stream), + Err(error) => response_dispatch_stream_error_response(error), + }; + } + + match state + .services + .responses + .complete(request_id.as_str(), codex_request, &model) + .await + { + Ok(body) => (StatusCode::OK, Json(body)).into_response(), + Err(ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore) => { + responses_no_available_accounts_response().into_response() + } + Err(ResponseDispatchError::Upstream(_)) => openai_error_response( + StatusCode::BAD_GATEWAY, + "Upstream Codex request failed", + "server_error", + "upstream_error", + ) + .into_response(), + Err(ResponseDispatchError::QuotaExhausted { + count, + upstream_error, + }) => openai_error_response( + StatusCode::PAYMENT_REQUIRED, + &format!( + "All accounts exhausted ({count} quota-exhausted). Codex upstream error: {upstream_error}" + ), + "server_error", + "upstream_error", + ) + .into_response(), + Err(ResponseDispatchError::RateLimited { + count, + upstream_error, + }) => openai_error_response( + StatusCode::TOO_MANY_REQUESTS, + &format!( + "All accounts exhausted ({count} rate-limited). Codex upstream error: {upstream_error}" + ), + "server_error", + "upstream_error", + ) + .into_response(), + Err(ResponseDispatchError::Expired { + count, + upstream_error, + }) => openai_error_response( + StatusCode::UNAUTHORIZED, + &format!( + "All accounts exhausted ({count} expired). Codex upstream error: {upstream_error}" + ), + "server_error", + "upstream_error", + ) + .into_response(), + Err(ResponseDispatchError::CloudflareChallenge { + count, + upstream_error, + }) => openai_error_response( + StatusCode::BAD_GATEWAY, + &format!( + "All accounts exhausted ({count} cloudflare-challenge). Codex upstream error: {upstream_error}" + ), + "server_error", + "upstream_error", + ) + .into_response(), + Err(ResponseDispatchError::CloudflarePathBlocked { + count, + upstream_error, + }) => openai_error_response( + StatusCode::BAD_GATEWAY, + &format!( + "All accounts exhausted ({count} cloudflare-path-block). Codex upstream error: {upstream_error}" + ), + "server_error", + "upstream_error", + ) + .into_response(), + Err(ResponseDispatchError::ModelUnsupported { + count, + upstream_error, + }) => openai_error_response( + StatusCode::BAD_REQUEST, + &format!( + "All accounts exhausted ({count} model-unsupported). Codex upstream error: {upstream_error}" + ), + "invalid_request_error", + "upstream_error", + ) + .into_response(), + Err( + ResponseDispatchError::InvalidSse(_) + | ResponseDispatchError::MissingCompleted + | ResponseDispatchError::EmptyUpstreamResponse, + ) => openai_error_response( + StatusCode::BAD_GATEWAY, + "Invalid upstream Codex response", + "server_error", + "invalid_upstream_response", + ) + .into_response(), + Err(ResponseDispatchError::Failed(_)) => openai_error_response( + StatusCode::BAD_GATEWAY, + "Upstream Codex response failed", + "server_error", + "upstream_error", + ) + .into_response(), + } +} + +async fn handle_compact_responses( + state: AppState, + request_id: RequestId, + headers: HeaderMap, + body: Bytes, +) -> Response { + if !authorize_client_api_key(&state, &headers).await { + return missing_client_api_key_response().into_response(); + } + + let Ok(openai_request) = serde_json::from_slice::(&body) else { + return openai_error_response( + StatusCode::BAD_REQUEST, + "Invalid responses request", + "invalid_request_error", + "invalid_request", + ) + .into_response(); + }; + let model = openai_request.model.trim().to_string(); + if model.is_empty() { + return openai_error_response( + StatusCode::BAD_REQUEST, + "Invalid responses request", + "invalid_request_error", + "invalid_request", + ) + .into_response(); + }; + let catalog = model_catalog_for_state(&state).await; + if !catalog.is_recognized_model_name(&model) { + return model_not_found_response().into_response(); + } + let compact_request = translate_response_to_compact(openai_request); + + match state + .services + .responses + .compact(request_id.as_str(), compact_request, &model) + .await + { + Ok(body) => (StatusCode::OK, Json(body)).into_response(), + Err(error) => response_dispatch_compact_error_response(error), + } +} + +fn force_openai_subagent(request: &mut OpenAiResponsesRequest, subagent: &str) { + let mut metadata = match request.client_metadata.take() { + Some(Value::Object(metadata)) => metadata, + _ => serde_json::Map::new(), + }; + metadata.insert( + OPENAI_SUBAGENT_METADATA_KEY.to_string(), + Value::String(subagent.to_string()), + ); + request.client_metadata = Some(Value::Object(metadata)); +} + +fn apply_responses_header_context(request: &mut OpenAiResponsesRequest, headers: &HeaderMap) { + fill_string_from_header(&mut request.turn_state, headers, "x-codex-turn-state"); + fill_string_from_header(&mut request.turn_metadata, headers, "x-codex-turn-metadata"); + fill_string_from_header(&mut request.beta_features, headers, "x-codex-beta-features"); + fill_string_from_header( + &mut request.include_timing_metrics, + headers, + "x-responsesapi-include-timing-metrics", + ); + fill_string_from_header(&mut request.version, headers, "version"); + fill_string_from_header(&mut request.codex_window_id, headers, "x-codex-window-id"); + fill_string_from_header( + &mut request.parent_thread_id, + headers, + "x-codex-parent-thread-id", + ); +} + +fn fill_string_from_header(field: &mut Option, headers: &HeaderMap, name: &str) { + if field + .as_deref() + .is_some_and(|value| !value.trim().is_empty()) + { + return; + } + *field = header_string(headers, name); +} + +fn openai_subagent_from_headers(headers: &HeaderMap) -> Option { + header_string(headers, OPENAI_SUBAGENT_HEADER).and_then(|value| { + if matches!( + value.as_str(), + "review" | "compact" | "memory_consolidation" | "collab_spawn" + ) { + Some(value) + } else { + None + } + }) +} + +fn header_string(headers: &HeaderMap, name: &str) -> Option { + headers + .get(name) + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) +} + +fn missing_client_api_key_response() -> (StatusCode, Json) { + openai_error_response( + StatusCode::UNAUTHORIZED, + "Missing client API key", + "invalid_request_error", + "invalid_api_key", + ) +} + +fn model_not_found_response() -> (StatusCode, Json) { + openai_error_response( + StatusCode::NOT_FOUND, + "Model not found", + "invalid_request_error", + "model_not_found", + ) +} + +fn responses_no_available_accounts_response() -> (StatusCode, Json) { + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({ + "type": "error", + "error": { + "type": "server_error", + "code": "no_available_accounts", + "message": "No available accounts. All accounts are expired or rate-limited.", + } + })), + ) +} + +fn responses_dispatch_error_response( + status: StatusCode, + message: &str, +) -> (StatusCode, Json) { + let (error_type, code) = if status == StatusCode::TOO_MANY_REQUESTS { + ("rate_limit_error", "rate_limit_exceeded") + } else if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) { + ("invalid_request_error", "authentication_error") + } else if status.is_client_error() { + ("invalid_request_error", "codex_api_error") + } else { + ("server_error", "codex_api_error") + }; + ( + status, + Json(json!({ + "type": "error", + "error": { + "type": error_type, + "code": code, + "message": message, + } + })), + ) +} + +fn openai_error_response( + status: StatusCode, + message: &str, + error_type: &str, + code: &str, +) -> (StatusCode, Json) { + ( + status, + Json(json!({ + "error": { + "message": message, + "type": error_type, + "code": code + } + })), + ) +} + +fn response_dispatch_compact_error_response(error: ResponseDispatchError) -> Response { + match error { + ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => { + responses_no_available_accounts_response().into_response() + } + ResponseDispatchError::Upstream(_) => responses_dispatch_error_response( + StatusCode::BAD_GATEWAY, + "Upstream Codex request failed", + ) + .into_response(), + ResponseDispatchError::QuotaExhausted { + count, + upstream_error, + } => responses_dispatch_error_response( + StatusCode::PAYMENT_REQUIRED, + &format!( + "All accounts exhausted ({count} quota-exhausted). Codex upstream error: {upstream_error}" + ), + ) + .into_response(), + ResponseDispatchError::RateLimited { + count, + upstream_error, + } => responses_dispatch_error_response( + StatusCode::TOO_MANY_REQUESTS, + &format!( + "All accounts exhausted ({count} rate-limited). Codex upstream error: {upstream_error}" + ), + ) + .into_response(), + ResponseDispatchError::Expired { + count, + upstream_error, + } => responses_dispatch_error_response( + StatusCode::UNAUTHORIZED, + &format!( + "All accounts exhausted ({count} expired). Codex upstream error: {upstream_error}" + ), + ) + .into_response(), + ResponseDispatchError::CloudflareChallenge { + count, + upstream_error, + } => responses_dispatch_error_response( + StatusCode::BAD_GATEWAY, + &format!( + "All accounts exhausted ({count} cloudflare-challenge). Codex upstream error: {upstream_error}" + ), + ) + .into_response(), + ResponseDispatchError::CloudflarePathBlocked { + count, + upstream_error, + } => responses_dispatch_error_response( + StatusCode::BAD_GATEWAY, + &format!( + "All accounts exhausted ({count} cloudflare-path-block). Codex upstream error: {upstream_error}" + ), + ) + .into_response(), + ResponseDispatchError::ModelUnsupported { + count, + upstream_error, + } => responses_dispatch_error_response( + StatusCode::BAD_REQUEST, + &format!( + "All accounts exhausted ({count} model-unsupported). Codex upstream error: {upstream_error}" + ), + ) + .into_response(), + ResponseDispatchError::InvalidSse(_) + | ResponseDispatchError::MissingCompleted + | ResponseDispatchError::EmptyUpstreamResponse => responses_dispatch_error_response( + StatusCode::BAD_GATEWAY, + "Invalid upstream Codex response", + ) + .into_response(), + ResponseDispatchError::Failed(_) => responses_dispatch_error_response( + StatusCode::BAD_GATEWAY, + "Upstream Codex response failed", + ) + .into_response(), + } +} + +fn event_stream_response(mut body: String) -> Response { + ensure_done_sse_frame(&mut body); + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream") + .header(CACHE_CONTROL, "no-cache") + .body(Body::from(body)) + .unwrap_or_else(|_| { + openai_error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "Failed to build stream response", + "server_error", + "stream_response_error", + ) + .into_response() + }) +} + +fn live_event_stream_response(stream: ResponseDispatchStream) -> Response { + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream") + .header(CACHE_CONTROL, "no-cache") + .body(Body::from_stream(stream.body)) + .unwrap_or_else(|_| { + openai_error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "Failed to build stream response", + "server_error", + "stream_response_error", + ) + .into_response() + }) +} + +fn response_dispatch_stream_error_response(error: ResponseDispatchError) -> Response { + let (status, message) = match error { + ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => ( + StatusCode::SERVICE_UNAVAILABLE, + "No active upstream account is available".to_string(), + ), + ResponseDispatchError::Upstream(_) => ( + StatusCode::BAD_GATEWAY, + "Upstream Codex request failed".to_string(), + ), + ResponseDispatchError::QuotaExhausted { + count, + upstream_error, + } => ( + StatusCode::PAYMENT_REQUIRED, + format!( + "All accounts exhausted ({count} quota-exhausted). Codex upstream error: {upstream_error}" + ), + ), + ResponseDispatchError::RateLimited { + count, + upstream_error, + } => ( + StatusCode::TOO_MANY_REQUESTS, + format!( + "All accounts exhausted ({count} rate-limited). Codex upstream error: {upstream_error}" + ), + ), + ResponseDispatchError::Expired { + count, + upstream_error, + } => ( + StatusCode::UNAUTHORIZED, + format!( + "All accounts exhausted ({count} expired). Codex upstream error: {upstream_error}" + ), + ), + ResponseDispatchError::CloudflareChallenge { + count, + upstream_error, + } => ( + StatusCode::BAD_GATEWAY, + format!( + "All accounts exhausted ({count} cloudflare-challenge). Codex upstream error: {upstream_error}" + ), + ), + ResponseDispatchError::CloudflarePathBlocked { + count, + upstream_error, + } => ( + StatusCode::BAD_GATEWAY, + format!( + "No accounts available. All accounts exhausted ({count} cloudflare-path-block). Codex upstream error: {upstream_error}" + ), + ), + ResponseDispatchError::ModelUnsupported { + count, + upstream_error, + } => ( + StatusCode::BAD_REQUEST, + format!( + "No accounts available. All accounts exhausted ({count} model-unsupported). Codex upstream error: {upstream_error}" + ), + ), + ResponseDispatchError::InvalidSse(_) + | ResponseDispatchError::MissingCompleted + | ResponseDispatchError::EmptyUpstreamResponse => ( + StatusCode::BAD_GATEWAY, + "Invalid upstream Codex response".to_string(), + ), + ResponseDispatchError::Failed(_) => ( + StatusCode::BAD_GATEWAY, + "Upstream Codex response failed".to_string(), + ), + }; + event_stream_response(response_failed_sse_event(status, &message)) +} + +fn ensure_done_sse_frame(body: &mut String) { + if body + .trim_end_matches(['\r', '\n']) + .ends_with("data: [DONE]") + { + return; + } + if !body.is_empty() && !body.ends_with('\n') { + body.push('\n'); + } + body.push_str(done_sse_frame()); +} + +fn response_failed_sse_event(status: StatusCode, message: &str) -> String { + let (error_type, code) = if status == StatusCode::TOO_MANY_REQUESTS { + ("rate_limit_error", "rate_limit_exceeded") + } else if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) { + ("invalid_request_error", "authentication_error") + } else if status.is_client_error() { + ("invalid_request_error", "codex_api_error") + } else { + ("server_error", "codex_api_error") + }; + let error = json!({ + "type": error_type, + "code": code, + "message": message, + }); + let data = json!({ + "type": "response.failed", + "response": { + "id": format!("resp_proxy_{}", uuid::Uuid::new_v4().simple()), + "status": "failed", + "error": error, + }, + "error": error, + }); + openai_sse_frame("response.failed", &data.to_string()) +} diff --git a/src/codex/serving/http/router.rs b/crates/server/src/openai_api/router.rs similarity index 71% rename from src/codex/serving/http/router.rs rename to crates/server/src/openai_api/router.rs index 71652630b..9614e8b99 100644 --- a/src/codex/serving/http/router.rs +++ b/crates/server/src/openai_api/router.rs @@ -1,16 +1,20 @@ +//! OpenAI API 路由。 + use axum::{ routing::{get, post}, Router, }; -use crate::runtime::state::AppState; +use codex_proxy_runtime::state::AppState; use super::{ chat::chat_completions, + diagnostics::{diagnostics, fingerprint, upstream}, models::{debug_models, model_catalog, model_detail, model_info, models}, responses::{compact_responses, responses, review_responses}, }; +/// 构造 OpenAI 兼容 API 路由。 pub fn router() -> Router { Router::new() .route("/v1/responses", post(responses)) @@ -22,4 +26,7 @@ pub fn router() -> Router { .route("/v1/models/{model_id}", get(model_detail)) .route("/v1/models/{model_id}/info", get(model_info)) .route("/debug/models", get(debug_models)) + .route("/debug/diagnostics", get(diagnostics)) + .route("/debug/fingerprint", get(fingerprint)) + .route("/debug/upstream", get(upstream)) } diff --git a/crates/server/src/openai_api/sse.rs b/crates/server/src/openai_api/sse.rs new file mode 100644 index 000000000..e5acc157c --- /dev/null +++ b/crates/server/src/openai_api/sse.rs @@ -0,0 +1,27 @@ +//! OpenAI SSE 封装。 + +use codex_proxy_core::protocol::codex::sse::encode_sse_event; + +/// 编码 OpenAI 兼容 SSE 事件帧。 +pub fn openai_sse_frame(event: &str, data: &str) -> String { + encode_sse_event(event, data) +} + +/// 编码 OpenAI 流结束标记。 +pub fn done_sse_frame() -> &'static str { + "data: [DONE]\n\n" +} + +#[cfg(test)] +mod tests { + use super::{done_sse_frame, openai_sse_frame}; + + #[test] + fn openai_sse_frame_should_encode_event_data_and_done_marker() { + assert_eq!( + openai_sse_frame("response.output_text.delta", "{\"delta\":\"hi\"}"), + "event: response.output_text.delta\ndata: {\"delta\":\"hi\"}\n\n" + ); + assert_eq!(done_sse_frame(), "data: [DONE]\n\n"); + } +} diff --git a/crates/server/src/router.rs b/crates/server/src/router.rs new file mode 100644 index 000000000..adae4581e --- /dev/null +++ b/crates/server/src/router.rs @@ -0,0 +1,31 @@ +//! 顶层 HTTP 路由。 + +use std::path::Path; + +use axum::{middleware, Router}; + +use codex_proxy_runtime::state::AppState; + +use crate::{ + admin_api, + middleware::{request_id::attach_request_id, trace::http_trace_layer}, + openai_api, +}; + +/// 默认前端构建产物目录。 +pub const DEFAULT_ASSET_DIST_DIR: &str = "web/dist"; + +/// 构造整个 HTTP 服务路由。 +pub fn router() -> Router { + router_with_assets(DEFAULT_ASSET_DIST_DIR) +} + +/// 使用指定前端构建产物目录构造整个 HTTP 服务路由。 +pub fn router_with_assets(dist_dir: impl AsRef) -> Router { + Router::new() + .merge(openai_api::router::router()) + .merge(admin_api::router()) + .fallback_service(codex_proxy_assets::router::spa_router(dist_dir)) + .layer(http_trace_layer()) + .layer(middleware::from_fn(attach_request_id)) +} diff --git a/crates/server/tests/admin_accounts_routes.rs b/crates/server/tests/admin_accounts_routes.rs new file mode 100644 index 000000000..d14b15002 --- /dev/null +++ b/crates/server/tests/admin_accounts_routes.rs @@ -0,0 +1,3634 @@ +use std::{collections::BTreeMap, fs, sync::Arc}; + +use async_trait::async_trait; +use axum::{ + body::{to_bytes, Body}, + http::{Request, StatusCode}, +}; +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use chrono::{Duration, Utc}; +use codex_proxy_adapters::sqlite::accounts::{NewAccount, SqliteAccountStore}; +use codex_proxy_core::{ + accounts::{model::AccountStatus, ports::AccountStore}, + auth::{ + oauth::{DeviceCode, OAuthError, RefreshFailure, TokenPair}, + ports::{OAuthClient, TokenRefresher}, + }, +}; +use codex_proxy_platform::{ + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, + UsageStatsConfig, WebSocketPoolConfig, + }, + crypto::SecretBox, + identity::ApiKeyHasher, + storage::connect_sqlite, +}; +use codex_proxy_runtime::state::AppState; +use codex_proxy_server::router; +use secrecy::{ExposeSecret, SecretString}; +use serde_json::{json, Value}; +use sqlx::SqlitePool; +use tokio::sync::Mutex; +use tower::util::ServiceExt; +use wiremock::{ + matchers::{header, method, path}, + Mock, MockServer, ResponseTemplate, +}; + +#[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; + 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("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(&pool) + .await + .unwrap(); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([63u8; 32]), + ApiKeyHasher::new([64u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/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["requestId"], "req_accounts_list"); + assert_eq!(body["data"][0]["id"], "acct_corrupt"); + assert_eq!(body["data"][0]["email"], "user@example.com"); +} + +#[tokio::test] +async fn admin_usage_stats_should_return_page_and_summary() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-usage.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + sqlx::query( + "insert into accounts (id, email, label, plan_type, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind("acct_usage") + .bind("usage@example.com") + .bind("primary") + .bind("plus") + .bind("cipher") + .bind("active") + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, image_input_tokens, image_output_tokens, image_request_count, image_request_failed_count, last_used_at) values (?, 3, 1, 21, 8, 5, 7, 2, 1, 0, ?)", + ) + .bind("acct_usage") + .bind("2026-06-18T00:10:00Z") + .execute(&pool) + .await + .unwrap(); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([73u8; 32]), + ApiKeyHasher::new([74u8; 32]), + ); + let app = router::router().with_state(state); + + let page_response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/usage-stats?limit=10") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_usage_page") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let page_status = page_response.status(); + let page_body = response_json(page_response).await; + + assert_eq!(page_status, StatusCode::OK); + assert_eq!(page_body["data"][0]["accountId"], "acct_usage"); + assert_eq!(page_body["data"][0]["email"], "usage@example.com"); + assert_eq!(page_body["data"][0]["requestCount"], 3); + assert_eq!(page_body["data"][0]["inputTokens"], 21); + + let summary_response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/usage-stats/summary") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_usage_summary") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let summary_status = summary_response.status(); + let summary_body = response_json(summary_response).await; + + assert_eq!(summary_status, StatusCode::OK); + assert_eq!(summary_body["data"]["accountCount"], 1); + assert_eq!(summary_body["data"]["requestCount"], 3); + assert_eq!(summary_body["data"]["outputTokens"], 8); +} + +#[tokio::test] +async fn admin_usage_stats_should_require_admin_session_cookie() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-usage-auth.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([125u8; 32]), + ApiKeyHasher::new([126u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/usage-stats") + .header("x-request-id", "req_usage_auth") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert_eq!(body["code"], 40101); + assert_eq!(body["requestId"], "req_usage_auth"); +} + +#[tokio::test] +async fn admin_usage_stats_should_cursor_page_account_usage() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-usage-cursor.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + seed_usage_account( + &pool, + "acct_a", + "a@example.com", + "primary", + "plus", + 3, + 0, + 12, + 5, + 1, + "2026-06-11T00:00:00Z", + ) + .await; + seed_usage_account( + &pool, + "acct_b", + "b@example.com", + "backup", + "free", + 2, + 1, + 7, + 3, + 2, + "2026-06-11T00:01:00Z", + ) + .await; + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([127u8; 32]), + ApiKeyHasher::new([128u8; 32]), + ); + let app = router::router().with_state(state); + + let first_response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/usage-stats?limit=1") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_usage_cursor") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let first_status = first_response.status(); + let first_body = response_json(first_response).await; + + assert_eq!(first_status, StatusCode::OK); + assert_eq!(first_body["code"], 200); + assert_eq!(first_body["requestId"], "req_usage_cursor"); + assert_eq!(first_body["data"].as_array().unwrap().len(), 1); + assert_eq!(first_body["data"][0]["accountId"], "acct_b"); + assert_eq!(first_body["data"][0]["requestCount"], 2); + assert_eq!(first_body["data"][0]["emptyResponseCount"], 1); + assert_eq!(first_body["data"][0]["inputTokens"], 7); + assert_eq!(first_body["page"]["limit"], 1); + let cursor = first_body["page"]["nextCursor"].as_str().unwrap(); + + let second_response = app + .oneshot( + Request::builder() + .method("GET") + .uri(format!("/api/admin/usage-stats?limit=1&cursor={cursor}")) + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let second_status = second_response.status(); + let second_body = response_json(second_response).await; + + assert_eq!(second_status, StatusCode::OK); + assert_eq!(second_body["data"][0]["accountId"], "acct_a"); + assert!(second_body["page"]["nextCursor"].is_null()); +} + +async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { + sqlx::query( + "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", + ) + .bind("admin_1") + .bind("hash") + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .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("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); +} + +#[expect( + clippy::too_many_arguments, + reason = "test fixture keeps usage rows explicit" +)] +async fn seed_usage_account( + pool: &SqlitePool, + id: &str, + email: &str, + label: &str, + plan_type: &str, + request_count: i64, + empty_response_count: i64, + input_tokens: i64, + output_tokens: i64, + cached_tokens: i64, + last_used_at: &str, +) { + sqlx::query( + "insert into accounts (id, email, label, plan_type, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?, 'active', ?, ?)", + ) + .bind(id) + .bind(email) + .bind(label) + .bind(plan_type) + .bind("encrypted") + .bind("2026-06-11T00:00:00Z") + .bind("2026-06-11T00:00:00Z") + .execute(pool) + .await + .unwrap(); + sqlx::query( + "insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)", + ) + .bind(id) + .bind(request_count) + .bind(empty_response_count) + .bind(input_tokens) + .bind(output_tokens) + .bind(cached_tokens) + .bind(last_used_at) + .execute(pool) + .await + .unwrap(); +} + +#[tokio::test] +async fn admin_accounts_lifecycle_should_update_and_delete_accounts() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-accounts-lifecycle.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + sqlx::query( + "insert into accounts (id, email, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)", + ) + .bind("acct_lifecycle") + .bind("life@example.com") + .bind("cipher") + .bind("active") + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(&pool) + .await + .unwrap(); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([75u8; 32]), + ApiKeyHasher::new([76u8; 32]), + ); + let app = router::router().with_state(state); + + let label = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/admin/accounts/acct_lifecycle/label") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(json!({"label": "primary"}).to_string())) + .unwrap(), + ) + .await + .unwrap(); + let label_status = label.status(); + let label_body = response_json(label).await; + + assert_eq!(label_status, StatusCode::OK); + assert_eq!(label_body["data"]["label"], "primary"); + + let status = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/admin/accounts/acct_lifecycle/status") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(json!({"status": "disabled"}).to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status_code = status.status(); + let status_body = response_json(status).await; + + assert_eq!(status_code, StatusCode::OK); + assert_eq!(status_body["data"]["status"], "disabled"); + + let deleted = app + .oneshot( + Request::builder() + .method("DELETE") + .uri("/api/admin/accounts/acct_lifecycle") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let delete_status = deleted.status(); + let delete_body = response_json(deleted).await; + + assert_eq!(delete_status, StatusCode::OK); + assert_eq!(delete_body["data"]["deleted"], true); +} + +#[tokio::test] +async fn admin_accounts_export_should_return_native_account_tokens() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-accounts-export.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([77u8; 32]); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_export".to_string(), + email: Some("export@example.com".to_string()), + account_id: Some("chatgpt_export".to_string()), + user_id: Some("user_export".to_string()), + label: Some("primary".to_string()), + plan_type: Some("plus".to_string()), + access_token: SecretString::new("access-export".to_string().into()), + refresh_token: Some(SecretString::new("refresh-export".to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + secret_box, + ApiKeyHasher::new([78u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/export?ids=acct_export") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_accounts_export") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["requestId"], "req_accounts_export"); + assert_eq!(body["data"]["sourceFormat"], "native"); + assert_eq!(body["data"]["accounts"][0]["id"], "acct_export"); + assert_eq!(body["data"]["accounts"][0]["token"], "access-export"); + assert_eq!( + body["data"]["accounts"][0]["refreshToken"], + "refresh-export" + ); +} + +#[tokio::test] +async fn admin_accounts_import_should_store_native_account_tokens() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-accounts-import.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([79u8; 32]); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool.clone(), + secret_box.clone(), + ApiKeyHasher::new([80u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/import") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_accounts_import") + .body(Body::from( + json!({ + "accounts": [{ + "id": "acct_import", + "email": "import@example.com", + "accountId": "chatgpt_import", + "userId": "user_import", + "label": "secondary", + "planType": "team", + "token": "access-import", + "refreshToken": "refresh-import", + "accessTokenExpiresAt": "2026-06-18T02:00:00Z", + "status": "disabled" + }] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let stored = SqliteAccountStore::new(pool, secret_box) + .get("acct_import") + .await + .unwrap() + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["requestId"], "req_accounts_import"); + assert_eq!(body["data"]["imported"], 1); + assert_eq!(body["data"]["skipped"], 0); + assert_eq!(body["data"]["sourceFormat"], "native"); + assert_eq!(stored.access_token.expose_secret(), "access-import"); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "refresh-import" + ); + assert_eq!(stored.status, AccountStatus::Disabled); +} + +#[tokio::test] +async fn admin_accounts_import_should_require_admin_session_cookie() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-accounts-import-auth.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([111u8; 32]), + ApiKeyHasher::new([112u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/import") + .header("content-type", "application/json") + .header("x-request-id", "req_accounts_import_auth") + .body(Body::from(r#"{"accounts":[]}"#)) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert_eq!(body["code"], 40101); + assert_eq!(body["requestId"], "req_accounts_import_auth"); +} + +#[tokio::test] +async fn admin_accounts_import_should_store_tokens_encrypted_and_list_sanitized_accounts() { + let (app, _state, pool, _dir, _secret_box) = + admin_accounts_test_app("admin-accounts-import-sanitized.sqlite", 113).await; + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/import") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_accounts_import_sanitized") + .body(Body::from( + json!({ + "accounts": [{ + "id": "acct_imported_sanitized", + "email": "user@example.com", + "accountId": "chatgpt-account", + "userId": "chatgpt-user", + "label": "primary", + "planType": "plus", + "token": "access-secret", + "refreshToken": "refresh-secret", + "status": "active" + }] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(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_sanitized") + .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("/api/admin/accounts?limit=10") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_accounts_import_list") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let list_status = list_response.status(); + let list_body = response_json(list_response).await; + + assert_eq!(list_status, StatusCode::OK); + assert_eq!(list_body["data"][0]["id"], "acct_imported_sanitized"); + assert_eq!(list_body["data"][0]["email"], "user@example.com"); + assert!(list_body["data"][0].get("token").is_none()); + assert!(list_body["data"][0].get("refreshToken").is_none()); + assert_eq!(list_body["page"]["limit"], 10); +} + +#[tokio::test] +async fn admin_accounts_import_should_reject_non_native_export_shape() { + let (app, _state, pool, _dir, _secret_box) = + admin_accounts_test_app("admin-accounts-import-external-shape.sqlite", 114).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/import") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_external_import") + .body(Body::from( + json!({ + "type": "external-data", + "version": 1, + "legacy": [], + "accounts": [{ + "platform": "openai", + "type": "oauth", + "credentials": { + "access_token": "Bearer external-access-secret", + "refresh_token": "rt_external", + "email": "team@example.com", + "chatgpt_account_id": "chatgpt-account", + "chatgpt_user_id": "chatgpt-user", + "plan_type": "team" + } + }] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let stored: (i64,) = sqlx::query_as("select count(*) from accounts") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(body["message"], "No importable accounts found"); + assert_eq!(stored.0, 0); +} + +#[tokio::test] +async fn admin_accounts_import_should_reject_native_payload_with_unknown_account_fields() { + let (app, _state, _pool, _dir, _secret_box) = + admin_accounts_test_app("admin-accounts-import-unknown-account.sqlite", 115).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/import") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_native_unknown_account") + .body(Body::from( + json!({ + "accounts": [{ + "id": "acct_native_unknown", + "token": "native-access-secret", + "refreshToken": "native-refresh-secret", + "email": "native@example.com", + "accountId": "native-account", + "userId": "native-user", + "label": "Native Unknown", + "planType": "plus", + "legacyField": "ignored-secret", + "status": "active" + }] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(body["message"], "No importable accounts found"); +} + +#[tokio::test] +async fn admin_accounts_import_should_reject_native_payload_with_unknown_container_fields() { + let (app, _state, _pool, _dir, _secret_box) = + admin_accounts_test_app("admin-accounts-import-unknown-container.sqlite", 116).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/import") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_native_unknown_container") + .body(Body::from( + json!({ + "accounts": [{ + "id": "acct_native_extra_container", + "token": "native-access-secret", + "refreshToken": "native-refresh-secret", + "email": "native@example.com", + "accountId": "native-account", + "userId": "native-user", + "label": "Native", + "planType": "plus", + "status": "active" + }], + "legacyContainer": true + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(body["message"], "No importable accounts found"); +} + +#[tokio::test] +async fn admin_accounts_export_should_return_native_accounts_with_tokens_and_filter_ids() { + let (app, _state, pool, _dir, secret_box) = + admin_accounts_test_app("admin-accounts-export-filter.sqlite", 117).await; + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_export_a".to_string(), + email: Some("export-a@example.com".to_string()), + account_id: Some("chatgpt-export-a".to_string()), + user_id: Some("user-export-a".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new("access-acct_export_a".to_string().into()), + refresh_token: Some(SecretString::new( + "refresh-acct_export_a".to_string().into(), + )), + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + seed_encrypted_account( + &pool, + secret_box, + NewAccount { + id: "acct_export_b".to_string(), + email: Some("export-b@example.com".to_string()), + account_id: Some("chatgpt-export-b".to_string()), + user_id: Some("user-export-b".to_string()), + label: None, + plan_type: Some("team".to_string()), + access_token: SecretString::new("access-acct_export_b".to_string().into()), + refresh_token: Some(SecretString::new( + "refresh-acct_export_b".to_string().into(), + )), + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/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("legacyField").is_none()); +} + +#[tokio::test] +async fn admin_accounts_export_should_reject_unsupported_external_format() { + let (app, _state, _pool, _dir, _secret_box) = + admin_accounts_test_app("admin-accounts-export-external.sqlite", 118).await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/export?format=external") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); +} + +#[tokio::test] +async fn admin_accounts_export_should_reject_unsupported_full_format() { + let (app, _state, _pool, _dir, _secret_box) = + admin_accounts_test_app("admin-accounts-export-full.sqlite", 119).await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/export?format=full") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); +} + +#[tokio::test] +async fn admin_account_cookies_should_store_encrypted_cookie_header() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-cookies.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([85u8; 32]); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_cookies".to_string(), + email: Some("cookies@example.com".to_string()), + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-cookies".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool.clone(), + secret_box, + ApiKeyHasher::new([86u8; 32]), + ); + let app = router::router().with_state(state); + + let set = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/acct_cookies/cookies") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({"cookies":"cf_clearance=clear-secret; __cf_bm=bm-secret"}).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let set_status = set.status(); + let set_body = response_json(set).await; + + assert_eq!(set_status, StatusCode::OK); + assert_eq!( + set_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 = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/acct_cookies/cookies") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let get_status = get.status(); + let get_body = response_json(get).await; + + assert_eq!(get_status, StatusCode::OK); + assert_eq!( + get_body["data"]["cookies"], + "__cf_bm=bm-secret; cf_clearance=clear-secret" + ); + + let delete = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri("/api/admin/accounts/acct_cookies/cookies") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let delete_status = delete.status(); + let delete_body = response_json(delete).await; + + assert_eq!(delete_status, StatusCode::OK); + assert_eq!(delete_body["data"]["deleted"], true); +} + +#[tokio::test] +async fn admin_account_cookies_should_reject_missing_account_and_empty_payload() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-cookies-invalid.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([87u8; 32]); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_cookie_invalid".to_string(), + email: Some("cookie-invalid@example.com".to_string()), + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-cookie-invalid".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + secret_box, + ApiKeyHasher::new([88u8; 32]), + ); + let app = router::router().with_state(state); + + let missing = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/missing/cookies") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(missing.status(), StatusCode::NOT_FOUND); + + let empty = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/acct_cookie_invalid/cookies") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(json!({"cookies": ""}).to_string())) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(empty.status(), StatusCode::BAD_REQUEST); +} + +#[tokio::test] +async fn admin_auth_login_start_should_return_pkce_auth_url_and_state() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-auth-login-start.sqlite"); + 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_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([95u8; 32]), + ApiKeyHasher::new([96u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/auth/login-start") + .header("cookie", "cpr_admin_session=session_1") + .header("host", "127.0.0.1:8080") + .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_status_should_require_admin_session_cookie() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-auth-status-auth.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([105u8; 32]), + ApiKeyHasher::new([106u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/auth/status") + .header("x-request-id", "req_auth_status_no_session") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert_eq!(body["code"], 40101); + assert_eq!(body["requestId"], "req_auth_status_no_session"); +} + +#[tokio::test] +async fn admin_auth_status_should_report_user_and_pool_summary_without_secrets() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-auth-status-summary.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([107u8; 32]); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_auth_active".to_string(), + email: Some("active-auth@example.com".to_string()), + account_id: Some("auth-active-account".to_string()), + user_id: Some("auth-active-user".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new("access-auth-active".to_string().into()), + refresh_token: Some(SecretString::new("refresh-auth-active".to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_auth_disabled".to_string(), + email: Some("disabled-auth@example.com".to_string()), + account_id: Some("auth-disabled-account".to_string()), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-auth-disabled".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Disabled, + }, + ) + .await; + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + secret_box, + ApiKeyHasher::new([108u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/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_eq!(body["data"]["user"]["email"], "active-auth@example.com"); + 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_usage_cookies_and_runtime_pool() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-auth-logout.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([109u8; 32]); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool.clone(), + secret_box.clone(), + ApiKeyHasher::new([110u8; 32]), + ); + let app = router::router().with_state(state.clone()); + let imported = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/import") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({ + "accounts": [{ + "id": "acct_auth_logout", + "email": "logout@example.com", + "planType": "plus", + "token": "access-auth-logout", + "refreshToken": "refresh-auth-logout", + "status": "active" + }] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(imported.status(), StatusCode::OK); + 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("/api/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("/api/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(); + let usage_count: (i64,) = sqlx::query_as("select count(*) from account_usage") + .fetch_one(&pool) + .await + .unwrap(); + let cookie_count: (i64,) = sqlx::query_as("select count(*) from account_cookies") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(account_count.0, 0); + assert_eq!(usage_count.0, 0); + assert_eq!(cookie_count.0, 0); + assert!(state + .services + .account_pool + .acquire("gpt-5.5", Utc::now()) + .await + .is_none()); +} + +#[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", 111, client) + .await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/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", 112, client) + .await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/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( + "device-account", + Some("device-user"), + Some("device@example.com"), + Some("plus"), + ); + 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", 113, client) + .await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/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 = SqliteAccountStore::new(pool.clone(), SecretBox::new([113u8; 32])) + .get( + state + .services + .account_pool + .acquire("gpt-5.5", Utc::now()) + .await + .unwrap() + .account + .id + .as_str(), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(stored.email.as_deref(), Some("device@example.com")); + assert_eq!(stored.account_id.as_deref(), Some("device-account")); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "device-refresh-secret" + ); +} + +#[tokio::test] +async fn admin_auth_login_start_should_use_configured_oauth_authorize_endpoint() { + let dir = tempfile::tempdir().unwrap(); + let db = dir + .path() + .join("admin-auth-login-start-configured-oauth.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + let mut config = test_config(url); + config.auth.oauth_client_id = "app_configured_client".to_string(); + config.auth.oauth_auth_endpoint = "https://auth.example.test/oauth/authorize".to_string(); + config.auth.oauth_token_endpoint = "https://auth.example.test/oauth/token".to_string(); + let state = AppState::with_pool_secret_and_api_key_hasher( + config, + pool, + SecretBox::new([114u8; 32]), + ApiKeyHasher::new([115u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/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_configured_oauth") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json(response).await; + let auth_url = body["data"]["authUrl"].as_str().unwrap(); + + assert!(auth_url.starts_with("https://auth.example.test/oauth/authorize?")); + assert!(auth_url.contains("client_id=app_configured_client")); +} + +#[tokio::test] +async fn admin_auth_code_relay_should_exchange_code_and_import_account() { + let token = test_jwt( + "pkce-account", + Some("pkce-user"), + Some("pkce@example.com"), + Some("plus"), + ); + 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("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", 116, client) + .await; + let start = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/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("/api/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(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + 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 acquired = state + .services + .account_pool + .acquire("gpt-5.5", Utc::now()) + .await + .unwrap() + .account; + 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", + 117, + client, + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/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); +} + +#[tokio::test] +async fn admin_auth_callback_should_exchange_code_and_redirect_to_return_host() { + let token = test_jwt( + "callback-account", + Some("callback-user"), + Some("callback@example.com"), + Some("plus"), + ); + 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", 118, client).await; + let start = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/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 count: (i64,) = sqlx::query_as("select count(*) from accounts where account_id = ?") + .bind("callback-account") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(count.0, 1); +} + +#[tokio::test] +async fn admin_account_health_check_should_skip_account_without_refresh_token() { + let upstream = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/usage")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "rate_limit": { + "primary_window": { + "used_percent": 1 + } + } + }))) + .expect(0) + .mount(&upstream) + .await; + + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-health-no-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([93u8; 32]); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_health".to_string(), + email: Some("health@example.com".to_string()), + account_id: Some("chatgpt-health".to_string()), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-health".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + let mut config = test_config(url); + config.api.base_url = upstream.uri(); + let state = AppState::with_pool_secret_and_api_key_hasher( + config, + pool, + secret_box, + ApiKeyHasher::new([94u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/health-check") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({ + "ids": ["acct_health"], + "concurrency": 1, + "staggerMs": 500 + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["data"]["summary"]["total"], 1); + assert_eq!(body["data"]["summary"]["alive"], 0); + assert_eq!(body["data"]["summary"]["dead"], 0); + assert_eq!(body["data"]["summary"]["skipped"], 1); + assert_eq!(body["data"]["results"][0]["id"], "acct_health"); + assert_eq!(body["data"]["results"][0]["result"], "skipped"); + assert_eq!(body["data"]["results"][0]["status"], "active"); + assert_eq!(body["data"]["results"][0]["error"], "no refresh token"); +} + +#[tokio::test] +async fn admin_accounts_health_check_should_refresh_oauth_without_touching_codex_backend() { + let upstream = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/usage")) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .mount(&upstream) + .await; + + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-health-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([124u8; 32]); + for (id, refresh_token) in [ + ("acct_health_alive", "refresh-acct_health_alive"), + ("acct_health_dead", "refresh-acct_health_dead"), + ] { + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: id.to_string(), + email: Some(format!("{id}@example.com")), + account_id: Some(format!("old-{id}")), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new(format!("old-access-{id}").into()), + refresh_token: Some(SecretString::new(refresh_token.to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + } + let calls = Arc::new(Mutex::new(Vec::new())); + let mut config = test_config(url); + config.api.base_url = upstream.uri(); + let state = AppState::with_pool_secret_api_key_hasher_and_token_refresher( + config, + pool.clone(), + secret_box.clone(), + ApiKeyHasher::new([124u8; 32]), + HealthCheckTokenRefresher { + calls: calls.clone(), + }, + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/health-check") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_health_refresh") + .body(Body::from( + json!({ + "ids": ["acct_health_alive", "acct_health_dead"], + "concurrency": 2, + "staggerMs": 500 + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + 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 mut refresh_calls = calls.lock().await.clone(); + refresh_calls.sort(); + assert_eq!( + refresh_calls, + vec![ + "refresh-acct_health_alive".to_string(), + "refresh-acct_health_dead".to_string(), + ] + ); + let serialized = serde_json::to_string(&body).unwrap(); + assert!(!serialized.contains("new-health-access")); + assert!(!serialized.contains("new-health-refresh")); + let dead_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_health_dead") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(dead_status.0, "expired"); +} + +#[tokio::test] +async fn admin_accounts_health_check_should_reject_unsupported_stagger_ms_field() { + let (app, _state, _pool, _dir, _secret_box) = + admin_accounts_test_app("admin-account-health-unsupported-field.sqlite", 125).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/health-check") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"stagger_ms":1000}"#)) + .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_account_quota_should_fetch_usage_store_quota_and_not_return_secrets() { + let upstream = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/usage")) + .and(header("authorization", "Bearer access-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, + "reset_at": 1770000400, + "limit_window_seconds": 3600 + } + } + }))) + .expect(1) + .mount(&upstream) + .await; + + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-quota.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([95u8; 32]); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_quota".to_string(), + email: Some("quota@example.com".to_string()), + account_id: Some("chatgpt-quota".to_string()), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-quota".to_string().into()), + refresh_token: Some(SecretString::new("refresh-quota".to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + let mut config = test_config(url); + config.api.base_url = upstream.uri(); + let state = AppState::with_pool_secret_and_api_key_hasher( + config, + pool.clone(), + secret_box, + ApiKeyHasher::new([96u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/acct_quota/quota") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_quota") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + + assert_eq!(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-quota")); + assert!(!serialized.contains("refresh-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-quota")); +} + +#[tokio::test] +async fn admin_account_quota_should_return_bad_gateway_when_usage_fetch_fails() { + let upstream = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/usage")) + .and(header("authorization", "Bearer access-quota-fail")) + .respond_with(ResponseTemplate::new(500).set_body_json(json!({ + "error": { + "message": "quota unavailable" + } + }))) + .expect(1) + .mount(&upstream) + .await; + + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-quota-fail.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([97u8; 32]); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_quota_fail".to_string(), + email: Some("quota-fail@example.com".to_string()), + account_id: Some("chatgpt-quota-fail".to_string()), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-quota-fail".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + let mut config = test_config(url); + config.api.base_url = upstream.uri(); + let state = AppState::with_pool_secret_and_api_key_hasher( + config, + pool, + secret_box, + ApiKeyHasher::new([98u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/acct_quota_fail/quota") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_quota_fail") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::BAD_GATEWAY); + assert_eq!(body["code"], 50201); + assert_eq!(body["message"], "Failed to fetch quota from Codex API"); + assert!(body["data"]["error"] + .as_str() + .is_some_and(|error| error.contains("quota unavailable"))); + assert_eq!(body["requestId"], "req_quota_fail"); +} + +#[tokio::test] +async fn admin_account_quota_should_reject_inactive_account_without_calling_upstream() { + let upstream = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/usage")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .expect(0) + .mount(&upstream) + .await; + + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-quota-inactive.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([99u8; 32]); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_quota_inactive".to_string(), + email: Some("quota-inactive@example.com".to_string()), + account_id: Some("chatgpt-quota-inactive".to_string()), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-quota-inactive".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Disabled, + }, + ) + .await; + let mut config = test_config(url); + config.api.base_url = upstream.uri(); + let state = AppState::with_pool_secret_and_api_key_hasher( + config, + pool.clone(), + secret_box, + ApiKeyHasher::new([100u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/acct_quota_inactive/quota") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_quota_inactive") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let quota_json: (Option,) = + sqlx::query_as("select quota_json from accounts where id = ?") + .bind("acct_quota_inactive") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(body["code"], 40901); + assert_eq!(body["message"], "Account is disabled, cannot query quota"); + assert_eq!(body["requestId"], "req_quota_inactive"); + assert!(quota_json.0.is_none()); +} + +#[tokio::test] +async fn admin_account_quota_should_return_store_error_when_quota_persistence_fails() { + let upstream = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/usage")) + .and(header("authorization", "Bearer access-quota-store-fail")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "plan_type": "plus", + "rate_limit": { + "allowed": true, + "limit_reached": false, + "primary_window": { + "used_percent": 40, + "reset_at": 1770000500, + "limit_window_seconds": 3600 + } + } + }))) + .expect(1) + .mount(&upstream) + .await; + + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-quota-store-fail.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([101u8; 32]); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_quota_store_fail".to_string(), + email: Some("quota-store-fail@example.com".to_string()), + account_id: Some("chatgpt-quota-store-fail".to_string()), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-quota-store-fail".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + sqlx::query( + "create trigger quota_write_denied before update of quota_json on accounts begin select raise(abort, 'quota write denied'); end", + ) + .execute(&pool) + .await + .unwrap(); + let mut config = test_config(url); + config.api.base_url = upstream.uri(); + let state = AppState::with_pool_secret_and_api_key_hasher( + config, + pool.clone(), + secret_box, + ApiKeyHasher::new([102u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/acct_quota_store_fail/quota") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_quota_store_fail") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let quota_json: (Option,) = + sqlx::query_as("select quota_json from accounts where id = ?") + .bind("acct_quota_store_fail") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(body["code"], 50001); + assert_eq!(body["message"], "Failed to store account quota"); + assert_eq!(body["requestId"], "req_quota_store_fail"); + assert!(quota_json.0.is_none()); +} + +#[tokio::test] +async fn admin_accounts_batch_status_should_update_found_accounts_and_report_invalid_requests() { + let (app, state, pool, _dir, secret_box) = + admin_accounts_test_app("admin-accounts-batch-status-route.sqlite", 120).await; + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_batch_status_a".to_string(), + email: Some("batch-a@example.com".to_string()), + account_id: Some("batch-a".to_string()), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-batch-a".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + seed_encrypted_account( + &pool, + secret_box, + NewAccount { + id: "acct_batch_status_b".to_string(), + email: Some("batch-b@example.com".to_string()), + account_id: Some("batch-b".to_string()), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-batch-b".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + let restored = state.restore_account_pool_from_repository().await.unwrap(); + assert_eq!(restored, 2); + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/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 invalid = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/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("/api/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_refresh_should_update_tokens_and_runtime_pool_without_returning_secrets() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-refresh-route.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([121u8; 32]); + let refreshed_access_token = test_jwt( + "refresh-account", + Some("refresh-user"), + Some("refresh@example.com"), + Some("plus"), + ); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_refresh_route".to_string(), + email: Some("old-refresh@example.com".to_string()), + account_id: Some("old-refresh-account".to_string()), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("old-access-token".to_string().into()), + refresh_token: Some(SecretString::new("old-refresh-token".to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + let state = AppState::with_pool_secret_api_key_hasher_and_token_refresher( + test_config(url), + pool.clone(), + secret_box.clone(), + ApiKeyHasher::new([122u8; 32]), + StaticTokenRefresher { + result: Ok(TokenPair { + access_token: refreshed_access_token.clone(), + refresh_token: Some("new-admin-refresh-rt".to_string()), + }), + calls: Arc::new(Mutex::new(Vec::new())), + }, + ); + let app = router::router().with_state(state.clone()); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/acct_refresh_route/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_route"); + assert_eq!(body["data"]["result"], "alive"); + assert_eq!(body["data"]["previousStatus"], "active"); + let serialized = serde_json::to_string(&body).unwrap(); + assert!(!serialized.contains(&refreshed_access_token)); + assert!(!serialized.contains("new-admin-refresh-rt")); + + let stored = SqliteAccountStore::new(pool.clone(), secret_box) + .get("acct_refresh_route") + .await + .unwrap() + .unwrap(); + assert_eq!(stored.access_token.expose_secret(), refreshed_access_token); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "new-admin-refresh-rt" + ); + let raw: (String, String) = sqlx::query_as( + "select access_token_cipher, refresh_token_cipher from accounts where id = ?", + ) + .bind("acct_refresh_route") + .fetch_one(&pool) + .await + .unwrap(); + assert!(raw.0.starts_with("v1:")); + assert!(!raw.0.contains("old-access-token")); + assert!(!raw.0.contains(&refreshed_access_token)); + assert!(raw.1.starts_with("v1:")); + assert!(!raw.1.contains("new-admin-refresh-rt")); + let acquired = state + .services + .account_pool + .acquire("gpt-5.5", Utc::now()) + .await + .unwrap() + .account; + assert_eq!(acquired.access_token, refreshed_access_token); +} + +#[tokio::test] +async fn admin_account_refresh_should_mark_invalid_refresh_token_as_expired() { + let dir = tempfile::tempdir().unwrap(); + let db = dir + .path() + .join("admin-account-refresh-invalid-route.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([123u8; 32]); + seed_encrypted_account( + &pool, + secret_box, + NewAccount { + id: "acct_refresh_invalid_route".to_string(), + email: Some("refresh-invalid@example.com".to_string()), + account_id: Some("refresh-invalid-account".to_string()), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("invalid-old-access".to_string().into()), + refresh_token: Some(SecretString::new( + "invalid-refresh-token".to_string().into(), + )), + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + let state = AppState::with_pool_secret_api_key_hasher_and_token_refresher( + test_config(url), + pool.clone(), + SecretBox::new([123u8; 32]), + ApiKeyHasher::new([124u8; 32]), + StaticTokenRefresher { + result: Err(RefreshFailure::InvalidGrant), + calls: Arc::new(Mutex::new(Vec::new())), + }, + ); + let restored = state.restore_account_pool_from_repository().await.unwrap(); + assert_eq!(restored, 1); + let app = router::router().with_state(state.clone()); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/acct_refresh_invalid_route/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 stored_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_refresh_invalid_route") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(stored_status.0, "expired"); + assert!(state + .services + .account_pool + .acquire("gpt-5.5", Utc::now()) + .await + .is_none()); +} + +#[tokio::test] +async fn admin_account_reset_usage_should_clear_local_counters_and_pool_last_used() { + let (app, state, pool, _dir, secret_box) = + admin_accounts_test_app("admin-account-reset-usage-route.sqlite", 125).await; + let window_reset_at = Utc::now() + Duration::minutes(5); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_reset_usage_route".to_string(), + email: Some("reset-usage@example.com".to_string()), + account_id: Some("reset-usage-account".to_string()), + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-reset-usage".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + sqlx::query( + "insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds, last_used_at) values (?, 7, 11, 13, 17, 5, 19, 23, 29, ?, ?, 300, ?)", + ) + .bind("acct_reset_usage_route") + .bind("2026-06-12T12:30:00Z") + .bind(window_reset_at.to_rfc3339()) + .bind("2026-06-12T12:00:00Z") + .execute(&pool) + .await + .unwrap(); + state.restore_account_pool_from_repository().await.unwrap(); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/acct_reset_usage_route/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_route"); + assert_eq!(body["data"]["reset"], true); + type ResetUsageRow = ( + i64, + i64, + i64, + i64, + i64, + i64, + i64, + i64, + Option, + Option, + ); + let usage: ResetUsageRow = sqlx::query_as( + "select request_count, input_tokens, output_tokens, cached_tokens, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_reset_at, last_used_at from account_usage where account_id = ?", + ) + .bind("acct_reset_usage_route") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!( + usage, + ( + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + Some(window_reset_at.to_rfc3339()), + None + ) + ); + let pool_account = SqliteAccountStore::new(pool, secret_box) + .get_pool_account("acct_reset_usage_route") + .await + .unwrap() + .unwrap(); + assert_eq!(pool_account.request_count, 0); + assert_eq!(pool_account.window_request_count, 0); + assert_eq!(pool_account.window_reset_at, Some(window_reset_at)); + assert!(pool_account.last_used_at.is_none()); +} + +#[tokio::test] +async fn admin_account_create_should_derive_claims_and_store_encrypted_tokens() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-create.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([91u8; 32]); + let token = test_jwt( + "jwt-account", + Some("jwt-user"), + Some("jwt@example.com"), + Some("team"), + ); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool.clone(), + secret_box.clone(), + ApiKeyHasher::new([92u8; 32]), + ); + let app = router::router().with_state(state.clone()); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/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": format!("Bearer {token}"), + "refreshToken": "manual-refresh-secret", + "status": "disabled" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + let created_id = body["data"]["id"].as_str().unwrap(); + assert_ne!(created_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 = SqliteAccountStore::new(pool.clone(), secret_box) + .get(created_id) + .await + .unwrap() + .unwrap(); + assert_eq!(stored.access_token.expose_secret(), &token); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "manual-refresh-secret" + ); + + let raw_cipher: (String,) = + sqlx::query_as("select access_token_cipher from accounts where id = ?") + .bind(created_id) + .fetch_one(&pool) + .await + .unwrap(); + assert!(!raw_cipher.0.contains(&token)); + let raw_refresh_cipher: (String,) = + sqlx::query_as("select refresh_token_cipher from accounts where id = ?") + .bind(created_id) + .fetch_one(&pool) + .await + .unwrap(); + assert!(raw_cipher.0.starts_with("v1:")); + assert!(raw_refresh_cipher.0.starts_with("v1:")); + assert!(!raw_refresh_cipher.0.contains("manual-refresh-secret")); + + let acquired = state + .services + .account_pool + .acquire("gpt-5.5", Utc::now()) + .await + .unwrap() + .account; + assert_eq!(acquired.id, created_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_create_should_reject_missing_invalid_expired_or_unbound_tokens() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-create-invalid.sqlite"); + 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_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([93u8; 32]), + ApiKeyHasher::new([94u8; 32]), + ); + let app = router::router().with_state(state); + + let cases = [ + ("missing tokens", json!({})), + ("invalid jwt", json!({"token": "not-a-jwt"})), + ( + "expired jwt", + json!({"token": test_jwt_with_exp( + Some("expired-account"), + Some("expired-user"), + Some("expired@example.com"), + Some("plus"), + 1_600_000_000, + )}), + ), + ( + "missing account claim", + json!({"token": test_jwt_with_exp( + None, + Some("claimless-user"), + Some("claimless@example.com"), + Some("free"), + 4_102_444_800, + )}), + ), + ]; + + for (name, payload) in cases { + let response = post_admin_account(&app, payload).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{name}"); + } +} + +#[tokio::test] +async fn admin_account_manual_create_refresh_only_should_exchange_rotate_and_sync_pool() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-refresh-only.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([95u8; 32]); + let token = test_jwt( + "rt-account", + Some("rt-user"), + Some("rt@example.com"), + Some("plus"), + ); + let calls = Arc::new(Mutex::new(Vec::new())); + let state = AppState::with_pool_secret_api_key_hasher_and_token_refresher( + test_config(url), + pool.clone(), + secret_box.clone(), + ApiKeyHasher::new([96u8; 32]), + StaticTokenRefresher { + result: Ok(TokenPair { + access_token: token.clone(), + refresh_token: Some("rotated-refresh".to_string()), + }), + calls: calls.clone(), + }, + ); + let app = router::router().with_state(state.clone()); + + let response = post_admin_account( + &app, + json!({ + "refreshToken": "initial-refresh", + "email": "caller@example.com", + "planType": "caller-plan" + }), + ) + .await; + let status = response.status(); + let body = response_json(response).await; + let created_id = body["data"]["id"].as_str().unwrap(); + + assert_eq!(status, StatusCode::OK); + 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 stored = SqliteAccountStore::new(pool, secret_box) + .get(created_id) + .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 + .account_pool + .acquire("gpt-5.5", Utc::now()) + .await + .unwrap() + .account; + assert_eq!(acquired.id, created_id); + assert_eq!(acquired.refresh_token.as_deref(), Some("rotated-refresh")); +} + +#[tokio::test] +async fn admin_account_manual_create_refresh_only_should_preserve_input_refresh_without_rotation() { + let dir = tempfile::tempdir().unwrap(); + let db = dir + .path() + .join("admin-account-refresh-preserve-input.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([97u8; 32]); + let token = test_jwt( + "rt-preserve-account", + Some("rt-preserve-user"), + Some("preserve@example.com"), + Some("free"), + ); + let state = AppState::with_pool_secret_api_key_hasher_and_token_refresher( + test_config(url), + pool.clone(), + secret_box.clone(), + ApiKeyHasher::new([98u8; 32]), + StaticTokenRefresher { + result: Ok(TokenPair { + access_token: token, + refresh_token: None, + }), + calls: Arc::new(Mutex::new(Vec::new())), + }, + ); + let app = router::router().with_state(state); + + let response = post_admin_account(&app, json!({"refreshToken": "preserved-refresh"})).await; + let body = response_json(response).await; + let created_id = body["data"]["id"].as_str().unwrap(); + + let stored = SqliteAccountStore::new(pool, secret_box) + .get(created_id) + .await + .unwrap() + .unwrap(); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "preserved-refresh" + ); +} + +#[tokio::test] +async fn admin_account_manual_create_should_update_existing_and_preserve_refresh() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-update-existing.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([99u8; 32]); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool.clone(), + secret_box.clone(), + ApiKeyHasher::new([100u8; 32]), + ); + let app = router::router().with_state(state.clone()); + + let first_token = test_jwt( + "team-account", + Some("same-user"), + Some("first@example.com"), + Some("free"), + ); + let first_response = post_admin_account( + &app, + json!({ + "token": first_token, + "refreshToken": "first-refresh" + }), + ) + .await; + 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( + "team-account", + Some("same-user"), + Some("second@example.com"), + Some("team"), + ); + let second_response = post_admin_account(&app, json!({"token": second_token})).await; + let status = second_response.status(); + let second_body = response_json(second_response).await; + + assert_eq!(status, StatusCode::OK); + 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 stored = SqliteAccountStore::new(pool, secret_box) + .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" + ); + + let acquired = state + .services + .account_pool + .acquire("gpt-5.5", Utc::now()) + .await + .unwrap() + .account; + 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 admin_account_manual_create_refresh_only_should_preserve_existing_refresh_without_rotation( +) { + let dir = tempfile::tempdir().unwrap(); + let db = dir + .path() + .join("admin-account-refresh-preserve-existing.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([101u8; 32]); + let refreshed_token = test_jwt( + "rt-existing-account", + Some("rt-existing-user"), + Some("second@example.com"), + Some("team"), + ); + let state = AppState::with_pool_secret_api_key_hasher_and_token_refresher( + test_config(url), + pool.clone(), + secret_box.clone(), + ApiKeyHasher::new([102u8; 32]), + StaticTokenRefresher { + result: Ok(TokenPair { + access_token: refreshed_token.clone(), + refresh_token: None, + }), + calls: Arc::new(Mutex::new(Vec::new())), + }, + ); + let app = router::router().with_state(state); + let first_token = test_jwt( + "rt-existing-account", + Some("rt-existing-user"), + Some("first@example.com"), + Some("free"), + ); + let first_response = post_admin_account( + &app, + json!({ + "token": first_token, + "refreshToken": "old-refresh" + }), + ) + .await; + 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 stored = SqliteAccountStore::new(pool, secret_box) + .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_auth_file_store_encrypted_and_sync_pool() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-import-cli.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([103u8; 32]); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool.clone(), + secret_box.clone(), + ApiKeyHasher::new([104u8; 32]), + ); + let app = router::router().with_state(state.clone()); + let codex_home = dir.path().join("codex-home"); + fs::create_dir_all(&codex_home).unwrap(); + let token = test_jwt( + "cli-account", + Some("cli-user"), + Some("cli@example.com"), + Some("plus"), + ); + 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("/api/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(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["requestId"], "req_import_cli"); + 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 = SqliteAccountStore::new(pool.clone(), secret_box) + .list(None, 10) + .await + .unwrap() + .items + .remove(0); + assert_ne!(stored.id, "caller-id"); + assert_eq!(stored.email.as_deref(), Some("cli@example.com")); + assert_eq!(stored.account_id.as_deref(), Some("cli-account")); + assert_eq!(stored.user_id.as_deref(), Some("cli-user")); + assert_eq!(stored.label, None); + assert_eq!(stored.plan_type.as_deref(), Some("plus")); + assert_eq!(stored.access_token.expose_secret(), &token); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "cli-refresh-secret" + ); + let raw: (String, String) = sqlx::query_as( + "select access_token_cipher, refresh_token_cipher from accounts where id = ?", + ) + .bind(&stored.id) + .fetch_one(&pool) + .await + .unwrap(); + assert!(raw.0.starts_with("v1:")); + assert!(!raw.0.contains(&token)); + assert!(raw.1.starts_with("v1:")); + assert!(!raw.1.contains("cli-refresh-secret")); + + let acquired = state + .services + .account_pool + .acquire("gpt-5.5", Utc::now()) + .await + .unwrap() + .account; + assert_eq!(acquired.id, stored.id); +} + +#[tokio::test] +async fn admin_account_quota_warnings_should_require_admin_session_cookie() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-quota-warnings-auth.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([87u8; 32]), + ApiKeyHasher::new([88u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/quota-warnings") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); +} + +#[tokio::test] +async fn admin_account_quota_warnings_should_return_threshold_matches_from_cached_quota() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-quota-warnings.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([89u8; 32]); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_warn".to_string(), + email: Some("warn@example.com".to_string()), + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-warn".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_quiet".to_string(), + email: Some("quiet@example.com".to_string()), + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-quiet".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + sqlx::query( + "update accounts set quota_json = ?, quota_fetched_at = ?, updated_at = ? where id = ?", + ) + .bind( + json!({ + "rate_limit": { + "used_percent": 85, + "reset_at": 1770000100 + }, + "secondary_rate_limit": { + "used_percent": 91, + "reset_at": 1770000200 + } + }) + .to_string(), + ) + .bind("2026-06-13T00:00:00Z") + .bind("2026-06-13T00:00:00Z") + .bind("acct_warn") + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "update accounts set quota_json = ?, quota_fetched_at = ?, updated_at = ? where id = ?", + ) + .bind( + json!({ + "rate_limit": { + "used_percent": 25, + "reset_at": 1770000300 + }, + "secondary_rate_limit": null + }) + .to_string(), + ) + .bind("2026-06-13T01:00:00Z") + .bind("2026-06-13T01:00:00Z") + .bind("acct_quiet") + .execute(&pool) + .await + .unwrap(); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + secret_box, + ApiKeyHasher::new([90u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/quota-warnings") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["data"]["updatedAt"], "2026-06-13T00:00:00+00:00"); + let warnings = body["data"]["warnings"].as_array().unwrap(); + assert_eq!(warnings.len(), 2); + assert!(warnings + .iter() + .all(|warning| warning["accountId"] == "acct_warn")); + assert!(warnings.iter().any(|warning| { + warning["window"] == "primary" + && warning["level"] == "warning" + && warning["usedPercent"] == 85.0 + && warning["resetAt"] == 1770000100 + })); + assert!(warnings.iter().any(|warning| { + warning["window"] == "secondary" + && warning["level"] == "critical" + && warning["usedPercent"] == 91.0 + && warning["resetAt"] == 1770000200 + })); +} + +#[tokio::test] +async fn admin_account_quota_warnings_should_ignore_invalid_and_below_threshold_snapshots() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-account-quota-warnings-edge.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([85u8; 32]); + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_invalid_quota".to_string(), + email: Some("invalid-quota@example.com".to_string()), + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-invalid-quota".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + seed_encrypted_account( + &pool, + secret_box.clone(), + NewAccount { + id: "acct_below_threshold".to_string(), + email: Some("below-threshold@example.com".to_string()), + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-below-threshold".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }, + ) + .await; + sqlx::query( + "update accounts set quota_json = ?, quota_fetched_at = ?, updated_at = ? where id = ?", + ) + .bind("{not valid json") + .bind("2026-06-13T00:00:00Z") + .bind("2026-06-13T00:00:00Z") + .bind("acct_invalid_quota") + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "update accounts set quota_json = ?, quota_fetched_at = ?, updated_at = ? where id = ?", + ) + .bind( + json!({ + "rate_limit": { + "used_percent": 79, + "reset_at": 1770000600 + }, + "secondary_rate_limit": { + "used_percent": 79, + "reset_at": 1770000700 + } + }) + .to_string(), + ) + .bind("2026-06-13T01:00:00Z") + .bind("2026-06-13T01:00:00Z") + .bind("acct_below_threshold") + .execute(&pool) + .await + .unwrap(); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + secret_box, + ApiKeyHasher::new([86u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/quota-warnings") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert!(body["data"]["warnings"].as_array().unwrap().is_empty()); + assert!(body["data"]["updatedAt"].is_null()); +} + +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 post_admin_account(app: &axum::Router, payload: Value) -> axum::response::Response { + app.clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/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 admin_accounts_test_app( + db_name: &str, + key_byte: u8, +) -> ( + axum::Router, + AppState, + SqlitePool, + tempfile::TempDir, + SecretBox, +) { + 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 secret_box = SecretBox::new([key_byte; 32]); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool.clone(), + secret_box.clone(), + ApiKeyHasher::new([key_byte; 32]), + ); + let app = router::router().with_state(state.clone()); + (app, state, pool, dir, secret_box) +} + +async fn admin_accounts_test_app_with_oauth_client( + db_name: &str, + key_byte: u8, + oauth_client: C, +) -> (axum::Router, AppState, SqlitePool, tempfile::TempDir) +where + C: OAuthClient + TokenRefresher + Clone, +{ + 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 = router::router().with_state(state.clone()); + (app, state, pool, dir) +} + +async fn seed_encrypted_account(pool: &SqlitePool, secret_box: SecretBox, account: NewAccount) { + SqliteAccountStore::new(pool.clone(), secret_box) + .insert(account) + .await + .unwrap(); +} + +fn test_jwt( + account_id: &str, + user_id: Option<&str>, + email: Option<&str>, + plan_type: Option<&str>, +) -> String { + test_jwt_with_exp(Some(account_id), user_id, email, plan_type, 4_102_444_800) +} + +fn test_jwt_with_exp( + account_id: Option<&str>, + user_id: Option<&str>, + email: Option<&str>, + plan_type: Option<&str>, + exp: i64, +) -> String { + let header = json!({"alg": "none", "typ": "JWT"}); + let payload = json!({ + "exp": exp, + "https://api.openai.com/auth": { + "chatgpt_account_id": account_id, + "chatgpt_user_id": user_id, + "chatgpt_plan_type": plan_type, + }, + "https://api.openai.com/profile": { + "email": email, + } + }); + format!("{}.{}.", jwt_part(&header), jwt_part(&payload)) +} + +#[derive(Clone)] +struct StaticTokenRefresher { + result: Result, + calls: Arc>>, +} + +#[derive(Clone)] +struct HealthCheckTokenRefresher { + calls: Arc>>, +} + +#[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 HealthCheckTokenRefresher { + async fn refresh(&self, refresh_token: &str) -> Result { + self.calls.lock().await.push(refresh_token.to_string()); + match refresh_token { + "refresh-acct_health_alive" => Ok(TokenPair { + access_token: test_jwt( + "health-alive-account", + Some("health-user"), + Some("health-alive@example.com"), + Some("plus"), + ), + refresh_token: Some("new-health-refresh".to_string()), + }), + "refresh-acct_health_dead" => Err(RefreshFailure::InvalidGrant), + _ => Err(RefreshFailure::Transport), + } + } +} + +#[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 jwt_part(value: &Value) -> String { + URL_SAFE_NO_PAD.encode(serde_json::to_vec(value).unwrap()) +} + +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(), + oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), + oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), + oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), + }, + 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, + }, + ws_pool: WebSocketPoolConfig::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + }, + } +} diff --git a/tests/admin/api_contract.rs b/crates/server/tests/admin_api_contract.rs similarity index 81% rename from tests/admin/api_contract.rs rename to crates/server/tests/admin_api_contract.rs index 1a807ff81..482dcc0d7 100644 --- a/tests/admin/api_contract.rs +++ b/crates/server/tests/admin_api_contract.rs @@ -3,15 +3,12 @@ use axum::{ http::StatusCode, response::{IntoResponse, Response}, }; +use codex_proxy_platform::json::Page; +use codex_proxy_server::admin_api::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse}; use serde_json::{json, Value}; -use codex_proxy_rs::{ - admin::api::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse}, - utils::pagination::Page, -}; - #[test] -fn admin_envelope_serializes_lower_camel_case_request_id() { +fn admin_envelope_should_serialize_lower_camel_case_request_id() { let body = AdminEnvelope::ok(json!({ "id": "acct_1" }), "req_1"); let value = serde_json::to_value(body).unwrap(); @@ -28,7 +25,7 @@ fn admin_envelope_serializes_lower_camel_case_request_id() { } #[test] -fn admin_page_envelope_exposes_items_as_data_with_page_metadata() { +fn admin_page_envelope_should_expose_items_as_data_with_page_metadata() { let page = Page { items: vec![json!({ "id": "evt_1" })], next_cursor: Some("cursor_1".to_string()), @@ -53,7 +50,7 @@ fn admin_page_envelope_exposes_items_as_data_with_page_metadata() { } #[test] -fn admin_response_keeps_http_status_separate_from_body_code() { +fn admin_response_should_keep_http_status_separate_from_body_code() { let body = AdminEnvelope::new(40101, "Admin login required", (), "req_1"); let response: Response = AdminResponse::new(StatusCode::UNAUTHORIZED, body).into_response(); @@ -61,7 +58,7 @@ fn admin_response_keeps_http_status_separate_from_body_code() { } #[test] -fn admin_error_body_uses_null_data() { +fn admin_error_body_should_use_null_data() { let body = AdminEnvelope::new(40101, "Admin login required", (), "req_1"); let value = serde_json::to_value(body).unwrap(); @@ -70,7 +67,7 @@ fn admin_error_body_uses_null_data() { } #[tokio::test] -async fn admin_error_into_response_uses_admin_envelope() { +async fn admin_error_into_response_should_use_admin_envelope() { let response: Response = AdminError::new( StatusCode::UNAUTHORIZED, 40101, diff --git a/tests/admin/client_keys_route.rs b/crates/server/tests/admin_client_keys_routes.rs similarity index 74% rename from tests/admin/client_keys_route.rs rename to crates/server/tests/admin_client_keys_routes.rs index 281e32df6..69edc78ea 100644 --- a/tests/admin/client_keys_route.rs +++ b/crates/server/tests/admin_client_keys_routes.rs @@ -1,26 +1,616 @@ use std::collections::BTreeMap; use axum::{ - body::Body, + body::{to_bytes, Body}, http::{Request, StatusCode}, }; -use serde_json::json; -use tower::ServiceExt; - -use codex_proxy_rs::{ +use codex_proxy_platform::{ config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, + UsageStatsConfig, WebSocketPoolConfig, }, - platform::crypto::SecretBox, - platform::identity::client_key::ApiKeyHasher, - platform::storage::db::connect_sqlite, - runtime::build_router, - runtime::state::AppState, + crypto::SecretBox, + identity::ApiKeyHasher, + storage::connect_sqlite, }; +use codex_proxy_runtime::state::AppState; +use codex_proxy_server::router; +use serde_json::{json, Value}; +use sqlx::SqlitePool; +use tower::util::ServiceExt; -use crate::support::{response_json, seed_admin_session}; +#[tokio::test] +async fn admin_client_keys_route_should_create_list_and_authorize_v1_requests() { + let (app, _dir) = admin_client_key_test_app("admin-client-key-create.sqlite").await; + + let rejected = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/models") + .header("authorization", "Bearer cpr_not_stored") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED); + + let create_response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/api-keys") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_api_key") + .body(Body::from(r#"{"name":"cursor"}"#)) + .unwrap(), + ) + .await + .unwrap(); + let create_status = create_response.status(); + let create_body = response_json(create_response).await; + + assert_eq!(create_status, StatusCode::OK); + assert!(create_body["data"]["plaintext"] + .as_str() + .is_some_and(|plaintext| plaintext.starts_with("cpr_"))); + assert_eq!(create_body["requestId"], "req_api_key"); + + let plaintext = create_body["data"]["plaintext"] + .as_str() + .unwrap() + .to_string(); + let list_response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/api-keys?limit=10") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let list_status = list_response.status(); + let list_body = response_json(list_response).await; + + assert_eq!(list_status, StatusCode::OK); + assert_eq!(list_body["data"][0]["name"], "cursor"); + assert!(list_body["data"][0].get("plaintext").is_none()); + assert!(list_body["data"][0].get("keyHash").is_none()); + + let models_response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/models") + .header("authorization", format!("Bearer {plaintext}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(models_response.status(), StatusCode::OK); +} + +#[tokio::test] +async fn admin_client_keys_status_should_disable_and_enable_authorization() { + let (app, _dir) = admin_client_key_test_app("admin-client-key-status.sqlite").await; + let (key_id, plaintext) = create_admin_client_key(&app, "status-key").await; + + let disabled = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri(format!("/api/admin/api-keys/{key_id}/status")) + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"status":"disabled"}"#)) + .unwrap(), + ) + .await + .unwrap(); + let disabled_status = disabled.status(); + let disabled_body = response_json(disabled).await; + + assert_eq!(disabled_status, StatusCode::OK); + assert_eq!(disabled_body["data"]["enabled"], false); + + let rejected = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/models") + .header("authorization", format!("Bearer {plaintext}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED); + + let enabled = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri(format!("/api/admin/api-keys/{key_id}/status")) + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"status":"active"}"#)) + .unwrap(), + ) + .await + .unwrap(); + let enabled_status = enabled.status(); + let enabled_body = response_json(enabled).await; + + assert_eq!(enabled_status, StatusCode::OK); + assert_eq!(enabled_body["data"]["enabled"], true); + + let accepted = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/models") + .header("authorization", format!("Bearer {plaintext}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(accepted.status(), StatusCode::OK); + + let invalid = app + .oneshot( + Request::builder() + .method("PATCH") + .uri(format!("/api/admin/api-keys/{key_id}/status")) + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"status":"expired"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(invalid.status(), StatusCode::BAD_REQUEST); +} + +#[tokio::test] +async fn admin_client_keys_delete_should_remove_authorization() { + let (app, _dir) = admin_client_key_test_app("admin-client-key-delete.sqlite").await; + let (key_id, plaintext) = create_admin_client_key(&app, "delete-key").await; + + let deleted = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri(format!("/api/admin/api-keys/{key_id}")) + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let delete_status = deleted.status(); + let delete_body = response_json(deleted).await; + + assert_eq!(delete_status, StatusCode::OK); + assert_eq!(delete_body["data"]["deleted"], true); + + let list_response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/api-keys?limit=10") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let list_body = response_json(list_response).await; + assert_eq!(list_body["data"].as_array().unwrap().len(), 0); + + let rejected = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/models") + .header("authorization", format!("Bearer {plaintext}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED); + + let missing = app + .oneshot( + Request::builder() + .method("DELETE") + .uri(format!("/api/admin/api-keys/{key_id}")) + .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_client_keys_label_should_update_clear_and_validate_label() { + let (app, _dir) = admin_client_key_test_app("admin-client-key-label.sqlite").await; + let (key_id, _plaintext) = create_admin_client_key(&app, "label-key").await; + + let renamed = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri(format!("/api/admin/api-keys/{key_id}/label")) + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"label":"automation"}"#)) + .unwrap(), + ) + .await + .unwrap(); + let renamed_status = renamed.status(); + let renamed_body = response_json(renamed).await; + + assert_eq!(renamed_status, StatusCode::OK); + assert_eq!(renamed_body["data"]["name"], "label-key"); + assert_eq!(renamed_body["data"]["label"], "automation"); + + let cleared = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri(format!("/api/admin/api-keys/{key_id}/label")) + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"label":null}"#)) + .unwrap(), + ) + .await + .unwrap(); + let cleared_status = cleared.status(); + let cleared_body = response_json(cleared).await; + + assert_eq!(cleared_status, StatusCode::OK); + assert_eq!(cleared_body["data"]["name"], "label-key"); + assert!(cleared_body["data"]["label"].is_null()); + + let too_long = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri(format!("/api/admin/api-keys/{key_id}/label")) + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(json!({ "label": "x".repeat(65) }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(too_long.status(), StatusCode::BAD_REQUEST); + + let missing = app + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/admin/api-keys/missing/label") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(r#"{"label":"automation"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(missing.status(), StatusCode::NOT_FOUND); +} + +#[tokio::test] +async fn admin_client_keys_batch_delete_should_remove_found_keys_and_report_missing_ids() { + let (app, _dir) = admin_client_key_test_app("admin-client-key-batch-delete.sqlite").await; + let (key_a, plaintext_a) = create_admin_client_key(&app, "batch-a").await; + let (key_b, plaintext_b) = create_admin_client_key(&app, "batch-b").await; + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/api-keys/batch-delete") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({ + "ids": [key_a, "ghost", key_b] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["data"]["deleted"], 2); + assert_eq!(body["data"]["notFound"], json!(["ghost"])); + + for plaintext in [plaintext_a, plaintext_b] { + let rejected = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/models") + .header("authorization", format!("Bearer {plaintext}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED); + } + + let empty = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/api-keys/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_client_keys_export_should_return_metadata_without_secret_material() { + let (app, _dir) = admin_client_key_test_app("admin-client-key-export.sqlite").await; + let (key_id, plaintext) = create_admin_client_key(&app, "export-key").await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri(format!("/api/admin/api-keys/export?ids={key_id}")) + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_api_key_export") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["requestId"], "req_api_key_export"); + assert_eq!(body["data"]["sourceFormat"], "rustLocalClientApiKeys"); + assert_eq!(body["data"]["rotationRequired"], true); + assert_eq!(body["data"]["apiKeys"][0]["id"], key_id); + assert_eq!(body["data"]["apiKeys"][0]["name"], "export-key"); + assert!(body["data"]["apiKeys"][0]["prefix"] + .as_str() + .is_some_and(|prefix| prefix.starts_with("cpr_"))); + assert!(body["data"]["apiKeys"][0].get("plaintext").is_none()); + assert!(body["data"]["apiKeys"][0].get("keyHash").is_none()); + assert!(body["data"].get("pepper").is_none()); + assert!(!body.to_string().contains(&plaintext)); +} + +#[tokio::test] +async fn admin_client_keys_import_should_rotate_exported_metadata_and_return_plaintext_once() { + let source_dir = tempfile::tempdir().unwrap(); + let source_db = source_dir + .path() + .join("admin-client-key-export-source.sqlite"); + let source_url = format!("sqlite://{}", source_db.display()); + let source_pool = connect_sqlite(&source_url).await.unwrap(); + seed_admin_session(&source_pool, "session_1").await; + let source_state = AppState::with_pool_secret_and_api_key_hasher( + test_config(source_url), + source_pool, + SecretBox::new([51u8; 32]), + ApiKeyHasher::new([52u8; 32]), + ); + let source_app = router::router().with_state(source_state); + let (source_key_id, source_plaintext) = + create_admin_client_key(&source_app, "rotated-key").await; + + let export_response = source_app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/api-keys/export") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(export_response.status(), StatusCode::OK); + let export_body = response_json(export_response).await; + + let target_dir = tempfile::tempdir().unwrap(); + let target_db = target_dir + .path() + .join("admin-client-key-import-target.sqlite"); + let target_url = format!("sqlite://{}", target_db.display()); + let target_pool = connect_sqlite(&target_url).await.unwrap(); + seed_admin_session(&target_pool, "session_1").await; + let target_state = AppState::with_pool_secret_and_api_key_hasher( + test_config(target_url), + target_pool, + SecretBox::new([61u8; 32]), + ApiKeyHasher::new([62u8; 32]), + ); + let target_app = router::router().with_state(target_state); + + let import_response = target_app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/api-keys/import") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_api_key_import") + .body(Body::from(export_body["data"].to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = import_response.status(); + let import_body = response_json(import_response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(import_body["requestId"], "req_api_key_import"); + assert_eq!(import_body["data"]["imported"], 1); + assert_eq!(import_body["data"]["skipped"], 0); + assert_eq!(import_body["data"]["rotated"], true); + assert_eq!(import_body["data"]["keys"][0]["sourceId"], source_key_id); + assert_eq!(import_body["data"]["keys"][0]["name"], "rotated-key"); + assert!(import_body["data"]["keys"][0].get("keyHash").is_none()); + let rotated_plaintext = import_body["data"]["keys"][0]["plaintext"] + .as_str() + .unwrap() + .to_string(); + assert!(rotated_plaintext.starts_with("cpr_")); + assert_ne!(rotated_plaintext, source_plaintext); + + let list_response = target_app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/api-keys?limit=10") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let list_body = response_json(list_response).await; + assert!(list_body["data"][0].get("plaintext").is_none()); + + let old_rejected = target_app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/models") + .header("authorization", format!("Bearer {source_plaintext}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(old_rejected.status(), StatusCode::UNAUTHORIZED); + + let new_accepted = target_app + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/models") + .header("authorization", format!("Bearer {rotated_plaintext}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(new_accepted.status(), StatusCode::OK); +} + +async fn admin_client_key_test_app(db_name: &str) -> (axum::Router, 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_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([43u8; 32]), + ApiKeyHasher::new([44u8; 32]), + ); + (router::router().with_state(state), dir) +} + +async fn create_admin_client_key(app: &axum::Router, name: &str) -> (String, String) { + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/api-keys") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(json!({ "name": name }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json(response).await; + ( + body["data"]["id"].as_str().unwrap().to_string(), + body["data"]["plaintext"].as_str().unwrap().to_string(), + ) +} + +async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { + sqlx::query( + "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", + ) + .bind("admin_1") + .bind("hash") + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .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("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); +} + +async fn response_json(response: axum::response::Response) -> Value { + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + serde_json::from_slice(&bytes).unwrap() +} fn test_config(database_url: String) -> AppConfig { AppConfig { @@ -68,12 +658,12 @@ fn test_config(database_url: String) -> AppConfig { tls: TlsConfig { force_http11: false, }, - ws_pool: Default::default(), + ws_pool: WebSocketPoolConfig::default(), admin: AdminConfig { session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, default_username: "admin".to_string(), default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, }, logging: LoggingConfig { directory: "logs".to_string(), @@ -84,554 +674,3 @@ fn test_config(database_url: String) -> AppConfig { }, } } - -#[tokio::test] -async fn admin_api_keys_should_create_list_and_authorize_v1_requests() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-api-keys.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_secret_and_api_key_hasher( - test_config(url), - pool, - SecretBox::new([41u8; 32]), - ApiKeyHasher::new([42u8; 32]), - )); - - let rejected = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models") - .header("authorization", "Bearer cpr_not_stored") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED); - - let create_response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/api-keys") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_api_key") - .body(Body::from(r#"{"name":"cursor"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(create_response.status(), StatusCode::OK); - let body = response_json(create_response).await; - let plaintext = body["data"]["plaintext"].as_str().unwrap().to_string(); - assert!(plaintext.starts_with("cpr_")); - assert_eq!(body["requestId"], "req_api_key"); - - let list_response = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/api-keys?limit=10") - .header("cookie", "cpr_admin_session=session_1") - .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]["name"], "cursor"); - assert!(body["data"][0].get("plaintext").is_none()); - assert!(body["data"][0].get("keyHash").is_none()); - - let models_response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models") - .header("authorization", format!("Bearer {plaintext}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(models_response.status(), StatusCode::OK); -} - -#[tokio::test] -async fn admin_api_key_status_should_disable_and_reenable_client_key_authorization() { - let (app, _dir) = admin_api_keys_test_app("admin-api-key-status.sqlite").await; - let (key_id, plaintext) = create_admin_api_key(&app, "session_1", "status-key").await; - - let disabled = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri(format!("/api/admin/api-keys/{key_id}/status")) - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"status":"disabled"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(disabled.status(), StatusCode::OK); - let body = response_json(disabled).await; - assert_eq!(body["data"]["enabled"], false); - - let rejected = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models") - .header("authorization", format!("Bearer {plaintext}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED); - - let enabled = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri(format!("/api/admin/api-keys/{key_id}/status")) - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"status":"active"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(enabled.status(), StatusCode::OK); - let body = response_json(enabled).await; - assert_eq!(body["data"]["enabled"], true); - - let accepted = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models") - .header("authorization", format!("Bearer {plaintext}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(accepted.status(), StatusCode::OK); - - let invalid = app - .oneshot( - Request::builder() - .method("PATCH") - .uri(format!("/api/admin/api-keys/{key_id}/status")) - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"status":"expired"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(invalid.status(), StatusCode::BAD_REQUEST); -} - -#[tokio::test] -async fn admin_api_key_delete_should_remove_client_key_authorization() { - let (app, _dir) = admin_api_keys_test_app("admin-api-key-delete.sqlite").await; - let (key_id, plaintext) = create_admin_api_key(&app, "session_1", "delete-key").await; - - let response = app - .clone() - .oneshot( - Request::builder() - .method("DELETE") - .uri(format!("/api/admin/api-keys/{key_id}")) - .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"]["deleted"], true); - - let list_response = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/api-keys?limit=10") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(list_response.status(), StatusCode::OK); - let body = response_json(list_response).await; - assert_eq!(body["data"].as_array().unwrap().len(), 0); - - let rejected = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models") - .header("authorization", format!("Bearer {plaintext}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED); - - let missing = app - .oneshot( - Request::builder() - .method("DELETE") - .uri(format!("/api/admin/api-keys/{key_id}")) - .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_api_key_label_should_update_and_clear_local_client_key_name() { - let (app, _dir) = admin_api_keys_test_app("admin-api-key-label.sqlite").await; - let (key_id, _plaintext) = create_admin_api_key(&app, "session_1", "label-key").await; - - let renamed = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri(format!("/api/admin/api-keys/{key_id}/label")) - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"label":"automation"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(renamed.status(), StatusCode::OK); - let body = response_json(renamed).await; - assert_eq!(body["data"]["name"], "label-key"); - assert_eq!(body["data"]["label"], "automation"); - - let cleared = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri(format!("/api/admin/api-keys/{key_id}/label")) - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"label":null}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(cleared.status(), StatusCode::OK); - let body = response_json(cleared).await; - assert_eq!(body["data"]["name"], "label-key"); - assert!(body["data"]["label"].is_null()); - - let too_long = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri(format!("/api/admin/api-keys/{key_id}/label")) - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(json!({ "label": "x".repeat(65) }).to_string())) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(too_long.status(), StatusCode::BAD_REQUEST); - - let missing = app - .oneshot( - Request::builder() - .method("PATCH") - .uri("/api/admin/api-keys/missing/label") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"label":"automation"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(missing.status(), StatusCode::NOT_FOUND); -} - -#[tokio::test] -async fn admin_api_keys_batch_delete_should_remove_found_keys_and_report_missing_ids() { - let (app, _dir) = admin_api_keys_test_app("admin-api-key-batch-delete.sqlite").await; - let (key_a, plaintext_a) = create_admin_api_key(&app, "session_1", "batch-a").await; - let (key_b, plaintext_b) = create_admin_api_key(&app, "session_1", "batch-b").await; - - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/api-keys/batch-delete") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - json!({ - "ids": [key_a, "ghost", key_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"])); - - for plaintext in [plaintext_a, plaintext_b] { - let rejected = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models") - .header("authorization", format!("Bearer {plaintext}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(rejected.status(), StatusCode::UNAUTHORIZED); - } - - let empty = app - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/api-keys/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); -} - -async fn admin_api_keys_test_app(db_name: &str) -> (axum::Router, 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 app = build_router(AppState::with_pool_secret_and_api_key_hasher( - test_config(url), - pool, - SecretBox::new([43u8; 32]), - ApiKeyHasher::new([44u8; 32]), - )); - (app, dir) -} - -async fn create_admin_api_key( - app: &axum::Router, - session_id: &str, - name: &str, -) -> (String, String) { - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/api-keys") - .header("content-type", "application/json") - .header("cookie", format!("cpr_admin_session={session_id}")) - .body(Body::from(json!({ "name": name }).to_string())) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - ( - body["data"]["id"].as_str().unwrap().to_string(), - body["data"]["plaintext"].as_str().unwrap().to_string(), - ) -} - -#[tokio::test] -async fn admin_api_keys_export_should_return_metadata_without_secret_material() { - let (app, _dir) = admin_api_keys_test_app("admin-api-key-export.sqlite").await; - let (key_id, plaintext) = create_admin_api_key(&app, "session_1", "export-key").await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri(format!("/api/admin/api-keys/export?ids={key_id}")) - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_api_key_export") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["requestId"], "req_api_key_export"); - assert_eq!(body["data"]["sourceFormat"], "rustLocalClientApiKeys"); - assert_eq!(body["data"]["rotationRequired"], true); - assert_eq!(body["data"]["apiKeys"][0]["id"], key_id); - assert_eq!(body["data"]["apiKeys"][0]["name"], "export-key"); - assert!(body["data"]["apiKeys"][0]["prefix"] - .as_str() - .unwrap() - .starts_with("cpr_")); - assert!(body["data"]["apiKeys"][0].get("plaintext").is_none()); - assert!(body["data"]["apiKeys"][0].get("keyHash").is_none()); - assert!(body["data"].get("pepper").is_none()); - assert!(!body.to_string().contains(&plaintext)); -} - -#[tokio::test] -async fn admin_api_keys_import_should_rotate_exported_metadata_and_return_new_plaintext_once() { - let source_dir = tempfile::tempdir().unwrap(); - let source_db = source_dir.path().join("admin-api-key-export-source.sqlite"); - let source_url = format!("sqlite://{}", source_db.display()); - let source_pool = connect_sqlite(&source_url).await.unwrap(); - seed_admin_session(&source_pool, "session_1").await; - let source_app = build_router(AppState::with_pool_secret_and_api_key_hasher( - test_config(source_url), - source_pool, - SecretBox::new([51u8; 32]), - ApiKeyHasher::new([52u8; 32]), - )); - let (source_key_id, source_plaintext) = - create_admin_api_key(&source_app, "session_1", "rotated-key").await; - - let export_response = source_app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/api-keys/export") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(export_response.status(), StatusCode::OK); - let export_body = response_json(export_response).await; - - let target_dir = tempfile::tempdir().unwrap(); - let target_db = target_dir.path().join("admin-api-key-import-target.sqlite"); - let target_url = format!("sqlite://{}", target_db.display()); - let target_pool = connect_sqlite(&target_url).await.unwrap(); - seed_admin_session(&target_pool, "session_1").await; - let target_app = build_router(AppState::with_pool_secret_and_api_key_hasher( - test_config(target_url), - target_pool, - SecretBox::new([61u8; 32]), - ApiKeyHasher::new([62u8; 32]), - )); - - let import_response = target_app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/api-keys/import") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_api_key_import") - .body(Body::from(export_body["data"].to_string())) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(import_response.status(), StatusCode::OK); - let import_body = response_json(import_response).await; - assert_eq!(import_body["requestId"], "req_api_key_import"); - assert_eq!(import_body["data"]["imported"], 1); - assert_eq!(import_body["data"]["skipped"], 0); - assert_eq!(import_body["data"]["rotated"], true); - assert_eq!(import_body["data"]["keys"][0]["sourceId"], source_key_id); - assert_eq!(import_body["data"]["keys"][0]["name"], "rotated-key"); - assert!(import_body["data"]["keys"][0].get("keyHash").is_none()); - let rotated_plaintext = import_body["data"]["keys"][0]["plaintext"] - .as_str() - .unwrap() - .to_string(); - assert!(rotated_plaintext.starts_with("cpr_")); - assert_ne!(rotated_plaintext, source_plaintext); - - let list_response = target_app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/api-keys?limit=10") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(list_response.status(), StatusCode::OK); - let list_body = response_json(list_response).await; - assert!(list_body["data"][0].get("plaintext").is_none()); - - let old_rejected = target_app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models") - .header("authorization", format!("Bearer {source_plaintext}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(old_rejected.status(), StatusCode::UNAUTHORIZED); - - let new_accepted = target_app - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models") - .header("authorization", format!("Bearer {rotated_plaintext}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(new_accepted.status(), StatusCode::OK); -} diff --git a/tests/admin/logs_route.rs b/crates/server/tests/admin_logs_routes.rs similarity index 50% rename from tests/admin/logs_route.rs rename to crates/server/tests/admin_logs_routes.rs index 0be88e044..7bd90ac4c 100644 --- a/tests/admin/logs_route.rs +++ b/crates/server/tests/admin_logs_routes.rs @@ -1,27 +1,402 @@ use std::collections::BTreeMap; use axum::{ - body::Body, + body::{to_bytes, Body}, http::{Request, StatusCode}, }; -use tower::ServiceExt; - -use codex_proxy_rs::{ - codex::events::{ - event::{EventLevel, EventLog}, - repository::EventLogRepository, - }, +use chrono::{Duration, Utc}; +use codex_proxy_adapters::sqlite::events::SqliteEventLogStore; +use codex_proxy_core::events::model::{EventLevel, EventLog}; +use codex_proxy_platform::{ config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, + UsageStatsConfig, WebSocketPoolConfig, }, - platform::storage::db::connect_sqlite, - runtime::build_router, - runtime::state::AppState, + crypto::SecretBox, + identity::ApiKeyHasher, + storage::connect_sqlite, }; +use codex_proxy_runtime::state::AppState; +use codex_proxy_server::router; +use serde_json::Value; +use sqlx::SqlitePool; +use tower::util::ServiceExt; -use crate::support::{response_json, seed_admin_session}; +#[tokio::test] +async fn admin_logs_should_require_admin_session_cookie() { + let (app, _store, _dir) = admin_logs_test_app("admin-logs-auth.sqlite").await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/logs") + .header("x-request-id", "req_logs_auth") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert_eq!(body["code"], 40101); + assert_eq!(body["requestId"], "req_logs_auth"); +} + +#[tokio::test] +async fn admin_logs_should_cursor_page_events_and_include_request_id() { + let (app, store, _dir) = admin_logs_test_app("admin-logs-cursor.sqlite").await; + let now = Utc::now(); + let mut older = EventLog::new("request", EventLevel::Info, "older"); + older.id = "log_older".to_string(); + older.created_at = now; + store.append(&older).await.unwrap(); + let mut newer = EventLog::new("request", EventLevel::Info, "newer"); + newer.id = "log_newer".to_string(); + newer.created_at = now + Duration::seconds(1); + store.append(&newer).await.unwrap(); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/logs?limit=1") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_logs_cursor") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["code"], 200); + assert_eq!(body["requestId"], "req_logs_cursor"); + assert_eq!(body["data"].as_array().unwrap().len(), 1); + assert_eq!(body["data"][0]["id"], "log_newer"); + assert_eq!(body["page"]["limit"], 1); + assert!(body["page"]["nextCursor"].is_string()); +} + +#[tokio::test] +async fn admin_logs_should_filter_and_cursor_page_events() { + let (app, store, _dir) = admin_logs_test_app("admin-logs.sqlite").await; + let mut matching = EventLog::new("request", EventLevel::Error, "upstream timeout"); + matching.id = "log_matching".to_string(); + matching.route = Some("/v1/responses".to_string()); + store.append(&matching).await.unwrap(); + store + .append(&EventLog::new( + "request", + EventLevel::Info, + "upstream timeout", + )) + .await + .unwrap(); + store + .append(&EventLog::new( + "account", + EventLevel::Error, + "upstream timeout", + )) + .await + .unwrap(); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/logs?kind=request&level=error&search=timeout&limit=1") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_logs_filter") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["requestId"], "req_logs_filter"); + assert_eq!(body["data"][0]["id"], "log_matching"); + assert_eq!(body["data"][0]["route"], "/v1/responses"); +} + +#[tokio::test] +async fn admin_logs_should_reject_unsupported_level_filter() { + let (app, _store, _dir) = admin_logs_test_app("admin-logs-invalid-level.sqlite").await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/logs?level=verbose") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_logs_invalid_level") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(body["code"], 40001); + assert_eq!(body["requestId"], "req_logs_invalid_level"); +} + +#[tokio::test] +async fn admin_logs_should_return_state_detail_and_clear_events() { + let (app, store, _dir) = admin_logs_test_app("admin-logs-state.sqlite").await; + let mut event = EventLog::new("request", EventLevel::Warn, "detail"); + event.id = "log_detail".to_string(); + event.request_id = Some("req_upstream".to_string()); + store.append(&event).await.unwrap(); + + let detail = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/logs/log_detail") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let detail_status = detail.status(); + let detail_body = response_json(detail).await; + + assert_eq!(detail_status, StatusCode::OK); + assert_eq!(detail_body["data"]["requestId"], "req_upstream"); + assert_eq!(detail_body["data"]["level"], "warn"); + + let state = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/logs/state") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let state_body = response_json(state).await; + assert_eq!(state_body["data"]["storedCount"], 1); + assert_eq!(state_body["data"]["enabled"], false); + + let cleared = app + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri("/api/admin/logs") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let cleared_body = response_json(cleared).await; + assert_eq!(cleared_body["data"]["cleared"], 1); + + let empty = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/logs?limit=50") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let empty_body = response_json(empty).await; + assert_eq!(empty_body["data"].as_array().unwrap().len(), 0); +} + +#[tokio::test] +async fn admin_logs_detail_should_return_not_found_for_missing_event() { + let (app, _store, _dir) = admin_logs_test_app("admin-logs-detail-missing.sqlite").await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/logs/missing") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_logs_missing") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::NOT_FOUND); + assert_eq!(body["code"], 40401); + assert_eq!(body["message"], "Log event not found"); + assert_eq!(body["requestId"], "req_logs_missing"); +} + +#[tokio::test] +async fn admin_logs_state_patch_should_require_admin_session_cookie() { + let (app, _store, _dir) = admin_logs_test_app("admin-logs-state-auth.sqlite").await; + + let response = app + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/admin/logs/state") + .header("content-type", "application/json") + .header("x-request-id", "req_logs_state_auth") + .body(Body::from( + serde_json::json!({ "enabled": true }).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert_eq!(body["code"], 40101); + assert_eq!(body["requestId"], "req_logs_state_auth"); +} + +#[tokio::test] +async fn admin_logs_should_update_state_and_trim_to_capacity() { + let (app, store, _dir) = admin_logs_test_app("admin-logs-update-state.sqlite").await; + for index in 0..3 { + let mut event = EventLog::new("request", EventLevel::Info, format!("event {index}")); + event.id = format!("log_capacity_{index}"); + event.created_at = Utc::now() + Duration::seconds(index); + store.append(&event).await.unwrap(); + } + + let response = app + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/admin/logs/state") + .header("cookie", "cpr_admin_session=session_1") + .header("content-type", "application/json") + .body(Body::from( + serde_json::json!({ + "enabled": true, + "capacity": 2, + "captureBody": true + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + assert_eq!(status, StatusCode::OK); + let body = response_json(response).await; + let remaining = store + .list(Default::default(), None, 10) + .await + .unwrap() + .items + .into_iter() + .map(|event| event.id) + .collect::>(); + + assert_eq!(body["data"]["enabled"], true); + assert_eq!(body["data"]["capacity"], 2); + assert_eq!(body["data"]["captureBody"], true); + assert_eq!(body["data"]["storedCount"], 2); + assert_eq!(remaining, vec!["log_capacity_2", "log_capacity_1"]); +} + +#[tokio::test] +async fn admin_logs_state_patch_should_reject_zero_capacity() { + let (app, _store, _dir) = admin_logs_test_app("admin-logs-zero-capacity.sqlite").await; + + let response = app + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/admin/logs/state") + .header("cookie", "cpr_admin_session=session_1") + .header("content-type", "application/json") + .header("x-request-id", "req_logs_zero_capacity") + .body(Body::from(serde_json::json!({ "capacity": 0 }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(body["code"], 40001); + assert_eq!(body["requestId"], "req_logs_zero_capacity"); +} + +async fn admin_logs_test_app( + db_name: &str, +) -> (axum::Router, SqliteEventLogStore, 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_secret_and_api_key_hasher( + test_config(url), + pool.clone(), + SecretBox::new([73u8; 32]), + ApiKeyHasher::new([74u8; 32]), + ); + ( + router::router().with_state(state), + SqliteEventLogStore::new(pool), + dir, + ) +} + +async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { + sqlx::query( + "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", + ) + .bind("admin_1") + .bind("hash") + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .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("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); +} + +async fn response_json(response: axum::response::Response) -> Value { + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + serde_json::from_slice(&bytes).unwrap() +} fn test_config(database_url: String) -> AppConfig { AppConfig { @@ -69,12 +444,12 @@ fn test_config(database_url: String) -> AppConfig { tls: TlsConfig { force_http11: false, }, - ws_pool: Default::default(), + ws_pool: WebSocketPoolConfig::default(), admin: AdminConfig { session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, default_username: "admin".to_string(), default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, }, logging: LoggingConfig { directory: "logs".to_string(), @@ -85,330 +460,3 @@ fn test_config(database_url: String) -> AppConfig { }, } } - -#[tokio::test] -async fn admin_logs_are_cursor_paginated_and_include_request_id() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-logs.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let repo = EventLogRepository::new(pool.clone()); - repo.insert(EventLog::new("request", EventLevel::Info, "first")) - .await - .unwrap(); - repo.insert(EventLog::new("request", EventLevel::Info, "second")) - .await - .unwrap(); - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/logs?limit=1") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_admin") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["code"], 200); - assert_eq!(body["requestId"], "req_admin"); - assert_eq!(body["data"].as_array().unwrap().len(), 1); - assert_eq!(body["page"]["limit"], 1); - assert!(body["page"]["nextCursor"].is_string()); -} - -#[tokio::test] -async fn admin_logs_should_filter_by_kind_level_and_search_text() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-logs-filter.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let repo = EventLogRepository::new(pool.clone()); - let mut matching = EventLog::new("request", EventLevel::Error, "upstream timeout"); - matching.id = "log_matching".to_string(); - matching.route = Some("/v1/responses".to_string()); - repo.insert(matching).await.unwrap(); - repo.insert(EventLog::new( - "request", - EventLevel::Info, - "upstream timeout", - )) - .await - .unwrap(); - repo.insert(EventLog::new( - "account", - EventLevel::Error, - "upstream timeout", - )) - .await - .unwrap(); - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/logs?kind=request&level=error&search=timeout") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_logs_filter") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["requestId"], "req_logs_filter"); - let items = body["data"].as_array().unwrap(); - assert_eq!(items.len(), 1); - assert_eq!(items[0]["id"], "log_matching"); - assert_eq!(items[0]["route"], "/v1/responses"); -} - -#[tokio::test] -async fn admin_logs_reject_missing_admin_session_cookie() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-logs.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/logs") - .header("x-request-id", "req_admin") - .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_admin"); -} - -#[tokio::test] -async fn admin_logs_state_should_return_runtime_logging_state_and_stored_count() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-logs-state.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let repo = EventLogRepository::new(pool.clone()); - repo.insert(EventLog::new("request", EventLevel::Info, "first")) - .await - .unwrap(); - repo.insert(EventLog::new("request", EventLevel::Warn, "second")) - .await - .unwrap(); - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/logs/state") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_logs_state") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["code"], 200); - assert_eq!(body["requestId"], "req_logs_state"); - assert_eq!(body["data"]["enabled"], false); - assert_eq!(body["data"]["captureBody"], false); - assert_eq!(body["data"]["capacity"], 2000); - assert_eq!(body["data"]["storedCount"], 2); -} - -#[tokio::test] -async fn admin_logs_state_patch_should_require_admin_session_cookie() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-logs-state-patch-auth.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .oneshot( - Request::builder() - .method("PATCH") - .uri("/api/admin/logs/state") - .header("content-type", "application/json") - .header("x-request-id", "req_logs_state_patch") - .body(Body::from( - serde_json::json!({ "enabled": true }).to_string(), - )) - .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_logs_state_patch"); -} - -#[tokio::test] -async fn admin_logs_state_patch_should_update_runtime_logging_state() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-logs-state-patch.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(test_config(url), pool)); - - let response = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri("/api/admin/logs/state") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_logs_state_patch") - .body(Body::from( - serde_json::json!({ - "enabled": true, - "captureBody": true, - "capacity": 512 - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["code"], 200); - assert_eq!(body["requestId"], "req_logs_state_patch"); - assert_eq!(body["data"]["enabled"], true); - assert_eq!(body["data"]["captureBody"], true); - assert_eq!(body["data"]["capacity"], 512); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/logs/state") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - let body = response_json(response).await; - assert_eq!(body["data"]["enabled"], true); - assert_eq!(body["data"]["captureBody"], true); - assert_eq!(body["data"]["capacity"], 512); -} - -#[tokio::test] -async fn admin_logs_detail_should_return_one_event_by_id() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-logs-detail.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let repo = EventLogRepository::new(pool.clone()); - let mut event = EventLog::new("request", EventLevel::Error, "detail"); - event.id = "log_detail".to_string(); - event.request_id = Some("req_upstream".to_string()); - repo.insert(event).await.unwrap(); - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/logs/log_detail") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_logs_detail") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["code"], 200); - assert_eq!(body["requestId"], "req_logs_detail"); - assert_eq!(body["data"]["id"], "log_detail"); - assert_eq!(body["data"]["requestId"], "req_upstream"); - assert_eq!(body["data"]["level"], "error"); - assert_eq!(body["data"]["message"], "detail"); -} - -#[tokio::test] -async fn admin_logs_clear_should_delete_stored_events() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-logs-clear.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let repo = EventLogRepository::new(pool.clone()); - repo.insert(EventLog::new("request", EventLevel::Info, "first")) - .await - .unwrap(); - repo.insert(EventLog::new("request", EventLevel::Info, "second")) - .await - .unwrap(); - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .clone() - .oneshot( - Request::builder() - .method("DELETE") - .uri("/api/admin/logs") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_logs_clear") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["code"], 200); - assert_eq!(body["requestId"], "req_logs_clear"); - assert_eq!(body["data"]["cleared"], 2); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/logs?limit=50") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let body = response_json(response).await; - assert_eq!(body["data"].as_array().unwrap().len(), 0); -} diff --git a/tests/admin/models_route.rs b/crates/server/tests/admin_models_routes.rs similarity index 63% rename from tests/admin/models_route.rs rename to crates/server/tests/admin_models_routes.rs index c259c0caa..357233569 100644 --- a/tests/admin/models_route.rs +++ b/crates/server/tests/admin_models_routes.rs @@ -1,37 +1,241 @@ use std::collections::BTreeMap; use axum::{ - body::Body, + body::{to_bytes, Body}, http::{Request, StatusCode}, }; +use codex_proxy_adapters::sqlite::{ + accounts::{NewAccount, SqliteAccountStore}, + models::ModelSnapshotRepository, +}; +use codex_proxy_core::{accounts::model::AccountStatus, gateway::fingerprint::Fingerprint}; +use codex_proxy_platform::{ + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, + UsageStatsConfig, WebSocketPoolConfig, + }, + crypto::SecretBox, + identity::ApiKeyHasher, + storage::connect_sqlite, +}; +use codex_proxy_runtime::state::AppState; +use codex_proxy_server::router; use secrecy::SecretString; -use serde_json::json; -use tower::ServiceExt; +use serde_json::{json, Value}; +use sqlx::SqlitePool; +use tower::util::ServiceExt; use wiremock::{ matchers::{method, path}, Mock, MockServer, ResponseTemplate, }; -use codex_proxy_rs::{ - codex::accounts::{ - model::{Account, AccountStatus}, - repository::{AccountRepository, NewAccount}, - }, - codex::gateway::fingerprint::model::Fingerprint, - codex::models::repository::ModelSnapshotRepository, - config::{ - AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, - QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, - }, - platform::crypto::SecretBox, - platform::identity::{client_key::ApiKeyHasher, client_key_repository::ClientApiKeyRepository}, - platform::storage::db::connect_sqlite, - runtime::build_router, - runtime::state::AppState, -}; +#[tokio::test] +async fn admin_refresh_models_should_require_admin_session() { + let upstream = MockServer::start().await; + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-models-auth.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url, upstream.uri()), + pool, + SecretBox::new([81u8; 32]), + ApiKeyHasher::new([82u8; 32]), + ); + let app = router::router().with_state(state); -use crate::support::{response_json, seed_admin_session}; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/refresh-models") + .header("x-request-id", "req_models_auth") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert_eq!(body["code"], 40101); + assert_eq!(body["requestId"], "req_models_auth"); +} + +#[tokio::test] +async fn admin_refresh_models_should_store_snapshots_for_distinct_active_plans() { + let upstream = MockServer::start().await; + let fingerprint = models_test_fingerprint(); + let expected_user_agent = fingerprint.user_agent(); + let expected_sec_ch_ua = fingerprint.sec_ch_ua(); + Mock::given(method("GET")) + .and(path("/codex/models")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "models": [{ + "slug": "gpt-refresh", + "display_name": "GPT Refresh", + "description": "Backend refreshed model", + "default_reasoning_level": "minimal", + "supported_reasoning_levels": [{ + "effort": "minimal", + "description": "Minimal" + }], + "contextWindow": 200000 + }] + }))) + .expect(2) + .mount(&upstream) + .await; + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-models-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([83u8; 32]); + let account_store = SqliteAccountStore::new(pool.clone(), secret_box.clone()); + seed_account(&account_store, "acct_plus_a", "plus", "access-plus-a").await; + seed_account(&account_store, "acct_plus_b", "plus", "access-plus-b").await; + seed_account(&account_store, "acct_team", "team", "access-team").await; + let state = AppState::with_pool_secret_api_key_hasher_and_fingerprint( + test_config(url, upstream.uri()), + pool.clone(), + secret_box, + ApiKeyHasher::new([84u8; 32]), + fingerprint, + ); + let generated = state + .services + .admin_client_keys + .create("catalog-reader") + .await + .unwrap(); + let app = router::router().with_state(state); + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/refresh-models") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_models_refresh") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let snapshots = ModelSnapshotRepository::new(pool) + .list_plan_snapshots() + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["requestId"], "req_models_refresh"); + assert_eq!(body["data"]["refreshedPlans"], 2); + assert_eq!(body["data"]["modelCount"], 2); + assert_eq!(body["data"]["failedPlans"], 0); + assert_eq!(snapshots.len(), 2); + assert_eq!(snapshots[0].models[0].id, "gpt-refresh"); + let requests = upstream.received_requests().await.unwrap(); + let model_requests = requests + .iter() + .filter(|request| request.url.path() == "/codex/models") + .collect::>(); + assert_eq!(model_requests.len(), 2); + assert!(model_requests.iter().all(|request| { + request + .headers + .get("user-agent") + .and_then(|value| value.to_str().ok()) + == Some(expected_user_agent.as_str()) + })); + assert!(model_requests.iter().all(|request| { + request + .headers + .get("sec-ch-ua") + .and_then(|value| value.to_str().ok()) + == Some(expected_sec_ch_ua.as_str()) + })); + + let catalog_response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/models/catalog") + .header("authorization", format!("Bearer {}", generated.plaintext)) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let catalog = response_json(catalog_response).await; + assert_eq!(catalog[0]["id"], "gpt-refresh"); + assert_eq!(catalog[0]["source"], "backend"); +} + +fn models_test_fingerprint() -> Fingerprint { + Fingerprint { + app_version: "27.333.444".to_string(), + build_number: "9002".to_string(), + chromium_version: "156".to_string(), + ..Fingerprint::default_for_tests() + } +} + +async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { + sqlx::query( + "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", + ) + .bind("admin_1") + .bind("hash") + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .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("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); +} + +async fn seed_account( + account_store: &SqliteAccountStore, + id: &str, + plan_type: &str, + access_token: &str, +) { + account_store + .insert(NewAccount { + id: id.to_string(), + email: None, + account_id: Some(format!("chatgpt-{id}")), + user_id: None, + label: None, + plan_type: Some(plan_type.to_string()), + access_token: SecretString::new(access_token.to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }) + .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_config(database_url: String, base_url: String) -> AppConfig { AppConfig { @@ -77,12 +281,12 @@ fn test_config(database_url: String, base_url: String) -> AppConfig { tls: TlsConfig { force_http11: false, }, - ws_pool: Default::default(), + ws_pool: WebSocketPoolConfig::default(), admin: AdminConfig { session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, default_username: "admin".to_string(), default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, }, logging: LoggingConfig { directory: "logs".to_string(), @@ -93,205 +297,3 @@ fn test_config(database_url: String, base_url: String) -> AppConfig { }, } } - -#[tokio::test] -async fn admin_refresh_models_should_require_admin_session_cookie() { - let server = MockServer::start().await; - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-models.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, server.uri()), - pool, - SecretBox::new([31u8; 32]), - )); - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/refresh-models") - .header("x-request-id", "req_refresh_models") - .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_refresh_models"); -} - -#[tokio::test] -async fn admin_refresh_models_should_fetch_snapshots_for_distinct_account_plans() { - let server = MockServer::start().await; - let fingerprint = models_test_fingerprint(); - let expected_user_agent = fingerprint.user_agent(); - let expected_sec_ch_ua = fingerprint.sec_ch_ua(); - Mock::given(method("GET")) - .and(path("/codex/models")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "models": [{ - "slug": "gpt-refresh", - "display_name": "GPT Refresh", - "description": "Backend refreshed model", - "default_reasoning_level": "minimal", - "supported_reasoning_levels": [{ - "effort": "minimal", - "description": "Minimal" - }], - "contextWindow": 200000 - }] - }))) - .expect(2) - .mount(&server) - .await; - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-models.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([32u8; 32]); - let hasher = ApiKeyHasher::new([33u8; 32]); - let generated = hasher.generate_client_api_key("test"); - ClientApiKeyRepository::new(pool.clone()) - .insert_generated("test", &generated) - .await - .unwrap(); - let account_repo = AccountRepository::new(pool.clone(), secret_box.clone()); - insert_account(&account_repo, "acct_plus_a", "plus", "access-plus-a").await; - insert_account(&account_repo, "acct_plus_b", "plus", "access-plus-b").await; - insert_account(&account_repo, "acct_team", "team", "access-team").await; - let state = AppState::with_pool_secret_api_key_hasher_and_fingerprint( - test_config(url, server.uri()), - pool.clone(), - secret_box, - hasher, - fingerprint, - ); - 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 - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/refresh-models") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_refresh_models") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["refreshedPlans"], 2); - assert_eq!(body["data"]["modelCount"], 2); - assert_eq!(body["data"]["failedPlans"], 0); - let snapshots = ModelSnapshotRepository::new(pool.clone()) - .list_plan_snapshots() - .await - .unwrap(); - assert_eq!(snapshots.len(), 2); - assert_eq!(snapshots[0].models[0].id, "gpt-refresh"); - let requests = server.received_requests().await.unwrap(); - let model_requests = requests - .iter() - .filter(|request| request.url.path() == "/codex/models") - .collect::>(); - assert_eq!(model_requests.len(), 2); - assert!(model_requests.iter().all(|request| { - request - .headers - .get("user-agent") - .and_then(|value| value.to_str().ok()) - == Some(expected_user_agent.as_str()) - })); - assert!(model_requests.iter().all(|request| { - request - .headers - .get("sec-ch-ua") - .and_then(|value| value.to_str().ok()) - == Some(expected_sec_ch_ua.as_str()) - })); - - let catalog_response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models/catalog") - .header("authorization", format!("Bearer {}", generated.plaintext)) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let catalog = response_json(catalog_response).await; - assert_eq!(catalog[0]["id"], "gpt-refresh"); - assert_eq!(catalog[0]["source"], "backend"); - - let acquired = state - .services - .accounts - .acquire_runtime_account("gpt-refresh") - .await - .unwrap(); - assert_ne!(acquired.plan_type.as_deref(), Some("free")); -} - -fn models_test_fingerprint() -> Fingerprint { - Fingerprint { - app_version: "27.333.444".to_string(), - build_number: "9002".to_string(), - chromium_version: "156".to_string(), - ..Fingerprint::default_for_tests() - } -} - -async fn insert_account( - repo: &AccountRepository, - id: &'static str, - plan_type: &'static str, - token: &'static str, -) { - repo.insert(NewAccount { - id: id.to_string(), - email: None, - account_id: Some(format!("{id}-chatgpt")), - user_id: None, - label: None, - plan_type: Some(plan_type.to_string()), - access_token: SecretString::new(token.to_string().into()), - refresh_token: Some(SecretString::new(format!("refresh-{id}").into())), - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); -} - -fn pool_account(id: &str, plan_type: &str) -> Account { - let mut account = Account::test(id, AccountStatus::Active); - account.plan_type = Some(plan_type.to_string()); - account -} diff --git a/tests/admin/session_login_route.rs b/crates/server/tests/admin_session_routes.rs similarity index 78% rename from tests/admin/session_login_route.rs rename to crates/server/tests/admin_session_routes.rs index 5a2263d86..85bc648e6 100644 --- a/tests/admin/session_login_route.rs +++ b/crates/server/tests/admin_session_routes.rs @@ -5,20 +5,154 @@ use axum::{ http::{Request, StatusCode}, }; use chrono::Utc; -use serde_json::Value; -use tower::ServiceExt; - -use codex_proxy_rs::{ +use codex_proxy_platform::{ config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, + UsageStatsConfig, WebSocketPoolConfig, }, - platform::identity::admin_session::hash_admin_password, - platform::storage::db::connect_sqlite, - runtime::build_router, - runtime::state::AppState, + crypto::SecretBox, + identity::{hash_admin_password, ApiKeyHasher}, + storage::connect_sqlite, }; +use codex_proxy_runtime::state::AppState; +use codex_proxy_server::router; +use serde_json::Value; +use sqlx::SqlitePool; +use tower::util::ServiceExt; + +#[tokio::test] +async fn admin_login_should_issue_http_only_session_cookie() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-login.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_user(&pool, "correct-password").await; + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([121u8; 32]), + ApiKeyHasher::new([122u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/login") + .header("content-type", "application/json") + .header("x-request-id", "req_login") + .body(Body::from(r#"{"password":"correct-password"}"#)) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let cookie = response + .headers() + .get("set-cookie") + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string); + + assert_eq!(status, StatusCode::OK); + let body = response_json(response).await; + let cookie = cookie.expect("login should set admin session cookie"); + assert!(cookie.starts_with("cpr_admin_session=")); + assert!(cookie.contains("HttpOnly")); + assert!(cookie.contains("SameSite=Lax")); + assert_eq!(body["code"], 200); + assert_eq!(body["requestId"], "req_login"); + assert!(body["data"]["expiresAt"].is_string()); + + let logs_response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/logs") + .header("cookie", cookie.split(';').next().unwrap()) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(logs_response.status(), StatusCode::OK); +} + +#[tokio::test] +async fn admin_login_should_reject_client_api_key_as_password_or_authorization() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-login-rejects-client-key.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_admin_user(&pool, "correct-password").await; + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([123u8; 32]), + ApiKeyHasher::new([124u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/login") + .header("content-type", "application/json") + .header("authorization", "Bearer cpr_not_an_admin_session") + .header("x-request-id", "req_login_bad") + .body(Body::from(r#"{"password":"cpr_not_an_admin_password"}"#)) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let has_cookie = response.headers().get("set-cookie").is_some(); + + assert_eq!(status, StatusCode::UNAUTHORIZED); + let body = response_json(response).await; + assert!(!has_cookie); + assert_eq!(body["code"], 40102); + assert_eq!(body["requestId"], "req_login_bad"); + + let logs_response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/logs") + .header("authorization", "Bearer cpr_not_an_admin_session") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(logs_response.status(), StatusCode::UNAUTHORIZED); +} + +async fn seed_admin_user(pool: &SqlitePool, password: &str) { + let now = Utc::now().to_rfc3339(); + let hash = hash_admin_password(password).unwrap(); + 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(); +} + +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_config(database_url: String) -> AppConfig { AppConfig { @@ -66,12 +200,12 @@ fn test_config(database_url: String) -> AppConfig { tls: TlsConfig { force_http11: false, }, - ws_pool: Default::default(), + ws_pool: WebSocketPoolConfig::default(), admin: AdminConfig { session_ttl_minutes: 60, + session_cleanup_interval_secs: 3600, default_username: "admin".to_string(), default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, }, logging: LoggingConfig { directory: "logs".to_string(), @@ -82,121 +216,3 @@ fn test_config(database_url: String) -> AppConfig { }, } } - -#[tokio::test] -async fn admin_login_issues_http_only_session_cookie() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-login.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_user(&pool, "correct-password").await; - let app = build_router(AppState::with_pool(test_config(url), pool.clone())); - - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/login") - .header("content-type", "application/json") - .header("x-request-id", "req_login") - .body(Body::from(r#"{"password":"correct-password"}"#)) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let cookie = response - .headers() - .get("set-cookie") - .unwrap() - .to_str() - .unwrap() - .to_string(); - assert!(cookie.starts_with("cpr_admin_session=")); - assert!(cookie.contains("HttpOnly")); - assert!(cookie.contains("SameSite=Lax")); - let session_cookie = cookie.split(';').next().unwrap().to_string(); - let body = response_json(response).await; - assert_eq!(body["code"], 200); - assert_eq!(body["requestId"], "req_login"); - assert!(body["data"]["expiresAt"].is_string()); - - let logs_response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/logs") - .header("cookie", session_cookie) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(logs_response.status(), StatusCode::OK); -} - -#[tokio::test] -async fn admin_login_rejects_client_api_key_as_password_or_authorization() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-login.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_user(&pool, "correct-password").await; - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/login") - .header("content-type", "application/json") - .header("authorization", "Bearer cpr_not_an_admin_session") - .header("x-request-id", "req_login_bad") - .body(Body::from(r#"{"password":"cpr_not_an_admin_password"}"#)) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - assert!(response.headers().get("set-cookie").is_none()); - let body = response_json(response).await; - assert_eq!(body["code"], 40102); - assert_eq!(body["requestId"], "req_login_bad"); - - let logs_response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/logs") - .header("authorization", "Bearer cpr_not_an_admin_session") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(logs_response.status(), StatusCode::UNAUTHORIZED); -} - -async fn seed_admin_user(pool: &sqlx::SqlitePool, password: &str) { - let now = Utc::now().to_rfc3339(); - let hash = hash_admin_password(password).unwrap(); - 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(); -} - -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.rs b/crates/server/tests/admin_settings_routes.rs similarity index 71% rename from tests/admin/settings_route.rs rename to crates/server/tests/admin_settings_routes.rs index 7695702d8..a9c80a17c 100644 --- a/tests/admin/settings_route.rs +++ b/crates/server/tests/admin_settings_routes.rs @@ -1,24 +1,344 @@ use std::{collections::BTreeMap, fs, path::Path}; use axum::{ - body::Body, + body::{to_bytes, Body}, http::{Request, StatusCode}, }; -use serde_json::json; -use tower::ServiceExt; - -use codex_proxy_rs::{ +use codex_proxy_platform::{ config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, + UsageStatsConfig, WebSocketPoolConfig, }, - platform::storage::db::connect_sqlite, - runtime::build_router, - runtime::state::AppState, + crypto::SecretBox, + identity::ApiKeyHasher, + storage::connect_sqlite, }; +use codex_proxy_runtime::state::AppState; +use codex_proxy_server::router; +use serde_json::{json, Value}; +use sqlx::SqlitePool; +use tower::util::ServiceExt; -use crate::support::{response_json, seed_admin_session}; +#[tokio::test] +async fn admin_settings_should_require_admin_session_cookie() { + let (app, _dir) = admin_settings_test_app("admin-settings-auth.sqlite").await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/settings") + .header("x-request-id", "req_settings") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert_eq!(body["code"], 40101); + assert_eq!(body["requestId"], "req_settings"); +} + +#[tokio::test] +async fn admin_diagnostics_should_return_runtime_summary_without_secrets() { + let (app, _dir) = admin_settings_test_app("admin-diagnostics.sqlite").await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/diagnostics") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_diagnostics") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["requestId"], "req_diagnostics"); + assert_eq!(body["data"]["status"], "ok"); + assert_eq!(body["data"]["runtime"]["packageName"], "codex-proxy-server"); + assert_eq!( + body["data"]["transport"]["backendBaseUrl"], + "https://chatgpt.com/backend-api" + ); + assert_eq!(body["data"]["settings"]["defaultModel"], "gpt-5.5"); + assert_eq!(body["data"]["accounts"]["pool"]["total"], 0); + assert!(!serde_json::to_string(&body).unwrap().contains("secret")); +} + +#[tokio::test] +async fn admin_settings_should_return_runtime_fields() { + let (app, _dir) = admin_settings_test_app("admin-settings-get.sqlite").await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/settings") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_settings_get") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["requestId"], "req_settings_get"); + assert_eq!( + body["data"], + json!({ + "defaultModel": "gpt-5.5", + "defaultReasoningEffort": "high", + "serviceTier": "flex", + "modelAliases": { "codex-fast": "gpt-5.5" }, + "refreshEnabled": true, + "refreshMarginSeconds": 240, + "refreshConcurrency": 2, + "maxConcurrentPerAccount": 4, + "requestIntervalMs": 50, + "rotationStrategy": "least_used", + "tierPriority": ["team", "plus"], + "quotaRefreshIntervalMinutes": 5, + "quotaWarningThresholds": { + "primary": [80, 90], + "secondary": [70, 95] + }, + "quotaSkipExhausted": true, + "logsEnabled": true, + "logsCapacity": 2000, + "logsCaptureBody": false, + "usageHistoryRetentionDays": 30 + }) + ); +} + +#[tokio::test] +async fn admin_settings_patch_should_require_admin_session_cookie() { + let (app, _dir) = admin_settings_test_app("admin-settings-patch-auth.sqlite").await; + + let response = app + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/admin/settings") + .header("content-type", "application/json") + .header("x-request-id", "req_settings_patch_auth") + .body(Body::from(r#"{"defaultModel":"gpt-5.5"}"#)) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert_eq!(body["code"], 40101); + assert_eq!(body["requestId"], "req_settings_patch_auth"); +} + +#[tokio::test] +async fn admin_settings_patch_should_persist_retained_fields_to_local_yaml() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-settings-patch.sqlite"); + let url = format!("sqlite://{}", db.display()); + let config = test_config(url); + write_config_yaml(dir.path(), &config); + let pool = connect_sqlite(&config.database.url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + let state = AppState::with_pool_secret_api_key_hasher_and_local_config_path( + config, + pool, + SecretBox::new([53u8; 32]), + ApiKeyHasher::new([54u8; 32]), + dir.path().join("local.yaml"), + ); + let app = router::router().with_state(state); + + let response = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/admin/settings") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_settings_patch") + .body(Body::from( + json!({ + "defaultModel": "gpt-6", + "defaultReasoningEffort": "medium", + "serviceTier": "priority", + "modelAliases": { "fast": "gpt-6-fast" }, + "refreshEnabled": false, + "refreshMarginSeconds": 180, + "refreshConcurrency": 3, + "maxConcurrentPerAccount": 5, + "requestIntervalMs": 125, + "rotationStrategy": "round_robin", + "tierPriority": ["pro", "plus"], + "quotaRefreshIntervalMinutes": 15, + "quotaWarningThresholds": { + "primary": [75, 90], + "secondary": [65, 95] + }, + "quotaSkipExhausted": false, + "logsEnabled": false, + "logsCapacity": 3000, + "logsCaptureBody": true, + "usageHistoryRetentionDays": 60 + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["requestId"], "req_settings_patch"); + assert_eq!(body["data"]["defaultModel"], "gpt-6"); + assert_eq!(body["data"]["rotationStrategy"], "round_robin"); + assert_eq!( + body["data"]["quotaWarningThresholds"]["primary"], + json!([75, 90]) + ); + + let reloaded = AppConfig::load_from_dir(dir.path()).unwrap(); + assert_eq!(reloaded.model.default_model, "gpt-6"); + assert_eq!( + reloaded.model.aliases.get("fast").map(String::as_str), + Some("gpt-6-fast") + ); + assert!(!reloaded.auth.refresh_enabled); + assert_eq!(reloaded.auth.rotation_strategy, "round_robin"); + assert_eq!(reloaded.quota.warning_thresholds.primary, vec![75, 90]); + assert_eq!(reloaded.usage_stats.history_retention_days, Some(60)); + + let get_response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/settings") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json(get_response).await; + assert_eq!(body["data"]["defaultModel"], "gpt-6"); +} + +#[tokio::test] +async fn admin_settings_patch_should_reject_unsupported_or_invalid_fields() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-settings-patch-invalid.sqlite"); + let url = format!("sqlite://{}", db.display()); + let config = test_config(url); + write_config_yaml(dir.path(), &config); + let pool = connect_sqlite(&config.database.url).await.unwrap(); + seed_admin_session(&pool, "session_1").await; + let state = AppState::with_pool_secret_api_key_hasher_and_local_config_path( + config, + pool, + SecretBox::new([53u8; 32]), + ApiKeyHasher::new([54u8; 32]), + dir.path().join("local.yaml"), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("PATCH") + .uri("/api/admin/settings") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .header("x-request-id", "req_settings_patch_invalid") + .body(Body::from( + json!({ + "unknownSetting": true, + "rotationStrategy": "random" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(body["code"], 40001); + assert_eq!(body["requestId"], "req_settings_patch_invalid"); + assert!(!dir.path().join("local.yaml").exists()); +} + +async fn admin_settings_test_app(db_name: &str) -> (axum::Router, 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_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([53u8; 32]), + ApiKeyHasher::new([54u8; 32]), + ); + (router::router().with_state(state), dir) +} + +async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { + sqlx::query( + "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", + ) + .bind("admin_1") + .bind("hash") + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .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("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); +} + +async fn response_json(response: axum::response::Response) -> Value { + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + serde_json::from_slice(&bytes).unwrap() +} + +fn write_config_yaml(path: &Path, config: &AppConfig) { + fs::write( + path.join("config.yaml"), + serde_yml::to_string(config).unwrap(), + ) + .unwrap(); +} fn test_config(database_url: String) -> AppConfig { let mut aliases = BTreeMap::new(); @@ -68,12 +388,12 @@ fn test_config(database_url: String) -> AppConfig { tls: TlsConfig { force_http11: false, }, - ws_pool: Default::default(), + ws_pool: WebSocketPoolConfig::default(), admin: AdminConfig { session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, default_username: "admin".to_string(), default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, }, logging: LoggingConfig { directory: "logs".to_string(), @@ -84,253 +404,3 @@ fn test_config(database_url: String) -> AppConfig { }, } } - -fn write_config_yaml(path: &Path, config: &AppConfig) { - fs::write( - path.join("config.yaml"), - serde_yml::to_string(config).unwrap(), - ) - .unwrap(); -} - -#[tokio::test] -async fn admin_settings_should_require_admin_session_cookie() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-settings.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/settings") - .header("x-request-id", "req_settings") - .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_settings"); -} - -#[tokio::test] -async fn admin_settings_patch_should_require_admin_session_cookie() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-settings-patch-auth.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .oneshot( - Request::builder() - .method("PATCH") - .uri("/api/admin/settings") - .header("content-type", "application/json") - .header("x-request-id", "req_settings_patch_auth") - .body(Body::from(r#"{"defaultModel":"gpt-5.5"}"#)) - .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_settings_patch_auth"); -} - -#[tokio::test] -async fn admin_settings_patch_should_persist_retained_fields_to_local_yaml() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-settings-patch.sqlite"); - let url = format!("sqlite://{}", db.display()); - let config = test_config(url); - write_config_yaml(dir.path(), &config); - let pool = connect_sqlite(&config.database.url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let app = build_router(AppState::with_pool_and_local_config_path( - config, - pool, - dir.path().join("local.yaml"), - )); - - let response = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri("/api/admin/settings") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_settings_patch") - .body(Body::from( - json!({ - "defaultModel": "gpt-6", - "defaultReasoningEffort": "medium", - "serviceTier": "priority", - "modelAliases": { "fast": "gpt-6-fast" }, - "refreshEnabled": false, - "refreshMarginSeconds": 180, - "refreshConcurrency": 3, - "maxConcurrentPerAccount": 5, - "requestIntervalMs": 125, - "rotationStrategy": "round_robin", - "tierPriority": ["pro", "plus"], - "quotaRefreshIntervalMinutes": 15, - "quotaWarningThresholds": { - "primary": [75, 90], - "secondary": [65, 95] - }, - "quotaSkipExhausted": false, - "logsEnabled": false, - "logsCapacity": 3000, - "logsCaptureBody": true, - "usageHistoryRetentionDays": 60 - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["code"], 200); - assert_eq!(body["requestId"], "req_settings_patch"); - assert_eq!(body["data"]["defaultModel"], "gpt-6"); - assert_eq!(body["data"]["rotationStrategy"], "round_robin"); - assert_eq!( - body["data"]["quotaWarningThresholds"]["primary"], - json!([75, 90]) - ); - - let reloaded = AppConfig::load_from_dir(dir.path()).unwrap(); - assert_eq!(reloaded.model.default_model, "gpt-6"); - assert_eq!( - reloaded.model.aliases.get("fast").map(String::as_str), - Some("gpt-6-fast") - ); - assert!(!reloaded.auth.refresh_enabled); - assert_eq!(reloaded.auth.rotation_strategy, "round_robin"); - assert_eq!(reloaded.quota.warning_thresholds.primary, vec![75, 90]); - assert_eq!(reloaded.usage_stats.history_retention_days, Some(60)); - - let get_response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/settings") - .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"]["defaultModel"], "gpt-6"); -} - -#[tokio::test] -async fn admin_settings_patch_should_reject_unsupported_or_invalid_fields() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-settings-patch-invalid.sqlite"); - let url = format!("sqlite://{}", db.display()); - let config = test_config(url); - write_config_yaml(dir.path(), &config); - let pool = connect_sqlite(&config.database.url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let app = build_router(AppState::with_pool_and_local_config_path( - config, - pool, - dir.path().join("local.yaml"), - )); - - let response = app - .oneshot( - Request::builder() - .method("PATCH") - .uri("/api/admin/settings") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_settings_patch_invalid") - .body(Body::from( - json!({ - "unknownSetting": true, - "rotationStrategy": "random" - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_REQUEST); - let body = response_json(response).await; - assert_eq!(body["code"], 40001); - assert_eq!(body["requestId"], "req_settings_patch_invalid"); - assert!(!dir.path().join("local.yaml").exists()); -} - -#[tokio::test] -async fn admin_settings_should_return_only_in_scope_runtime_fields() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-settings.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(test_config(url), pool)); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/settings") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_settings") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["code"], 200); - assert_eq!(body["requestId"], "req_settings"); - assert_eq!( - body["data"], - json!({ - "defaultModel": "gpt-5.5", - "defaultReasoningEffort": "high", - "serviceTier": "flex", - "modelAliases": { "codex-fast": "gpt-5.5" }, - "refreshEnabled": true, - "refreshMarginSeconds": 240, - "refreshConcurrency": 2, - "maxConcurrentPerAccount": 4, - "requestIntervalMs": 50, - "rotationStrategy": "least_used", - "tierPriority": ["team", "plus"], - "quotaRefreshIntervalMinutes": 5, - "quotaWarningThresholds": { - "primary": [80, 90], - "secondary": [70, 95] - }, - "quotaSkipExhausted": true, - "logsEnabled": true, - "logsCapacity": 2000, - "logsCaptureBody": false, - "usageHistoryRetentionDays": 30 - }) - ); -} diff --git a/crates/server/tests/assets_routes.rs b/crates/server/tests/assets_routes.rs new file mode 100644 index 000000000..f6d08037e --- /dev/null +++ b/crates/server/tests/assets_routes.rs @@ -0,0 +1,155 @@ +use axum::{ + body::{to_bytes, Body}, + http::{header, Request, StatusCode}, +}; +use codex_proxy_platform::{ + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, + UsageStatsConfig, WebSocketPoolConfig, + }, + crypto::SecretBox, + identity::ApiKeyHasher, + storage::connect_sqlite, +}; +use codex_proxy_runtime::state::AppState; +use codex_proxy_server::router; +use serde_json::Value; +use tower::ServiceExt; + +#[tokio::test] +async fn server_router_should_serve_frontend_assets_without_shadowing_api_routes() { + let dir = tempfile::tempdir().expect("temp dir"); + let dist = dir.path().join("dist"); + std::fs::create_dir_all(dist.join("assets")).expect("assets dir should be created"); + std::fs::write( + dist.join("index.html"), + "
Codex Proxy Admin
", + ) + .expect("index should be written"); + std::fs::write(dist.join("assets").join("app.js"), "window.__asset = true;") + .expect("asset should be written"); + let db = dir.path().join("assets-routes.sqlite"); + let database_url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&database_url).await.expect("sqlite pool"); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(database_url), + pool, + SecretBox::new([61u8; 32]), + ApiKeyHasher::new([62u8; 32]), + ); + let app = router::router_with_assets(&dist).with_state(state); + + let index = app.clone().oneshot(request("/")).await.expect("index"); + let asset = app + .clone() + .oneshot(request("/assets/app.js")) + .await + .expect("asset"); + let api = app + .oneshot(request("/api/admin/settings")) + .await + .expect("api response"); + + assert_eq!(index.status(), StatusCode::OK); + assert_eq!(asset.status(), StatusCode::OK); + assert_eq!(api.status(), StatusCode::UNAUTHORIZED); + assert_eq!( + asset.headers().get(header::CACHE_CONTROL), + Some(&header::HeaderValue::from_static( + "public, max-age=31536000, immutable" + )) + ); + assert_body_contains(index, "Codex Proxy Admin").await; + assert_body_contains(asset, "__asset").await; + assert_eq!(response_json(api).await["code"], Value::from(40101)); +} + +fn request(uri: &str) -> Request { + Request::builder() + .uri(uri) + .body(Body::empty()) + .expect("request should build") +} + +async fn assert_body_contains(response: axum::response::Response, expected: &str) { + let bytes = to_bytes(response.into_body(), usize::MAX) + .await + .expect("body should collect"); + let text = std::str::from_utf8(&bytes).expect("body should be utf8"); + + assert!( + text.contains(expected), + "body should contain `{expected}`, got `{text}`" + ); +} + +async fn response_json(response: axum::response::Response) -> Value { + let bytes = to_bytes(response.into_body(), usize::MAX) + .await + .expect("body should collect"); + serde_json::from_slice(&bytes).expect("body should be json") +} + +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: Default::default(), + }, + auth: AuthConfig { + refresh_margin_seconds: 240, + refresh_enabled: true, + refresh_concurrency: 2, + max_concurrent_per_account: 4, + request_interval_ms: 50, + rotation_strategy: "least_used".to_string(), + tier_priority: Vec::new(), + oauth_client_id: "app_id".to_string(), + oauth_auth_endpoint: "https://auth.invalid/authorize".to_string(), + oauth_token_endpoint: "https://auth.invalid/token".to_string(), + }, + 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: Some(30), + }, + 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, + }, + ws_pool: WebSocketPoolConfig::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + }, + } +} diff --git a/tests/runtime/http_trace.rs b/crates/server/tests/http_trace_middleware.rs similarity index 91% rename from tests/runtime/http_trace.rs rename to crates/server/tests/http_trace_middleware.rs index efd59a5db..8d779009a 100644 --- a/tests/runtime/http_trace.rs +++ b/crates/server/tests/http_trace_middleware.rs @@ -10,11 +10,10 @@ use axum::{ routing::get, Router, }; +use codex_proxy_server::middleware::{request_id::attach_request_id, trace::http_trace_layer}; use tower::ServiceExt; use tracing_subscriber::fmt::MakeWriter; -use codex_proxy_rs::{platform::http::request_id::attach_request_id, runtime::http_trace_layer}; - #[derive(Clone, Default)] struct SharedLogBuffer { inner: Arc>>, @@ -84,8 +83,8 @@ async fn http_trace_should_include_request_id_and_completion_fields() { let output = logs.content(); assert!( output.contains("req_trace") - && output.contains("收到 HTTP 请求") - && output.contains("HTTP 请求完成") + && output.contains("received HTTP request") + && output.contains("completed HTTP request") && output.contains("\"status\":204") && output.contains("latency_ms"), "unexpected trace output: {output}" diff --git a/crates/server/tests/openai_chat_upstream.rs b/crates/server/tests/openai_chat_upstream.rs new file mode 100644 index 000000000..05e6acb8b --- /dev/null +++ b/crates/server/tests/openai_chat_upstream.rs @@ -0,0 +1,8579 @@ +use std::{ + collections::BTreeMap, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, Mutex, + }, + time::{Duration as StdDuration, Instant}, +}; + +use axum::{ + body::{to_bytes, Body, Bytes}, + http::{Request, StatusCode}, +}; +use chrono::{DateTime, Duration, Utc}; +use codex_proxy_adapters::sqlite::{ + cookies::SqliteCookieStore, + events::{EventLogFilter, SqliteEventLogStore}, + session_affinity::SqliteSessionAffinityStore, +}; +use codex_proxy_core::{ + events::model::EventLevel, gateway::conversation::build_conversation_identity, + serving::affinity::SessionAffinityEntry, +}; +use codex_proxy_platform::{ + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, + UsageStatsConfig, WebSocketPoolConfig, + }, + crypto::SecretBox, + identity::ApiKeyHasher, + storage::connect_sqlite, +}; +use codex_proxy_runtime::state::AppState; +use codex_proxy_server::router; +use futures::{SinkExt, StreamExt}; +use secrecy::SecretString; +use serde_json::{json, Value}; +use sqlx::SqlitePool; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + sync::oneshot, + time::timeout, +}; +use tokio_tungstenite::{ + accept_async, accept_hdr_async, + tungstenite::{ + handshake::server::{Request as WsRequest, Response as WsResponse}, + Message, + }, +}; +use tower::util::ServiceExt; +use wiremock::{ + matchers::{header, method, path}, + Mock, MockServer, ResponseTemplate, +}; + +const TEST_INSTALLATION_ID: &str = "b4f9d503-07b1-457b-a0da-87e6836b1c43"; + +const CHAT_SUCCESS_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"hello"} + +event: response.completed +data: {"response":{"id":"resp_1","output_text":"hello","usage":{"input_tokens":9,"output_tokens":3}}} + +"#; + +const RESPONSES_SUCCESS_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"response hello"} + +event: response.completed +data: {"response":{"id":"resp_response_1","object":"response","status":"completed","usage":{"input_tokens":5,"output_tokens":2}}} + +"#; + +const RESPONSES_COMPLETED_USAGE_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"usage hello"} + +event: response.completed +data: {"response":{"id":"resp_usage","object":"response","status":"completed","output_text":"usage hello","usage":{"input_tokens":7,"output_tokens":4,"input_tokens_details":{"cached_tokens":2}}}} + +"#; + +const RESPONSES_COMPLETED_IMAGE_USAGE_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"image done"} + +event: response.completed +data: {"response":{"id":"resp_image","object":"response","status":"completed","output_text":"image done","usage":{"input_tokens":7,"output_tokens":4,"input_tokens_details":{"cached_tokens":2}},"tool_usage":{"image_gen":{"input_tokens":31,"output_tokens":9}}}} + +"#; + +const RESPONSES_EMPTY_COMPLETED_SSE: &str = r#"event: response.completed +data: {"response":{"id":"resp_empty","object":"response","status":"completed"}} + +"#; + +const RESPONSES_TEXT_DELTAS_COMPLETED_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"hello "} + +event: response.output_text.delta +data: {"delta":"from deltas"} + +event: response.completed +data: {"response":{"id":"resp_text","object":"response","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":3}}} + +"#; + +const RESPONSES_DONE_ITEM_COMPLETED_SSE: &str = r#"event: response.output_item.done +data: {"item":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"from done item"}]}} + +event: response.completed +data: {"response":{"id":"resp_item","object":"response","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":3}}} + +"#; + +const RESPONSES_TUPLE_OBJECT_SSE: &str = r#"event: response.output_text.delta +data: {"type":"response.output_text.delta","delta":"{\"point\":{\"0\":1,\"1\":2}}"} + +event: response.completed +data: {"response":{"id":"resp_tuple","object":"response","status":"completed","output_text":"{\"point\":{\"0\":0,\"1\":0}}","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"{\"point\":{\"0\":1,\"1\":2}}"}]}],"usage":{"input_tokens":1,"output_tokens":1}}} + +"#; + +const RESPONSES_AFTER_5XX_RETRY_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"retry hello"} + +event: response.completed +data: {"response":{"id":"resp_after_5xx_retry","object":"response","status":"completed","usage":{"input_tokens":4,"output_tokens":2}}} + +"#; + +const RESPONSES_AFTER_402_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"quota fallback hello"} + +event: response.completed +data: {"response":{"id":"resp_after_402","object":"response","status":"completed","usage":{"input_tokens":6,"output_tokens":2}}} + +"#; + +const RESPONSES_FAILED_QUOTA_SSE: &str = r#"event: response.failed +data: {"response":{"id":"resp_sse_quota_failed","object":"response","status":"failed","error":{"code":"quota_exceeded","message":"quota exhausted"}}} + +"#; + +const RESPONSES_FAILED_AUTH_SSE: &str = r#"event: response.failed +data: {"response":{"id":"resp_sse_auth_failed","object":"response","status":"failed","error":{"code":"token_invalid","message":"token revoked"}}} + +"#; + +const RESPONSES_AFTER_401_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"auth fallback hello"} + +event: response.completed +data: {"response":{"id":"resp_after_401","object":"response","status":"completed","usage":{"input_tokens":3,"output_tokens":1}}} + +"#; + +const RESPONSES_AFTER_403_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"banned fallback hello"} + +event: response.completed +data: {"response":{"id":"resp_after_403","object":"response","status":"completed","usage":{"input_tokens":5,"output_tokens":1}}} + +"#; + +const RESPONSES_STREAM_USAGE_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"stream hello"} + +event: response.completed +data: {"response":{"id":"resp_stream_usage","object":"response","status":"completed","usage":{"input_tokens":3,"output_tokens":5}}} + +"#; + +const RESPONSES_STREAM_AFTER_429_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"rate fallback hello"} + +event: response.completed +data: {"response":{"id":"resp_stream_after_429","object":"response","status":"completed","usage":{"input_tokens":4,"output_tokens":1}}} + +"#; + +const RESPONSES_STREAM_AFTER_5XX_RETRY_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"retry stream hello"} + +event: response.completed +data: {"response":{"id":"resp_stream_after_5xx_retry","object":"response","status":"completed","usage":{"input_tokens":7,"output_tokens":2}}} + +"#; + +const RESPONSES_STREAM_AFTER_402_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"quota stream fallback hello"} + +event: response.completed +data: {"response":{"id":"resp_stream_after_402","object":"response","status":"completed","usage":{"input_tokens":6,"output_tokens":3}}} + +"#; + +const RESPONSES_STREAM_AFTER_401_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"auth stream fallback hello"} + +event: response.completed +data: {"response":{"id":"resp_stream_after_401","object":"response","status":"completed","usage":{"input_tokens":4,"output_tokens":2}}} + +"#; + +const RESPONSES_STREAM_AFTER_403_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"banned stream fallback hello"} + +event: response.completed +data: {"response":{"id":"resp_stream_after_403","object":"response","status":"completed","usage":{"input_tokens":5,"output_tokens":2}}} + +"#; + +const RESPONSES_AFTER_CLOUDFLARE_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"cloudflare fallback hello"} + +event: response.completed +data: {"response":{"id":"resp_after_cloudflare","object":"response","status":"completed","usage":{"input_tokens":5,"output_tokens":2}}} + +"#; + +const RESPONSES_STREAM_AFTER_CLOUDFLARE_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"cloudflare stream fallback hello"} + +event: response.completed +data: {"response":{"id":"resp_stream_after_cloudflare","object":"response","status":"completed","usage":{"input_tokens":6,"output_tokens":2}}} + +"#; + +const RESPONSES_AFTER_MODEL_UNSUPPORTED_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"model fallback hello"} + +event: response.completed +data: {"response":{"id":"resp_after_model_unsupported","object":"response","status":"completed","usage":{"input_tokens":4,"output_tokens":2}}} + +"#; + +const RESPONSES_STREAM_AFTER_MODEL_UNSUPPORTED_SSE: &str = r#"event: response.output_text.delta +data: {"delta":"model stream fallback hello"} + +event: response.completed +data: {"response":{"id":"resp_stream_after_model_unsupported","object":"response","status":"completed","usage":{"input_tokens":5,"output_tokens":2}}} + +"#; + +const RESPONSES_FAILED_MODEL_UNSUPPORTED_SSE: &str = r#"event: response.failed +data: {"response":{"id":"resp_model_unsupported_failed","object":"response","status":"failed","error":{"code":"model_not_supported","message":"Model gpt-5.5 is not supported on this account plan"}}} + +"#; + +const WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY: &str = include_str!( + "../../../tests/fixtures/responses/websocket/completed_with_reasoning_replay.json" +); +const WEBSOCKET_HISTORY_RATE_LIMITED: &str = + include_str!("../../../tests/fixtures/responses/websocket/history_rate_limited.json"); +const WEBSOCKET_RATE_LIMITED: &str = + include_str!("../../../tests/fixtures/responses/websocket/rate_limited.json"); +const WEBSOCKET_TOKEN_REVOKED: &str = + include_str!("../../../tests/fixtures/responses/websocket/token_revoked.json"); +const WEBSOCKET_FIRST_ACCOUNT_LIMITED: &str = + include_str!("../../../tests/fixtures/responses/websocket/first_account_limited.json"); +const WEBSOCKET_SECOND_ACCOUNT_LIMITED: &str = + include_str!("../../../tests/fixtures/responses/websocket/second_account_limited.json"); +const WEBSOCKET_INVALID_ENCRYPTED_CONTENT: &str = + include_str!("../../../tests/fixtures/responses/websocket/invalid_encrypted_content.json"); +const WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND: &str = + include_str!("../../../tests/fixtures/responses/websocket/previous_response_not_found.json"); +const WEBSOCKET_UNANSWERED_FUNCTION_CALL: &str = + include_str!("../../../tests/fixtures/responses/websocket/unanswered_function_call.json"); +const REASONING_REPLAY_REQUEST_GOLDEN: &str = + include_str!("../../../tests/fixtures/responses/golden/reasoning_replay_request.json"); + +#[tokio::test] +async fn chat_completions_should_dispatch_to_codex_and_return_openai_response() { + 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")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(CHAT_SUCCESS_SSE), + ) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [ + {"role": "system", "content": "You are concise."}, + {"role": "user", "content": "Say hello"} + ] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["object"], "chat.completion"); + assert_eq!(body["model"], "gpt-5.5"); + assert_eq!(body["choices"][0]["message"]["content"], "hello"); + assert_eq!(body["usage"]["prompt_tokens"], 9); + assert_eq!(body["usage"]["completion_tokens"], 3); + let usage: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(usage, (1, 9, 3)); +} + +#[tokio::test] +async fn chat_completions_should_forward_runtime_installation_id_to_codex() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("x-codex-installation-id", TEST_INSTALLATION_ID)) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(CHAT_SUCCESS_SSE), + ) + .mount(&server) + .await; + + let (app, api_key, _dir) = + test_app_with_account_and_installation_id(server.uri(), TEST_INSTALLATION_ID.to_string()) + .await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [ + {"role": "user", "content": "Say hello"} + ] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); +} + +#[tokio::test] +async fn chat_completions_should_dispatch_from_restored_runtime_account_pool() { + 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")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(CHAT_SUCCESS_SSE), + ) + .mount(&server) + .await; + + let (app, api_key, _dir) = + test_app_with_restored_pool_then_disabled_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [ + {"role": "user", "content": "Say hello"} + ] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); +} + +#[tokio::test] +async fn chat_completions_should_fallback_to_next_account_after_rate_limit() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("retry-after", "120") + .set_body_json(json!({ + "error": { + "message": "rate limited", + "resets_in_seconds": 120 + } + })), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(CHAT_SUCCESS_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "Say hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let authorizations = received_authorizations(&server).await; + + assert_eq!(status, StatusCode::OK, "requests: {authorizations:?}"); + assert_eq!(body["choices"][0]["message"]["content"], "hello"); + let persisted_quota_state: (i64, Option) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", + ) + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(persisted_quota_state.0, 1); + let cooldown_until = + chrono::DateTime::parse_from_rfc3339(persisted_quota_state.1.as_deref().unwrap()) + .unwrap() + .with_timezone(&Utc); + assert!(cooldown_until > Utc::now()); +} + +#[tokio::test] +async fn chat_completions_should_return_rate_limit_error_when_429_fallback_is_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(429).set_body_json(json!({ + "error": {"message": "rate limited"} + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "Say hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let quota_state: (i64, Option) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + let message = body["error"]["message"].as_str().unwrap_or_default(); + + assert_eq!(status, StatusCode::TOO_MANY_REQUESTS); + assert!(message.contains("All accounts exhausted (1 rate-limited)")); + assert!(message.contains("rate limited")); + assert_eq!(body["error"]["code"], "upstream_error"); + assert_eq!(quota_state.0, 1); + assert!(quota_state.1.is_some()); +} + +#[tokio::test] +async fn chat_completions_should_fallback_after_http_model_unsupported() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": { + "code": "model_not_supported", + "message": "Model gpt-5.5 is not supported on this account plan" + } + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(CHAT_SUCCESS_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "Say hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let primary_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["choices"][0]["message"]["content"], "hello"); + assert_eq!(primary_status.0, "active"); +} + +#[tokio::test] +async fn chat_completions_should_return_model_unsupported_error_when_fallback_is_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": { + "code": "model_not_supported", + "message": "Model gpt-5.5 is not supported on this account plan" + } + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "Say hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let message = body["error"]["message"].as_str().unwrap_or_default(); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert!(message.contains("All accounts exhausted (1 model-unsupported)")); + assert!(message.contains("model_not_supported")); + assert_eq!(body["error"]["code"], "upstream_error"); + assert_eq!(account_status.0, "active"); +} + +#[tokio::test] +async fn chat_completions_should_mark_expired_after_401_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(401).set_body_json(json!({ + "error": {"code": "token_revoked", "message": "token revoked"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(CHAT_SUCCESS_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "Say hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["choices"][0]["message"]["content"], "hello"); + assert_eq!(account_status.0, "expired"); +} + +#[tokio::test] +async fn chat_completions_should_return_auth_error_when_401_fallback_is_exhausted() { + 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", "message": "token revoked"} + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "Say hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let message = body["error"]["message"].as_str().unwrap_or_default(); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert!(message.contains("All accounts exhausted (1 expired)")); + assert!(message.contains("token_revoked")); + assert_eq!(body["error"]["code"], "upstream_error"); + assert_eq!(account_status.0, "expired"); +} + +#[tokio::test] +async fn chat_completions_should_mark_banned_when_401_says_account_deactivated() { + 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": {"message": "account deactivated"} + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "Say hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + assert_eq!(account_status.0, "banned"); +} + +#[tokio::test] +async fn chat_completions_should_mark_banned_after_403_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(403).set_body_json(json!({ + "error": {"message": "account banned"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(CHAT_SUCCESS_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "Say hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["choices"][0]["message"]["content"], "hello"); + assert_eq!(account_status.0, "banned"); +} + +#[tokio::test] +async fn chat_completions_should_cool_down_cloudflare_403_and_fallback() { + let started_at = Utc::now(); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(403).set_body_string( + "Just a moment...cf_chl challenge", + )) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(CHAT_SUCCESS_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + cookie_store + .set_cookie_header("acct_primary", "cf_clearance=old") + .await + .unwrap(); + cookie_store + .set_cookie_header("acct_secondary", "cf_clearance=keep") + .await + .unwrap(); + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "Say hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let primary_state: (String, Option) = + sqlx::query_as("select status, cloudflare_cooldown_until from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + let cooldown_until = chrono::DateTime::parse_from_rfc3339(primary_state.1.as_deref().unwrap()) + .unwrap() + .with_timezone(&Utc); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["choices"][0]["message"]["content"], "hello"); + assert_eq!(primary_state.0, "active"); + assert!(cooldown_until > started_at); + assert_eq!( + cookie_store + .cookie_header("acct_primary", "chatgpt.com") + .await + .unwrap(), + None + ); + assert_eq!( + cookie_store + .cookie_header("acct_secondary", "chatgpt.com") + .await + .unwrap() + .as_deref(), + Some("cf_clearance=keep") + ); +} + +#[tokio::test] +async fn chat_completions_should_mark_quota_exhausted_after_402_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(402).set_body_json(json!({ + "error": {"message": "quota reached"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(CHAT_SUCCESS_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "Say hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["choices"][0]["message"]["content"], "hello"); + assert_eq!(account_status.0, "quota_exhausted"); +} + +#[tokio::test] +async fn chat_completions_should_return_quota_error_when_402_fallback_is_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(402).set_body_json(json!({ + "error": {"message": "quota reached"} + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "messages": [{"role": "user", "content": "Say hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + let message = body["error"]["message"].as_str().unwrap_or_default(); + + assert_eq!(status, StatusCode::PAYMENT_REQUIRED); + assert!(message.contains("All accounts exhausted (1 quota-exhausted)")); + assert!(message.contains("quota reached")); + assert_eq!(body["error"]["code"], "upstream_error"); + assert_eq!(account_status.0, "quota_exhausted"); +} + +#[tokio::test] +async fn responses_should_dispatch_to_codex_and_return_completed_response() { + 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")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_SUCCESS_SSE), + ) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_response_1"); + assert_eq!(body["output_text"], "response hello"); + assert_eq!(body["usage"]["input_tokens"], 5); + assert_eq!(body["usage"]["output_tokens"], 2); +} + +#[tokio::test] +async fn responses_should_reject_invalid_json_without_upstream_request() { + let server = MockServer::start().await; + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from("{")) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let requests = server.received_requests().await.unwrap(); + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(body["error"]["code"], "invalid_request"); + assert!(requests.is_empty()); +} + +#[tokio::test] +async fn responses_should_reject_non_object_json_without_upstream_request() { + let server = MockServer::start().await; + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from("[]")) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let requests = server.received_requests().await.unwrap(); + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(body["error"]["code"], "invalid_request"); + assert!(requests.is_empty()); +} + +#[tokio::test] +async fn responses_should_return_no_available_accounts_error_when_no_accounts_are_available() { + let server = MockServer::start().await; + let (app, api_key, _dir) = test_app_without_accounts(server.uri()).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(body["type"], "error"); + assert_eq!(body["error"]["type"], "server_error"); + assert_eq!(body["error"]["code"], "no_available_accounts"); +} + +#[tokio::test] +async fn responses_should_honor_explicit_http_sse_transport() { + 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(RESPONSES_STREAM_USAGE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_text(response).await; + let requests = server.received_requests().await.unwrap(); + let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(body.contains("event: response.output_text.delta")); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].method.as_str(), "POST"); + assert_eq!(upstream_body["model"], "gpt-5.5"); + assert!(upstream_body.get("use_websocket").is_none()); +} + +#[tokio::test] +async fn responses_should_stagger_same_account_requests_before_sending_upstream() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let request_times = Arc::new(Mutex::new(Vec::new())); + let request_times_for_server = Arc::clone(&request_times); + let (first_seen_tx, first_seen_rx) = oneshot::channel(); + let (release_first_tx, release_first_rx) = oneshot::channel(); + let upstream = tokio::spawn(async move { + let (mut first_socket, _) = listener.accept().await.unwrap(); + request_times_for_server + .lock() + .unwrap() + .push(Instant::now()); + first_seen_tx.send(()).unwrap(); + read_http_request(&mut first_socket).await; + + let (mut second_socket, _) = listener.accept().await.unwrap(); + request_times_for_server + .lock() + .unwrap() + .push(Instant::now()); + read_http_request(&mut second_socket).await; + write_http_sse_response(&mut second_socket, RESPONSES_COMPLETED_USAGE_SSE).await; + + release_first_rx.await.unwrap(); + write_http_sse_response(&mut first_socket, RESPONSES_COMPLETED_USAGE_SSE).await; + }); + + let (app, api_key, _pool, _dir) = test_app_with_account_pool_config(base_url, |config| { + config.auth.max_concurrent_per_account = 2; + config.auth.request_interval_ms = 300; + }) + .await; + let first_app = app.clone(); + let first_api_key = api_key.clone(); + let first_response = tokio::spawn(async move { + first_app + .oneshot(responses_http_sse_request( + &first_api_key, + "req_stagger_first", + )) + .await + .unwrap() + }); + first_seen_rx.await.unwrap(); + + let second = app + .clone() + .oneshot(responses_http_sse_request(&api_key, "req_stagger_second")) + .await + .unwrap(); + assert_eq!(second.status(), StatusCode::OK); + release_first_tx.send(()).unwrap(); + let first = first_response.await.unwrap(); + assert_eq!(first.status(), StatusCode::OK); + upstream.await.unwrap(); + + let times = request_times.lock().unwrap(); + assert_eq!(times.len(), 2); + let elapsed = times[1].duration_since(times[0]); + assert!( + elapsed >= StdDuration::from_millis(180), + "second upstream request was sent too early: {elapsed:?}" + ); +} + +#[tokio::test] +async fn responses_should_use_websocket_upstream_by_default_while_serving_sse() { + let (base_url, upstream) = spawn_single_websocket_completed_upstream("resp_ws_default").await; + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "generate": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let captured = upstream.await.unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("event: response.completed")); + assert!(body.contains("\"id\":\"resp_ws_default\"")); + assert_eq!(captured.payload["type"], "response.create"); + assert_eq!(captured.payload["model"], "gpt-5.5"); + assert_eq!(captured.payload["generate"], false); + assert!(captured.payload.get("previous_response_id").is_none()); + assert!(captured.payload["prompt_cache_key"] + .as_str() + .is_some_and(|value| value.starts_with("cp_"))); + assert!(captured.payload["client_metadata"]["x-codex-installation-id"].is_string()); + assert!(captured.payload["client_metadata"]["x-codex-window-id"] + .as_str() + .is_some_and(|value| value.starts_with("cp_") && value.ends_with(":0"))); +} + +#[tokio::test] +async fn responses_should_ignore_camel_case_use_websocket_field() { + let (base_url, upstream) = + spawn_single_websocket_completed_upstream("resp_ws_camel_case").await; + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "useWebSocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let captured = upstream.await.unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(captured.payload["type"], "response.create"); + assert!(captured.payload.get("useWebSocket").is_none()); + assert!(captured.payload.get("use_websocket").is_none()); +} + +#[tokio::test] +async fn responses_websocket_should_stream_first_frame_before_terminal_event() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let (first_frame_tx, first_frame_rx) = oneshot::channel(); + let (terminal_tx, terminal_rx) = oneshot::channel(); + let upstream = 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 payload = serde_json::from_str::(&message.into_text().unwrap()).unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "first websocket frame" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + first_frame_tx.send(()).unwrap(); + terminal_rx.await.unwrap(); + websocket + .send(Message::Text( + response_completed_websocket_message("resp_ws_streaming").into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + payload + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + let response_task = tokio::spawn(async move { + app.oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": true, + "use_websocket": true + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap() + }); + + first_frame_rx.await.unwrap(); + let response = timeout(StdDuration::from_millis(250), response_task) + .await + .expect("websocket response should be returned after the first non-error frame") + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let first_chunk = timeout( + StdDuration::from_secs(1), + first_response_body_chunk(response), + ) + .await + .expect("downstream should receive a websocket-backed SSE chunk before completion") + .expect("response body should produce a chunk"); + terminal_tx.send(()).unwrap(); + let payload = upstream.await.unwrap(); + + assert!(first_chunk.contains("event: response.output_text.delta")); + assert_eq!(payload["type"], "response.create"); +} + +#[tokio::test] +async fn responses_websocket_stream_should_synthesize_response_failed_when_closed_before_terminal() +{ + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = 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 payload = serde_json::from_str::(&message.into_text().unwrap()).unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "partial before websocket close" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + payload + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": true, + "use_websocket": true + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_text(response).await; + let payload = upstream.await.unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(payload["type"], "response.create"); + assert!(body.contains("event: response.output_text.delta")); + assert!(body.contains("event: response.failed")); + assert!(body.contains("stream_disconnected")); +} + +#[tokio::test] +async fn responses_websocket_should_reuse_connection_for_recorded_conversation() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let first_message = websocket.next().await.unwrap().unwrap(); + let first_payload = serde_json::from_str::(&first_message.into_text().unwrap()) + .expect("first websocket payload should be json"); + websocket + .send(Message::Text( + response_completed_websocket_message("resp_pool_first").into(), + )) + .await + .unwrap(); + + loop { + tokio::select! { + message = websocket.next() => { + match message { + Some(Ok(message)) if message.is_text() => { + let second_payload = serde_json::from_str::( + &message.into_text().unwrap(), + ) + .expect("second websocket payload should be json"); + websocket + .send(Message::Text( + response_completed_websocket_message("resp_pool_second").into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + break (true, first_payload, second_payload); + } + Some(_) => continue, + None => { + let second_payload = accept_successful_websocket_response( + &listener, + "resp_pool_second", + ) + .await; + break (false, first_payload, second_payload); + } + } + } + accepted = listener.accept() => { + let (stream, _) = accepted.unwrap(); + let mut second_websocket = accept_async(stream).await.unwrap(); + let second_message = second_websocket.next().await.unwrap().unwrap(); + let second_payload = serde_json::from_str::( + &second_message.into_text().unwrap(), + ) + .expect("second websocket payload should be json"); + second_websocket + .send(Message::Text( + response_completed_websocket_message("resp_pool_second").into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + break (false, first_payload, second_payload); + } + } + } + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let first_response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{ + "role": "user", + "content": "reuse this upstream websocket" + }] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + assert!(response_text(first_response) + .await + .contains("\"id\":\"resp_pool_first\"")); + + let second_response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "previous_response_id": "resp_pool_first" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + assert!(response_text(second_response) + .await + .contains("\"id\":\"resp_pool_second\"")); + let (reused_connection, first_payload, second_payload) = upstream.await.unwrap(); + + assert!(reused_connection, "second request opened a new websocket"); + assert_eq!( + second_payload["prompt_cache_key"], first_payload["prompt_cache_key"], + "pooled websocket reuse should stay on the recorded conversation key" + ); + assert_eq!(second_payload["previous_response_id"], "resp_pool_first"); +} + +#[tokio::test] +async fn responses_websocket_should_not_reuse_connection_when_pool_is_disabled() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let first_message = websocket.next().await.unwrap().unwrap(); + let first_payload = serde_json::from_str::(&first_message.into_text().unwrap()) + .expect("first websocket payload should be json"); + websocket + .send(Message::Text( + response_completed_websocket_message("resp_disabled_pool_first").into(), + )) + .await + .unwrap(); + + loop { + tokio::select! { + message = websocket.next() => { + match message { + Some(Ok(message)) if message.is_text() => { + let second_payload = serde_json::from_str::( + &message.into_text().unwrap(), + ) + .expect("second websocket payload should be json"); + websocket + .send(Message::Text( + response_completed_websocket_message( + "resp_disabled_pool_second", + ) + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + break (true, first_payload, second_payload); + } + Some(_) => continue, + None => { + let second_payload = accept_successful_websocket_response( + &listener, + "resp_disabled_pool_second", + ) + .await; + break (false, first_payload, second_payload); + } + } + } + accepted = listener.accept() => { + let (stream, _) = accepted.unwrap(); + let mut second_websocket = accept_async(stream).await.unwrap(); + let second_message = second_websocket.next().await.unwrap().unwrap(); + let second_payload = serde_json::from_str::( + &second_message.into_text().unwrap(), + ) + .expect("second websocket payload should be json"); + second_websocket + .send(Message::Text( + response_completed_websocket_message("resp_disabled_pool_second") + .into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + break (false, first_payload, second_payload); + } + } + } + }); + let (app, api_key, _pool, _dir) = test_app_with_account_pool_config(base_url, |config| { + config.ws_pool.enabled = false; + }) + .await; + + let first_response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{ + "role": "user", + "content": "do not reuse this upstream websocket" + }] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + assert!(response_text(first_response) + .await + .contains("\"id\":\"resp_disabled_pool_first\"")); + + let second_response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "previous_response_id": "resp_disabled_pool_first" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + assert!(response_text(second_response) + .await + .contains("\"id\":\"resp_disabled_pool_second\"")); + let (reused_connection, first_payload, second_payload) = upstream.await.unwrap(); + + assert!( + !reused_connection, + "disabled pool reused the upstream websocket" + ); + assert_eq!( + second_payload["prompt_cache_key"], first_payload["prompt_cache_key"], + "disabling the pool must not change the recorded conversation key" + ); + assert_eq!( + second_payload["previous_response_id"], + "resp_disabled_pool_first" + ); +} + +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn responses_websocket_stream_should_record_metadata_turn_state_for_continuation() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + let mut first_websocket = accept_async(first_stream).await.unwrap(); + let first_message = first_websocket.next().await.unwrap().unwrap(); + let first_payload = serde_json::from_str::(&first_message.into_text().unwrap()) + .expect("first websocket payload should be json"); + first_websocket + .send(Message::Text( + json!({ + "type": "response.metadata", + "headers": { + "x-codex-turn-state": "turn-from-metadata" + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + first_websocket + .send(Message::Text( + response_completed_websocket_message("resp_metadata_turn_state").into(), + )) + .await + .unwrap(); + first_websocket.close(None).await.unwrap(); + + let (second_stream, _) = listener.accept().await.unwrap(); + let second_headers = Arc::new(Mutex::new(Vec::new())); + let second_headers_for_callback = Arc::clone(&second_headers); + let mut second_websocket = accept_hdr_async( + second_stream, + move |request: &WsRequest, response: WsResponse| { + *second_headers_for_callback.lock().unwrap() = request_headers(request); + Ok(response) + }, + ) + .await + .unwrap(); + let second_message = second_websocket.next().await.unwrap().unwrap(); + let second_payload = serde_json::from_str::(&second_message.into_text().unwrap()) + .expect("second websocket payload should be json"); + second_websocket + .send(Message::Text( + response_completed_websocket_message("resp_metadata_turn_state_next").into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + let second_headers = second_headers.lock().unwrap().clone(); + (first_payload, second_payload, second_headers) + }); + let (app, api_key, _pool, _dir) = test_app_with_account_pool_config(base_url, |config| { + config.ws_pool.enabled = false; + }) + .await; + + let first_response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": true, + "use_websocket": true + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + let first_body = response_text(first_response).await; + assert!(first_body.contains("\"id\":\"resp_metadata_turn_state\"")); + assert!(!first_body.contains("response.metadata")); + + let second_response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "previous_response_id": "resp_metadata_turn_state" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + assert!(response_text(second_response) + .await + .contains("\"id\":\"resp_metadata_turn_state_next\"")); + let (first_payload, second_payload, second_headers) = upstream.await.unwrap(); + + assert!(first_payload.get("previous_response_id").is_none()); + assert_eq!( + second_payload["previous_response_id"], + "resp_metadata_turn_state" + ); + assert_eq!( + captured_header(&second_headers, "x-codex-turn-state"), + Some("turn-from-metadata") + ); +} + +#[tokio::test] +async fn responses_websocket_should_implicitly_resume_full_history_with_reasoning_replay() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let first_payload = send_websocket_response_and_capture_payload( + &mut websocket, + WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY.trim().to_string(), + ) + .await; + let second_payload = accept_followup_websocket_response( + &listener, + &mut websocket, + websocket_completed_response("resp_implicit_resume_second", 4, 1), + ) + .await; + let _ = websocket.close(None).await; + (first_payload, second_payload) + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let first_response = app + .clone() + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "answer briefly", + "stream": false, + "input": [{"role": "user", "content": "remember this"}] + }), + )) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + + let second_response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "answer briefly", + "stream": false, + "input": [ + {"role": "user", "content": "remember this"}, + {"role": "assistant", "content": "cached answer"}, + {"role": "user", "content": "continue"} + ] + }), + )) + .await + .unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + + let (first_payload, mut second_payload) = upstream.await.unwrap(); + assert!(first_payload["prompt_cache_key"].as_str().is_some()); + assert_eq!( + second_payload["previous_response_id"], + "resp_implicit_resume_first" + ); + assert_eq!( + second_payload["prompt_cache_key"], + first_payload["prompt_cache_key"] + ); + assert_eq!( + second_payload["input"][0]["encrypted_content"], + "enc_reasoning_replay" + ); + assert_eq!(second_payload["input"][1]["content"], "continue"); + assert_eq!(second_payload["input"].as_array().unwrap().len(), 2); + + second_payload + .as_object_mut() + .unwrap() + .remove("prompt_cache_key"); + second_payload + .as_object_mut() + .unwrap() + .remove("client_metadata"); + let expected: Value = serde_json::from_str(REASONING_REPLAY_REQUEST_GOLDEN).unwrap(); + assert_eq!(second_payload, expected); +} + +#[tokio::test] +async fn responses_websocket_should_not_implicitly_resume_unmatched_function_call_output() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let first_payload = send_websocket_response_and_capture_payload( + &mut websocket, + websocket_completed_function_call_response("resp_call_first", "call_expected"), + ) + .await; + let second_payload = accept_followup_websocket_response( + &listener, + &mut websocket, + websocket_completed_response("resp_call_mismatch_second", 4, 1), + ) + .await; + let _ = websocket.close(None).await; + (first_payload, second_payload) + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let first_response = app + .clone() + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "use the lookup tool", + "stream": false, + "input": [{"role": "user", "content": "call the lookup tool"}] + }), + )) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + + let second_response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "use the lookup tool", + "stream": false, + "input": [ + {"role": "user", "content": "call the lookup tool"}, + { + "type": "function_call", + "call_id": "call_expected", + "name": "lookup", + "arguments": "{}" + }, + { + "type": "function_call_output", + "call_id": "call_missing", + "output": "tool output" + } + ] + }), + )) + .await + .unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + + let (first_payload, second_payload) = upstream.await.unwrap(); + assert!(first_payload["prompt_cache_key"].as_str().is_some()); + assert!(second_payload.get("previous_response_id").is_none()); + assert_eq!(second_payload["input"].as_array().unwrap().len(), 3); + assert_eq!(second_payload["input"][2]["call_id"], "call_missing"); +} + +#[tokio::test] +async fn responses_websocket_should_implicitly_resume_after_sqlite_affinity_restore() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let server_base_url = base_url.clone(); + let upstream = tokio::spawn(async move { + let first_payload = accept_websocket_response_with_message( + &listener, + WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY.trim().to_string(), + ) + .await; + let second_payload = accept_websocket_response_with_message( + &listener, + websocket_completed_response("resp_restored_implicit_resume", 4, 1), + ) + .await; + (first_payload, second_payload) + }); + let (app, api_key, pool, dir) = test_app_with_account_pool_config(base_url, |_| {}).await; + + let first_response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "answer briefly", + "stream": false, + "input": [{"role": "user", "content": "remember this"}] + }), + )) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + + let db = dir.path().join("openai-record-affinity.sqlite"); + let restored_state = AppState::with_pool_secret_api_key_hasher_and_installation_id( + test_config(format!("sqlite://{}", db.display()), server_base_url), + pool.clone(), + SecretBox::new([83u8; 32]), + ApiKeyHasher::new([84u8; 32]), + TEST_INSTALLATION_ID.to_string(), + ); + assert_eq!( + restored_state + .restore_account_pool_from_repository() + .await + .unwrap(), + 1 + ); + assert_eq!( + restored_state + .restore_session_affinity_from_repository_now() + .await + .unwrap(), + 1 + ); + let restored_app = router::router().with_state(restored_state); + + let second_response = restored_app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "answer briefly", + "stream": false, + "input": [ + {"role": "user", "content": "remember this"}, + {"role": "assistant", "content": "cached answer"}, + {"role": "user", "content": "continue"} + ] + }), + )) + .await + .unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + + let (first_payload, second_payload) = upstream.await.unwrap(); + assert_eq!( + second_payload["prompt_cache_key"], + first_payload["prompt_cache_key"] + ); + assert_eq!( + second_payload["previous_response_id"], + "resp_implicit_resume_first" + ); + assert_eq!(second_payload["input"].as_array().unwrap().len(), 1); + assert_eq!(second_payload["input"][0]["content"], "continue"); +} + +#[tokio::test] +async fn responses_websocket_pool_should_be_evicted_after_admin_account_status_cycle() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let _first_payload = send_websocket_response_and_capture_payload( + &mut websocket, + websocket_completed_response("resp_pool_status_first", 4, 1), + ) + .await; + + tokio::select! { + message = websocket.next() => { + match message { + Some(Ok(message)) if message.is_text() => { + websocket + .send(Message::Text( + websocket_completed_response("resp_pool_status_second", 3, 1).into(), + )) + .await + .unwrap(); + let _ = websocket.close(None).await; + true + } + _ => { + accept_websocket_response_with_message( + &listener, + websocket_completed_response("resp_pool_status_second", 3, 1), + ) + .await; + false + } + } + } + accepted = listener.accept() => { + let (stream, _) = accepted.unwrap(); + let mut second_websocket = accept_async(stream).await.unwrap(); + send_websocket_response_and_capture_payload( + &mut second_websocket, + websocket_completed_response("resp_pool_status_second", 3, 1), + ) + .await; + second_websocket.close(None).await.unwrap(); + let _ = websocket.close(None).await; + false + } + } + }); + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(base_url).await; + seed_openai_admin_session(&pool, "session_status_cycle").await; + + let first_response = app + .clone() + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "prompt_cache_key": "status-cycle" + }), + )) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + assert!(response_text(first_response) + .await + .contains("\"id\":\"resp_pool_status_first\"")); + + update_admin_account_status(&app, "acct_chat", "disabled").await; + update_admin_account_status(&app, "acct_chat", "active").await; + + let second_response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "previous_response_id": "resp_pool_status_first" + }), + )) + .await + .unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + assert!(response_text(second_response) + .await + .contains("\"id\":\"resp_pool_status_second\"")); + + let reused_connection = upstream.await.unwrap(); + assert!( + !reused_connection, + "admin status lifecycle should evict the old pooled websocket" + ); +} + +#[tokio::test] +async fn responses_websocket_should_not_implicitly_resume_self_contained_function_call_replay() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let first_payload = send_websocket_response_and_capture_payload( + &mut websocket, + websocket_completed_function_call_response("resp_self_contained_first", "call_self"), + ) + .await; + let second_payload = accept_followup_websocket_response( + &listener, + &mut websocket, + websocket_completed_response("resp_self_contained_second", 4, 1), + ) + .await; + let _ = websocket.close(None).await; + (first_payload, second_payload) + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let first_response = app + .clone() + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "use the lookup tool", + "stream": false, + "input": [{"role": "user", "content": "call the lookup tool"}] + }), + )) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + + let second_response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "use the lookup tool", + "stream": false, + "input": [ + {"role": "user", "content": "call the lookup tool"}, + { + "type": "function_call", + "call_id": "call_self", + "name": "lookup", + "arguments": "{}" + }, + { + "type": "function_call_output", + "call_id": "call_self", + "output": "tool output" + } + ] + }), + )) + .await + .unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + + let (first_payload, second_payload) = upstream.await.unwrap(); + assert!(first_payload["prompt_cache_key"].as_str().is_some()); + assert!(second_payload.get("previous_response_id").is_none()); + assert_eq!(second_payload["input"].as_array().unwrap().len(), 3); + assert_eq!(second_payload["input"][2]["call_id"], "call_self"); +} + +#[tokio::test] +async fn responses_websocket_should_not_implicitly_resume_across_codex_windows() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let first_payload = send_websocket_response_and_capture_payload( + &mut websocket, + WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY.trim().to_string(), + ) + .await; + let second_payload = accept_followup_websocket_response( + &listener, + &mut websocket, + websocket_completed_response("resp_window_b", 8, 2), + ) + .await; + let _ = websocket.close(None).await; + (first_payload, second_payload) + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let first_response = app + .clone() + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "answer briefly", + "prompt_cache_key": "shared-variant-session", + "codexWindowId": "window-a", + "stream": false, + "input": [{"role": "user", "content": "remember this"}] + }), + )) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + + let second_response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "answer briefly", + "prompt_cache_key": "shared-variant-session", + "codexWindowId": "window-b", + "stream": false, + "input": [ + {"role": "user", "content": "remember this"}, + {"role": "assistant", "content": "cached answer"}, + {"role": "user", "content": "continue in another window"} + ] + }), + )) + .await + .unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + + let (first_payload, second_payload) = upstream.await.unwrap(); + assert_eq!( + second_payload["prompt_cache_key"], + first_payload["prompt_cache_key"] + ); + assert!(second_payload.get("previous_response_id").is_none()); + assert_eq!(second_payload["input"].as_array().unwrap().len(), 3); +} + +#[tokio::test] +async fn responses_websocket_should_evict_reasoning_replay_after_invalid_encrypted_content() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let first_payload = send_websocket_response_and_capture_payload( + &mut websocket, + WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY.trim().to_string(), + ) + .await; + let invalid_payload = accept_followup_websocket_response( + &listener, + &mut websocket, + WEBSOCKET_INVALID_ENCRYPTED_CONTENT.trim().to_string(), + ) + .await; + let _ = websocket.close(None).await; + let retried_payload = accept_websocket_response_with_message( + &listener, + websocket_completed_response("resp_after_replay_eviction", 4, 1), + ) + .await; + (first_payload, invalid_payload, retried_payload) + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let first_response = app + .clone() + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "answer briefly", + "stream": false, + "input": [{"role": "user", "content": "remember this"}] + }), + )) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + + let second_response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "answer briefly", + "stream": false, + "input": [ + {"role": "user", "content": "remember this"}, + {"role": "assistant", "content": "cached answer"}, + {"role": "user", "content": "continue"} + ] + }), + )) + .await + .unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + + let (first_payload, invalid_payload, retried_payload) = upstream.await.unwrap(); + assert_eq!( + invalid_payload["previous_response_id"], + "resp_implicit_resume_first" + ); + assert_eq!( + invalid_payload["input"][0]["encrypted_content"], + "enc_reasoning_replay" + ); + assert_eq!( + retried_payload["prompt_cache_key"], + first_payload["prompt_cache_key"] + ); + assert!(retried_payload.get("previous_response_id").is_none()); + let retried_input = retried_payload["input"].as_array().unwrap(); + assert!(retried_input + .iter() + .all(|item| item.get("encrypted_content").is_none())); + assert_eq!(retried_input.last().unwrap()["content"], "continue"); +} + +#[tokio::test] +async fn responses_websocket_should_restore_full_history_when_implicit_resume_previous_response_is_missing( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let first_payload = send_websocket_response_and_capture_payload( + &mut websocket, + WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY.trim().to_string(), + ) + .await; + let implicit_payload = accept_followup_websocket_response( + &listener, + &mut websocket, + WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(), + ) + .await; + let _ = websocket.close(None).await; + let restored_payload = accept_websocket_response_with_message( + &listener, + websocket_completed_response("resp_implicit_resume_restored", 10, 2), + ) + .await; + (first_payload, implicit_payload, restored_payload) + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let first_response = app + .clone() + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "answer briefly", + "stream": false, + "input": [{"role": "user", "content": "remember this"}] + }), + )) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + + let second_response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "answer briefly", + "stream": false, + "input": [ + {"role": "user", "content": "remember this"}, + {"role": "assistant", "content": "cached answer"}, + {"role": "user", "content": "continue"} + ] + }), + )) + .await + .unwrap(); + let body = response_json(second_response).await; + assert_eq!(body["id"], "resp_implicit_resume_restored"); + + let (_first_payload, implicit_payload, restored_payload) = upstream.await.unwrap(); + assert_eq!( + implicit_payload["previous_response_id"], + "resp_implicit_resume_first" + ); + assert_eq!(implicit_payload["input"].as_array().unwrap().len(), 2); + assert!(restored_payload.get("previous_response_id").is_none()); + assert_eq!(restored_payload["input"].as_array().unwrap().len(), 3); + assert_eq!(restored_payload["input"][0]["role"], "user"); + assert_eq!(restored_payload["input"][1]["role"], "assistant"); + assert_eq!(restored_payload["input"][2]["content"], "continue"); +} + +#[tokio::test] +async fn responses_websocket_should_route_previous_response_id_to_recorded_account() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (_reused_connection, first_payload, second_payload) = + accept_two_successful_websocket_responses_with_authorization( + &listener, + "Bearer access-primary", + "resp_affinity_first", + "resp_affinity_second", + ) + .await; + (first_payload, second_payload) + }); + let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await; + + let first_response = app + .clone() + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "instructions": "answer briefly", + "input": [{ + "role": "user", + "content": "keep this conversation on the same account" + }] + }), + )) + .await + .unwrap(); + assert_eq!(first_response.status(), StatusCode::OK); + assert!(response_text(first_response) + .await + .contains("\"id\":\"resp_affinity_first\"")); + let stored_affinity: (String, String, String, Option, String) = sqlx::query_as( + "select account_id, conversation_id, function_call_ids_json, input_tokens, expires_at from session_affinities where response_id = ?", + ) + .bind("resp_affinity_first") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(stored_affinity.0, "acct_primary"); + assert!(!stored_affinity.1.is_empty()); + assert_eq!(stored_affinity.2, "[]"); + assert_eq!(stored_affinity.3, Some(3)); + assert!(!stored_affinity.4.is_empty()); + + let second_response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "previous_response_id": "resp_affinity_first" + }), + )) + .await + .unwrap(); + assert_eq!(second_response.status(), StatusCode::OK); + let second_body = response_text(second_response).await; + assert!( + second_body.contains("\"id\":\"resp_affinity_second\""), + "{second_body}" + ); + let (first_payload, second_payload) = upstream.await.unwrap(); + + assert_ne!(first_payload["prompt_cache_key"], Value::Null); + assert_eq!( + second_payload["previous_response_id"], + "resp_affinity_first" + ); + assert_eq!( + second_payload["prompt_cache_key"], first_payload["prompt_cache_key"], + "previous_response_id should inherit the recorded conversation identity" + ); +} + +#[tokio::test] +async fn responses_websocket_non_stream_previous_response_not_found_should_strip_history_and_retry_same_account( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let first_payload = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-secret", + WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(), + ) + .await; + let second_payload = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-secret", + websocket_completed_response("resp_after_history_strip", 3, 1), + ) + .await; + (first_payload, second_payload) + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "previous_response_id": "resp_missing" + }), + )) + .await + .unwrap(); + let body = response_json(response).await; + let (first_payload, second_payload) = upstream.await.unwrap(); + + assert_eq!(body["id"], "resp_after_history_strip"); + assert_eq!(first_payload["previous_response_id"], "resp_missing"); + assert!(second_payload.get("previous_response_id").is_none()); +} + +#[tokio::test] +async fn responses_websocket_stream_previous_response_not_found_should_strip_history_and_retry_same_account( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let first_payload = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-secret", + WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(), + ) + .await; + let second_payload = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-secret", + websocket_completed_response("resp_stream_after_history_strip", 3, 1), + ) + .await; + (first_payload, second_payload) + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "stream": true, + "previous_response_id": "resp_missing" + }), + )) + .await + .unwrap(); + let body = response_text(response).await; + let (first_payload, second_payload) = upstream.await.unwrap(); + + assert!(body.contains("\"id\":\"resp_stream_after_history_strip\"")); + assert_eq!(first_payload["previous_response_id"], "resp_missing"); + assert!(second_payload.get("previous_response_id").is_none()); +} + +#[tokio::test] +async fn responses_websocket_non_stream_unanswered_function_call_should_strip_history_and_retry_same_account( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let first_payload = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-secret", + WEBSOCKET_UNANSWERED_FUNCTION_CALL.trim().to_string(), + ) + .await; + let second_payload = accept_websocket_response_with_authorization_and_message( + &listener, + "Bearer access-secret", + websocket_completed_response("resp_after_function_call_strip", 3, 1), + ) + .await; + (first_payload, second_payload) + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "previous_response_id": "resp_with_call" + }), + )) + .await + .unwrap(); + let body = response_json(response).await; + let (first_payload, second_payload) = upstream.await.unwrap(); + + assert_eq!(body["id"], "resp_after_function_call_strip"); + assert_eq!(first_payload["previous_response_id"], "resp_with_call"); + assert!(second_payload.get("previous_response_id").is_none()); +} + +#[tokio::test] +async fn responses_websocket_previous_response_id_should_retry_fallback_account_after_429() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + reject_next_websocket_upgrade( + &listener, + "Bearer access-primary", + 429, + "Too Many Requests", + Some(77), + WEBSOCKET_HISTORY_RATE_LIMITED, + ) + .await; + accept_successful_websocket_response_with_authorization( + &listener, + "Bearer access-secondary", + "resp_history_fallback", + ) + .await + }); + let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "previous_response_id": "resp_prev" + }), + )) + .await + .unwrap(); + let body = response_text(response).await; + let fallback_payload = upstream.await.unwrap(); + let secondary_usage: (i64,) = + sqlx::query_as("select request_count from account_usage where account_id = ?") + .bind("acct_secondary") + .fetch_one(&pool) + .await + .unwrap(); + + assert!(body.contains("\"id\":\"resp_history_fallback\"")); + assert_eq!(fallback_payload["previous_response_id"], "resp_prev"); + assert_eq!(secondary_usage.0, 1); +} + +#[tokio::test] +async fn responses_websocket_non_stream_previous_response_id_should_retry_fallback_account_after_429( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + reject_next_websocket_upgrade( + &listener, + "Bearer access-primary", + 429, + "Too Many Requests", + Some(77), + WEBSOCKET_HISTORY_RATE_LIMITED, + ) + .await; + accept_successful_websocket_response_with_authorization( + &listener, + "Bearer access-secondary", + "resp_history_fallback_non_stream", + ) + .await + }); + let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "previous_response_id": "resp_prev" + }), + )) + .await + .unwrap(); + let body = response_json(response).await; + let fallback_payload = upstream.await.unwrap(); + let secondary_usage: (i64,) = + sqlx::query_as("select request_count from account_usage where account_id = ?") + .bind("acct_secondary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(body["id"], "resp_history_fallback_non_stream"); + assert_eq!(fallback_payload["previous_response_id"], "resp_prev"); + assert_eq!(secondary_usage.0, 1); +} + +#[tokio::test] +async fn responses_websocket_without_history_should_mark_expired_after_fallback_401() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + reject_next_websocket_upgrade( + &listener, + "Bearer access-primary", + 429, + "Too Many Requests", + Some(30), + WEBSOCKET_RATE_LIMITED, + ) + .await; + reject_next_websocket_upgrade( + &listener, + "Bearer access-secondary", + 401, + "Unauthorized", + None, + WEBSOCKET_TOKEN_REVOKED, + ) + .await; + }); + let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "stream": true, + "use_websocket": true + }), + )) + .await + .unwrap(); + let body = response_text(response).await; + upstream.await.unwrap(); + let secondary_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_secondary") + .fetch_one(&pool) + .await + .unwrap(); + let primary_usage: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert!(body.contains("event: response.failed")); + assert!(body.contains("\"type\":\"invalid_request_error\"")); + assert!(body.contains("\"code\":\"authentication_error\"")); + assert!(body.contains("All accounts exhausted")); + assert!(body.contains("token_revoked")); + assert_eq!(secondary_status.0, "expired"); + assert_eq!(primary_usage, (1, 0, 0)); +} + +#[tokio::test] +async fn responses_websocket_without_history_should_return_rate_limit_stream_error_when_fallback_accounts_exhausted( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + reject_next_websocket_upgrade( + &listener, + "Bearer access-primary", + 429, + "Too Many Requests", + Some(11), + WEBSOCKET_FIRST_ACCOUNT_LIMITED, + ) + .await; + reject_next_websocket_upgrade( + &listener, + "Bearer access-secondary", + 429, + "Too Many Requests", + Some(22), + WEBSOCKET_SECOND_ACCOUNT_LIMITED, + ) + .await; + }); + let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "stream": true, + "use_websocket": true + }), + )) + .await + .unwrap(); + let body = response_text(response).await; + upstream.await.unwrap(); + let primary_usage: (i64,) = + sqlx::query_as("select request_count from account_usage where account_id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + let secondary_usage: (i64,) = + sqlx::query_as("select request_count from account_usage where account_id = ?") + .bind("acct_secondary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_response_failed_stream( + &body, + "rate_limit_error", + "rate_limit_exceeded", + &[ + "All accounts exhausted (2 rate-limited)", + "second account limited", + ], + ); + assert_eq!(primary_usage.0, 1); + assert_eq!(secondary_usage.0, 1); +} + +#[tokio::test] +async fn responses_websocket_response_failed_quota_should_retry_fallback_account() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + let mut first_websocket = + accept_websocket_with_authorization(first_stream, "Bearer access-primary").await; + let _first_payload = send_websocket_response_and_capture_payload( + &mut first_websocket, + json!({ + "type": "response.failed", + "response": { + "id": "resp_ws_quota_failed", + "error": { + "code": "insufficient_quota", + "message": "quota exhausted" + } + } + }) + .to_string(), + ) + .await; + + let (second_stream, _) = listener.accept().await.unwrap(); + let mut second_websocket = + accept_websocket_with_authorization(second_stream, "Bearer access-secondary").await; + send_websocket_response_and_capture_payload( + &mut second_websocket, + websocket_completed_response("resp_after_ws_quota", 3, 1), + ) + .await; + second_websocket.close(None).await.unwrap(); + }); + let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": true + }), + )) + .await + .unwrap(); + let body = response_json(response).await; + upstream.await.unwrap(); + let primary_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(body["id"], "resp_after_ws_quota"); + assert_eq!(primary_status.0, "quota_exhausted"); +} + +#[tokio::test] +async fn responses_websocket_without_history_should_return_quota_stream_error_when_402_has_no_fallback( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + reject_next_websocket_upgrade( + &listener, + "Bearer access-secret", + 402, + "Payment Required", + None, + r#"{"error":{"message":"quota reached"}}"#, + ) + .await; + }); + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "stream": true, + "use_websocket": true + }), + )) + .await + .unwrap(); + let body = response_text(response).await; + upstream.await.unwrap(); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_response_failed_stream( + &body, + "invalid_request_error", + "codex_api_error", + &[ + "All accounts exhausted (1 quota-exhausted)", + "quota reached", + ], + ); + assert_eq!(account_status.0, "quota_exhausted"); +} + +#[tokio::test] +async fn responses_websocket_without_history_should_return_model_unsupported_stream_error_when_no_fallback( +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + reject_next_websocket_upgrade( + &listener, + "Bearer access-secret", + 400, + "Bad Request", + None, + r#"{"error":{"code":"model_not_available","message":"Model gpt-5.5 is not available on this account plan"}}"#, + ) + .await; + }); + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "stream": true, + "use_websocket": true + }), + )) + .await + .unwrap(); + let body = response_text(response).await; + upstream.await.unwrap(); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_response_failed_stream( + &body, + "invalid_request_error", + "codex_api_error", + &[ + "No accounts available", + "model_not_available", + "not available", + ], + ); + assert_eq!(account_status.0, "active"); +} + +#[tokio::test] +async fn responses_websocket_with_history_should_return_path_block_stream_error_when_no_fallback() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + reject_next_websocket_upgrade( + &listener, + "Bearer access-secret", + 404, + "Not Found", + None, + "", + ) + .await; + }); + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + + let response = app + .oneshot(responses_json_request( + &api_key, + json!({ + "model": "gpt-5.5", + "input": [], + "stream": true, + "previous_response_id": "resp_prev" + }), + )) + .await + .unwrap(); + let body = response_text(response).await; + upstream.await.unwrap(); + + assert_response_failed_stream( + &body, + "server_error", + "codex_api_error", + &["No accounts available", "Cloudflare path-block"], + ); +} + +#[tokio::test] +async fn responses_should_convert_tuple_schema_before_upstream() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_TUPLE_OBJECT_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from(tuple_response_request_body(false))) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + + assert_eq!(status, StatusCode::OK); + let requests = server.received_requests().await.unwrap(); + let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + let schema = &upstream_body["text"]["format"]["schema"]; + assert!(schema["properties"]["point"].get("prefixItems").is_none()); + assert!(schema["properties"]["point"].get("items").is_none()); + assert_eq!( + schema["properties"]["point"], + json!({ + "type": "object", + "properties": { + "0": {"type": "number"}, + "1": {"type": "number"} + }, + "required": ["0", "1"], + "additionalProperties": false + }) + ); +} + +#[tokio::test] +async fn responses_should_reconvert_tuple_schema_output_for_client() { + 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(RESPONSES_TUPLE_OBJECT_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from(tuple_response_request_body(false))) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_tuple"); + assert_eq!(body["output_text"], "{\"point\":[1,2]}"); + assert_eq!(body["output"][0]["content"][0]["text"], "{\"point\":[1,2]}"); +} + +#[tokio::test] +async fn responses_stream_should_reconvert_tuple_schema_output_for_client() { + 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(RESPONSES_TUPLE_OBJECT_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from(tuple_response_request_body(true))) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_text(response).await; + + assert_eq!(status, StatusCode::OK); + assert!(body.contains(r#""delta":"{\"point\":[1,2]}""#)); + assert!(body.contains(r#""output_text":"{\"point\":[1,2]}""#)); + assert!(!body.contains(r#""point\":{\"0":1,"1":2}"#)); + assert!(body.ends_with("data: [DONE]\n\n")); +} + +#[tokio::test] +async fn responses_should_forward_parity_fields_context_headers_and_account_scoped_identity() { + 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(RESPONSES_SUCCESS_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {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", "false") + .header("version", "header-version") + .header("x-codex-window-id", "window-header") + .header("x-codex-parent-thread-id", "parent-header") + .header("x-openai-subagent", "review") + .body(Body::from( + json!({ + "model": "gpt-5.5-high-fast", + "stream": false, + "use_websocket": false, + "input": [], + "prompt_cache_key": "pcache", + "client_metadata": {"safe": "yes", "drop": 42}, + "turnState": "turn-body", + "turnMetadata": "meta-body", + "betaFeatures": "beta-body", + "includeTimingMetrics": "true", + "version": "2026-06-12", + "codexWindowId": "window-body", + "parentThreadId": "parent-body", + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_response_1"); + let requests = server.received_requests().await.unwrap(); + let upstream = requests + .iter() + .find(|request| request.url.path() == "/codex/responses") + .expect("responses upstream request should be sent"); + let upstream_body: Value = serde_json::from_slice(&upstream.body).unwrap(); + let identity = build_conversation_identity(Some("pcache"), Some("window-body"), "acct_chat"); + let conversation_id = identity + .conversation_id + .as_deref() + .expect("conversation identity should be scoped"); + let window_id = identity + .window_id + .as_deref() + .expect("window identity should be scoped"); + let upstream_header = |name: &str| { + upstream + .headers + .get(name) + .and_then(|value| value.to_str().ok()) + }; + + assert_eq!(upstream_body["model"], "gpt-5.5"); + assert_eq!(upstream_body["service_tier"], "priority"); + assert_eq!( + upstream_body["reasoning"], + json!({"summary": "auto", "effort": "high"}) + ); + assert_eq!( + upstream_body["include"], + json!(["reasoning.encrypted_content"]) + ); + assert_eq!(upstream_body["prompt_cache_key"], conversation_id); + assert_eq!(upstream_body["client_metadata"]["safe"], "yes"); + assert_eq!( + upstream_body["client_metadata"]["x-openai-subagent"], + "review" + ); + assert_eq!( + upstream_body["client_metadata"]["x-codex-installation-id"], + TEST_INSTALLATION_ID + ); + assert_eq!( + upstream_body["client_metadata"]["x-codex-window-id"], + window_id + ); + assert_eq!( + upstream_body["client_metadata"]["x-codex-turn-metadata"], + "meta-body" + ); + assert_eq!( + upstream_body["client_metadata"]["x-codex-parent-thread-id"], + "parent-body" + ); + assert_eq!(upstream_header("session_id"), Some(conversation_id)); + assert_eq!(upstream_header("x-codex-window-id"), Some(window_id)); + assert_eq!(upstream_header("x-codex-turn-state"), Some("turn-body")); + assert_eq!(upstream_header("x-codex-turn-metadata"), Some("meta-body")); + assert_eq!(upstream_header("x-codex-beta-features"), Some("beta-body")); + assert_eq!( + upstream_header("x-responsesapi-include-timing-metrics"), + Some("true") + ); + assert_eq!(upstream_header("version"), Some("2026-06-12")); + assert_eq!( + upstream_header("x-codex-parent-thread-id"), + Some("parent-body") + ); + assert_eq!(upstream_header("x-openai-subagent"), Some("review")); + for local_field in [ + "turnState", + "turnMetadata", + "betaFeatures", + "includeTimingMetrics", + "version", + "codexWindowId", + "parentThreadId", + "use_websocket", + ] { + assert!(upstream_body.get(local_field).is_none()); + } +} + +#[tokio::test] +async fn responses_should_preserve_non_empty_include_when_reasoning_defaults_apply() { + 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(RESPONSES_SUCCESS_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5-high", + "stream": false, + "use_websocket": false, + "input": [], + "include": ["file_search_call.results"], + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let requests = server.received_requests().await.unwrap(); + let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!( + upstream_body["reasoning"], + json!({"summary": "auto", "effort": "high"}) + ); + assert_eq!( + upstream_body["include"], + json!(["file_search_call.results"]) + ); +} + +#[tokio::test] +async fn responses_should_sanitize_reasoning_and_compaction_input_before_upstream() { + 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(RESPONSES_COMPLETED_USAGE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "stream": false, + "use_websocket": false, + "input": [ + { + "type": "reasoning", + "id": "rs_1", + "status": "completed", + "summary": [ + {"type": "summary_text", "text": "valid summary"}, + {"type": "ignored", "text": "drop"} + ], + "encrypted_content": "enc_reasoning", + "content": [ + {"type": "reasoning_text", "text": "valid reasoning"}, + {"type": "ignored", "text": "drop"} + ], + "extra": "drop" + }, + { + "type": "reasoning", + "id": "", + "summary": [{"type": "summary_text", "text": "drop"}] + }, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": "enc_compaction", + "extra": "drop" + }, + {"type": "compaction", "id": "cmp_drop"}, + {"type": "message", "role": "user", "content": "keep me", "extra": 42} + ] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let requests = server.received_requests().await.unwrap(); + let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!( + upstream_body["input"], + json!([ + { + "type": "reasoning", + "id": "rs_1", + "summary": [{"type": "summary_text", "text": "valid summary"}], + "status": "completed", + "encrypted_content": "enc_reasoning", + "content": [{"type": "reasoning_text", "text": "valid reasoning"}] + }, + { + "type": "compaction", + "encrypted_content": "enc_compaction", + "id": "cmp_1" + }, + {"type": "message", "role": "user", "content": "keep me", "extra": 42} + ]) + ); +} + +#[tokio::test] +async fn 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(RESPONSES_TEXT_DELTAS_COMPLETED_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json(response).await; + + assert_eq!(body["id"], "resp_text"); + assert_eq!(body["output"][0]["role"], "assistant"); + assert_eq!(body["output"][0]["content"][0]["text"], "hello from deltas"); + assert_eq!(body["output_text"], "hello from deltas"); +} + +#[tokio::test] +async fn 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(RESPONSES_DONE_ITEM_COMPLETED_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json(response).await; + + assert_eq!(body["id"], "resp_item"); + assert_eq!(body["output"][0]["content"][0]["text"], "from done item"); + assert_eq!(body["output_text"], "from done item"); +} + +#[tokio::test] +async fn responses_review_route_should_force_review_subagent_upstream() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("x-openai-subagent", "review")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_SUCCESS_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses/review") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_response_1"); + let requests = server.received_requests().await.unwrap(); + let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap(); + assert_eq!( + upstream_body["client_metadata"]["x-openai-subagent"], + "review" + ); +} + +#[tokio::test] +async fn responses_compact_should_post_json_to_codex_compact_upstream() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses/compact")) + .and(header("authorization", "Bearer access-secret")) + .and(header("chatgpt-account-id", "chatgpt-account")) + .and(header("content-type", "application/json")) + .and(header("openai-beta", "responses_websockets=2026-02-06")) + .and(header("x-openai-internal-codex-residency", "us")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "compacted"}] + } + ] + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses/compact") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5-fast", + "instructions": "compress the session", + "input": [ + {"role": "user", "content": "hello"}, + { + "type": "reasoning", + "id": "rs_1", + "status": "completed", + "summary": [{"type": "summary_text", "text": "kept"}], + "ignored": "drop" + }, + {"type": "compaction", "encrypted_content": "enc_compact"}, + {"type": "compaction", "id": "drop_missing_encrypted"} + ], + "tools": [{"type": "function", "name": "lookup"}], + "parallel_tool_calls": false, + "reasoning": {"effort": "high", "summary": "auto", "extra": "drop"}, + "text": { + "format": { + "type": "json_schema", + "name": "Compact", + "schema": {"type": "object"}, + "strict": true + } + }, + "stream": true, + "use_websocket": false, + "store": true, + "prompt_cache_key": "must_not_forward" + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["output"][0]["content"][0]["text"], "compacted"); + let requests = server.received_requests().await.unwrap(); + let compact_request = requests + .iter() + .find(|request| request.url.path() == "/codex/responses/compact") + .expect("compact upstream request should be sent"); + assert_ne!( + compact_request + .headers + .get("accept") + .and_then(|value| value.to_str().ok()), + Some("text/event-stream") + ); + assert!(compact_request + .headers + .get("x-codex-installation-id") + .is_some()); + let upstream_body: Value = serde_json::from_slice(&compact_request.body).unwrap(); + assert_eq!(upstream_body["model"], "gpt-5.5"); + assert_eq!(upstream_body["instructions"], "compress the session"); + assert_eq!(upstream_body["parallel_tool_calls"], false); + assert_eq!( + upstream_body["reasoning"], + json!({"effort": "high", "summary": "auto"}) + ); + assert_eq!( + upstream_body["tools"], + json!([{"type": "function", "name": "lookup"}]) + ); + assert_eq!(upstream_body["text"]["format"]["type"], "json_schema"); + assert!(upstream_body.get("stream").is_none()); + assert!(upstream_body.get("store").is_none()); + assert!(upstream_body.get("prompt_cache_key").is_none()); + assert_eq!(upstream_body["input"].as_array().unwrap().len(), 3); + assert!(upstream_body["input"][1].get("ignored").is_none()); + assert_eq!( + upstream_body["input"][2]["encrypted_content"], + "enc_compact" + ); +} + +#[tokio::test] +async fn responses_compact_should_return_rate_limit_error_when_fallback_is_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses/compact")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(429).set_body_json(json!({ + "error": { + "code": "rate_limit_exceeded", + "message": "compact quota reached" + } + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses/compact") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let message = body["error"]["message"].as_str().unwrap_or_default(); + + assert_eq!(status, StatusCode::TOO_MANY_REQUESTS); + assert_eq!(body["type"], "error"); + assert_eq!(body["error"]["type"], "rate_limit_error"); + assert_eq!(body["error"]["code"], "rate_limit_exceeded"); + assert!(message.contains("All accounts exhausted (1 rate-limited)")); + assert!(message.contains("compact quota reached")); +} + +#[tokio::test] +async fn responses_compact_should_return_responses_error_when_no_accounts_are_available() { + let server = MockServer::start().await; + let (app, api_key, _dir) = test_app_without_accounts(server.uri()).await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses/compact") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "hello"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(body["type"], "error"); + assert_eq!(body["error"]["type"], "server_error"); + assert_eq!(body["error"]["code"], "no_available_accounts"); +} + +#[tokio::test] +async fn responses_should_use_imported_account_record_usage_cookie_and_event_log() { + 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")) + .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(RESPONSES_COMPLETED_USAGE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .header("x-request-id", "req_non_stream_usage_log") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_usage"); + 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_chat") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(usage, (1, 7, 4, 2)); + let cookie_header = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])) + .cookie_header("acct_chat", "chatgpt.com") + .await + .unwrap(); + assert_eq!(cookie_header.as_deref(), Some("cf_clearance=new")); + let event = latest_response_event_log(&pool).await; + let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap(); + assert_eq!(event.level, "info"); + assert_eq!( + event.request_id.as_deref(), + Some("req_non_stream_usage_log") + ); + assert_eq!(event.account_id.as_deref(), Some("acct_chat")); + assert_eq!(event.status_code, Some(200)); + assert_eq!(metadata["stream"], false); + assert_eq!(metadata["usage"]["inputTokens"], 7); +} + +#[tokio::test] +async fn responses_should_skip_event_log_when_logging_disabled() { + 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(RESPONSES_COMPLETED_USAGE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .header("x-request-id", "req_non_stream_no_log") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response_event_log_count(&pool).await, 0); +} + +#[tokio::test] +async fn responses_should_record_image_generation_usage_when_tool_succeeds() { + 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(RESPONSES_COMPLETED_IMAGE_USAGE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false, + "tools": [{"type": "image_generation"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let usage: (i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64) = + sqlx::query_as( + "select request_count, input_tokens, output_tokens, cached_tokens, image_input_tokens, image_output_tokens, image_request_count, image_request_failed_count, window_image_input_tokens, window_image_output_tokens, window_image_request_count, window_image_request_failed_count from account_usage where account_id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(usage, (1, 7, 4, 2, 31, 9, 1, 0, 31, 9, 1, 0)); +} + +#[tokio::test] +async fn responses_should_record_failed_image_generation_attempt_when_tool_has_no_output() { + 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(RESPONSES_COMPLETED_USAGE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false, + "tools": [{"type": "image_generation"}] + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let usage: (i64, i64, i64, i64) = sqlx::query_as( + "select image_request_count, image_request_failed_count, window_image_request_count, window_image_request_failed_count from account_usage where account_id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(usage, (0, 1, 0, 1)); +} + +#[tokio::test] +async fn responses_should_scope_upstream_cookie_by_codex_response_path() { + 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(RESPONSES_COMPLETED_USAGE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + cookie_store + .capture_set_cookie( + "acct_chat", + "cf_clearance=root; Domain=.chatgpt.com; Path=/", + ) + .await + .unwrap(); + cookie_store + .capture_set_cookie( + "acct_chat", + "cf_clearance=codex; Domain=.chatgpt.com; Path=/codex", + ) + .await + .unwrap(); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let requests = server.received_requests().await.unwrap(); + let cookie_header = requests + .iter() + .find(|request| request.url.path() == "/codex/responses") + .and_then(|request| request.headers.get("cookie")) + .and_then(|value| value.to_str().ok()); + assert_eq!(cookie_header, Some("cf_clearance=codex; cf_clearance=root")); +} + +#[tokio::test] +async fn responses_should_passively_cache_rate_limit_headers() { + let server = MockServer::start().await; + let reset_at = Utc::now().timestamp() + 300; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .insert_header("x-codex-primary-used-percent", "100") + .insert_header("x-codex-primary-window-minutes", "5") + .insert_header("x-codex-primary-reset-at", reset_at.to_string()) + .set_body_string(RESPONSES_COMPLETED_USAGE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + sqlx::query( + r#"update accounts set quota_json = '{"credits":{"has_credits":true,"unlimited":false,"balance":12}}', quota_verify_required = 1 where id = ?"#, + ) + .bind("acct_chat") + .execute(&pool) + .await + .unwrap(); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let stored: (String, i64, Option, i64) = sqlx::query_as( + "select quota_json, quota_limit_reached, quota_cooldown_until, quota_verify_required from accounts where id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + let quota: Value = serde_json::from_str(&stored.0).unwrap(); + assert_eq!(quota["rate_limit"]["limit_reached"], true); + assert_eq!(quota["rate_limit"]["reset_at"], reset_at); + assert_eq!(quota["credits"]["balance"], 12); + assert_eq!(stored.1, 1); + assert!(stored.2.is_some()); + assert_eq!(stored.3, 0); + let window: (i64, i64, i64, i64, Option, Option, Option) = + sqlx::query_as( + "select window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds from account_usage where account_id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(window.0, 1); + assert_eq!(window.1, 7); + assert_eq!(window.2, 4); + assert_eq!(window.3, 2); + assert!(window.4.is_some()); + assert_eq!( + DateTime::parse_from_rfc3339(window.5.as_deref().unwrap()) + .unwrap() + .timestamp(), + reset_at + ); + assert_eq!(window.6, Some(300)); +} + +#[tokio::test] +async fn responses_should_record_empty_response_attempts() { + 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(RESPONSES_EMPTY_COMPLETED_SSE), + ) + .expect(3) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + let usage: (i64, i64, i64, i64) = sqlx::query_as( + "select request_count, empty_response_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(usage, (3, 3, 0, 0)); +} + +#[tokio::test] +async fn responses_should_fallback_to_next_account_after_rate_limit() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("retry-after", "120") + .set_body_json(json!({ + "error": { + "message": "rate limited", + "resets_in_seconds": 120 + } + })), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_SUCCESS_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_response_1"); + let primary_usage: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + let secondary_usage: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_secondary") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(primary_usage, (1, 0, 0)); + assert_eq!(secondary_usage, (1, 5, 2)); +} + +#[tokio::test] +async fn responses_should_mark_expired_after_401_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(401).set_body_json(json!({ + "error": {"code": "token_revoked", "message": "token revoked"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_AFTER_401_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_401"); + assert_eq!(account_status.0, "expired"); +} + +#[tokio::test] +async fn responses_should_recover_auth_failure_from_sse_failed_event() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_FAILED_AUTH_SSE), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_AFTER_401_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_401"); + assert_eq!(account_status.0, "expired"); +} + +#[tokio::test] +async fn responses_should_return_auth_error_when_401_fallback_is_exhausted() { + 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", "message": "token revoked"} + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let message = body["error"]["message"].as_str().unwrap_or_default(); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert!(message.contains("All accounts exhausted (1 expired)")); + assert!(message.contains("token_revoked")); + assert_eq!(body["error"]["code"], "upstream_error"); + assert_eq!(account_status.0, "expired"); +} + +#[tokio::test] +async fn responses_should_mark_banned_when_401_says_account_deactivated() { + 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": {"message": "account deactivated"} + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + assert_eq!(account_status.0, "banned"); +} + +#[tokio::test] +async fn responses_should_mark_banned_after_403_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(403).set_body_json(json!({ + "error": {"message": "account banned"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_AFTER_403_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_403"); + assert_eq!(account_status.0, "banned"); +} + +#[tokio::test] +async fn responses_should_return_cloudflare_challenge_error_when_403_fallback_is_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(403).set_body_string( + "Just a moment...cf_chl challenge", + )) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + cookie_store + .set_cookie_header("acct_chat", "cf_clearance=old") + .await + .unwrap(); + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let message = body["error"]["message"].as_str().unwrap_or_default(); + let primary_state: (String, Option) = + sqlx::query_as("select status, cloudflare_cooldown_until from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::BAD_GATEWAY); + assert!(message.contains("All accounts exhausted (1 cloudflare-challenge)")); + assert!(message.contains("Cloudflare challenge")); + assert_eq!(body["error"]["code"], "upstream_error"); + assert_eq!(primary_state.0, "active"); + assert!(primary_state.1.is_some()); + assert_eq!( + cookie_store + .cookie_header("acct_chat", "chatgpt.com") + .await + .unwrap(), + None + ); +} + +#[tokio::test] +async fn responses_should_cool_down_cloudflare_403_and_fallback() { + let started_at = Utc::now(); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(403).set_body_string( + "Just a moment...cf_chl challenge", + )) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_AFTER_CLOUDFLARE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + cookie_store + .set_cookie_header("acct_primary", "cf_clearance=old") + .await + .unwrap(); + cookie_store + .set_cookie_header("acct_secondary", "cf_clearance=keep") + .await + .unwrap(); + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let primary_state: (String, Option) = + sqlx::query_as("select status, cloudflare_cooldown_until from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + let cooldown_until = chrono::DateTime::parse_from_rfc3339(primary_state.1.as_deref().unwrap()) + .unwrap() + .with_timezone(&Utc); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_cloudflare"); + assert_eq!(primary_state.0, "active"); + assert!(cooldown_until > started_at); + assert_eq!( + cookie_store + .cookie_header("acct_primary", "chatgpt.com") + .await + .unwrap(), + None + ); + assert_eq!( + cookie_store + .cookie_header("acct_secondary", "chatgpt.com") + .await + .unwrap() + .as_deref(), + Some("cf_clearance=keep") + ); +} + +#[tokio::test] +async fn responses_should_clear_cookies_after_cloudflare_path_block_404_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(404)) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_AFTER_CLOUDFLARE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + cookie_store + .set_cookie_header("acct_primary", "cf_clearance=old") + .await + .unwrap(); + cookie_store + .set_cookie_header("acct_secondary", "cf_clearance=keep") + .await + .unwrap(); + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let primary_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_cloudflare"); + assert_eq!(primary_status.0, "active"); + assert_eq!( + cookie_store + .cookie_header("acct_primary", "chatgpt.com") + .await + .unwrap(), + None + ); + assert_eq!( + cookie_store + .cookie_header("acct_secondary", "chatgpt.com") + .await + .unwrap() + .as_deref(), + Some("cf_clearance=keep") + ); +} + +#[tokio::test] +async fn responses_should_disable_account_after_three_cloudflare_path_block_404s() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(404)) + .expect(3) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + cookie_store + .set_cookie_header("acct_chat", "cf_clearance=old") + .await + .unwrap(); + for _ in 0..3 { + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let message = body["error"]["message"].as_str().unwrap_or_default(); + + assert_eq!(status, StatusCode::BAD_GATEWAY); + assert!(message.contains("cloudflare-path-block")); + } + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(account_status.0, "disabled"); + assert_eq!( + cookie_store + .cookie_header("acct_chat", "chatgpt.com") + .await + .unwrap(), + None + ); +} + +#[tokio::test] +async fn responses_should_fallback_after_http_model_unsupported() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": { + "code": "model_not_supported", + "message": "Model gpt-5.5 is not supported on this account plan" + } + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_AFTER_MODEL_UNSUPPORTED_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let primary_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_model_unsupported"); + assert_eq!(primary_status.0, "active"); +} + +#[tokio::test] +async fn responses_should_fallback_after_sse_model_unsupported() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_FAILED_MODEL_UNSUPPORTED_SSE), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_AFTER_MODEL_UNSUPPORTED_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_model_unsupported"); +} + +#[tokio::test] +async fn responses_stream_should_fallback_after_sse_model_unsupported() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_FAILED_MODEL_UNSUPPORTED_SSE), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_STREAM_AFTER_MODEL_UNSUPPORTED_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("resp_stream_after_model_unsupported")); +} + +#[tokio::test] +async fn responses_stream_should_return_model_unsupported_error_when_fallback_is_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": { + "code": "model_not_available", + "message": "Model gpt-5.5 is not available on this account plan" + } + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("event: response.failed")); + assert!(body.contains("All accounts exhausted (1 model-unsupported)")); + assert!(body.contains("model_not_available")); + assert_eq!(account_status.0, "active"); +} + +#[tokio::test] +async fn responses_should_strip_history_after_previous_response_not_found() { + let (base_url, upstream) = + spawn_websocket_failure_then_websocket_success_upstream(response_failed_websocket_message( + "resp_previous_missing_failed", + "previous_response_not_found", + "Previous response with id resp_stale was not found", + )) + .await; + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "previous_response_id": "resp_stale", + "turnState": "turn-stale", + "input": [{"role": "user", "content": "Continue"}], + "stream": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let captured = upstream.await.unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_history_recovery"); + assert_eq!( + captured_header(&captured.first_ws_headers, "x-codex-turn-state"), + Some("turn-stale") + ); + assert!(captured_header(&captured.second_ws_headers, "x-codex-turn-state").is_none()); + assert_eq!( + captured.first_ws_payload["previous_response_id"], + "resp_stale" + ); + assert!(captured + .second_ws_payload + .get("previous_response_id") + .is_none()); +} + +#[tokio::test] +async fn responses_should_strip_history_after_unanswered_function_call() { + let (base_url, upstream) = + spawn_websocket_failure_then_websocket_success_upstream(response_failed_websocket_message( + "resp_unanswered_call_failed", + "invalid_request", + "No tool output found for function call call_1", + )) + .await; + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "previous_response_id": "resp_stale", + "turnState": "turn-stale", + "input": [{"role": "user", "content": "Continue"}], + "stream": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let captured = upstream.await.unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_history_recovery"); + assert!(captured_header(&captured.second_ws_headers, "x-codex-turn-state").is_none()); + assert!(captured + .second_ws_payload + .get("previous_response_id") + .is_none()); +} + +#[tokio::test] +async fn responses_should_strip_history_after_websocket_invalid_encrypted_reasoning_replay() { + let (base_url, upstream) = + spawn_websocket_failure_then_websocket_success_upstream(response_failed_websocket_message( + "resp_invalid_reasoning_failed", + "invalid_encrypted_content", + "Invalid encrypted content in reasoning replay", + )) + .await; + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "previous_response_id": "resp_stale", + "turnState": "turn-stale", + "input": [{"role": "user", "content": "Continue"}], + "stream": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let captured = upstream.await.unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_history_recovery"); + assert!(captured_header(&captured.second_ws_headers, "x-codex-turn-state").is_none()); + assert!(captured + .second_ws_payload + .get("previous_response_id") + .is_none()); +} + +#[tokio::test] +async fn responses_should_strip_history_after_sse_invalid_encrypted_reasoning_replay() { + let (base_url, upstream) = + spawn_websocket_failure_then_websocket_success_upstream(response_failed_websocket_message( + "resp_invalid_reasoning_failed", + "invalid_encrypted_content", + "Invalid encrypted content in reasoning replay", + )) + .await; + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "previous_response_id": "resp_stale", + "turnState": "turn-stale", + "input": [{"role": "user", "content": "Continue"}], + "stream": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let captured = upstream.await.unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_history_recovery"); + assert!(captured_header(&captured.second_ws_headers, "x-codex-turn-state").is_none()); + assert!(captured + .second_ws_payload + .get("previous_response_id") + .is_none()); +} + +#[tokio::test] +async fn responses_stream_should_strip_history_after_previous_response_not_found() { + let (base_url, upstream) = + spawn_websocket_failure_then_websocket_success_upstream(response_failed_websocket_message( + "resp_previous_missing_failed", + "previous_response_not_found", + "Previous response with id resp_stale was not found", + )) + .await; + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "previous_response_id": "resp_stale", + "turnState": "turn-stale", + "input": [{"role": "user", "content": "Continue"}], + "stream": true + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_text(response).await; + let captured = upstream.await.unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(body.contains("resp_after_history_recovery")); + assert_eq!( + captured_header(&captured.first_ws_headers, "x-codex-turn-state"), + Some("turn-stale") + ); + assert_eq!( + captured.first_ws_payload["previous_response_id"], + "resp_stale" + ); + assert!(captured_header(&captured.second_ws_headers, "x-codex-turn-state").is_none()); + assert!(captured + .second_ws_payload + .get("previous_response_id") + .is_none()); +} + +#[tokio::test] +async fn responses_stream_should_strip_history_after_websocket_invalid_encrypted_reasoning_replay() +{ + let (base_url, upstream) = + spawn_websocket_failure_then_websocket_success_upstream(response_failed_websocket_message( + "resp_invalid_reasoning_failed", + "invalid_encrypted_content", + "Invalid encrypted content in reasoning replay", + )) + .await; + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "previous_response_id": "resp_stale", + "turnState": "turn-stale", + "input": [{"role": "user", "content": "Continue"}], + "stream": true + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_text(response).await; + let captured = upstream.await.unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(body.contains("resp_after_history_recovery")); + assert_eq!( + captured_header(&captured.first_ws_headers, "x-codex-turn-state"), + Some("turn-stale") + ); + assert_eq!( + captured.first_ws_payload["previous_response_id"], + "resp_stale" + ); + assert!(captured_header(&captured.second_ws_headers, "x-codex-turn-state").is_none()); + assert!(captured + .second_ws_payload + .get("previous_response_id") + .is_none()); +} + +#[tokio::test] +async fn responses_stream_should_proxy_sse_and_record_usage() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_STREAM_USAGE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let usage: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("event: response.completed")); + assert!(body.contains("resp_stream_usage")); + assert!( + body.ends_with("data: [DONE]\n\n"), + "stream responses should terminate clients, body was {body:?}" + ); + assert_eq!(usage, (1, 3, 5)); +} + +#[tokio::test] +async fn responses_stream_should_close_http_sse_upstream_when_client_disconnects() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + read_http_request(&mut socket).await; + write_chunked_http_sse_headers(&mut socket).await; + write_http_chunk( + &mut socket, + b"event: response.output_text.delta\ndata: {\"delta\":\"hello\"}\n\n", + ) + .await; + socket.flush().await.unwrap(); + + timeout( + StdDuration::from_secs(2), + wait_for_http_sse_upstream_disconnect(&mut socket), + ) + .await + .is_ok() + }); + + let (app, api_key, _dir) = test_app_with_account(base_url).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let mut body = response.into_body().into_data_stream(); + let first_chunk = timeout(StdDuration::from_secs(1), body.next()) + .await + .expect("first SSE chunk should arrive before disconnect") + .expect("stream should yield a first chunk") + .expect("chunk should be readable"); + assert!(String::from_utf8(first_chunk.to_vec()) + .unwrap() + .contains("event: response.output_text.delta")); + + drop(body); + assert!( + upstream.await.unwrap(), + "dropping the downstream stream should close the HTTP SSE upstream socket" + ); +} + +#[tokio::test] +async fn responses_stream_should_record_event_log_after_completed_stream() { + 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") + .insert_header("retry-after", "7") + .insert_header("x-ratelimit-limit-requests", "99") + .set_body_string(RESPONSES_STREAM_USAGE_SSE), + ) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .header("x-request-id", "req_stream_completed_log") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "stream with log"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let body = response_text(response).await; + let event = latest_response_event_log(&pool).await; + let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap(); + + assert!(body.contains("resp_stream_usage")); + assert_eq!(event.level, "info"); + assert_eq!( + event.request_id.as_deref(), + Some("req_stream_completed_log") + ); + assert_eq!(event.account_id.as_deref(), Some("acct_chat")); + assert_eq!(event.status_code, Some(200)); + assert_eq!(metadata["stream"], true); + assert_eq!(metadata["completed"], true); + assert_eq!(metadata["responseId"], "resp_stream_usage"); + assert_eq!(metadata["usage"]["inputTokens"], 3); + assert_eq!(metadata["usage"]["outputTokens"], 5); + assert_rate_limit_header(&metadata, "retry-after", "7"); + assert_rate_limit_header(&metadata, "x-ratelimit-limit-requests", "99"); + assert!(metadata.get("requestBody").is_none()); + assert!(metadata.get("responseBody").is_none()); +} + +#[tokio::test] +async fn responses_stream_should_preserve_body_metadata_when_capture_body_enabled() { + 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(RESPONSES_STREAM_USAGE_SSE), + ) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = + test_app_with_account_pool_and_logging_capture_body(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .header("x-request-id", "req_stream_capture_body_log") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "stream body capture"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let body = response_text(response).await; + let event = latest_response_event_log(&pool).await; + let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap(); + + assert!(body.contains("resp_stream_usage")); + assert_eq!(event.level, "info"); + assert_eq!( + event.request_id.as_deref(), + Some("req_stream_capture_body_log") + ); + assert_eq!(metadata["requestBody"]["model"], "gpt-5.5"); + assert_eq!(metadata["requestBody"]["stream"], true); + assert!(metadata["responseBody"] + .as_str() + .is_some_and(|body| body.contains("resp_stream_usage"))); +} + +#[tokio::test] +async fn responses_stream_should_record_event_log_after_late_disconnect() { + let first_frame = r#"event: response.output_text.delta +data: {"delta":"partial before logged disconnect"} + +"#; + let (base_url, first_chunk_sent, close_upstream) = + spawn_chunked_sse_upstream_then_clean_close_with_headers( + first_frame, + &[ + ("retry-after", "11"), + ("x-codex-primary-used-percent", "88"), + ], + ) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(base_url).await; + let response_task = tokio::spawn(async move { + app.oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .header("x-request-id", "req_stream_disconnect_log") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "stream disconnect log"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap() + }); + first_chunk_sent.await.unwrap(); + let response = response_task.await.unwrap(); + close_upstream.send(()).unwrap(); + let body = response_text(response).await; + let event = latest_response_event_log(&pool).await; + let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap(); + + assert!(body.contains("stream_disconnected")); + assert_eq!(event.level, "error"); + assert_eq!( + event.request_id.as_deref(), + Some("req_stream_disconnect_log") + ); + assert_eq!(event.account_id.as_deref(), Some("acct_chat")); + assert_eq!(event.status_code, Some(502)); + assert_eq!(metadata["stream"], true); + assert_eq!(metadata["failed"], true); + assert_eq!(metadata["upstreamCode"], "stream_disconnected"); + assert_rate_limit_header(&metadata, "retry-after", "11"); + assert_rate_limit_header(&metadata, "x-codex-primary-used-percent", "88"); +} + +#[tokio::test] +async fn responses_stream_should_forward_first_chunk_before_upstream_completes() { + let first_frame = r#"event: response.output_text.delta +data: {"delta":"live stream hello"} + +"#; + let final_frame = r#"event: response.completed +data: {"response":{"id":"resp_live_stream","object":"response","status":"completed","usage":{"input_tokens":3,"output_tokens":4}}} + +"#; + let (base_url, first_chunk_sent, finish_upstream) = + spawn_chunked_sse_upstream(first_frame, final_frame).await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(base_url).await; + let response_task = tokio::spawn(async move { + app.oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap() + }); + + first_chunk_sent.await.unwrap(); + let response = timeout(StdDuration::from_millis(300), response_task) + .await + .expect("stream response should be returned before upstream completes") + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let mut body_stream = response.into_body().into_data_stream(); + let first_chunk = timeout(StdDuration::from_millis(300), body_stream.next()) + .await + .expect("first proxied SSE chunk should arrive before upstream completes") + .unwrap() + .unwrap(); + let first_chunk = String::from_utf8(first_chunk.to_vec()).unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(first_chunk.contains("live stream hello")); + assert!(!first_chunk.contains("resp_live_stream")); + + finish_upstream.send(()).unwrap(); + let mut rest = Vec::new(); + while let Some(chunk) = body_stream.next().await { + rest.extend_from_slice(&chunk.unwrap()); + } + let rest = String::from_utf8(rest).unwrap(); + let usage: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert!(rest.contains("resp_live_stream")); + assert!( + rest.ends_with("data: [DONE]\n\n"), + "stream responses should terminate clients, body was {rest:?}" + ); + assert_eq!(usage, (1, 3, 4)); +} + +#[tokio::test] +async fn responses_stream_should_emit_failed_event_after_upstream_read_error_once_downstream_started( +) { + let first_frame = r#"event: response.output_text.delta +data: {"delta":"partial before transport failure"} + +"#; + let (base_url, first_chunk_sent, close_upstream) = + spawn_chunked_sse_upstream_then_abrupt_close(first_frame).await; + + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + let response_task = tokio::spawn(async move { + app.oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Start then fail"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap() + }); + + first_chunk_sent.await.unwrap(); + let response = timeout(StdDuration::from_millis(300), response_task) + .await + .expect("stream response should be returned before upstream closes") + .unwrap(); + let mut body_stream = response.into_body().into_data_stream(); + let first_chunk = timeout(StdDuration::from_millis(300), body_stream.next()) + .await + .expect("first proxied SSE chunk should arrive before upstream closes") + .unwrap() + .unwrap(); + assert!(String::from_utf8(first_chunk.to_vec()) + .unwrap() + .contains("partial before transport failure")); + + close_upstream.send(()).unwrap(); + let rest = collect_stream_body(body_stream).await; + + assert!(rest.contains("event: response.failed")); + assert!(rest.contains("stream_disconnected")); + assert!(rest.ends_with("data: [DONE]\n\n")); +} + +#[tokio::test] +async fn responses_stream_should_emit_failed_event_when_upstream_closes_without_completed() { + let first_frame = r#"event: response.output_text.delta +data: {"delta":"partial before clean close"} + +"#; + let (base_url, first_chunk_sent, close_upstream) = + spawn_chunked_sse_upstream_then_clean_close(first_frame).await; + + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await; + let response_task = tokio::spawn(async move { + app.oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Start then close"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap() + }); + + first_chunk_sent.await.unwrap(); + let response = timeout(StdDuration::from_millis(300), response_task) + .await + .expect("stream response should be returned before upstream closes") + .unwrap(); + let mut body_stream = response.into_body().into_data_stream(); + let first_chunk = timeout(StdDuration::from_millis(300), body_stream.next()) + .await + .expect("first proxied SSE chunk should arrive before upstream closes") + .unwrap() + .unwrap(); + assert!(String::from_utf8(first_chunk.to_vec()) + .unwrap() + .contains("partial before clean close")); + + close_upstream.send(()).unwrap(); + let rest = collect_stream_body(body_stream).await; + + assert!(rest.contains("event: response.failed")); + assert!(rest.contains("stream_disconnected")); + assert!(rest.ends_with("data: [DONE]\n\n")); +} + +#[tokio::test] +async fn responses_stream_should_fallback_to_next_account_after_rate_limit() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("retry-after", "90") + .set_body_json(json!({ + "error": {"message": "rate limited"} + })), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_STREAM_AFTER_429_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let primary_quota_state: (i64, Option) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", + ) + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + let primary_usage: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + let secondary_usage: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_secondary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("resp_stream_after_429")); + assert_eq!(primary_quota_state.0, 1); + assert!(primary_quota_state.1.is_some()); + assert_eq!(primary_usage, (1, 0, 0)); + assert_eq!(secondary_usage, (1, 4, 1)); +} + +#[tokio::test] +async fn responses_stream_should_mark_expired_after_401_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(401).set_body_json(json!({ + "error": {"code": "token_revoked", "message": "token revoked"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_STREAM_AFTER_401_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + let secondary_usage: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_secondary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("resp_stream_after_401")); + assert_eq!(account_status.0, "expired"); + assert_eq!(secondary_usage, (1, 4, 2)); +} + +#[tokio::test] +async fn responses_stream_should_return_auth_error_when_401_fallback_is_exhausted() { + 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", "message": "token revoked"} + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("event: response.failed")); + assert!(body.contains("All accounts exhausted (1 expired)")); + assert!(body.contains("token_revoked")); + assert_eq!(account_status.0, "expired"); +} + +#[tokio::test] +async fn responses_stream_should_mark_banned_after_403_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(403).set_body_json(json!({ + "error": {"message": "account banned"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_STREAM_AFTER_403_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("resp_stream_after_403")); + assert_eq!(account_status.0, "banned"); +} + +#[tokio::test] +async fn responses_stream_should_cool_down_cloudflare_403_and_fallback() { + let started_at = Utc::now(); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(403).set_body_string( + "Just a moment...cf_chl challenge", + )) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_STREAM_AFTER_CLOUDFLARE_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + cookie_store + .set_cookie_header("acct_primary", "cf_clearance=old") + .await + .unwrap(); + cookie_store + .set_cookie_header("acct_secondary", "cf_clearance=keep") + .await + .unwrap(); + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let primary_state: (String, Option) = + sqlx::query_as("select status, cloudflare_cooldown_until from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + let cooldown_until = chrono::DateTime::parse_from_rfc3339(primary_state.1.as_deref().unwrap()) + .unwrap() + .with_timezone(&Utc); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("resp_stream_after_cloudflare")); + assert_eq!(primary_state.0, "active"); + assert!(cooldown_until > started_at); + assert_eq!( + cookie_store + .cookie_header("acct_primary", "chatgpt.com") + .await + .unwrap(), + None + ); + assert_eq!( + cookie_store + .cookie_header("acct_secondary", "chatgpt.com") + .await + .unwrap() + .as_deref(), + Some("cf_clearance=keep") + ); +} + +#[tokio::test] +async fn responses_stream_should_return_cloudflare_path_block_error_when_404_fallback_is_exhausted() +{ + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(404)) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + cookie_store + .set_cookie_header("acct_chat", "cf_clearance=old") + .await + .unwrap(); + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("event: response.failed")); + assert!(body.contains("All accounts exhausted (1 cloudflare-path-block)")); + assert!(body.contains("Cloudflare path-block")); + assert_eq!( + cookie_store + .cookie_header("acct_chat", "chatgpt.com") + .await + .unwrap(), + None + ); +} + +#[tokio::test] +async fn responses_should_retry_same_account_after_5xx_before_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(500).set_body_json(json!({ + "error": {"message": "temporary upstream failure"} + }))) + .up_to_n_times(2) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_AFTER_5XX_RETRY_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let authorizations = received_authorizations(&server).await; + let primary_requests = authorizations + .iter() + .filter(|authorization| authorization.as_str() == "Bearer access-primary") + .count(); + let secondary_requests = authorizations + .iter() + .filter(|authorization| authorization.as_str() == "Bearer access-secondary") + .count(); + + assert_eq!(status, StatusCode::OK, "requests: {authorizations:?}"); + assert_eq!(body["id"], "resp_after_5xx_retry"); + assert_eq!(primary_requests, 3, "requests: {authorizations:?}"); + assert_eq!(secondary_requests, 0, "requests: {authorizations:?}"); +} + +#[tokio::test] +async fn responses_stream_should_retry_same_account_after_5xx_before_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(500).set_body_json(json!({ + "error": {"message": "temporary upstream failure"} + }))) + .up_to_n_times(2) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_STREAM_AFTER_5XX_RETRY_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, _pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let authorizations = received_authorizations(&server).await; + let primary_requests = authorizations + .iter() + .filter(|authorization| authorization.as_str() == "Bearer access-primary") + .count(); + let secondary_requests = authorizations + .iter() + .filter(|authorization| authorization.as_str() == "Bearer access-secondary") + .count(); + + assert_eq!(status, StatusCode::OK, "requests: {authorizations:?}"); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("resp_stream_after_5xx_retry")); + assert_eq!(primary_requests, 3, "requests: {authorizations:?}"); + assert_eq!(secondary_requests, 0, "requests: {authorizations:?}"); +} + +#[tokio::test] +async fn responses_should_mark_quota_exhausted_after_402_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(402).set_body_json(json!({ + "error": {"message": "quota exhausted"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_AFTER_402_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + let failed_usage: Option<(i64,)> = + sqlx::query_as("select request_count from account_usage where account_id = ?") + .bind("acct_primary") + .fetch_optional(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_402"); + assert_eq!(account_status.0, "quota_exhausted"); + assert_eq!(failed_usage.map(|row| row.0).unwrap_or_default(), 1); +} + +#[tokio::test] +async fn responses_stream_should_mark_quota_exhausted_after_402_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with(ResponseTemplate::new(402).set_body_json(json!({ + "error": {"message": "quota exhausted"} + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_STREAM_AFTER_402_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + let secondary_usage: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + ) + .bind("acct_secondary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("resp_stream_after_402")); + assert_eq!(account_status.0, "quota_exhausted"); + assert_eq!(secondary_usage, (1, 6, 3)); +} + +#[tokio::test] +async fn responses_should_return_quota_error_when_402_fallback_is_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(402).set_body_json(json!({ + "error": {"message": "quota reached"} + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + let message = body["error"]["message"].as_str().unwrap_or_default(); + + assert_eq!(status, StatusCode::PAYMENT_REQUIRED); + assert!(message.contains("All accounts exhausted (1 quota-exhausted)")); + assert!(message.contains("quota reached")); + assert_eq!(body["error"]["code"], "upstream_error"); + assert_eq!(account_status.0, "quota_exhausted"); +} + +#[tokio::test] +async fn responses_should_return_rate_limit_error_when_429_fallback_is_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(429).set_body_json(json!({ + "error": {"message": "rate limited"} + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let quota_state: (i64, Option) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + let message = body["error"]["message"].as_str().unwrap_or_default(); + + assert_eq!(status, StatusCode::TOO_MANY_REQUESTS); + assert!(message.contains("All accounts exhausted (1 rate-limited)")); + assert!(message.contains("rate limited")); + assert_eq!(body["error"]["code"], "upstream_error"); + assert_eq!(quota_state.0, 1); + assert!(quota_state.1.is_some()); +} + +#[tokio::test] +async fn responses_stream_should_return_rate_limit_error_when_429_fallback_is_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(429).set_body_json(json!({ + "error": {"message": "rate limited"} + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let quota_state: (i64, Option) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("event: response.failed")); + assert!(body.contains("rate_limit_exceeded")); + assert!(body.contains("All accounts exhausted (1 rate-limited)")); + assert!(body.contains("rate limited")); + assert_eq!(quota_state.0, 1); + assert!(quota_state.1.is_some()); +} + +#[tokio::test] +async fn responses_stream_should_return_quota_error_when_402_fallback_is_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(402).set_body_json(json!({ + "error": {"message": "quota reached"} + }))) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": true, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let body = response_text(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("event: response.failed")); + assert!(body.contains("All accounts exhausted (1 quota-exhausted)")); + assert!(body.contains("quota reached")); + assert_eq!(account_status.0, "quota_exhausted"); +} + +#[tokio::test] +async fn responses_should_classify_sse_quota_failure_and_fallback() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-primary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_FAILED_QUOTA_SSE), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secondary")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_AFTER_402_SSE), + ) + .expect(1) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["id"], "resp_after_402"); + assert_eq!(account_status.0, "quota_exhausted"); +} + +#[tokio::test] +async fn responses_should_record_request_count_when_5xx_retries_are_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .and(header("authorization", "Bearer access-secret")) + .respond_with(ResponseTemplate::new(500).set_body_json(json!({ + "error": {"message": "temporary upstream failure"} + }))) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let authorizations = received_authorizations(&server).await; + let failed_usage: Option<(i64,)> = + sqlx::query_as("select request_count from account_usage where account_id = ?") + .bind("acct_chat") + .fetch_optional(&pool) + .await + .unwrap(); + + assert_eq!(status, StatusCode::BAD_GATEWAY); + assert_eq!(authorizations.len(), 3, "requests: {authorizations:?}"); + assert_eq!(failed_usage.map(|row| row.0).unwrap_or_default(), 1); +} + +#[tokio::test] +async fn responses_should_prefer_session_affinity_account_for_previous_response() { + let (base_url, upstream) = spawn_single_websocket_completed_upstream("resp_affinity_ws").await; + let (app, api_key, _dir) = test_app_with_two_accounts_and_affinity(base_url).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "previous_response_id": "resp_previous", + "input": [{"role": "user", "content": "Continue"}], + "stream": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let captured = upstream.await.unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + captured_header(&captured.headers, "authorization"), + Some("Bearer access-affinity") + ); + assert_eq!( + captured_header(&captured.headers, "chatgpt-account-id"), + Some("chatgpt-affinity") + ); + assert_eq!(captured.payload["previous_response_id"], "resp_previous"); +} + +#[tokio::test] +async fn responses_with_previous_response_id_should_use_websocket_and_configured_pool() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut first_websocket = accept_async(first_stream).await.unwrap(); + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_websocket + .send(Message::Text( + response_completed_websocket_message("resp_runtime_pool_first").into(), + )) + .await + .unwrap(); + + tokio::select! { + second_message = first_websocket.next() => { + let _second_message = second_message.unwrap().unwrap(); + first_websocket + .send(Message::Text( + response_completed_websocket_message("resp_runtime_pool_second").into(), + )) + .await + .unwrap(); + first_websocket.close(None).await.unwrap(); + } + accepted = listener.accept() => { + let (second_stream, _) = accepted.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut second_websocket = accept_async(second_stream).await.unwrap(); + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + response_completed_websocket_message("resp_runtime_pool_second").into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + first_websocket.close(None).await.unwrap(); + } + } + }); + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let first = app + .clone() + .oneshot(responses_previous_request( + &api_key, + "Continue from pooled runtime websocket", + )) + .await + .unwrap(); + let second = app + .oneshot(responses_previous_request( + &api_key, + "Continue again from pooled runtime websocket", + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(first.status(), StatusCode::OK); + assert_eq!(second.status(), StatusCode::OK); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn responses_stream_with_previous_response_id_should_forward_websocket_chunks_before_completion( +) { + let (base_url, first_chunk_sent_rx, finish_tx, upstream) = + spawn_chunked_websocket_upstream().await; + let (app, api_key, _dir) = test_app_with_account(base_url).await; + + let response = app + .oneshot(responses_previous_stream_request( + &api_key, + "Continue as a WebSocket stream", + )) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + first_chunk_sent_rx + .await + .expect("upstream should send the first websocket event"); + let first_chunk = timeout( + StdDuration::from_secs(1), + first_response_body_chunk(response), + ) + .await + .expect("downstream should receive a websocket-backed SSE chunk before upstream completes") + .expect("response body should produce a chunk"); + finish_tx + .send(()) + .expect("test should be able to finish upstream websocket"); + let captured = upstream.await.unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(first_chunk.contains("event: response.output_text.delta")); + assert_eq!( + captured_header(&captured.headers, "authorization"), + Some("Bearer access-secret") + ); + assert_eq!( + captured.payload["previous_response_id"], + "resp_runtime_pool_previous" + ); +} + +#[tokio::test] +async fn responses_stream_with_previous_response_id_should_record_websocket_audit_metadata() { + let (base_url, first_chunk_sent_rx, finish_tx, upstream) = + spawn_chunked_websocket_upstream().await; + let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(base_url).await; + + let response = app + .oneshot(responses_previous_stream_request( + &api_key, + "Continue as a logged WebSocket stream", + )) + .await + .unwrap(); + first_chunk_sent_rx + .await + .expect("upstream should send the first websocket event"); + finish_tx + .send(()) + .expect("test should be able to finish upstream websocket"); + let body = response_text(response).await; + let captured = upstream.await.unwrap(); + let event = latest_response_event_log(&pool).await; + let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap(); + + assert!(body.contains("resp_live_websocket_stream")); + assert_eq!( + captured.payload["previous_response_id"], + "resp_runtime_pool_previous" + ); + assert_eq!(event.level, "info"); + assert_eq!(metadata["stream"], true); + assert_eq!(metadata["transport"], "websocket"); + assert_eq!(metadata["usage"]["inputTokens"], 3); + assert_eq!(metadata["usage"]["outputTokens"], 1); + assert_rate_limit_header(&metadata, "x-codex-primary-used-percent", "44"); + assert_rate_limit_header(&metadata, "x-codex-primary-window-minutes", "5"); +} + +#[tokio::test] +async fn responses_should_record_session_affinity_for_completed_response() { + 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")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(RESPONSES_SUCCESS_SSE), + ) + .mount(&server) + .await; + + let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [{"role": "user", "content": "Say hello"}], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let records = SqliteSessionAffinityStore::new(pool) + .list_active(Utc::now()) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(records.len(), 1); + assert_eq!(records[0].response_id, "resp_response_1"); + assert_eq!(records[0].entry.account_id, "acct_chat"); +} + +async fn test_app_with_account(base_url: String) -> (axum::Router, String, tempfile::TempDir) { + test_app_with_account_and_installation_id(base_url, TEST_INSTALLATION_ID.to_string()).await +} + +async fn test_app_with_account_and_installation_id( + base_url: String, + installation_id: String, +) -> (axum::Router, String, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("openai-chat-upstream.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let secret_box = SecretBox::new([83u8; 32]); + let hasher = ApiKeyHasher::new([84u8; 32]); + let api_key = insert_client_api_key(&pool, &hasher).await; + insert_account(&pool, &secret_box).await; + let state = AppState::with_pool_secret_api_key_hasher_and_installation_id( + test_config(url, base_url), + pool, + secret_box, + hasher, + installation_id, + ); + state + .restore_account_pool_from_repository() + .await + .expect("account pool should restore"); + (router::router().with_state(state), api_key, dir) +} + +async fn test_app_with_account_and_pool( + base_url: String, +) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { + test_app_with_account_pool_config(base_url, |_| {}).await +} + +async fn test_app_without_accounts(base_url: String) -> (axum::Router, String, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("openai-responses-no-accounts.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let secret_box = SecretBox::new([83u8; 32]); + let hasher = ApiKeyHasher::new([84u8; 32]); + let api_key = insert_client_api_key(&pool, &hasher).await; + let state = AppState::with_pool_secret_api_key_hasher_and_installation_id( + test_config(url, base_url), + pool, + secret_box, + hasher, + TEST_INSTALLATION_ID.to_string(), + ); + state + .restore_account_pool_from_repository() + .await + .expect("empty account pool should restore"); + (router::router().with_state(state), api_key, dir) +} + +async fn test_app_with_account_pool_and_logging( + base_url: String, +) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { + test_app_with_account_pool_config(base_url, |config| { + config.logging.enabled = true; + }) + .await +} + +async fn test_app_with_account_pool_and_logging_capture_body( + base_url: String, +) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { + test_app_with_account_pool_config(base_url, |config| { + config.logging.enabled = true; + config.logging.capture_body = true; + }) + .await +} + +async fn test_app_with_account_pool_config( + base_url: String, + configure: impl FnOnce(&mut AppConfig), +) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("openai-record-affinity.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let secret_box = SecretBox::new([83u8; 32]); + let hasher = ApiKeyHasher::new([84u8; 32]); + let api_key = insert_client_api_key(&pool, &hasher).await; + insert_account(&pool, &secret_box).await; + let mut config = test_config(url, base_url); + configure(&mut config); + let state = AppState::with_pool_secret_api_key_hasher_and_installation_id( + config, + pool.clone(), + secret_box, + hasher, + TEST_INSTALLATION_ID.to_string(), + ); + state + .restore_account_pool_from_repository() + .await + .expect("account pool should restore"); + (router::router().with_state(state), api_key, pool, dir) +} + +async fn test_app_with_restored_pool_then_disabled_account( + base_url: String, +) -> (axum::Router, String, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("openai-chat-restored-pool.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let secret_box = SecretBox::new([83u8; 32]); + let hasher = ApiKeyHasher::new([84u8; 32]); + let api_key = insert_client_api_key(&pool, &hasher).await; + insert_account(&pool, &secret_box).await; + let state = AppState::with_pool_secret_api_key_hasher_and_installation_id( + test_config(url, base_url), + pool.clone(), + secret_box, + hasher, + TEST_INSTALLATION_ID.to_string(), + ); + state + .restore_account_pool_from_repository() + .await + .expect("account pool should restore"); + sqlx::query("update accounts set status = 'disabled' where id = ?") + .bind("acct_chat") + .execute(&pool) + .await + .unwrap(); + (router::router().with_state(state), api_key, dir) +} + +async fn test_app_with_two_accounts_and_affinity( + base_url: String, +) -> (axum::Router, String, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("openai-responses-affinity.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let secret_box = SecretBox::new([83u8; 32]); + let hasher = ApiKeyHasher::new([84u8; 32]); + let api_key = insert_client_api_key(&pool, &hasher).await; + insert_named_account( + &pool, + &secret_box, + "acct_a", + "access-default", + "chatgpt-default", + ) + .await; + insert_named_account( + &pool, + &secret_box, + "acct_z", + "access-affinity", + "chatgpt-affinity", + ) + .await; + let now = Utc::now(); + SqliteSessionAffinityStore::new(pool.clone()) + .upsert( + "resp_previous", + &SessionAffinityEntry { + account_id: "acct_z".to_string(), + conversation_id: "conv_affinity".to_string(), + turn_state: Some("turn_affinity".to_string()), + instructions_hash: None, + input_tokens: None, + function_call_ids: Vec::new(), + variant_hash: None, + created_at: now, + }, + Duration::hours(4), + ) + .await + .unwrap(); + let state = AppState::with_pool_secret_api_key_hasher_and_installation_id( + test_config(url, base_url), + pool, + secret_box, + hasher, + TEST_INSTALLATION_ID.to_string(), + ); + state + .restore_account_pool_from_repository() + .await + .expect("account pool should restore"); + state + .restore_session_affinity_from_repository(now + Duration::minutes(1)) + .await + .expect("session affinity should restore"); + (router::router().with_state(state), api_key, dir) +} + +async fn test_app_with_two_accounts( + base_url: String, +) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("openai-responses-fallback.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let secret_box = SecretBox::new([83u8; 32]); + let hasher = ApiKeyHasher::new([84u8; 32]); + let api_key = insert_client_api_key(&pool, &hasher).await; + insert_named_account( + &pool, + &secret_box, + "acct_primary", + "access-primary", + "chatgpt-primary", + ) + .await; + insert_named_account( + &pool, + &secret_box, + "acct_secondary", + "access-secondary", + "chatgpt-secondary", + ) + .await; + let state = AppState::with_pool_secret_api_key_hasher_and_installation_id( + test_config(url, base_url), + pool.clone(), + secret_box, + hasher, + TEST_INSTALLATION_ID.to_string(), + ); + state + .restore_account_pool_from_repository() + .await + .expect("account pool should restore"); + (router::router().with_state(state), api_key, pool, dir) +} + +async fn insert_client_api_key(pool: &SqlitePool, hasher: &ApiKeyHasher) -> String { + let generated = hasher.generate_client_api_key("test"); + sqlx::query( + "insert into client_api_keys (id, name, prefix, key_hash, enabled, created_at) values (?, ?, ?, ?, 1, ?)", + ) + .bind("key_test") + .bind("test") + .bind(&generated.prefix) + .bind(&generated.key_hash) + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); + generated.plaintext +} + +async fn seed_openai_admin_session(pool: &SqlitePool, session_id: &str) { + sqlx::query( + "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", + ) + .bind("admin_openai") + .bind("hash") + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); + sqlx::query( + "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", + ) + .bind(session_id) + .bind("admin_openai") + .bind("2999-01-01T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); +} + +async fn update_admin_account_status(app: &axum::Router, account_id: &str, status: &str) { + let response = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri(format!("/api/admin/accounts/{account_id}/status")) + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_status_cycle") + .body(Body::from(json!({"status": status}).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); +} + +async fn insert_account(pool: &SqlitePool, secret_box: &SecretBox) { + let access_token = secret_box + .encrypt(&SecretString::new("access-secret".to_string().into())) + .unwrap(); + sqlx::query( + "insert into accounts (id, email, account_id, user_id, access_token_cipher, access_token_expires_at, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind("acct_chat") + .bind("user@example.com") + .bind("chatgpt-account") + .bind("chatgpt-user") + .bind(access_token) + .bind("2100-01-01T00:00:00Z") + .bind("active") + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); +} + +async fn insert_named_account( + pool: &SqlitePool, + secret_box: &SecretBox, + id: &str, + access_token_plaintext: &str, + chatgpt_account_id: &str, +) { + let access_token = secret_box + .encrypt(&SecretString::new( + access_token_plaintext.to_string().into(), + )) + .unwrap(); + sqlx::query( + "insert into accounts (id, email, account_id, user_id, access_token_cipher, access_token_expires_at, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(id) + .bind(format!("{id}@example.com")) + .bind(chatgpt_account_id) + .bind(format!("user-{id}")) + .bind(access_token) + .bind("2100-01-01T00:00:00Z") + .bind("active") + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); +} + +fn tuple_response_request_body(stream: bool) -> String { + json!({ + "model": "gpt-5.5", + "stream": stream, + "use_websocket": false, + "input": [], + "text": { + "format": { + "type": "json_schema", + "name": "TupleAnswer", + "schema": { + "type": "object", + "properties": { + "point": { + "type": "array", + "prefixItems": [ + {"type": "number"}, + {"type": "number"} + ], + "items": false + } + }, + "required": ["point"] + }, + "strict": true + } + } + }) + .to_string() +} + +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() +} + +struct CapturedWebSocketRequest { + headers: Vec<(String, String)>, + payload: Value, +} + +struct HistoryRecoveryCapture { + first_ws_headers: Vec<(String, String)>, + first_ws_payload: Value, + second_ws_headers: Vec<(String, String)>, + second_ws_payload: Value, +} + +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn spawn_single_websocket_completed_upstream( + response_id: &'static str, +) -> (String, tokio::task::JoinHandle) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let captured_headers = Arc::new(Mutex::new(Vec::new())); + let captured_headers_for_callback = Arc::clone(&captured_headers); + let mut websocket = accept_hdr_async( + stream, + move |request: &WsRequest, mut response: WsResponse| { + *captured_headers_for_callback.lock().unwrap() = request_headers(request); + response + .headers_mut() + .insert("x-ratelimit-limit-requests", "55".parse().unwrap()); + response + .headers_mut() + .insert("x-codex-primary-used-percent", "44".parse().unwrap()); + Ok(response) + }, + ) + .await + .unwrap(); + let message = websocket.next().await.unwrap().unwrap(); + let payload = serde_json::from_str::(&message.into_text().unwrap()) + .expect("websocket payload should be json"); + websocket + .send(Message::Text( + response_completed_websocket_message(response_id).into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + let headers = captured_headers.lock().unwrap().clone(); + CapturedWebSocketRequest { headers, payload } + }); + (base_url, server) +} + +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn spawn_websocket_failure_then_websocket_success_upstream( + failure: String, +) -> (String, tokio::task::JoinHandle) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + let first_headers = Arc::new(Mutex::new(Vec::new())); + let first_headers_for_callback = Arc::clone(&first_headers); + let mut first_websocket = accept_hdr_async( + first_stream, + move |request: &WsRequest, response: WsResponse| { + *first_headers_for_callback.lock().unwrap() = request_headers(request); + Ok(response) + }, + ) + .await + .unwrap(); + let first_message = first_websocket.next().await.unwrap().unwrap(); + let first_ws_payload = serde_json::from_str::(&first_message.into_text().unwrap()) + .expect("websocket payload should be json"); + first_websocket + .send(Message::Text(failure.into())) + .await + .unwrap(); + drop(first_websocket); + + let (second_stream, _) = listener.accept().await.unwrap(); + let second_headers = Arc::new(Mutex::new(Vec::new())); + let second_headers_for_callback = Arc::clone(&second_headers); + let mut second_websocket = accept_hdr_async( + second_stream, + move |request: &WsRequest, response: WsResponse| { + *second_headers_for_callback.lock().unwrap() = request_headers(request); + Ok(response) + }, + ) + .await + .unwrap(); + let second_message = second_websocket.next().await.unwrap().unwrap(); + let second_ws_payload = serde_json::from_str::(&second_message.into_text().unwrap()) + .expect("second websocket payload should be json"); + second_websocket + .send(Message::Text( + response_completed_websocket_message("resp_after_history_recovery").into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + let first_ws_headers = first_headers.lock().unwrap().clone(); + let second_ws_headers = second_headers.lock().unwrap().clone(); + + HistoryRecoveryCapture { + first_ws_headers, + first_ws_payload, + second_ws_headers, + second_ws_payload, + } + }); + (base_url, server) +} + +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn spawn_chunked_websocket_upstream() -> ( + String, + oneshot::Receiver<()>, + oneshot::Sender<()>, + tokio::task::JoinHandle, +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let (first_chunk_sent_tx, first_chunk_sent_rx) = oneshot::channel(); + let (finish_tx, finish_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let captured_headers = Arc::new(Mutex::new(Vec::new())); + let captured_headers_for_callback = Arc::clone(&captured_headers); + let mut websocket = + accept_hdr_async(stream, move |request: &WsRequest, response: WsResponse| { + *captured_headers_for_callback.lock().unwrap() = request_headers(request); + Ok(response) + }) + .await + .unwrap(); + let message = websocket.next().await.unwrap().unwrap(); + let payload = serde_json::from_str::(&message.into_text().unwrap()) + .expect("websocket payload should be json"); + websocket + .send(Message::Text( + json!({ + "type": "codex.rate_limits", + "rate_limits": { + "primary": { + "used_percent": 44.0, + "window_minutes": 5 + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "live websocket hello" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + let _ = first_chunk_sent_tx.send(()); + let _ = finish_rx.await; + websocket + .send(Message::Text( + response_completed_websocket_message("resp_live_websocket_stream").into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + let headers = captured_headers.lock().unwrap().clone(); + CapturedWebSocketRequest { headers, payload } + }); + (base_url, first_chunk_sent_rx, finish_tx, server) +} + +async fn accept_successful_websocket_response(listener: &TcpListener, response_id: &str) -> Value { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let message = websocket.next().await.unwrap().unwrap(); + let payload = serde_json::from_str::(&message.into_text().unwrap()) + .expect("websocket payload should be json"); + websocket + .send(Message::Text( + response_completed_websocket_message(response_id).into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + payload +} + +async fn accept_successful_websocket_response_with_authorization( + listener: &TcpListener, + expected_authorization: &'static str, + response_id: &str, +) -> Value { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_websocket_with_authorization(stream, expected_authorization).await; + let payload = send_websocket_response_and_capture_payload( + &mut websocket, + websocket_completed_response(response_id, 3, 1), + ) + .await; + websocket.close(None).await.unwrap(); + payload +} + +async fn accept_two_successful_websocket_responses_with_authorization( + listener: &TcpListener, + expected_authorization: &'static str, + first_response_id: &str, + second_response_id: &str, +) -> (bool, Value, Value) { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_websocket_with_authorization(stream, expected_authorization).await; + let first_payload = send_websocket_response_and_capture_payload( + &mut websocket, + websocket_completed_response(first_response_id, 3, 1), + ) + .await; + + loop { + tokio::select! { + message = websocket.next() => { + match message { + Some(Ok(message)) if message.is_text() => { + let second_payload = serde_json::from_str::( + &message.into_text().unwrap(), + ) + .expect("second websocket payload should be json"); + websocket + .send(Message::Text( + websocket_completed_response(second_response_id, 3, 1).into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + break (true, first_payload, second_payload); + } + Some(_) => continue, + None => { + let second_payload = accept_successful_websocket_response_with_authorization( + listener, + expected_authorization, + second_response_id, + ) + .await; + break (false, first_payload, second_payload); + } + } + } + accepted = listener.accept() => { + let (stream, _) = accepted.unwrap(); + let mut second_websocket = + accept_websocket_with_authorization(stream, expected_authorization).await; + let second_payload = send_websocket_response_and_capture_payload( + &mut second_websocket, + websocket_completed_response(second_response_id, 3, 1), + ) + .await; + second_websocket.close(None).await.unwrap(); + break (false, first_payload, second_payload); + } + } + } +} + +async fn accept_websocket_response_with_authorization_and_message( + listener: &TcpListener, + expected_authorization: &'static str, + response_message: String, +) -> Value { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_websocket_with_authorization(stream, expected_authorization).await; + let payload = + send_websocket_response_and_capture_payload(&mut websocket, response_message).await; + websocket.close(None).await.unwrap(); + payload +} + +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn accept_websocket_with_authorization( + stream: TcpStream, + expected_authorization: &'static str, +) -> tokio_tungstenite::WebSocketStream { + 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) + ); + Ok(response) + }) + .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")); + let lower_request = request.to_ascii_lowercase(); + assert!( + lower_request.contains(&format!( + "authorization: {}", + expected_authorization.to_ascii_lowercase() + )), + "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(); +} + +async fn accept_websocket_response_with_message( + listener: &TcpListener, + response_message: String, +) -> Value { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let payload = + send_websocket_response_and_capture_payload(&mut websocket, response_message).await; + websocket.close(None).await.unwrap(); + payload +} + +async fn accept_followup_websocket_response( + listener: &TcpListener, + websocket: &mut tokio_tungstenite::WebSocketStream, + response_message: String, +) -> Value { + tokio::select! { + message = websocket.next() => { + match message { + Some(Ok(message)) if message.is_text() => { + let payload = serde_json::from_str::(&message.into_text().unwrap()) + .expect("websocket payload should be json"); + websocket + .send(Message::Text(response_message.into())) + .await + .unwrap(); + payload + } + _ => accept_websocket_response_with_message(listener, response_message).await, + } + } + accepted = listener.accept() => { + let (stream, _) = accepted.unwrap(); + let mut followup = accept_async(stream).await.unwrap(); + let payload = + send_websocket_response_and_capture_payload(&mut followup, response_message).await; + followup.close(None).await.unwrap(); + payload + } + } +} + +async fn send_websocket_response_and_capture_payload( + websocket: &mut tokio_tungstenite::WebSocketStream, + response_message: String, +) -> Value { + let message = websocket.next().await.unwrap().unwrap(); + let payload = serde_json::from_str::(&message.into_text().unwrap()) + .expect("websocket payload should be json"); + websocket + .send(Message::Text(response_message.into())) + .await + .unwrap(); + payload +} + +fn request_headers(request: &WsRequest) -> Vec<(String, String)> { + request + .headers() + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.as_str().to_string(), value.to_string())) + }) + .collect() +} + +fn captured_header<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> { + headers + .iter() + .find(|(header_name, _)| header_name.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_str()) +} + +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() +} + +fn assert_response_failed_stream( + body: &str, + expected_error_type: &str, + expected_code: &str, + expected_fragments: &[&str], +) { + assert!(body.contains("event: response.failed")); + assert!( + body.contains(&format!("\"type\":\"{expected_error_type}\"")) + || body.contains(&format!("\"type\": \"{expected_error_type}\"")), + "missing error type {expected_error_type} in {body}" + ); + assert!( + body.contains(&format!("\"code\":\"{expected_code}\"")) + || body.contains(&format!("\"code\": \"{expected_code}\"")), + "missing error code {expected_code} in {body}" + ); + for fragment in expected_fragments { + assert!( + body.contains(fragment), + "missing fragment {fragment:?} in {body}" + ); + } +} + +fn responses_http_sse_request(api_key: &str, request_id: &str) -> Request { + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .header("x-request-id", request_id) + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap() +} + +fn responses_json_request(api_key: &str, body: Value) -> Request { + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap() +} + +fn responses_previous_request(api_key: &str, content: &str) -> Request { + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "previous_response_id": "resp_runtime_pool_previous", + "input": [{"role": "user", "content": content}], + "stream": false + }) + .to_string(), + )) + .unwrap() +} + +fn response_completed_websocket_message(response_id: &str) -> String { + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "object": "response", + "status": "completed", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4 + } + } + }) + .to_string() +} + +fn websocket_completed_response( + response_id: &str, + input_tokens: u64, + output_tokens: u64, +) -> String { + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "object": "response", + "status": "completed", + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens + } + } + }) + .to_string() +} + +fn websocket_completed_function_call_response(response_id: &str, call_id: &str) -> String { + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "object": "response", + "status": "completed", + "output": [{ + "type": "function_call", + "id": format!("fc_{call_id}"), + "call_id": call_id, + "name": "lookup", + "arguments": "{}" + }], + "usage": { + "input_tokens": 6, + "output_tokens": 2, + "total_tokens": 8 + } + } + }) + .to_string() +} + +fn responses_previous_stream_request(api_key: &str, content: &str) -> Request { + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "previous_response_id": "resp_runtime_pool_previous", + "input": [{"role": "user", "content": content}], + "stream": true + }) + .to_string(), + )) + .unwrap() +} + +fn response_failed_websocket_message(response_id: &str, code: &str, message: &str) -> String { + json!({ + "type": "response.failed", + "response": { + "id": response_id, + "object": "response", + "status": "failed", + "error": { + "code": code, + "message": message + } + } + }) + .to_string() +} + +async fn first_response_body_chunk(response: axum::response::Response) -> Option { + let mut body_stream = response.into_body().into_data_stream(); + let chunk = body_stream.next().await?.ok()?; + String::from_utf8(chunk.to_vec()).ok() +} + +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() +} + +struct ResponseEventLog { + level: String, + request_id: Option, + account_id: Option, + status_code: Option, + metadata_json: String, +} + +async fn latest_response_event_log(pool: &SqlitePool) -> ResponseEventLog { + let page = SqliteEventLogStore::new(pool.clone()) + .list( + EventLogFilter { + kind: Some("v1.response".to_string()), + ..EventLogFilter::default() + }, + None, + 1, + ) + .await + .unwrap(); + let event = page + .items + .into_iter() + .next() + .expect("expected a v1.response event log"); + ResponseEventLog { + level: event_level_name(event.level).to_string(), + request_id: event.request_id, + account_id: event.account_id, + status_code: event.status_code, + metadata_json: event.metadata.to_string(), + } +} + +async fn response_event_log_count(pool: &SqlitePool) -> i64 { + let (count,): (i64,) = sqlx::query_as("select count(*) from event_logs where kind = ?") + .bind("v1.response") + .fetch_one(pool) + .await + .unwrap(); + count +} + +fn event_level_name(level: EventLevel) -> &'static str { + match level { + EventLevel::Debug => "debug", + EventLevel::Info => "info", + EventLevel::Warn => "warn", + EventLevel::Error => "error", + } +} + +fn assert_rate_limit_header(metadata: &Value, name: &str, value: &str) { + let headers = metadata["rateLimitHeaders"] + .as_array() + .expect("rateLimitHeaders should be an array"); + let expected_name = name.to_ascii_lowercase(); + assert!( + headers.iter().any(|entry| { + let Some(entry) = entry.as_array() else { + return false; + }; + let Some(header_name) = entry.first().and_then(Value::as_str) else { + return false; + }; + let Some(header_value) = entry.get(1).and_then(Value::as_str) else { + return false; + }; + header_name.eq_ignore_ascii_case(&expected_name) && header_value == value + }), + "expected {name}: {value} in rateLimitHeaders, got {headers:?}" + ); +} + +async fn spawn_chunked_sse_upstream( + first_frame: &'static str, + final_frame: &'static str, +) -> (String, oneshot::Receiver<()>, oneshot::Sender<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let (first_chunk_sent_tx, first_chunk_sent_rx) = oneshot::channel(); + let (finish_tx, finish_rx) = oneshot::channel(); + + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + read_http_request(&mut socket).await; + socket + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\nconnection: close\r\n\r\n", + ) + .await + .unwrap(); + write_http_chunk(&mut socket, first_frame.as_bytes()).await; + socket.flush().await.unwrap(); + let _ = first_chunk_sent_tx.send(()); + let _ = finish_rx.await; + write_http_chunk(&mut socket, final_frame.as_bytes()).await; + socket.write_all(b"0\r\n\r\n").await.unwrap(); + socket.flush().await.unwrap(); + }); + + (base_url, first_chunk_sent_rx, finish_tx) +} + +async fn spawn_chunked_sse_upstream_then_abrupt_close( + first_frame: &'static str, +) -> (String, oneshot::Receiver<()>, oneshot::Sender<()>) { + spawn_chunked_sse_upstream_then_close(first_frame, ChunkedSseCloseMode::Abrupt).await +} + +async fn spawn_chunked_sse_upstream_then_clean_close( + first_frame: &'static str, +) -> (String, oneshot::Receiver<()>, oneshot::Sender<()>) { + spawn_chunked_sse_upstream_then_close(first_frame, ChunkedSseCloseMode::Clean).await +} + +async fn spawn_chunked_sse_upstream_then_clean_close_with_headers( + first_frame: &'static str, + headers: &'static [(&'static str, &'static str)], +) -> (String, oneshot::Receiver<()>, oneshot::Sender<()>) { + spawn_chunked_sse_upstream_then_close_with_headers( + first_frame, + ChunkedSseCloseMode::Clean, + headers, + ) + .await +} + +enum ChunkedSseCloseMode { + Abrupt, + Clean, +} + +async fn spawn_chunked_sse_upstream_then_close( + first_frame: &'static str, + close_mode: ChunkedSseCloseMode, +) -> (String, oneshot::Receiver<()>, oneshot::Sender<()>) { + spawn_chunked_sse_upstream_then_close_with_headers(first_frame, close_mode, &[]).await +} + +async fn spawn_chunked_sse_upstream_then_close_with_headers( + first_frame: &'static str, + close_mode: ChunkedSseCloseMode, + headers: &'static [(&'static str, &'static str)], +) -> (String, oneshot::Receiver<()>, oneshot::Sender<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let (first_chunk_sent_tx, first_chunk_sent_rx) = oneshot::channel(); + let (close_tx, close_rx) = oneshot::channel(); + + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + read_http_request(&mut socket).await; + let extra_headers = headers + .iter() + .map(|(name, value)| format!("{name}: {value}\r\n")) + .collect::(); + let response_head = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\nconnection: close\r\n{extra_headers}\r\n" + ); + socket.write_all(response_head.as_bytes()).await.unwrap(); + write_http_chunk(&mut socket, first_frame.as_bytes()).await; + socket.flush().await.unwrap(); + let _ = first_chunk_sent_tx.send(()); + let _ = close_rx.await; + if matches!(close_mode, ChunkedSseCloseMode::Clean) { + socket.write_all(b"0\r\n\r\n").await.unwrap(); + socket.flush().await.unwrap(); + } + }); + + (base_url, first_chunk_sent_rx, close_tx) +} + +async fn collect_stream_body(mut body_stream: S) -> String +where + S: futures::Stream> + Unpin, + E: std::fmt::Debug, +{ + let mut body = Vec::new(); + while let Some(chunk) = body_stream.next().await { + let chunk = chunk.expect("late upstream failures should be converted into SSE frames"); + body.extend_from_slice(&chunk); + } + String::from_utf8(body).unwrap() +} + +async fn read_http_request(socket: &mut TcpStream) { + let mut buffer = Vec::new(); + let mut chunk = [0u8; 1024]; + while !buffer.windows(4).any(|window| window == b"\r\n\r\n") { + let read = socket.read(&mut chunk).await.unwrap(); + if read == 0 { + return; + } + buffer.extend_from_slice(&chunk[..read]); + } + + let Some(header_end) = buffer + .windows(4) + .position(|window| window == b"\r\n\r\n") + .map(|position| position + 4) + else { + return; + }; + let headers = String::from_utf8_lossy(&buffer[..header_end]); + let content_length = headers.lines().find_map(|line| { + let (name, value) = line.split_once(':')?; + if name.eq_ignore_ascii_case("content-length") { + value.trim().parse::().ok() + } else { + None + } + }); + let Some(content_length) = content_length else { + return; + }; + let already_read_body = buffer.len().saturating_sub(header_end); + let remaining = content_length.saturating_sub(already_read_body); + if remaining > 0 { + let mut discard = vec![0u8; remaining]; + socket.read_exact(&mut discard).await.unwrap(); + } +} + +async fn write_http_sse_response(socket: &mut TcpStream, body: &str) { + socket + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .await + .unwrap(); +} + +async fn write_chunked_http_sse_headers(socket: &mut TcpStream) { + socket + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n", + ) + .await + .unwrap(); +} + +async fn write_http_chunk(socket: &mut TcpStream, bytes: &[u8]) { + socket + .write_all(format!("{:X}\r\n", bytes.len()).as_bytes()) + .await + .unwrap(); + socket.write_all(bytes).await.unwrap(); + socket.write_all(b"\r\n").await.unwrap(); +} + +async fn wait_for_http_sse_upstream_disconnect(socket: &mut TcpStream) { + let mut buffer = [0u8; 1024]; + loop { + match socket.read(&mut buffer).await { + Ok(0) | Err(_) => return, + Ok(_) => {} + } + } +} + +async fn received_authorizations(server: &MockServer) -> Vec { + server + .received_requests() + .await + .unwrap() + .into_iter() + .filter_map(|request| { + request + .headers + .get("authorization") + .and_then(|value| value.to_str().ok()) + .map(str::to_string) + }) + .collect() +} + +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(), + oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), + oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), + oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), + }, + 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, + }, + ws_pool: WebSocketPoolConfig::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + }, + } +} diff --git a/crates/server/tests/openai_diagnostics_routes.rs b/crates/server/tests/openai_diagnostics_routes.rs new file mode 100644 index 000000000..726989b7f --- /dev/null +++ b/crates/server/tests/openai_diagnostics_routes.rs @@ -0,0 +1,323 @@ +use std::collections::BTreeMap; + +use axum::{ + body::{to_bytes, Body}, + http::{Request, StatusCode}, +}; +use chrono::{Duration, Utc}; +use codex_proxy_core::{accounts::model::AccountStatus, gateway::fingerprint::Fingerprint}; +use codex_proxy_platform::{ + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, + UsageStatsConfig, WebSocketPoolConfig, + }, + crypto::SecretBox, + identity::ApiKeyHasher, + storage::connect_sqlite, +}; +use codex_proxy_runtime::state::AppState; +use codex_proxy_server::router; +use secrecy::SecretString; +use serde_json::Value; +use tower::util::ServiceExt; +use wiremock::{ + matchers::{header, method, path}, + Mock, MockServer, ResponseTemplate, +}; + +use codex_proxy_adapters::sqlite::accounts::{NewAccount, SqliteAccountStore}; + +#[tokio::test] +async fn debug_diagnostics_should_return_runtime_pool_transport_paths_and_fingerprint() { + let (app, _dir) = test_app_with_account( + "https://chatgpt.test/backend-api".to_string(), + diagnostics_fingerprint(), + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/debug/diagnostics") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["status"], "ok"); + assert_eq!(body["runtime"]["packageName"], "codex-proxy-server"); + assert_eq!(body["paths"]["config"], "config.yaml"); + assert_eq!(body["paths"]["localConfig"], "local.yaml"); + assert!(body["paths"]["databaseUrl"] + .as_str() + .unwrap() + .starts_with("sqlite://")); + assert_eq!( + body["transport"]["backendBaseUrl"], + "https://chatgpt.test/backend-api" + ); + assert_eq!( + body["transport"]["fingerprint"]["userAgent"], + "Codex Desktop/27.100.200 (linux; x64)" + ); + assert_eq!(body["accounts"]["pool"]["total"], 1); + assert_eq!(body["accounts"]["pool"]["active"], 1); + assert_eq!(body["accounts"]["capacity"]["maxConcurrentPerAccount"], 3); + assert_eq!(body["accounts"]["capacity"]["totalSlots"], 3); + assert_eq!(body["accounts"]["capacity"]["availableSlots"], 3); + assert!(!serde_json::to_string(&body) + .unwrap() + .contains("access-secret")); +} + +#[tokio::test] +async fn debug_fingerprint_should_return_runtime_fingerprint_summary() { + let (app, _dir) = test_app_with_account( + "https://chatgpt.test/backend-api".to_string(), + diagnostics_fingerprint(), + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/debug/fingerprint") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["source"], "runtime"); + assert_eq!(body["originator"], "Codex Desktop"); + assert_eq!(body["appVersion"], "27.100.200"); + assert_eq!(body["buildNumber"], "9001"); + assert_eq!(body["chromiumVersion"], "147"); + assert_eq!(body["userAgent"], "Codex Desktop/27.100.200 (linux; x64)"); +} + +#[tokio::test] +async fn debug_upstream_should_probe_codex_models_endpoint_without_returning_secrets() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/codex/models")) + .and(header("originator", "Codex Desktop")) + .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({ + "error": { + "message": "missing or invalid token" + } + }))) + .expect(1) + .mount(&server) + .await; + let (app, _dir) = test_app_with_account(server.uri(), diagnostics_fingerprint()).await; + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/debug/upstream") + .header("x-request-id", "req_debug_probe") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::OK); + assert_eq!(body["target"], "codexModels"); + assert_eq!(body["backendBaseUrl"], server.uri()); + assert_eq!(body["reachable"], true); + assert_eq!(body["statusCode"], 401); + assert_eq!(body["authorization"], "rejected"); + let endpoint = body["endpoint"].as_str().unwrap(); + assert!(endpoint.contains("/codex/models")); + assert!(endpoint.contains("client_version=27.100.200")); + assert!(!serde_json::to_string(&body) + .unwrap() + .contains("access-secret")); + + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].headers.get("x-client-request-id").is_none()); +} + +#[tokio::test] +async fn debug_endpoints_should_reject_forwarded_remote_requests_without_probe() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/codex/models")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let (app, _dir) = test_app_with_account(server.uri(), diagnostics_fingerprint()).await; + + let diagnostics = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/debug/diagnostics") + .header("x-forwarded-for", "203.0.113.10") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let fingerprint = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri("/debug/fingerprint") + .header("x-real-ip", "203.0.113.20") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let upstream = app + .oneshot( + Request::builder() + .method("GET") + .uri("/debug/upstream") + .header("x-forwarded-for", "203.0.113.50") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(diagnostics.status(), StatusCode::FORBIDDEN); + assert_eq!(fingerprint.status(), StatusCode::FORBIDDEN); + assert_eq!(upstream.status(), StatusCode::FORBIDDEN); +} + +async fn test_app_with_account( + base_url: String, + fingerprint: Fingerprint, +) -> (axum::Router, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("openai-diagnostics.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let secret_box = SecretBox::new([77u8; 32]); + let store = SqliteAccountStore::new(pool.clone(), secret_box.clone()); + store + .insert(NewAccount { + id: "acct_1".to_string(), + email: Some("acct@example.test".to_string()), + account_id: Some("chatgpt-account".to_string()), + user_id: Some("user_1".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new("access-secret".to_string().into()), + refresh_token: Some(SecretString::new("refresh-secret".to_string().into())), + access_token_expires_at: Some(Utc::now() + Duration::hours(1)), + status: AccountStatus::Active, + }) + .await + .unwrap(); + let state = AppState::with_pool_secret_api_key_hasher_and_fingerprint( + test_config(url, base_url), + pool, + secret_box, + ApiKeyHasher::new([79u8; 32]), + fingerprint, + ); + state.restore_account_pool_from_repository().await.unwrap(); + (router::router().with_state(state), dir) +} + +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 diagnostics_fingerprint() -> Fingerprint { + Fingerprint { + originator: "Codex Desktop".to_string(), + app_version: "27.100.200".to_string(), + build_number: "9001".to_string(), + platform: "linux".to_string(), + arch: "x64".to_string(), + chromium_version: "147".to_string(), + user_agent_template: "Codex Desktop/{version} ({platform}; {arch})".to_string(), + default_headers: Fingerprint::default_headers(), + header_order: Fingerprint::default_header_order(), + } +} + +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(), + oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), + oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), + oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), + }, + 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, + }, + ws_pool: WebSocketPoolConfig::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + }, + } +} diff --git a/crates/server/tests/openai_models_auth.rs b/crates/server/tests/openai_models_auth.rs new file mode 100644 index 000000000..d5c7e7bdf --- /dev/null +++ b/crates/server/tests/openai_models_auth.rs @@ -0,0 +1,160 @@ +use std::collections::BTreeMap; + +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use tower::util::ServiceExt; + +use codex_proxy_platform::{ + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, + UsageStatsConfig, WebSocketPoolConfig, + }, + crypto::SecretBox, + identity::ApiKeyHasher, + storage::connect_sqlite, +}; +use codex_proxy_runtime::state::AppState; +use codex_proxy_server::router; +use sqlx::SqlitePool; + +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(), + oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), + oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), + oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), + }, + 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, + }, + ws_pool: WebSocketPoolConfig::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + }, + } +} + +#[tokio::test] +async fn models_route_should_reject_unknown_client_api_key() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("openai-models-auth.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([91u8; 32]), + ApiKeyHasher::new([92u8; 32]), + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/models") + .header("authorization", "Bearer cpr_not_stored") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); +} + +#[tokio::test] +async fn models_route_should_accept_stored_client_api_key() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("openai-models-auth-valid.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let hasher = ApiKeyHasher::new([93u8; 32]); + let plaintext = insert_client_api_key(&pool, &hasher).await; + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([94u8; 32]), + hasher, + ); + let app = router::router().with_state(state); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/v1/models") + .header("authorization", format!("Bearer {plaintext}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); +} + +async fn insert_client_api_key(pool: &SqlitePool, hasher: &ApiKeyHasher) -> String { + let generated = hasher.generate_client_api_key("test"); + sqlx::query( + "insert into client_api_keys (id, name, prefix, key_hash, enabled, created_at) values (?, ?, ?, ?, 1, ?)", + ) + .bind("key_test") + .bind("test") + .bind(&generated.prefix) + .bind(&generated.key_hash) + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); + generated.plaintext +} diff --git a/crates/server/tests/openai_responses_routes.rs b/crates/server/tests/openai_responses_routes.rs new file mode 100644 index 000000000..c6f2c5bfc --- /dev/null +++ b/crates/server/tests/openai_responses_routes.rs @@ -0,0 +1,244 @@ +use std::collections::BTreeMap; + +use axum::{ + body::{to_bytes, Body}, + http::{Request, StatusCode}, +}; +use codex_proxy_platform::{ + config::{ + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, + QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, + UsageStatsConfig, WebSocketPoolConfig, + }, + crypto::SecretBox, + identity::ApiKeyHasher, + storage::connect_sqlite, +}; +use codex_proxy_runtime::state::AppState; +use codex_proxy_server::router; +use serde_json::Value; +use sqlx::SqlitePool; +use tower::util::ServiceExt; + +#[tokio::test] +async fn responses_route_should_reject_missing_client_api_key() { + let (app, _key, _dir) = test_app_with_client_api_key().await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .body(Body::from(r#"{"model":"gpt-5.5","input":[]}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); +} + +#[tokio::test] +async fn responses_route_should_reject_unknown_models_with_openai_error() { + let (app, api_key, _dir) = test_app_with_client_api_key().await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from(r#"{"model":"unknown-model","input":[]}"#)) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = response_json(response).await; + + assert_eq!(status, StatusCode::NOT_FOUND); + assert_eq!(body["error"]["code"], "model_not_found"); +} + +#[tokio::test] +async fn responses_route_should_terminate_generated_stream_errors_with_done_marker() { + let (app, api_key, _dir) = test_app_with_client_api_key().await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":true}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let body = String::from_utf8(bytes.to_vec()).unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(body.contains("event: response.failed")); + assert!( + body.ends_with("data: [DONE]\n\n"), + "generated stream errors should terminate clients, body was {body:?}" + ); +} + +#[tokio::test] +async fn responses_route_should_default_omitted_stream_to_sse() { + let (app, api_key, _dir) = test_app_with_client_api_key().await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from(r#"{"model":"gpt-5.5","input":[]}"#)) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(); + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let body = String::from_utf8(bytes.to_vec()).unwrap(); + + assert_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("event: response.failed")); + assert!(body.ends_with("data: [DONE]\n\n")); +} + +#[tokio::test] +async fn chat_completions_route_should_reject_empty_messages_with_bad_request() { + let (app, api_key, _dir) = test_app_with_client_api_key().await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from(r#"{"model":"gpt-5.5","messages":[]}"#)) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); +} + +async fn test_app_with_client_api_key() -> (axum::Router, String, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("openai-responses-routes.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let hasher = ApiKeyHasher::new([95u8; 32]); + let plaintext = insert_client_api_key(&pool, &hasher).await; + let state = AppState::with_pool_secret_and_api_key_hasher( + test_config(url), + pool, + SecretBox::new([96u8; 32]), + hasher, + ); + (router::router().with_state(state), plaintext, dir) +} + +async fn insert_client_api_key(pool: &SqlitePool, hasher: &ApiKeyHasher) -> String { + let generated = hasher.generate_client_api_key("test"); + sqlx::query( + "insert into client_api_keys (id, name, prefix, key_hash, enabled, created_at) values (?, ?, ?, ?, 1, ?)", + ) + .bind("key_test") + .bind("test") + .bind(&generated.prefix) + .bind(&generated.key_hash) + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); + generated.plaintext +} + +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_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(), + oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), + oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), + oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), + }, + 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, + }, + ws_pool: WebSocketPoolConfig::default(), + admin: AdminConfig { + session_ttl_minutes: 1440, + session_cleanup_interval_secs: 3600, + default_username: "admin".to_string(), + default_password: "admin".to_string(), + }, + logging: LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled: false, + capacity: 2_000, + capture_body: false, + }, + } +} diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml new file mode 100644 index 000000000..bb040de8b --- /dev/null +++ b/crates/xtask/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "xtask" +version.workspace = true +edition.workspace = true +license.workspace = true +publish.workspace = true +rust-version.workspace = true + +[dependencies] + +[lints] +workspace = true diff --git a/crates/xtask/src/build_web.rs b/crates/xtask/src/build_web.rs new file mode 100644 index 000000000..1b596d7f2 --- /dev/null +++ b/crates/xtask/src/build_web.rs @@ -0,0 +1,10 @@ +//! 前端构建命令。 + +use crate::{repository_root, run_command, XtaskResult}; + +/// 构建 web 前端资源。 +pub fn run() -> XtaskResult { + let web_dir = repository_root()?.join("web"); + run_command("pnpm", &["install", "--frozen-lockfile"], &web_dir)?; + run_command("pnpm", &["build"], &web_dir) +} diff --git a/crates/xtask/src/check_architecture.rs b/crates/xtask/src/check_architecture.rs new file mode 100644 index 000000000..aaed5d244 --- /dev/null +++ b/crates/xtask/src/check_architecture.rs @@ -0,0 +1,9 @@ +//! 架构检查命令。 + +use crate::{repository_root, run_command, XtaskResult}; + +/// 运行架构测试。 +pub fn run() -> XtaskResult { + let root = repository_root()?; + run_command("cargo", &["test", "--test", "architecture"], &root) +} diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs new file mode 100644 index 000000000..96f8fdf16 --- /dev/null +++ b/crates/xtask/src/main.rs @@ -0,0 +1,60 @@ +mod build_web; +mod check_architecture; +mod release; + +use std::{ + error::Error, + io, + path::{Path, PathBuf}, + process::{Command, ExitCode}, +}; + +type XtaskResult = Result<(), Box>; + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + eprintln!("xtask failed: {error}"); + ExitCode::FAILURE + } + } +} + +fn run() -> XtaskResult { + match std::env::args().nth(1).as_deref() { + Some("build-web") => build_web::run(), + Some("check-architecture") => check_architecture::run(), + Some("release") => release::run(), + _ => { + print_usage(); + Ok(()) + } + } +} + +fn print_usage() { + eprintln!("usage: cargo xtask "); +} + +pub(crate) fn repository_root() -> Result { + Path::new(env!("CARGO_MANIFEST_DIR")) + .parent() + .and_then(Path::parent) + .map(Path::to_path_buf) + .ok_or_else(|| io::Error::other("failed to derive repository root")) +} + +pub(crate) fn run_command(program: &str, args: &[&str], current_dir: &Path) -> XtaskResult { + eprintln!("running: {} {}", program, args.join(" ")); + let status = Command::new(program) + .args(args) + .current_dir(current_dir) + .status()?; + + if status.success() { + Ok(()) + } else { + Err(io::Error::other(format!("command `{}` exited with status {status}", program)).into()) + } +} diff --git a/crates/xtask/src/release.rs b/crates/xtask/src/release.rs new file mode 100644 index 000000000..ae55a27b9 --- /dev/null +++ b/crates/xtask/src/release.rs @@ -0,0 +1,25 @@ +//! 发布打包命令。 + +use crate::{build_web, repository_root, run_command, XtaskResult}; + +/// 执行发布打包流程。 +pub fn run() -> XtaskResult { + let root = repository_root()?; + run_command("cargo", &["fmt", "--all", "--", "--check"], &root)?; + run_command("cargo", &["test", "--workspace", "--all-targets"], &root)?; + run_command( + "cargo", + &[ + "clippy", + "--workspace", + "--all-targets", + "--all-features", + "--locked", + "--", + "-D", + "warnings", + ], + &root, + )?; + build_web::run() +} diff --git a/docs/OPENAI_CODEX_PARITY_AUDIT.md b/docs/OPENAI_CODEX_PARITY_AUDIT.md index 916c2871b..280538ad1 100644 --- a/docs/OPENAI_CODEX_PARITY_AUDIT.md +++ b/docs/OPENAI_CODEX_PARITY_AUDIT.md @@ -1,14 +1,14 @@ # OpenAI Codex 原版链路对齐审计 -日期：2026-06-14 +日期：2026-06-19 范围：对比原版 Node.js 仓库 `/home/zyy/桌面/Codes/codex-proxy`，审计 Rust 版的安全链路、请求链路、响应链路。这里的“代理”不包含 IP 代理、VPN、`HttpsProxyAgent`、本地代理探测等网络代理能力，这部分明确不移植。 ## 结论 -当前 Rust 版已补齐关键安全模拟字段、默认 WebSocket 上游请求、稳定 `prompt_cache_key`、`previous_response_id` 账号亲和性、WS 首帧错误分类、WS 实时边收边转发、账号 + conversation 维度的 WS 连接复用。IP 代理/VPN 能力明确不移植。 +当前 Rust 版已补齐关键安全模拟字段、默认 WebSocket 上游请求、稳定 `prompt_cache_key`、`previous_response_id` 账号亲和性、WS 首帧错误分类、WS 实时边收边转发、账号 + conversation 维度的 WS 连接复用、implicit resume、strip-and-retry、reasoning replay。IP 代理/VPN 能力明确不移植。 -本轮已补齐请求发送细节和账号生命周期差异：自动 Cookie 捕获白名单、`Max-Age` 优先级、账号刷新/管理状态变化后的 WS pool 驱逐、`request_interval_ms` 发送前 stagger、`least_used` 的 reset 缺失排序。仍未达到 100% 对齐的项目是 Responses implicit resume/reasoning replay 状态机。 +本轮已补齐请求发送细节和账号生命周期差异：自动 Cookie 捕获白名单、`Max-Age` 优先级、账号刷新/管理状态变化后的 WS pool 驱逐、`request_interval_ms` 发送前 stagger、`least_used` 的 reset 缺失排序。当前没有保留的 OpenAI/Codex 请求链路迁移阻塞项；IP 代理/VPN 仍是明确非目标。 ## 1. 安全链路和指纹 @@ -103,7 +103,9 @@ Rust 状态： Rust 状态： - 到位：显式 `previous_response_id` 会强制 WebSocket 并优先回到原账号。 -- 未对齐：implicit resume、strip-and-retry、reasoning replay 尚未完整迁移，应单独设计和测试。 +- 到位：implicit resume 会基于已记录 affinity/replay 条目续接同账号、同 conversation、同变体历史。 +- 到位：strip-and-retry 覆盖 `previous_response_not_found`、unanswered function call、invalid encrypted reasoning replay，并会在需要时恢复完整历史。 +- 到位：reasoning replay 缓存、记录、驱逐和 SQLite affinity restore 后续接均由 crate-local tests 覆盖。 ## 9. 非目标：IP 代理/VPN @@ -111,6 +113,7 @@ Rust 状态： ## 当前验证点 -- `tests/codex_gateway/websocket.rs` 覆盖默认 WS、forced HTTP SSE、WS 安全字段、握手错误、首帧错误分类、HTTP SSE 降级。 -- `tests/codex_serving/responses_websocket.rs` 覆盖默认 WS 上游、派生 `prompt_cache_key`、`client_metadata` 安全字段、`previous_response_id` 原账号亲和性、WS 首帧实时返回、同 conversation 连接复用、限流/刷新路径。 -- `tests/codex_serving/responses_http_sse.rs` 覆盖显式 HTTP SSE 兼容路径和 body/header 安全字段。 +- `crates/adapters/tests/codex.rs` 覆盖默认 WS、opening/header/payload 捕获、WS 安全字段、握手错误、首帧错误分类、deflate、ping、pool、idle timeout、rate-limit/metadata 捕获。 +- `crates/server/tests/openai_chat_upstream.rs` 覆盖默认 WS 上游、派生 `prompt_cache_key`、`client_metadata` 安全字段、`previous_response_id` 原账号亲和性、WS 首帧实时返回、同 conversation 连接复用、限流/刷新路径、implicit resume、strip-and-retry、reasoning replay。 +- `crates/core/tests/protocol.rs` 覆盖 WebSocket 纯协议转换、错误分类、Responses payload 顺序、事件形状校验和 replay 相关元数据抽取。 +- `crates/server/tests/openai_diagnostics_routes.rs` 覆盖 `/debug/diagnostics`、`/debug/fingerprint`、`/debug/upstream` 本地诊断和 secret redaction。 diff --git a/docs/OPENAI_REQUEST_ALIGNMENT.md b/docs/OPENAI_REQUEST_ALIGNMENT.md index de2f20263..fe862ecb6 100644 --- a/docs/OPENAI_REQUEST_ALIGNMENT.md +++ b/docs/OPENAI_REQUEST_ALIGNMENT.md @@ -3,15 +3,15 @@ 本文档记录了 `codex-proxy-rs` (Rust) 与 `codex-proxy` (Node.js) 在 OpenAI 请求链路上的对比审查和修复工作。 ## 审查日期 -2026-06-13 +2026-06-19 ## 当前结论 -Rust 版已经补齐 OpenAI/Codex 主请求链路的核心路径：默认 WebSocket 上游、`previous_response_id` 账号亲和性、稳定 `prompt_cache_key`、关键安全 header、`client_metadata`、WS frame 到 SSE 的响应转换、usage/rate-limit 记录和 HTTP SSE fallback。 +Rust 版已经补齐 OpenAI/Codex 主请求链路的核心路径：默认 WebSocket 上游、`previous_response_id` 账号亲和性、稳定 `prompt_cache_key`、关键安全 header、`client_metadata`、WS frame 到 SSE 的响应转换、usage/rate-limit 记录、HTTP SSE fallback、implicit resume、strip-and-retry、reasoning replay。 本轮改造后，直接影响 OpenAI 交互语义和请求发送行为的差异已补齐：自动 Cookie 捕获白名单、`Max-Age` 优先级、账号生命周期后的 WS pool 驱逐、`request_interval_ms` 发送前 stagger、`least_used` reset 缺失排序。 -仍不能宣称与 Node.js 原版 100% 一致，因为 Responses implicit resume、strip-and-retry、reasoning replay 这类更细的状态机尚未完整迁移。 +IP 代理/VPN、本地代理探测和 `HttpsProxyAgent` 不属于 Rust 版迁移目标。除此之外，当前审计没有保留的 OpenAI/Codex 请求链路迁移阻塞项。 ## 本轮已对齐项 @@ -51,15 +51,16 @@ Rust 版已经补齐 OpenAI/Codex 主请求链路的核心路径：默认 WebSoc - `(Some(a), Some(b))` 且不相等时比较 reset - 任一侧缺失 reset 时继续比较 `request_count` -## 仍未完全对齐项 +## 历史未对齐项（已关闭） ### Responses 隐式续接和 reasoning replay 状态机 -原版 shared handler 还包含 implicit resume、strip-and-retry、reasoning replay 等更细的 Responses 状态机。Rust 当前主请求链路已对齐默认 WS 和显式 `previous_response_id`，但这些边缘续接状态机尚未完整迁移。 +原版 shared handler 包含 implicit resume、strip-and-retry、reasoning replay 等更细的 Responses 状态机。这个缺口已经在当前 Rust workspace 中关闭： -对齐要求： -- 单独设计并测试 implicit resume/reasoning replay -- 不在完成前把 OpenAI/Codex 链路标记为 100% 一致 +- `crates/core/src/serving/implicit_resume.rs` 保存纯策略； +- `crates/core/src/serving/reasoning_replay.rs` 保存 replay 缓存策略； +- `crates/runtime/src/services.rs` 在 Responses 非流式、HTTP SSE 流式、WebSocket 路径中接入恢复、驱逐和历史还原； +- `crates/server/tests/openai_chat_upstream.rs` 覆盖 WebSocket implicit resume、SQLite 恢复后续接、跨 window 拒绝、自包含 function-call replay 拒绝、未匹配 function-call 输出拒绝、invalid encrypted reasoning replay 后驱逐、previous_response_not_found / unanswered_function_call strip-and-retry、SSE/WebSocket invalid reasoning replay strip-and-retry。 ## 已修复的问题 @@ -236,9 +237,7 @@ test codex::gateway::installation::tests::test_generate_and_persist ... ok 4. ✅ 实现 installation ID 持久化 5. ✅ 所有测试通过 6. ✅ 编译无错误 - -仍需完成： -1. implicit resume/reasoning replay 状态机 +7. ✅ implicit resume / strip-and-retry / reasoning replay 状态机已迁移并由 crate-local 测试覆盖 ## 下一步建议 @@ -246,11 +245,10 @@ test codex::gateway::installation::tests::test_generate_and_persist ... ok 2. **监控**：观察 `session_id` 和 `installation_id` 在日志中的出现，确保正确传递 3. **性能验证**：确认 prompt cache 命中率是否提升（session affinity 的主要效果） 4. **兼容性测试**：验证与真实 Codex Desktop 的 `~/.codex/installation_id` 共享是否正常 -5. **状态机补齐**：单独迁移 implicit resume/reasoning replay，并补充原版行为级测试 ## 参考 - Node.js 实现：`/home/zyy/桌面/Codes/codex-proxy/src/proxy/codex-api.ts` -- Rust 实现：`/home/zyy/Codes/codex-proxy-rs/src/codex/` +- Rust 实现：`/home/zyy/Codes/codex-proxy-rs/crates/{core,adapters,runtime,server}/` - Installation ID 规范：Node.js `src/proxy/installation-id.ts` - Identity 派生规范：Node.js `buildAccountScopedIdentity()` 方法 diff --git a/docs/api.md b/docs/api.md index baba52cec..8f9bde266 100644 --- a/docs/api.md +++ b/docs/api.md @@ -96,10 +96,7 @@ Success response: "paths": { "config": "config.yaml", "localConfig": "local.yaml", - "databaseUrl": "sqlite://data/codex-proxy.sqlite", - "logsDirectory": "logs", - "masterKeyFile": "data/master.key", - "apiKeyPepperFile": "data/api-key-pepper.key" + "databaseUrl": "sqlite://.runtime/data/codex-proxy-rs.sqlite" }, "transport": { "backendBaseUrl": "https://chatgpt.com/backend-api", diff --git a/docs/architecture-audit.md b/docs/architecture-audit.md index 8eaa8eccd..04eafca86 100644 --- a/docs/architecture-audit.md +++ b/docs/architecture-audit.md @@ -1,378 +1,6796 @@ -# 🔍 codex-proxy-rs 架构审计报告 +# Architecture Migration Audit -## 项目概览 +Date: 2026-06-19 -**项目规模：** -- 📁 111 个 Rust 源文件 -- 📝 ~17,609 行代码 -- 🧪 142 个测试（单元测试 + 集成测试） -- ⚡ 259 个 async 函数 +This document records the current audit state for the workspace architecture +defined in `docs/architecture.md`. -**架构分层：** +## Current Verdict + +The architecture migration is no longer blocked by empty source files, empty +source directories, empty Rust function/module shells, unwired `xtask`, root +behavior tests left under `tests/`, or an unmapped baseline behavior-test suite. + +Current source-level audit result: + +- `codegraph sync . && codegraph status .` was rerun on 2026-06-19 and reports + the index up to date: 285 files, 4,332 nodes, 13,241 edges. +- `tests/` now contains root architecture tests only; behavior tests are under + `crates/*/tests`. +- Baseline root behavior-test migration is now guarded for all 46 old files + that contain `#[test]` / `#[tokio::test]`; the 8 old root `mod` aggregators + are not counted as behavior suites. +- Current test inventory after the latest migration pass: 495 test attributes + across 65 Rust test files, including 43 crate-local integration test files. +- `find tests crates src -type f \( -name '*.rs' -o -name '*.toml' -o -name + '*.md' \) -empty` returned no source/document/manifest empty files. +- `find tests crates src -type d -empty ...` returned no empty source/test + directories. +- Placeholder/empty-shell scan in `crates src tests` only matched the + architecture guard's own sample strings in + `tests/architecture/placeholder_implementations.rs`. +- `crates/xtask` is wired and intentionally thin: it dispatches + `build-web`, `check-architecture`, and `release` to real commands. + +The current completion state is therefore: source migration placement and +obvious scaffold cleanup are complete by current evidence. The broad +verification gate below has passed, and older OpenAI/Codex parity documents have +been reconciled to the current crate-local coverage. + +## Audit Inputs + +- Baseline before the migration squash: + `f07442a28b0c186bd22b598a53cc4856ab4b2445` +- Current audited head: + `2aa27475723474469d442efab690f70b8f1f0fc9` +- CodeGraph: + - `codegraph sync . && codegraph status .` run during this audit. + - Result: synced after the latest diagnostics/trace/log-rotation/test-mapping + migration pass. + - Status after latest sync: 285 files, 4,332 nodes, 13,241 edges. +- Empty local directories: + - Follow-up `find tests crates src ... -type d -empty` found no remaining + empty source/test directories. + +## 2026-06-19 Current Full Migration Audit Pass + +Status: complete; source/test placement gaps found in this pass have been +implemented, and the full verification gate has passed. + +New gaps found and migrated in this pass: + +- Old `tests/codex_serving/diagnostics_route.rs` behavior was only partially + represented in crate-local tests. Added + `crates/server/tests/openai_diagnostics_routes.rs` and implemented: + `/debug/diagnostics` path/fingerprint/capacity fields, `/debug/fingerprint`, + `/debug/upstream`, local-only guards, and secret redaction. +- Old `tests/runtime/http_trace.rs` behavior had no clear crate-local target. + Added `crates/server/tests/http_trace_middleware.rs`, implemented + `server::middleware::trace::http_trace_layer()`, and mounted it in the + server router. +- Old `tests/platform/log_rotation.rs` rolling appender behavior had no + crate-local test. Added `crates/platform/tests/log_rotation.rs`. +- `tests/architecture/test_migration.rs` now maps all 46 old behavior files to + current crate-local test files, instead of only a high-risk subset. + +Targeted verification run in this pass: + +- `cargo test -p codex-proxy-server --test openai_diagnostics_routes -- --nocapture` + -> 4 passed. +- `cargo test -p codex-proxy-server --test http_trace_middleware -- --nocapture` + -> 1 passed. +- `cargo test -p codex-proxy-platform --test log_rotation -- --nocapture` + -> 2 passed. +- `cargo test --test architecture test_migration -- --nocapture` + -> 3 passed. +- `cargo fmt --all -- --check` -> passed. +- `cargo check --workspace --tests` -> passed. +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` + -> passed after removing a `runtime` boundary violation in the diagnostics + probe and factoring the HTTP trace layer type. +- `cargo test --workspace --tests` -> passed. +- `cargo run -p xtask -- check-architecture` -> passed; 27 architecture + tests passed. +- `cargo doc --workspace --no-deps` -> passed. +- `cargo run -p xtask -- release` -> passed: + `cargo fmt --all -- --check`, `cargo test --workspace --all-targets`, + `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings`, + `pnpm install --frozen-lockfile`, and `pnpm build`. +- Final `codegraph sync . && codegraph status .` -> index up to date with 285 + files, 4,332 nodes, and 13,241 edges. +- Final empty-file and empty-directory scans across `tests crates src` returned + no entries. +- Final placeholder/empty-shell scan across `tests crates src` only matched the + architecture guard's own sample strings. + +Current remaining work before claiming final migration completion: none found +by this pass. The migration is complete by current source, test, architecture, +documentation, and release-gate evidence. + +Note: the long slice log below is retained as historical evidence. Earlier +"Blocking Findings" entries may describe states that were later closed by +subsequent slices and by this current audit pass. + +## Implementation Progress + +### 2026-06-18 Slice 1: Core Installation ID IO Boundary + +Status: implemented and verified in this slice. + +Changes: + +- Added an architecture guard so `crates/core/src` fails tests if it uses + filesystem or environment/path IO patterns: `std::fs`, `std::env`, `dirs::`, + `std::path`, `PathBuf`, or `Path::`. +- Reworked `crates/core/src/gateway/installation.rs` into pure UUID rules: + `generate_installation_id()` and `parse_installation_id()`. +- Moved Codex Desktop/data-dir installation ID file resolution and persistence + to `platform::storage`. +- Removed `dirs` and `tempfile` from `codex-proxy-core` dependencies. +- Added crate-local tests for the old lookup order: + Codex Desktop file, data-dir file, then generate-and-persist. + +Verification run so far: + +- `cargo test --test architecture source_imports_should_respect_layer_boundaries` +- `cargo test -p codex-proxy-platform storage::paths::tests::resolve_installation_id` +- `cargo test -p codex-proxy-core gateway::installation::tests` +- `cargo fmt --all -- --check` +- `cargo test --test architecture` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `cargo test --workspace --all-targets` +- `codegraph sync .` + +Note: + +- The first full workspace test attempt exhausted disk space while writing + `target/`. `cargo clean` removed 72.1 GiB of build artifacts, then + `cargo test --workspace --all-targets` was rerun successfully. + +Remaining after this slice: + +- Runtime startup/request wiring still needs to call the new + `platform::storage` installation ID API where the real startup path is + restored. **Closed later in Slice 7.** +- This only closes the direct `core` IO boundary violation; it does not close + the server entrypoint, xtask, runtime task, WebSocket, serving dispatch, or + deleted-test migration findings. + +### 2026-06-18 Slice 2: Placeholder Guard, Xtask Wiring, And Server Main + +Status: implemented and targeted verification passed in this slice. + +Changes: + +- Added `tests/architecture/placeholder_implementations.rs` to fail when Rust + sources contain machine-detectable placeholder markers: + - `fn main() {}` + - `not wired yet` + - `后续承载` +- Replaced `crates/server/src/main.rs` empty entrypoint with a real startup + path that: + - loads `config.yaml`/local config through `platform::config`; + - initializes logging when enabled; + - connects SQLite through `platform::storage`; + - loads/creates the master secret box and API key pepper; + - constructs `runtime::state::AppState`; + - serves the Axum router with `ctrl_c` graceful shutdown. +- Wired `cargo run -p xtask -- check-architecture` to actually execute + `cargo test --test architecture`. +- Wired `cargo run -p xtask -- build-web` to run + `pnpm install --frozen-lockfile` and `pnpm build` in `web/`. +- Wired `cargo run -p xtask -- release` to run Rust format/test/clippy gates + and then the web build command. +- Replaced root crate docs that described future placeholder behavior with + current responsibility descriptions. + +Verification run so far: + +- `cargo test --test architecture rust_sources_should_not_keep_placeholder_markers` +- `cargo check -p codex-proxy-server --all-targets` +- `cargo check -p xtask --all-targets` +- `cargo test --test architecture` +- `cargo run -p xtask -- check-architecture` +- `cargo run -p xtask -- build-web` +- `cargo run -p xtask -- release` + +Note: + +- The first `xtask build-web` run failed on npm registry `ECONNRESET`/fetch + errors while `pnpm install --frozen-lockfile` was retrying. The command was + restarted and then completed successfully, including `pnpm build`. + +Remaining after this slice: + +- Server startup is no longer an empty function, but it still does not prove the + full baseline startup parity: account restoration, session-affinity + restoration, remaining background task parity details, and installation ID + request wiring still need explicit migration and tests. Installation ID + request wiring was closed later in Slice 7. +- `xtask release` is now wired and has been run successfully once. It should be + kept in the final verification gate after the remaining migration slices + stabilize. + +### 2026-06-18 Slice 3: Runtime Cleanup Task Wiring + +Status: cookie/session cleanup task slices and migrated-task coordinator wiring +implemented; targeted verification passed. + +Changes: + +- Added `SqliteCookieStore::cleanup_expired(now)` to persistently delete + expired account cookies. +- Added `SqliteAdminSessionStore::cleanup_expired_sessions(now)` to persistently + delete expired admin sessions. +- Implemented `runtime::tasks::cookie_cleanup::CookieCleanupTask` with: + - default 5-minute interval; + - `start()` returning the shared `SchedulerHandle`; + - `cleanup_once()` and testable `cleanup_once_at(now)`. +- Implemented `runtime::tasks::session_cleanup::SessionCleanupTask` with: + - configured interval; + - `start()` returning the shared `SchedulerHandle`; + - `cleanup_once()` and testable `cleanup_once_at(now)`. +- Added crate-local runtime tests proving cleanup deletes only expired rows and + leaves unexpired rows intact. +- Added `BackgroundTaskCoordinator::task_names()` so runtime task wiring can be + regression-tested without depending on private handle internals. +- Added runtime task stores to `Services` so background tasks are assembled from + the same SQLite adapters that power request/admin services. +- Wired `start_background_tasks(&AppState)` to start the migrated + `cookie_cleanup`, `session_cleanup`, and `model_refresh` tasks. +- Wired the server entrypoint to own the background task coordinator and shut it + down after Axum graceful shutdown returns. +- Added an architecture guard requiring the server entrypoint to start runtime + background tasks and shut the coordinator down. + +Verification run so far: + +- `cargo test -p codex-proxy-runtime --test tasks cleanup_task_should_delete_only_expired` +- `cargo test -p codex-proxy-runtime --test tasks start_background_tasks_should_register_migrated_runtime_tasks` +- `cargo test --test architecture server_entrypoint_should_own_background_task_lifecycle` + +Remaining after this slice: + +- `token_refresh.rs` and `quota_refresh.rs` are still one-line modules. +- `start_background_tasks(&AppState)` now wires the migrated cleanup/model + tasks, but it still does not wire token refresh, quota refresh, or fingerprint + update. +- Startup now owns the coordinator lifetime for the wired tasks, but full + baseline parity still depends on migrating the remaining schedulers. + +### 2026-06-18 Slice 4: Token And Quota Refresh Task Migration + +Status: token/quota task minimum runtime loops implemented and targeted +verification passed. + +Changes: + +- Added `core::accounts::jwt::jwt_expiration(token)` so runtime token refresh + can persist refreshed access-token expiry without duplicating JWT parsing. +- Implemented `runtime::tasks::token_refresh::TokenRefreshTask` with: + - core `RefreshScheduler` and `TokenRefresher` port usage; + - SQLite account scanning via `SqliteAccountStore`; + - refreshed access-token persistence through `update_from_claims`; + - status-only persistence for scheduler status transitions; + - periodic `start()` integration with `SchedulerHandle`. +- Added a runtime integration test proving a due account refresh writes the new + access token, preserves the refresh token when upstream omits it, stores the + new JWT expiry, and keeps the account active. +- Added `SqliteAccountStore::update_quota_json()` and `get_quota_json()` based + on the old quota repository behavior. +- Migrated the old usage-to-quota snapshot normalization into + `core::serving::quota::quota_from_usage()`. +- Implemented `runtime::tasks::quota_refresh::QuotaRefreshTask` with: + - scan of active `quota_limit_reached` accounts; + - Codex usage fetch through the existing `CodexBackendClient`; + - normalized quota JSON persistence in SQLite; + - in-process minimum refresh interval tracking for the periodic loop. +- Added runtime quota refresh test using a local mocked Codex usage endpoint. +- Extended `RuntimeConfig` with auth/quota/admin task settings. +- Exposed the assembled Codex backend client from `Services` so runtime tasks + reuse the same adapter instance. +- Wired `start_background_tasks(&AppState)` to include `token_refresh` when + auth refresh is enabled, and `quota_refresh` using configured quota interval. + +Verification run so far: + +- `cargo test -p codex-proxy-core accounts::jwt::tests::jwt_expiration_should_return_exp_as_utc_datetime` +- `cargo test -p codex-proxy-runtime --test token_refresh token_refresh_task_should_persist_refreshed_access_token_and_keep_refresh_token` +- `cargo test -p codex-proxy-adapters --test account_repository account_repository_should_update_quota_json_and_fetched_at` +- `cargo test -p codex-proxy-runtime --test quota_refresh quota_refresh_task_should_fetch_usage_for_quota_locked_accounts_and_store_quota` +- `cargo test -p codex-proxy-runtime --test tasks start_background_tasks_should_register_migrated_runtime_tasks` + +Remaining after this slice: + +- Token refresh is functional but not yet at full old parity: no per-account + timer map, in-flight tracking, crash recovery states, exponential retry loop, + or refresh lease store usage. +- Quota refresh is functional for active quota-locked accounts, but admin + explicit quota route parity and all old quota edge-case tests are still + incomplete. +- Server startup still needs account/session-affinity restoration. + +### 2026-06-18 Slice 5: Fingerprint Update Task Wiring + +Status: implemented and targeted coordinator verification passed. + +Changes: + +- Added `FingerprintRepository` to runtime repository/task-store assembly. +- Stored the runtime fingerprint in `Services` so background tasks use the same + version/build values as the request clients. +- Added default runtime constants for the Codex Desktop appcast URL and optional + `data/extracted-fingerprint.json` path. +- Wired `start_background_tasks(&AppState)` to start `fingerprint_update` with + the assembled repository and current runtime fingerprint version/build. +- Extended the runtime coordinator test to require `fingerprint_update` in the + task list. + +Verification run so far: + +- `cargo test -p codex-proxy-runtime --test tasks start_background_tasks_should_register_migrated_runtime_tasks` + +Remaining after this slice: + +- Fingerprint update is wired, but the broader old fingerprint update tests + still need to be audited against current crate-local coverage. +- Server startup still needs account/session-affinity restoration. + +### 2026-06-18 Slice 6: WebSocket Adapter Shell Cleanup + +Status: placeholder unit structs removed; targeted adapter and architecture +verification passed. + +Changes: + +- Replaced `CodexWebSocketConnection` unit struct with a connection description + that stores the WebSocket endpoint and ordered opening headers. +- Added `CodexWebSocketConnection::opening_audit_snapshot()` so opening header + order can be audited through the core `OpeningAuditSnapshot` type. +- Replaced `CodexWebSocketPool` unit struct with a concrete pool policy holding + `max_per_account` and `max_age`. +- Added pool policy helpers for connection capacity and age-based recycling. +- Extended the placeholder architecture guard to reject the WebSocket adapter + unit structs if they regress. +- Added adapter tests for connection header-order preservation and pool policy. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test codex websocket_` +- `cargo test --test architecture rust_sources_should_not_keep_placeholder_markers` +- `cargo test --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` + +Remaining after this slice: + +- This removes the obvious WebSocket adapter shells, but it is not full + WebSocket parity. Real `tokio-tungstenite` connect/send/receive, pooling of + live sockets, deflate handling, and the large old WebSocket behavior tests are + still incomplete. + +### 2026-06-18 Slice 7: Installation ID Request Context Wiring + +Status: implemented and targeted verification passed. + +Changes: + +- Added `AppState::with_pool_secret_api_key_hasher_and_installation_id()` and a + lower-level fingerprint-aware constructor so tests and startup can inject the + runtime installation ID without touching `core` filesystem rules. +- Updated the server entrypoint to resolve the data directory through + `platform::storage`, load/create the installation ID there, and inject it into + runtime state. +- Stored the optional installation ID in `Services` and threaded it through: + - OpenAI Chat Completions dispatch; + - OpenAI Responses dispatch; + - admin account health probes; + - admin model refresh requests; + - background model refresh; + - background quota refresh. +- Extended the pure `core::gateway::ports::CodexModelCatalogRequest` with an + optional `installation_id` field so model catalog fetches can carry runtime + request context without introducing IO into `core`. +- Updated the Codex adapter so auxiliary JSON requests (`usage`, model catalog, + connectivity probes) include `x-codex-installation-id` when context provides + it, matching the existing Responses/compact request behavior. +- Added a server route regression test proving `/v1/chat/completions` forwards + the runtime installation ID to the mocked Codex upstream. +- Added an architecture guard requiring the server entrypoint to load the + platform installation ID and inject it into runtime state. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test codex codex_backend_client_usage_should_use_original_auxiliary_headers` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_forward_runtime_installation_id_to_codex` +- `cargo test -p codex-proxy-adapters --test codex` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo test -p codex-proxy-core --test models` +- `cargo test --test architecture` + +Remaining after this slice: + +- Server startup still needs account/session-affinity restoration. +- Full serving dispatch/fallback/recovery parity remains incomplete. +- Full live WebSocket transport/pooling/parity remains incomplete, although + pure transport selection plus opening/payload audit snapshots are restored. + +### 2026-06-18 Slice 8: Session Affinity Store And Startup Restore + +Status: core/store/runtime restore path implemented and targeted verification +passed. + +Changes: + +- Replaced the thin `core::serving::affinity` placeholder with pure + `SessionAffinityEntry`, `StoredSessionAffinity`, and `SessionAffinityMap` + types. +- Implemented active-record restore, previous-response account lookup, + conversation lookup, turn-state lookup, variant-aware latest-response lookup, + and forget/size helpers in the pure core map. +- Implemented `SqliteSessionAffinityStore` over the existing + `session_affinities` schema: + - upsert response affinity with TTL-derived expiry; + - list active records; + - delete expired records; + - parse persisted timestamps and function-call ID JSON. +- Added the session affinity store to runtime repository assembly. +- Added `RuntimeSessionAffinityService` with an in-memory core map guarded by + `tokio::sync::RwLock` and restore/lookup/record APIs. +- Added `AppState::restore_session_affinity_from_repository()` and + `restore_session_affinity_from_repository_now()`. +- Wired the server entrypoint to restore session affinity mappings from SQLite + before background task startup and serving. +- Added an architecture guard requiring server startup to restore session + affinity. + +Verification run so far: + +- `cargo test -p codex-proxy-core --test session_affinity` +- `cargo test -p codex-proxy-adapters --test session_affinity` +- `cargo test -p codex-proxy-runtime --test session_affinity` +- `cargo check -p codex-proxy-server --all-targets` +- `cargo test --test architecture` + +Remaining after this slice: + +- Startup still needs active account pool restoration into a runtime pool. + **Closed later in Slice 9.** +- The runtime OpenAI dispatch path still does not use the session affinity map + for preferred-account selection or record new response affinities; that + belongs with the broader serving dispatch/fallback/recovery migration. +- Expired session affinity cleanup is available in the store but is not yet + wired into a background task. + +### 2026-06-18 Slice 9: Account Pool Startup Restore + +Status: implemented and targeted verification passed. + +Changes: + +- Added `RuntimeAccountPoolService` that owns a core `AccountPool` behind a + `tokio::sync::Mutex` and restores its contents from the existing + `AccountStore` port. +- Mapped runtime account pool options from application config: + max concurrent requests per account, stale slot TTL, rotation strategy, + quota-limited skipping, and tier priority. +- Exposed account acquisition/release helpers on the runtime account pool + service for later serving dispatch integration. +- Added `Services::account_pool` and + `AppState::restore_account_pool_from_repository()`. +- Wired the server entrypoint to restore the runtime account pool from SQLite + before restoring session affinity, starting background tasks, and serving. +- Added an architecture guard requiring server startup to restore the runtime + account pool. +- Added a runtime integration test proving an active persisted account is loaded + into the runtime pool and can be acquired. + +Verification run so far: + +- `cargo test -p codex-proxy-runtime --test account_pool_restore` +- `cargo check -p codex-proxy-server --all-targets` +- `cargo test --test architecture` + +Remaining after this slice: + +- Runtime OpenAI Chat/Responses dispatch still directly scans the account store + and must be moved to the restored account pool. **Closed later in Slice 10.** +- Session affinity is restored but is not yet used for preferred-account + selection or recording new response affinities. +- Serving fallback/recovery, quota transitions, implicit resume, and stream + audit parity remain incomplete. + +### 2026-06-18 Slice 10: Dispatch Uses Restored Account Pool + +Status: implemented and targeted verification passed. + +Changes: + +- Reworked `ChatDispatchService` and `ResponseDispatchService` to acquire + upstream accounts from `RuntimeAccountPoolService` instead of scanning + `AccountStore` for the first active account. +- Released the account pool slot after the upstream Codex request returns, + preserving the core pool's concurrency accounting for subsequent dispatches. +- Updated server upstream route fixtures to restore the account pool before + exercising OpenAI routes. +- Added a regression test proving `/v1/chat/completions` dispatches from the + restored runtime pool even if the SQLite row is changed after startup restore. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + +Remaining after this slice: + +- Session affinity is restored but is not yet used to set + `preferred_account_id` for previous-response requests. **Closed later in + Slice 11.** +- New response affinities are not yet recorded after successful upstream + responses. +- Fallback/recovery, quota transitions, implicit resume, and stream audit parity + remain incomplete. + +### 2026-06-18 Slice 11: Session Affinity Preferred Account Dispatch + +Status: implemented and targeted verification passed. + +Changes: + +- Extended OpenAI Responses request translation to preserve + `previous_response_id` in the Codex request. +- Updated `ResponseDispatchService` to use restored session affinity state for + previous-response requests: + - look up preferred account ID and pass it to `AccountAcquireRequest`; + - restore prompt cache key from the affinity conversation ID when absent; + - restore turn state from affinity when the incoming request lacks it. +- Added a server regression test with two restored accounts proving + `/v1/responses` chooses the account bound to `previous_response_id` instead of + the default pool winner. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_prefer_session_affinity_account_for_previous_response` + +Remaining after this slice: + +- New response affinities are not yet recorded after successful upstream + responses. **Closed later in Slice 12.** +- Fallback/recovery, quota transitions, implicit resume, and stream audit parity + remain incomplete. + +### 2026-06-18 Slice 12: Session Affinity Recording + +Status: implemented and targeted verification passed. + +Changes: + +- Updated `ResponseDispatchService` to record a session affinity entry after a + successful `response.completed` upstream result. +- Persisted the completed response ID, selected account ID, conversation ID, + upstream turn state, input token count when available, and created timestamp + through `RuntimeSessionAffinityService`. +- Kept affinity recording best-effort: persistence failures are logged and do + not turn an otherwise successful upstream response into a 5xx. +- Added a server regression test proving a successful `/v1/responses` call + writes the completed response ID to SQLite session affinity storage. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_record_session_affinity_for_completed_response` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo fmt --all -- --check` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Fallback/recovery, dirty quota transitions, implicit resume edge cases, and + stream audit parity remain incomplete. +- Expired session affinity cleanup is available in the store but is not yet + wired into a background task. **Closed later in Slice 13.** + +### 2026-06-18 Slice 13: Session Affinity Cleanup Task Wiring + +Status: implemented and targeted verification passed. + +Changes: + +- Added `runtime::tasks::session_affinity_cleanup::SessionAffinityCleanupTask` + with: + - periodic `start()` integration through the shared `SchedulerHandle`; + - `cleanup_once()` and testable `cleanup_once_at(now)`; + - cleanup through the existing `SqliteSessionAffinityStore::delete_expired`. +- Added the SQLite session-affinity store to `BackgroundTaskStores` so + background tasks use runtime-assembled repositories. +- Wired `start_background_tasks(&AppState)` to register + `session_affinity_cleanup` after admin session cleanup. +- Added runtime task coverage proving only expired session-affinity rows are + deleted and the coordinator registers the new task. +- Updated `docs/architecture.md` and architecture source-tree whitelists so the + new task file is part of the precise runtime task layout rather than an + out-of-band source file. + +Verification run so far: + +- `cargo test -p codex-proxy-runtime --test tasks session_affinity_cleanup_task_should_delete_only_expired_affinities` +- `cargo test -p codex-proxy-runtime --test tasks start_background_tasks_should_register_migrated_runtime_tasks` +- `cargo test -p codex-proxy-runtime --test tasks` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Fallback/recovery, dirty quota transitions, implicit resume edge cases, and + stream audit parity remain incomplete. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Admin explicit quota route parity and old quota edge-case tests remain + incomplete. + +### 2026-06-18 Slice 14: Admin Explicit Account Quota Route + +Status: implemented and targeted verification passed. + +Changes: + +- Migrated the old explicit `GET /api/admin/accounts/{account_id}/quota` + behavior into the new `server`/`runtime` split. +- Added `AdminAccountService::account_quota()` to: + - load the selected active account from SQLite; + - call the Codex usage endpoint through the assembled `CodexBackendClient`; + - normalize the usage response with `core::serving::quota::quota_from_usage`; + - persist the normalized quota snapshot via `SqliteAccountStore`; + - return normalized `quota` plus upstream `raw` usage without exposing stored + token secrets. +- Added the admin HTTP handler and router entry for + `/api/admin/accounts/{account_id}/quota`. +- Added a server regression test migrated from the old quota suite proving the + route calls upstream with the selected account token, stores quota JSON, and + does not return access/refresh tokens. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_account_quota_should_fetch_usage_store_quota_and_not_return_secrets` +- `cargo test -p codex-proxy-server --test admin_accounts_routes` +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `codegraph sync .` + +Remaining after this slice: + +- Old quota edge-case tests still need migration: inactive account errors, + upstream failure mapping, quota persistence failure behavior, and quota + warning edge cases beyond the current cached-warning test. +- Fallback/recovery, dirty quota transitions, implicit resume edge cases, and + stream audit parity remain incomplete. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. + +### 2026-06-18 Slice 15: Admin Account Quota Upstream Failure Shape + +Status: implemented and targeted verification passed. + +Changes: + +- Added a server regression test for the old explicit quota route behavior when + Codex usage fetch fails. +- Reworked the account quota HTTP handler to use a quota-specific error + response so upstream fetch failures return: + - HTTP `502 Bad Gateway`; + - admin code `50201`; + - message `Failed to fetch quota from Codex API`; + - `data.error` containing the public upstream failure detail. +- Kept standard admin errors (`401`, `404`, inactive account conflicts, and + persistence failures) on the normal admin error path. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_account_quota_should_return_bad_gateway_when_usage_fetch_fails` +- `cargo test -p codex-proxy-server --test admin_accounts_routes` +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. +- Fallback/recovery, dirty quota transitions, implicit resume edge cases, and + stream audit parity remain incomplete. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. + +### 2026-06-18 Slice 16: Responses Rate-Limit Account Fallback + +Status: implemented and targeted verification passed. + +Changes: + +- Added a server regression test migrated from the old upstream fallback suite + proving non-streaming `/v1/responses` retries the next restored runtime-pool + account after the first account returns HTTP `429` with `retry-after`. +- Added `RuntimeAccountPoolService::mark_quota_limited_until()` so runtime + dispatch can update the in-memory core pool instead of keeping fallback state + in the HTTP/server layer. +- Reworked `ResponseDispatchService::complete()` to: + - acquire accounts through the restored runtime pool with an exclusion list; + - call Codex using the selected account; + - release the selected account slot after each attempt; + - on upstream `429`, mark that account quota-limited until the retry-after + cooldown and retry with the next available account; + - preserve the existing completed-response session affinity recording path + for the account that actually succeeds. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_fallback_to_next_account_after_rate_limit` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Fallback/recovery remains incomplete: streaming Responses fallback, Chat + fallback, 5xx same-account retry, exhausted-account error aggregation, + auth/cloudflare recovery, implicit resume, stream audit, and persisted dirty + quota transitions still need migration. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 17: Chat Rate-Limit Account Fallback + +Status: implemented and targeted verification passed. + +Changes: + +- Added a server regression test proving `/v1/chat/completions` retries the + next restored runtime-pool account after the first account returns HTTP `429` + with `retry-after`. +- Reused the runtime rate-limit fallback path for Chat dispatch: + - acquire accounts through `RuntimeAccountPoolService`; + - release the selected account slot after each attempt; + - on upstream `429`, mark the account quota-limited until the retry-after + cooldown and exclude it from the next acquisition. +- Extracted a shared runtime helper for constructing Codex response requests + with selected account context so Chat and Responses do not drift in auxiliary + headers, turn state, installation ID, or account ID handling. +- Stabilized the OpenAI upstream route test fixtures by moving + `access_token_expires_at` from a same-day timestamp to a far-future timestamp; + the fixed timestamp had started expiring during the audit run and caused + account-pool acquisition to fail before any upstream request. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_fallback_to_next_account_after_rate_limit` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Fallback/recovery remains incomplete: streaming Responses fallback, 5xx + same-account retry, exhausted-account error aggregation, + auth/cloudflare recovery, implicit resume, stream audit, and persisted dirty + quota transitions still need migration. The HTTP 429 cooldown persistence + path was closed later in Slice 18. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 18: Persisted Quota Cooldown Transition + +Status: implemented and targeted verification passed. + +Changes: + +- Extended the existing Chat rate-limit fallback regression test to assert that + the first account is persisted as quota-limited after an upstream HTTP `429`: + `quota_limit_reached = 1` and `quota_cooldown_until` is a future timestamp. +- Added `AccountStore::mark_quota_limited_until()` to the core account-store + port so runtime dispatch can persist quota state without depending on a + concrete SQLite adapter. +- Implemented the new port in `SqliteAccountStore`, updating + `quota_limit_reached`, clearing `quota_verify_required`, writing + `quota_cooldown_until`, and refreshing `updated_at`. +- Updated `RuntimeAccountPoolService::mark_quota_limited_until()` to write the + quota cooldown to storage and still update the in-memory account pool. Storage + errors are logged and do not prevent the current request from retrying another + account. +- Updated the runtime task fake account store for the expanded port. + +TDD evidence: + +- Before the implementation, the targeted test failed as expected at the new + SQLite assertion with `quota_limit_reached` still `0`. +- After the implementation, the same test passed. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_fallback_to_next_account_after_rate_limit` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- HTTP 429 cooldown persistence is restored for the shared runtime fallback + path used by non-streaming Chat and Responses dispatch, but fallback/recovery + remains incomplete: streaming Responses fallback, 5xx same-account retry, + exhausted-account error aggregation, auth/cloudflare recovery, implicit + resume, and stream audit still need migration. The non-streaming Responses + 5xx same-account retry path was closed later in Slice 19. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 19: Responses 5xx Same-Account Retry + +Status: implemented and targeted verification passed. + +Changes: + +- Migrated the old upstream fallback behavior test proving non-streaming + `/v1/responses` retries the same account after transient upstream HTTP 5xx + failures before falling back to another account. +- Added a Responses-specific runtime helper that retries the selected account + for upstream HTTP 5xx responses up to two additional attempts. +- Kept HTTP 429 handling on the existing quota-cooldown fallback path; 429 still + marks the account quota-limited and acquires the next eligible account. +- Kept Chat dispatch unchanged in this slice; the migrated old behavior covered + non-streaming Responses. + +TDD evidence: + +- Before implementation, the migrated test failed with HTTP `502` and only one + request to the primary account. +- After implementation, the same test passed and confirmed three primary + requests with zero secondary requests. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_retry_same_account_after_5xx_before_fallback` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Fallback/recovery remains incomplete: streaming Responses fallback, streaming + 5xx retry, exhausted-account error aggregation, auth/cloudflare recovery, + implicit resume, and stream audit still need migration. +- Persisted per-request usage for selected accounts is not yet restored. + **Closed for request_count in Slice 20; token deltas remain incomplete.** +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 20: Persist Selected-Account Request Count + +Status: implemented and targeted verification passed. + +Changes: + +- Migrated the old fallback test expectation that an account selected for a + failed non-streaming `/v1/responses` request still records one persisted + request in `account_usage`. +- Added `AccountStore::record_usage_delta()` plus a default + `record_request()` helper to the core account-store port so runtime account + acquisition can persist request-count usage through the same abstraction used + by the account pool. +- Implemented `record_usage_delta()` in `SqliteAccountStore` using the existing + `record_usage(..., UsageDelta::default())` semantics for request-count usage. +- Updated `RuntimeAccountPoolService::acquire_with()` and `acquire()` so every + selected account records one persisted external request. Persistence failures + are logged and do not block dispatch. +- Updated the runtime fake account store for the expanded port. + +TDD evidence: + +- Before implementation, the migrated test failed with `request_count = 0`. +- After implementation, the same test passed with `request_count = 1` and three + upstream attempts on the selected account. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_record_request_count_when_5xx_retries_are_exhausted` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo test -p codex-proxy-runtime --test account_pool_restore` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Request-count persistence is restored for selected accounts, including failed + requests and multi-account fallback attempts, but persisted token deltas from + successful Responses/Chat completions are still incomplete. **Closed for + non-streaming Responses in Slice 21; Chat and streaming paths remain + incomplete.** +- The old exhausted-account HTTP/status/error aggregation shape is still not + restored; current OpenAI upstream failures still map through the current 502 + handler. +- Fallback/recovery remains incomplete: streaming Responses fallback, streaming + 5xx retry, auth/cloudflare recovery, implicit resume, and stream audit still + need migration. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 21: Persist Responses Token Usage Delta + +Status: implemented and targeted verification passed. + +Changes: + +- Extended the migrated Responses HTTP 429 fallback test with the old usage + assertions: + - the primary 429 account records one request and zero tokens; + - the successful secondary account records one request plus the completed + response input/output token usage. +- Extended `core::accounts::usage::AccountUsageDelta` with cached and image + token fields so runtime can pass token deltas through a core-owned type. +- Reworked SQLite `UsageDelta` and `RECORD_USAGE_SQL` so `request_count` is a + bound delta rather than a fixed `1`; this allows token-only updates with + `request_count = 0`. +- Added conversion from core `AccountUsageDelta` into SQLite `UsageDelta` with + saturating `u64 -> i64` conversion. +- Updated `RuntimeAccountPoolService` to persist successful Responses token + usage after `response.completed` is parsed and to sync the in-memory account + pool window token counters. + +TDD evidence: + +- Before implementation, the extended 429 fallback test failed with the + secondary account usage row `(1, 0, 0)` instead of `(1, 5, 2)`. +- After implementation, the same test passed. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_fallback_to_next_account_after_rate_limit` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo test -p codex-proxy-adapters --test account_repository` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Persisted token usage is restored for non-streaming Responses completion + paths, but Chat completion token persistence and streaming Responses token + persistence remain incomplete. **Chat completion token persistence closed in + Slice 22.** +- Empty-response counters, image request success/failure counters, and old + exhausted-account HTTP/status/error aggregation are still not restored. +- Fallback/recovery remains incomplete: streaming Responses fallback, streaming + 5xx retry, auth/cloudflare recovery, implicit resume, and stream audit still + need migration. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 22: Persist Chat Token Usage Delta + +Status: implemented and targeted verification passed. + +Changes: + +- Extended the Chat Completions upstream regression test to assert that a + successful `/v1/chat/completions` request persists `account_usage` with one + request plus the completed response input/output token usage. +- Updated `ChatDispatchService` to keep the selected account ID from the + successful dispatch attempt and call the shared + `RuntimeAccountPoolService::record_token_usage()` after the Chat SSE response + is parsed successfully. +- Kept token persistence after successful SSE parsing so invalid or empty + upstream responses do not record completed token deltas. + +TDD evidence: + +- Before implementation, the extended Chat test failed with persisted usage + `(1, 0, 0)` instead of `(1, 9, 3)`. +- After implementation, the same test passed. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_dispatch_to_codex_and_return_openai_response` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Persisted token usage is restored for non-streaming Responses and Chat + completion paths, but streaming Responses token persistence remains + incomplete. +- Empty-response counters, image request success/failure counters, and old + exhausted-account HTTP/status/error aggregation are still not restored. +- Fallback/recovery remains incomplete: streaming Responses fallback, streaming + 5xx retry, auth/cloudflare recovery, implicit resume, and stream audit still + need migration. Non-streaming Responses HTTP 402 quota fallback was closed in + Slice 23. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 23: Responses 402 Quota Exhausted Fallback + +Status: implemented and targeted verification passed. + +Changes: + +- Migrated the old upstream fallback behavior proving non-streaming + `/v1/responses` treats upstream HTTP `402` as quota exhaustion for the + selected account, marks that account `quota_exhausted`, and retries the next + available account. +- Added `AccountStore::set_status()` to the core account-store port. +- Implemented the new status port in `SqliteAccountStore` by reusing the + existing status update path. +- Added `RuntimeAccountPoolService::set_status()` to persist status changes and + update the in-memory account pool, logging persistence failures without + panicking. +- Updated Responses dispatch to handle upstream HTTP `402` by marking the + selected account `QuotaExhausted`, excluding it, and continuing the fallback + acquire loop. + +TDD evidence: + +- Before implementation, the migrated 402 fallback test failed with HTTP `502` + instead of retrying the secondary account. +- After implementation, the same test passed and confirmed the primary account + was persisted as `quota_exhausted` with one failed request recorded. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_mark_quota_exhausted_after_402_and_fallback` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo test -p codex-proxy-adapters --test account_repository` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Non-streaming Responses HTTP 402 quota fallback is restored, but SSE + `response.failed` quota classification, streaming quota fallback, and no + fallback exhausted-account error aggregation remain incomplete. Non-streaming + SSE quota failure fallback was closed in Slice 24. +- Auth/cloudflare recovery, implicit resume, and stream audit still need + migration. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 24: Responses SSE Quota Failure Fallback + +Status: implemented and targeted verification passed. + +Changes: + +- Migrated the old non-streaming Responses SSE failure behavior proving + `response.failed` with a quota error marks the selected account + `quota_exhausted` and retries the next available account. +- Moved Responses SSE collection into the per-account attempt loop so fallback + decisions can inspect `CollectedResponse::Failed` before the dispatch result + is finalized. +- Added quota-failure classification for SSE failures using known upstream + codes (`quota_exceeded`, `insufficient_quota`) and quota-bearing messages. +- Preserved the existing error path for non-quota SSE failures by breaking out + of the attempt loop and returning `ResponseDispatchError::Failed`. + +TDD evidence: + +- Before implementation, the migrated SSE quota failure test failed with HTTP + `502` instead of retrying the secondary account. +- After implementation, the same test passed and confirmed the primary account + was persisted as `quota_exhausted`. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_classify_sse_quota_failure_and_fallback` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Non-streaming Responses HTTP 402 and SSE quota failure fallback are restored, + but streaming quota fallback and no-fallback exhausted-account error + aggregation remain incomplete. +- Chat 402 quota fallback remains incomplete. **Closed for multi-account HTTP + 402 fallback in Slice 25; no-fallback aggregation remains incomplete.** +- Auth/cloudflare recovery, implicit resume, and stream audit still need + migration. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 25: Chat 402 Quota Exhausted Fallback + +Status: implemented and targeted verification passed. + +Changes: + +- Added Chat Completions coverage proving a primary account that returns + upstream HTTP `402` is marked `quota_exhausted` and excluded before retrying + the next restored runtime-pool account. +- Reused the runtime status migration added for Responses quota fallback, so + Chat and Responses persist and update the in-memory account pool consistently. +- Kept this slice scoped to multi-account fallback; old no-fallback Chat quota + error aggregation remains open. + +TDD evidence: + +- Before implementation, the new Chat 402 fallback test failed with HTTP `502`. +- After implementation, the same test passed and confirmed the primary account + was persisted as `quota_exhausted`. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_mark_quota_exhausted_after_402_and_fallback` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Multi-account Chat and non-streaming Responses quota fallback paths are + restored for HTTP 402, but no-fallback exhausted-account error aggregation is + still incomplete. Chat HTTP 402 no-fallback aggregation was closed in Slice + 26. +- Streaming quota fallback, auth/cloudflare recovery, implicit resume, and + stream audit still need migration. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 26: Chat 402 No-Fallback Quota Error + +Status: implemented and targeted verification passed. + +Changes: + +- Migrated the old Chat no-fallback quota behavior for a single account that + returns upstream HTTP `402`. +- Added `ChatDispatchError::QuotaExhausted` so dispatch can preserve the number + of quota-exhausted accounts and the last upstream error body instead of + collapsing the result to generic `NoActiveAccount`. +- Updated the Chat dispatch acquire loop to return the quota-exhausted aggregate + once all retried accounts have been excluded by quota exhaustion. +- Mapped the new Chat error to an OpenAI-compatible HTTP `402 Payment Required` + response with the old `All accounts exhausted (N quota-exhausted)` wording + and `upstream_error` code. + +TDD evidence: + +- Before implementation, the migrated test failed with HTTP `503` instead of + `402`. +- After implementation, the same test passed and confirmed the account was + persisted as `quota_exhausted`. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_return_quota_error_when_402_fallback_is_exhausted` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Chat HTTP 402 no-fallback aggregation is restored, but analogous Responses + no-fallback quota aggregation remains incomplete. **Closed for Responses HTTP + 402 in Slice 27; streaming/SSE no-fallback aggregation remains incomplete.** +- Streaming quota fallback, auth/cloudflare recovery, implicit resume, and + stream audit still need migration. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 27: Responses 402 No-Fallback Quota Error + +Status: implemented and targeted verification passed. + +Changes: + +- Added non-streaming Responses coverage for a single account that returns + upstream HTTP `402` with no fallback account available. +- Added `ResponseDispatchError::QuotaExhausted`, mirroring the Chat aggregate + error, to preserve the quota-exhausted account count and last upstream error + body. +- Updated the Responses dispatch acquire loop to return the quota aggregate + after all candidates are excluded by HTTP 402 or quota-classified SSE + failures. +- Mapped the new Responses error to HTTP `402 Payment Required` with the old + `All accounts exhausted (N quota-exhausted)` message and `upstream_error` + code. + +TDD evidence: + +- Before implementation, the migrated Responses no-fallback test failed with + HTTP `503` instead of `402`. +- After implementation, the same test passed and confirmed the account was + persisted as `quota_exhausted`. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_return_quota_error_when_402_fallback_is_exhausted` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Chat and non-streaming Responses HTTP 402 no-fallback aggregation are + restored, but streaming quota fallback and SSE no-fallback aggregation remain + incomplete. +- Chat HTTP 429 no-fallback aggregation remains incomplete. **Closed in Slice + 28; analogous Responses 429 no-fallback aggregation remains incomplete.** +- Auth/cloudflare recovery, implicit resume, and stream audit still need + migration. +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 28: Chat 429 No-Fallback Rate-Limit Error + +Status: implemented and targeted verification passed. + +Changes: + +- Migrated the old Chat no-fallback rate-limit behavior for a single account + that returns upstream HTTP `429`. +- Added `ChatDispatchError::RateLimited` so Chat dispatch can preserve the count + of rate-limited accounts and the last upstream error body instead of + collapsing the result to generic `NoActiveAccount`. +- Updated the Chat dispatch loop to return the rate-limited aggregate after all + candidates are excluded by HTTP 429 cooldown. +- Mapped the new Chat error to an OpenAI-compatible HTTP `429 Too Many + Requests` response with the old `All accounts exhausted (N rate-limited)` + wording and `upstream_error` code. +- The test also asserts the account's persisted quota cooldown state is updated. + +TDD evidence: + +- Before implementation, the migrated test failed with HTTP `503` instead of + `429`. +- After implementation, the same test passed and confirmed the account was + marked quota-limited with a cooldown timestamp. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_return_rate_limit_error_when_429_fallback_is_exhausted` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Chat HTTP 429 no-fallback aggregation is restored, but analogous Responses + HTTP 429 no-fallback aggregation remains incomplete. **Closed in Slice 29.** +- Streaming quota/rate-limit fallback, auth/cloudflare recovery, implicit + resume, and stream audit still need migration. **The buffered HTTP SSE + `stream: true` fallback/usage path is partially restored in Slice 30; live + chunk-by-chunk streaming, WebSocket stream parity, and stream audit still + remain.** +- Token refresh still lacks the old per-account timer, in-flight, retry, and + crash-recovery behavior. +- Old quota edge-case tests still need migration: inactive account route + assertions, quota persistence failure behavior, and quota warning edge cases + beyond the current cached-warning test. + +### 2026-06-18 Slice 29: Responses 429 No-Fallback Rate-Limit Error + +Status: implemented and targeted verification passed. + +Changes: + +- Migrated the old Responses no-fallback rate-limit behavior for a single + account that returns upstream HTTP `429`. +- Added `ResponseDispatchError::RateLimited`, mirroring the Chat aggregate + error restored in Slice 28. +- Updated `ResponseDispatchService::complete()` to count HTTP 429 accounts, + preserve the last upstream error body, and return the rate-limit aggregate + after all candidates are excluded by cooldown. +- Mapped the new Responses error to an OpenAI-compatible HTTP `429 Too Many + Requests` response with the old `All accounts exhausted (N rate-limited)` + wording and `upstream_error` code. +- The regression test also asserts the account is persisted as quota-limited + with a cooldown timestamp. + +TDD evidence: + +- Before implementation, + `responses_should_return_rate_limit_error_when_429_fallback_is_exhausted` + failed with HTTP `503` instead of `429`. +- After implementation, the same test passed and confirmed the persisted + cooldown state. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_return_rate_limit_error_when_429_fallback_is_exhausted` +- `codegraph sync .` + +Remaining after this slice: + +- Chat and Responses HTTP 429 no-fallback aggregation are restored for + non-streaming paths. +- Buffered HTTP SSE `stream: true` fallback and usage persistence still needed + restoration. **Partially closed in Slice 30.** +- Auth/cloudflare recovery, implicit resume, live stream audit, WebSocket stream + parity, and token-refresh old parity still need migration. + +### 2026-06-18 Slice 30: Buffered Responses Stream Fallback And Usage + +Status: implemented and targeted verification passed. + +Changes: + +- Added a `ResponseDispatchService::stream()` path for OpenAI + `/v1/responses` requests with `stream: true`. +- The stream path now returns `text/event-stream` instead of collecting Codex + SSE into a JSON Responses object. +- Restored the HTTP SSE fallback behaviors covered by the migrated tests: + - completed SSE usage is persisted to `account_usage`; + - upstream HTTP `429` marks the first account quota-limited, excludes it, and + retries the next account; + - transient upstream HTTP 5xx responses retry the same account before + falling back; + - upstream HTTP `402` marks the first account `QuotaExhausted`, excludes it, + and retries the next account. +- The stream path also classifies upstream SSE `response.failed` quota failures + for fallback using the same collector already restored for non-streaming + Responses. +- Added server-side SSE error encoding for dispatch errors on `stream: true`. + +Scope note: + +- This is a buffered HTTP SSE compatibility path: the proxy returns the Codex + SSE body with the correct OpenAI streaming content type after the upstream + response has been read. The old live chunk-by-chunk forwarding, + premature-close handling, heartbeat behavior, tuple stream reconversion, + WebSocket streaming, and stream audit still need separate migration. + +TDD evidence: + +- The first RED run showed three new stream tests failing because the handler + returned JSON instead of `text/event-stream`; the 5xx test was then tightened + because it had passed by only checking the JSON body. +- The corrected RED run failed all four stream tests on the missing + `text/event-stream` response type. +- After implementation, all four tests passed: + `responses_stream_should_proxy_sse_and_record_usage`, + `responses_stream_should_fallback_to_next_account_after_rate_limit`, + `responses_stream_should_retry_same_account_after_5xx_before_fallback`, and + `responses_stream_should_mark_quota_exhausted_after_402_and_fallback`. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Buffered HTTP SSE `stream: true` fallback and usage persistence are partially + restored, but true live streaming and stream audit are still incomplete. +- WebSocket transport/pool streaming parity remains incomplete. +- Stream no-fallback edge cases, auth/cloudflare recovery, implicit resume, + and token-refresh old parity still need migration. **HTTP 401/token-invalid + account recovery is partially restored in Slice 31; Cloudflare/path-block + recovery remains incomplete.** + +### 2026-06-18 Slice 31: Auth Failure Account Recovery + +Status: implemented and targeted verification passed. + +Changes: + +- Restored account-level recovery for upstream HTTP `401` authentication + failures: + - Chat Completions marks the failed account `expired`, excludes it, and + retries the next account. + - non-streaming Responses marks the failed account `expired`, excludes it, + and retries the next account. + - buffered Responses `stream: true` marks the failed account `expired`, + excludes it, and retries the next account. +- Restored non-streaming Responses recovery for upstream SSE + `response.failed` auth/token failures such as `token_invalid`. +- Persisted status changes through `RuntimeAccountPoolService::set_status()`, + so SQLite and the in-memory pool stay consistent. + +TDD evidence: + +- Before implementation, `cargo test ... 401` failed as expected: + Chat and non-streaming Responses returned HTTP `502` instead of falling back, + and the stream case did not produce the fallback account usage row. +- Before implementation, `cargo test ... sse_failed_event` failed with HTTP + `502` instead of falling back from the token-invalid SSE failure. +- After implementation, both targeted filters passed. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream 401` +- `cargo test -p codex-proxy-server --test openai_chat_upstream sse_failed_event` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- HTTP 401 and token-invalid SSE account recovery are restored for the covered + Chat/Responses paths. +- Cloudflare/path-block recovery, auth no-fallback aggregate/error parity, + live stream audit, WebSocket stream parity, implicit resume, and token-refresh + old parity still need migration. **Auth no-fallback aggregate/error parity is + restored for the covered Chat/Responses paths in Slice 32.** + +### 2026-06-18 Slice 32: Auth No-Fallback Expired Aggregation + +Status: implemented and targeted verification passed. + +Changes: + +- Added expired-account aggregate errors to Chat and Responses dispatch: + `ChatDispatchError::Expired` and `ResponseDispatchError::Expired`. +- Dispatch now counts HTTP `401`/token-invalid accounts, preserves the last + upstream error body, and returns an expired aggregate after all candidates are + excluded. +- Mapped non-streaming Chat and Responses expired aggregates to HTTP `401 + Unauthorized` with `All accounts exhausted (N expired)` wording and + `upstream_error` code. +- Mapped Responses `stream: true` expired aggregates to a `response.failed` + SSE event with authentication error semantics. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-server --test openai_chat_upstream auth_error_when_401` + failed for all three new tests: + - Chat and non-streaming Responses returned HTTP `503` instead of `401`. + - Responses `stream: true` returned a no-active stream error instead of the + expired aggregate. +- After implementation, the same filter passed. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream auth_error_when_401` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- HTTP 401/token-invalid fallback and no-fallback aggregation are restored for + the covered Chat/Responses paths. +- Cloudflare/path-block recovery, banned/deactivated account status edge cases, + live stream audit, WebSocket stream parity, implicit resume, and token-refresh + old parity still need migration. **HTTP 401 deactivated/banned status + classification is restored in Slice 33.** + +### 2026-06-18 Slice 33: Deactivated Auth Marks Banned + +Status: implemented and targeted verification passed. + +Changes: + +- Added auth failure account-status classification for upstream HTTP `401` + bodies and SSE auth failures. +- Auth failures containing `account_deactivated`, `account deactivated`, + `account has been deactivated`, `deactivated`, or `banned` now set + `AccountStatus::Banned`; other token auth failures still set + `AccountStatus::Expired`. +- Covered both Chat Completions and non-streaming Responses HTTP `401` + deactivation bodies. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-server --test openai_chat_upstream deactivated` + failed because both accounts were persisted as `expired` instead of `banned`. +- After implementation, the same filter passed. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream deactivated` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- HTTP 401 token-invalid/deactivated fallback and no-fallback status handling + are restored for the covered Chat/Responses paths. +- Cloudflare/path-block recovery, live stream audit, WebSocket stream parity, + implicit resume, and token-refresh old parity still need migration. + +### 2026-06-18 Slice 34: HTTP 403 Banned Fallback + +Status: implemented and targeted verification passed. + +Changes: + +- Restored account-level fallback for upstream HTTP `403` bodies that signal a + banned or deactivated account. +- Chat Completions, non-streaming Responses, and buffered Responses + `stream: true` now mark the failed account `banned`, exclude it, and retry + the next eligible account. +- HTTP `403` bodies without a banned/deactivated signal still use the normal + upstream-error path; Cloudflare challenge/path-block recovery remains a + separate migration gap. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-server --test openai_chat_upstream after_403` + failed because Chat and non-streaming Responses returned HTTP `502` instead + of falling back, and the buffered stream response did not contain the fallback + account response id. +- After implementation, the same filter passed. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream after_403` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- HTTP 401 token-invalid/deactivated and HTTP 403 banned/deactivated account + fallback are restored for the covered Chat/Responses paths. +- Cloudflare challenge/path-block recovery, live stream audit, WebSocket stream + parity, implicit resume, and token-refresh old parity still need migration. + +### 2026-06-18 Slice 35: Cloudflare Challenge And Path-Block Recovery + +Status: implemented and targeted verification passed. + +Changes: + +- Added `core::accounts::cloudflare::CloudflarePathBlockTracker` for pure + per-account path-block counting, including the old 3-hit disable threshold + and 1-hour stale counter window. +- Extended `AccountStore` and `SqliteAccountStore` with + `set_cloudflare_cooldown_until()` so runtime can persist Cloudflare + cooldowns without reaching around the port boundary. +- Added a shared runtime Cloudflare recovery helper used by Chat Completions, + non-streaming Responses, and buffered Responses `stream: true`: + - reads stored `chatgpt.com` cookies into upstream Codex requests; + - clears the failing account's cookies after Cloudflare challenge/path-block; + - sets `cloudflare_cooldown_until` after HTTP `403` Cloudflare challenge; + - counts HTTP `404` empty-body path-blocks and disables the account after + three recent hits; + - excludes the failed account and retries the next eligible account. +- Added Cloudflare aggregate dispatch errors and mapped them to HTTP `502` + JSON errors for non-streaming requests and `response.failed` SSE events for + buffered streaming requests. +- Updated the exact architecture source whitelist for + `crates/core/src/accounts/cloudflare.rs`. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-server --test openai_chat_upstream cloudflare` + failed 7/7 newly added tests: fallback cases returned HTTP `502`, cooldowns + were not written, cookies were not cleared, path-block errors were generic, + and the 3-hit disabled transition did not occur. +- After implementation, the same filter passed 7/7. +- The first post-implementation architecture run failed because + `crates/core/src/accounts/cloudflare.rs` was not yet listed in the exact + architecture shape; `docs/architecture.md` and the architecture test + whitelist were updated, then the architecture suite passed. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream cloudflare` + (RED before implementation; failed 7/7 as expected) +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream cloudflare` +- `cargo fmt --all` +- `cargo test -p codex-proxy-core cloudflare` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` (failed once on the new source whitelist) +- `cargo test --test architecture` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream cloudflare` +- `cargo test -p codex-proxy-core cloudflare` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Cloudflare challenge/path-block recovery is restored for the covered + non-streaming Chat, non-streaming Responses, and buffered Responses + `stream: true` paths. +- The old live chunk-by-chunk stream behavior, WebSocket stream parity, stream + audit, implicit resume, and token-refresh old parity still need migration. + +### 2026-06-18 Slice 36: Responses Model-Unsupported Fallback + +Status: implemented and targeted verification passed. + +Changes: + +- Restored Responses account fallback for model-plan incompatibility signals: + - upstream HTTP `400` bodies with `model_not_supported`, + `model_not_available`, `not supported`, or `not available`; + - upstream SSE `response.failed` events with the same model unsupported + signals. +- Non-streaming Responses and buffered Responses `stream: true` now exclude the + failed account and retry the next account once without changing account + status. +- Added model-unsupported aggregate errors for exhausted fallback: + non-streaming requests map to HTTP `400`, and buffered streaming requests map + to a `response.failed` SSE event. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-server --test openai_chat_upstream model_unsupported` + failed 4/4 newly added tests: HTTP model unsupported returned `502`, SSE + model unsupported did not retry the next account, buffered stream did not + contain the fallback response id, and exhausted fallback returned a generic + stream error. +- After implementation, the same filter passed 4/4. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream model_unsupported` + (RED before implementation; failed 4/4 as expected) +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream model_unsupported` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream model_unsupported` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Responses model-unsupported fallback is restored for the covered + non-streaming and buffered `stream: true` paths. +- Chat model-unsupported fallback, implicit resume/history stripping, live + stream audit, WebSocket stream parity, and token-refresh old parity still + need migration. + +### 2026-06-18 Slice 37: Responses History Recovery + +Status: implemented and targeted verification passed. + +Changes: + +- Added `turnState` / `turn_state` parsing to OpenAI Responses requests and + propagation to the Codex `x-codex-turn-state` upstream header. +- Restored one-shot Responses request history recovery for upstream SSE + `response.failed` signals: + - `previous_response_not_found`; + - messages containing `Previous response ... not found`; + - messages containing `No tool output found for function call`. +- On recovery, runtime clears `previous_response_id`, `turn_state`, + `turn_metadata`, and non-explicit `prompt_cache_key`, then retries the request + once instead of switching accounts or changing account status. +- Covered non-streaming Responses and buffered Responses `stream: true`. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-server --test openai_chat_upstream strip_history` + failed 3/3 newly added tests because only one upstream request was sent. +- After implementation and fixing the same-account WireMock fixtures to use the + old `up_to_n_times(1)` pattern, the same filter passed 3/3. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream strip_history` + (RED before implementation; failed 3/3 as expected) +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream strip_history` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream strip_history` +- `cargo test -p codex-proxy-core --test protocol openai_response_request_should_translate_to_codex_request` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Remaining after this slice: + +- Responses history recovery is restored for the covered non-streaming and + buffered `stream: true` SSE failure paths. +- HTTP upstream history errors, invalid encrypted reasoning replay recovery, + Chat model-unsupported fallback, live stream audit, WebSocket stream parity, + and token-refresh old parity still need migration. + +### 2026-06-18 Slice 38: Chat Model Fallback And Invalid Reasoning Recovery + +Status: implemented and targeted verification passed. + +Changes: + +- Restored Chat Completions account fallback for upstream HTTP model-plan + incompatibility signals such as `model_not_supported`, + `model_not_available`, `not supported`, and `not available`. +- Chat model-unsupported fallback now excludes the failed account, retries one + alternate account once, preserves the failed account status, and maps + exhausted fallback to OpenAI-style HTTP `400`. +- Extended Responses history recovery to recognize invalid encrypted reasoning + replay signals: + - upstream HTTP error bodies containing `invalid_encrypted_content`; + - equivalent `invalid` + `encrypted` + `content` text; + - upstream SSE `response.failed` events carrying the same signals. +- Invalid encrypted reasoning replay now bypasses generic same-account 5xx + retry, clears `previous_response_id`, `turn_state`, `turn_metadata`, and + non-explicit `prompt_cache_key`, then retries the stripped request once. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-server --test openai_chat_upstream model_unsupported` + failed the 2 newly added Chat tests: fallback returned `502` instead of + retrying, and exhausted fallback returned `502` instead of HTTP `400`. +- Before implementation, + `cargo test -p codex-proxy-server --test openai_chat_upstream invalid_encrypted_reasoning` + failed 3/3 newly added tests: SSE paths sent only one upstream request, and + the HTTP 503 path retried without stripping `x-codex-turn-state`. +- After implementation, both filters passed. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream model_unsupported` + (RED before implementation; failed the 2 new Chat tests as expected) +- `cargo test -p codex-proxy-server --test openai_chat_upstream invalid_encrypted_reasoning` + (RED before implementation; failed 3/3 as expected) +- `cargo test -p codex-proxy-server --test openai_chat_upstream model_unsupported` +- `cargo test -p codex-proxy-server --test openai_chat_upstream invalid_encrypted_reasoning` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo fmt --all` +- `cargo check --workspace --all-targets` +- `cargo fmt --all -- --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- `codegraph sync . && codegraph status .` + (synced 14 changed files; index up to date with 274 files, 3,733 nodes, and + 10,666 edges before this documentation-only update) +- `cargo test --test architecture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo fmt --all -- --check` +- `git diff --check` + +Remaining after this slice: + +- Chat HTTP model-unsupported fallback is restored for the covered + non-streaming Chat path. +- Responses invalid encrypted reasoning replay recovery is restored for covered + HTTP error, non-streaming SSE failure, and buffered `stream: true` SSE failure + paths. +- Live stream audit, WebSocket stream parity, implicit resume/reasoning replay + cache parity, and token-refresh old parity still need migration. + +### 2026-06-19 Slice 39: Token Refresh Lease Coordination + +Status: implemented and targeted verification passed. + +Changes: + +- Implemented `SqliteRefreshLeaseStore` behavior on top of the existing + `account_refresh_leases` table: + - acquire creates a lease when absent; + - acquire is blocked while another owner holds an unexpired lease; + - acquire succeeds after expiry or for the same owner; + - release only succeeds for the current owner. +- Added the refresh lease store to runtime repository assembly and + `BackgroundTaskStores`. +- Wired the background token refresh task to use refresh leases before invoking + the upstream refresher. +- Token refresh scans now skip a due account when another owner holds a live + lease, and do not call the refresher for that account. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-adapters --test refresh_leases` failed to compile + because `SqliteRefreshLeaseStore` had no `try_acquire` or `release` methods. +- Before implementation, + `cargo test -p codex-proxy-runtime --test token_refresh` failed to compile + because the lease store methods and `TokenRefreshTask::with_refresh_lease_store` + did not exist. +- After implementation, both filters passed. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test refresh_leases` + (RED before implementation; missing lease API as expected) +- `cargo test -p codex-proxy-runtime --test token_refresh` + (RED before implementation; missing lease API/task wiring as expected) +- `cargo test -p codex-proxy-adapters --test refresh_leases` +- `cargo test -p codex-proxy-runtime --test token_refresh` +- `codegraph sync .` +- `cargo test -p codex-proxy-runtime --test tasks` +- `cargo fmt --all` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo test -p codex-proxy-adapters --test refresh_leases` +- `cargo test -p codex-proxy-runtime --test token_refresh` +- `cargo fmt --all -- --check` +- `git diff --check` + +Remaining after this slice: + +- Token refresh now uses the SQLite refresh lease table for cross-process + single-owner coordination on scanned accounts. +- Token refresh still lacks the old per-account timer map, explicit in-flight + tracking, scheduled `refreshing` crash recovery, exponential retry loop, and + delayed recovery scheduling. +- Live stream audit, WebSocket stream parity, and implicit resume/reasoning + replay cache parity still need migration. + +### 2026-06-19 Slice 40: Token Refresh Refreshing Recovery + +Status: implemented and targeted verification passed. + +Changes: + +- Token refresh scans now distinguish refresh candidates before taking a + refresh lease, so non-due accounts are skipped without being marked + `refreshing`. +- Due active accounts are persisted as `refreshing` before invoking the + upstream token refresher, matching the old scheduled refresh state transition. +- Accounts found persisted as `refreshing` are treated as restart/crash + recovery candidates and refreshed immediately even when their access token is + not inside the normal pre-expiry margin. +- Transport failures during a refresh attempt now restore the persisted account + status to `active` instead of leaving the account stuck in `refreshing`. +- Exposed a read-only core `RefreshScheduler::should_refresh_account_at` helper + so runtime can reuse the domain refresh policy before it mutates persisted + state. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-runtime --test token_refresh` failed the 3 newly + added tests: + - a due active account was still observed as `Active` when the refresher ran; + - a persisted `Refreshing` account was skipped with `summary.refreshed = 0`; + - a transport-failure attempt was still observed as `Active` instead of + `Refreshing`. +- After implementation, the same test target passed 5/5. + +Verification run so far: + +- `cargo test -p codex-proxy-runtime --test token_refresh` + (RED before implementation; failed 3/5 as expected) +- `cargo test -p codex-proxy-runtime --test token_refresh` +- `codegraph sync .` +- `cargo test -p codex-proxy-runtime --test tasks` +- `cargo fmt --all` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo fmt --all -- --check` +- `git diff --check` + +Remaining after this slice: + +- Token refresh now covers SQLite lease coordination, refresh-before-call + `refreshing` persistence, restart/crash recovery from persisted `refreshing`, + and active restoration after a transport failure. +- Token refresh still lacks the old per-account timer map, explicit in-flight + tracking, exponential retry loop, and delayed recovery scheduling. +- Live stream audit, WebSocket stream parity, and implicit resume/reasoning + replay cache parity still need migration. + +### 2026-06-19 Slice 41: Token Refresh Retry Confirmation + +Status: implemented and targeted verification passed. + +Changes: + +- Added token-refresh retry delays with production defaults matching the old + scheduler's five-attempt exponential retry shape: + `5s`, `15s`, `45s`, and `135s` between attempts, capped at `300s`. +- Added a configurable `TokenRefreshTask::with_retry_delays(...)` hook so + tests can run the retry loop with zero delays without pausing Tokio's global + clock. +- Transport failures now retry before the account is restored to `active` and + counted as failed. +- A transient transport failure followed by a successful refresh now persists + the refreshed token and reports the account as refreshed. +- Permanent refresh failures represented by terminal account statuses are now + confirmed twice before the final status is persisted, matching the old + invalid-grant/permanent-failure confirmation behavior. +- Accounts without a refresh token still skip the retry loop and are marked + `expired` immediately. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-runtime --test token_refresh` first failed because + the new tests needed a missing `TokenRefreshTask::with_retry_delays(...)` + hook. +- Before adding the retry implementation, the behavior expectations were also + missing: transport failure made only one refresher call, transient success + after a retry was not reached, and invalid-grant was persisted after one hit + instead of two. +- After implementation, the same test target passed 7/7. + +Verification run so far: + +- `cargo test -p codex-proxy-runtime --test token_refresh` + (RED before implementation; missing retry-delay hook as expected) +- `cargo test -p codex-proxy-runtime --test token_refresh` +- `codegraph sync .` +- `cargo test -p codex-proxy-runtime --test tasks` +- `cargo fmt --all` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo test -p codex-proxy-runtime --test token_refresh` +- `cargo fmt --all -- --check` +- `git diff --check` + +Remaining after this slice: + +- Token refresh now covers SQLite lease coordination, refresh-before-call + `refreshing` persistence, restart/crash recovery from persisted `refreshing`, + active restoration after retry exhaustion, transient retry success, and + two-hit permanent failure confirmation. +- Token refresh still lacks the old per-account timer map, explicit in-flight + tracking, and delayed recovery scheduling after retry exhaustion. +- Live stream audit, WebSocket stream parity, and implicit resume/reasoning + replay cache parity still need migration. + +### 2026-06-19 Slice 42: Token Refresh In-Flight Guard + +Status: implemented and targeted verification passed. + +Changes: + +- Added process-local in-flight tracking to `TokenRefreshTask`. +- Concurrent scans for the same account now skip the duplicate in-process + refresh attempt instead of treating the persisted `refreshing` status as a + crash recovery signal while the first refresh is still running. +- The in-flight guard is acquired before the SQLite refresh lease, so duplicate + scans from the same task owner cannot acquire and then release another + in-process refresh attempt's lease. +- The in-flight marker is released on lease-skip, lease-error, refresh success, + refresh failure, and status-update paths. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-runtime --test token_refresh token_refresh_task_should_skip_duplicate_in_flight_refresh` + failed because the second concurrent scan blocked inside the refresher instead + of returning a skipped summary. +- After implementation, the filtered test passed, and the full token-refresh + test target passed 8/8. + +Verification run so far: + +- `cargo test -p codex-proxy-runtime --test token_refresh token_refresh_task_should_skip_duplicate_in_flight_refresh` + (RED before implementation; duplicate refresh waited in the refresher) +- `cargo test -p codex-proxy-runtime --test token_refresh token_refresh_task_should_skip_duplicate_in_flight_refresh` +- `cargo test -p codex-proxy-runtime --test token_refresh` +- `codegraph sync .` +- `cargo test -p codex-proxy-runtime --test tasks` +- `cargo fmt --all` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo test -p codex-proxy-runtime --test token_refresh` +- `cargo fmt --all -- --check` +- `git diff --check` + +Remaining after this slice: + +- Token refresh now covers SQLite lease coordination, process-local in-flight + protection, refresh-before-call `refreshing` persistence, restart/crash + recovery from persisted `refreshing`, retry exhaustion restoration, + transient retry success, and two-hit permanent failure confirmation. +- Token refresh still lacks the old per-account timer map and delayed recovery + scheduling after retry exhaustion. +- Live stream audit, WebSocket stream parity, and implicit resume/reasoning + replay cache parity still need migration. + +### 2026-06-19 Slice 43: Token Refresh Delayed Recovery + +Status: implemented and targeted verification passed. + +Changes: + +- Added process-local delayed recovery tracking for token refresh retry + exhaustion. +- After all transport retry attempts are exhausted, the task restores the + account to `active`, records a 10-minute recovery window, and skips scans for + that account until the window expires. +- Once the recovery window expires, the next scan clears the delay and attempts + refresh again. +- Successful refreshes, terminal status updates, and missing-refresh-token + expiry clear any pending delayed recovery entry for that account. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-runtime --test token_refresh token_refresh_task_should_delay_recovery_after_retry_exhaustion` + failed because a scan 5 minutes after retry exhaustion did not skip the + account. +- After implementation, the filtered test passed, and the full token-refresh + test target passed 9/9. + +Verification run so far: + +- `cargo test -p codex-proxy-runtime --test token_refresh token_refresh_task_should_delay_recovery_after_retry_exhaustion` + (RED before implementation; recovery-window skip missing) +- `cargo test -p codex-proxy-runtime --test token_refresh token_refresh_task_should_delay_recovery_after_retry_exhaustion` +- `cargo test -p codex-proxy-runtime --test token_refresh` +- `codegraph sync .` +- `cargo test -p codex-proxy-runtime --test tasks` +- `cargo fmt --all` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo test -p codex-proxy-runtime --test token_refresh` +- `cargo fmt --all -- --check` +- `git diff --check` + +Remaining after this slice: + +- Token refresh now covers SQLite lease coordination, process-local in-flight + protection, refresh-before-call `refreshing` persistence, restart/crash + recovery from persisted `refreshing`, retry exhaustion delayed recovery, + transient retry success, and two-hit permanent failure confirmation. +- Token refresh still lacks the old per-account timer map; the current + implementation remains scan-driven rather than scheduling one timer per + account at `exp - margin`. +- Live stream audit, WebSocket stream parity, and implicit resume/reasoning + replay cache parity still need migration. + +### 2026-06-19 Slice 44: Responses Stream No-Fallback Coverage + +Status: migrated tests added; existing implementation passed. + +Changes: + +- Added crate-local server tests for Responses `stream: true` no-fallback + exhausted-account aggregation: + - HTTP 429 rate-limit exhaustion returns `text/event-stream` with + `response.failed`, `rate_limit_exceeded`, and the upstream error message; + - HTTP 402 quota exhaustion returns `text/event-stream` with + `response.failed`, the quota-exhausted aggregate message, and the upstream + error message. +- The tests also verify the persisted account side effects: + - HTTP 429 sets quota cooldown state; + - HTTP 402 persists `quota_exhausted`. + +TDD / coverage evidence: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_return_` + passed after adding the migrated tests, showing the behavior was already + restored and the gap was missing crate-local coverage rather than missing + production code. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_return_` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo fmt --all` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_return_` +- `cargo fmt --all -- --check` +- `git diff --check` + +Remaining after this slice: + +- Responses stream no-fallback HTTP 429 and HTTP 402 aggregate error coverage is + now migrated. +- True live chunk-by-chunk stream proxying, live stream usage audit, and + WebSocket stream parity still need migration. +- Token refresh still lacks the old per-account timer map. + +### 2026-06-19 Slice 45: Admin Quota Edge-Case Coverage + +Status: migrated tests added; existing implementation passed. + +Changes: + +- Added crate-local server tests for explicit admin quota edge cases from the + old admin account quota suite: + - inactive accounts return HTTP `409` with code `40901` and do not call the + upstream usage endpoint; + - quota persistence failures after a successful upstream usage fetch return + HTTP `500` with code `50001`; + - quota warnings require an admin session cookie; + - invalid cached quota JSON and below-threshold quota snapshots are ignored + without producing warnings or an `updatedAt` timestamp. +- Existing happy-path and upstream-failure quota route coverage remains in the + same crate-local test target. + +TDD / coverage evidence: + +- The first new edge-case run exposed a bad test input: the secondary + `used_percent` was `89`, above the configured `80` warning threshold. After + correcting the fixture to stay below threshold, the new edge-case tests passed. +- No production code changes were needed for these old quota route semantics; + the gap was missing migrated coverage. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_account_quota` + (first run failed because the new below-threshold fixture used an above-threshold + secondary quota value) +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_account_quota` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test admin_accounts_routes` +- `cargo fmt --all` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_account_quota` +- `cargo fmt --all -- --check` +- `git diff --check` + +Remaining after this slice: + +- Explicit admin quota route happy path, upstream failure shape, inactive-account + rejection, persistence-failure response, quota-warning auth, and quota-warning + invalid/below-threshold edge coverage are migrated. +- Broader old quota/task behavior still needs final parity review, especially + quota refresh task and usage/quota helper edge cases outside the explicit + admin route. +- True live chunk-by-chunk stream proxying, live stream usage audit, WebSocket + stream parity, and token refresh per-account timer parity still need migration. + +### 2026-06-19 Slice 46: Fingerprint Runtime Startup Coverage + +Status: migrated tests added; existing implementation passed. + +Changes: + +- Added runtime task coverage proving `start_fingerprint_update_task(...)` + performs the initial appcast check and writes the `auto_update` fingerprint + row through `FingerprintRepository`. +- Added coverage for matching current appcast version/build: the task checks + appcast but does not persist an `auto_update` row when no update is available. +- Added coverage documenting the current `repository: None` semantics: the task + still checks appcast but has no persistence target. + +TDD / coverage evidence: + +- `cargo test -p codex-proxy-runtime --test tasks fingerprint_update_task` + passed after adding the migrated tests, showing the runtime wrapper already + drove the adapter behavior and the gap was missing crate-local coverage. + +Verification run so far: + +- `cargo test -p codex-proxy-runtime --test tasks fingerprint_update_task` +- `codegraph sync .` +- `cargo test -p codex-proxy-runtime --test tasks fingerprint_update_task` + +Remaining after this slice: + +- Fingerprint update startup/update behavior from the old updater tests is now + covered at adapter and runtime-wrapper levels. +- True live chunk-by-chunk stream proxying, live stream usage audit, WebSocket + stream parity, and token refresh per-account timer parity still need migration. + +### 2026-06-19 Slice 47: Quota Refresh Request Spacing + +Status: implemented and targeted verification passed. + +Changes: + +- Restored the old quota refresh task's request staggering behavior for + multiple quota-locked accounts. +- Added `QuotaRefreshTask::default_request_spacing()` with the old 3-second + default and `with_request_spacing(...)` for fast tests. +- The refresh loop now sleeps between consecutive quota refresh candidates when + the configured spacing is non-zero. + +TDD evidence: + +- Before implementation, `cargo test -p codex-proxy-runtime --test quota_refresh` + failed to compile because `default_request_spacing()` and + `with_request_spacing(...)` did not exist. +- The first implementation still failed the stagger behavior test because it + used `Filter::size_hint()` to detect a remaining candidate; the lower bound + was `0`, so the second request was sent immediately. Switching to + `peekable()` made the test pass. + +Verification run so far: + +- `cargo test -p codex-proxy-runtime --test quota_refresh` + (RED: missing request spacing API) +- `cargo test -p codex-proxy-runtime --test quota_refresh` + (failed once: second request was not delayed) +- `cargo test -p codex-proxy-runtime --test quota_refresh` +- `codegraph sync .` +- `cargo test -p codex-proxy-runtime --test quota_refresh` + +Remaining after this slice: + +- The quota refresh task now covers the old scan interval shape, per-account + minimum refresh window, active quota-locked filtering, usage fetch/persist + path, and inter-account request spacing. +- Token refresh still lacks old per-account timer parity. +- True live chunk-by-chunk stream proxying, live stream usage audit, and + WebSocket stream parity still need migration. + +### 2026-06-19 Slice 48: Assets Static Router And Server Fallback + +Status: implemented and targeted verification passed. + +Changes: + +- Replaced empty `crates/assets` modules with real static asset behavior: + - cache policy helper for SPA HTML vs fingerprinted `/assets/*` files; + - security headers including `x-content-type-options`, `x-frame-options`, + `referrer-policy`, and CSP; + - content type selection for common frontend artifacts; + - Axum SPA router serving `/`, `/assets/{*path}`, and deep-link fallback. +- Added crate-local assets tests for header policy, index serving, SPA fallback, + and cached asset serving. +- Added `codex-proxy-assets` as a server dependency and wired the assets router + as the final fallback behind the OpenAI/admin API routes. +- Added server integration coverage proving `/` and `/assets/*` serve frontend + files while `/api/admin/settings` still resolves to the admin API handler. +- Removed one redundant `serde_json::json!` clone in the Responses stream error + mapping after the final workspace clippy gate exposed it. + +TDD evidence: + +- Before implementation, `cargo test -p codex-proxy-assets` failed because the + assets crate had no `cache_control_for_path(...)` or `spa_router(...)`; it + also exposed a missing `tokio` test dependency. +- After implementing the assets crate, its tests passed. +- Before server wiring, `cargo test -p codex-proxy-server --test assets_routes` + failed because `router_with_assets(...)` did not exist. +- After adding the server dependency and fallback wiring, the server assets + integration test passed. + +Verification run so far: + +- `cargo test -p codex-proxy-assets` (RED: missing API/test dependency) +- `cargo test -p codex-proxy-assets` +- `cargo test -p codex-proxy-server --test assets_routes` + (RED: missing `router_with_assets(...)`) +- `cargo test -p codex-proxy-server --test assets_routes` +- `codegraph sync .` +- `cargo test -p codex-proxy-assets` +- `cargo test -p codex-proxy-server --test assets_routes` +- `cargo fmt --all` +- `cargo test -p codex-proxy-runtime --test tasks` +- `cargo test -p codex-proxy-runtime --test quota_refresh` +- `cargo test -p codex-proxy-assets` +- `cargo test -p codex-proxy-server --test assets_routes` +- `cargo test -p codex-proxy-server --test admin_settings_routes` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo fmt --all -- --check` +- `git diff --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` + (failed once on an existing redundant `json!` clone in + `crates/server/src/openai_api/responses.rs`) +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `cargo test -p codex-proxy-server --test assets_routes` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo fmt --all -- --check` +- `git diff --check` + +Additional source/empty-folder audit in this slice: + +- `rg 'fn ... {}'` over `crates`, `src`, and `tests` found no empty source + functions beyond the architecture placeholder test's own marker string. +- Placeholder scan for `todo!`, `unimplemented!`, `not wired yet`, + `后续承载`, `placeholder`, and `stub` found only architecture guard/doc + references in Rust sources. +- Empty-directory scan found only `web/node_modules/.vite-temp`, an ignored + frontend dependency cache directory, not a migration leftover. + +Remaining after this slice: + +- `assets` is no longer an empty scaffold and is wired into the server router. +- True live chunk-by-chunk stream proxying, live stream usage audit, WebSocket + stream parity, and token refresh per-account timer parity still need migration. + +### 2026-06-19 Slice 49: Token Refresh Per-Account Timers + +Status: implemented and targeted verification passed. + +Changes: + +- Restored an account-level token refresh timer map in + `crates/runtime/src/tasks/token_refresh.rs`, replacing the scan-only loop for + normal background scheduling. +- The scheduler now computes the planned refresh instant from + `access_token_expires_at`/JWT exp and stores one timer per account. +- Already-due active or persisted `refreshing` accounts refresh immediately + during scheduling, matching the old `schedule_one`/crash-recovery behavior + instead of leaving a zero-delay stale timer behind. +- Timer callbacks use the planned trigger timestamp when asking the core + refresh policy whether an account is due. This preserves deterministic tests + and avoids skipping refresh if wall-clock time moves backward after a timer + was scheduled. +- Successful scheduled refreshes immediately schedule the next per-account + timer from the newly persisted token expiry, restoring the old + `schedule_next_refresh` behavior. +- Existing in-flight, SQLite lease, retry, permanent-failure confirmation, and + delayed-recovery logic remains shared by scan-triggered and timer-triggered + refreshes. + +TDD evidence: + +- Before implementation, `cargo test -p codex-proxy-runtime --test token_refresh timer -- --nocapture` + failed: + - `token_refresh_task_should_fire_per_account_timer_at_refresh_time` timed out + because the timer callback recomputed due-ness from real `Utc::now()` rather + than the planned trigger timestamp. + - `token_refresh_task_should_reschedule_next_timer_after_scheduled_refresh` + failed with `left: 0, right: 1`, proving a successful scheduled refresh did + not leave the next account-level timer installed. +- The first implementation attempt exposed a `Send` future compile error from + the recursive shape `timer -> refresh -> schedule next -> timer`; boxing the + timer scheduling future keeps the recursion out of the concrete async type. +- A second run exposed a JWT-vs-stored-expiry precision issue in the new + sub-second timer test. The scheduler now keeps stored + `access_token_expires_at` as the primary expiry source and only writes the new + JWT exp into the refreshed account snapshot before scheduling the next timer. + +Verification run so far: + +- `codegraph sync .` +- `cargo test -p codex-proxy-runtime --test token_refresh timer -- --nocapture` + (RED: timed-out due timer and missing next timer) +- `cargo test -p codex-proxy-runtime --test token_refresh timer -- --nocapture` + (compile failed once on mutable `updated`, then once on recursive async + `Send`, then once on JWT/stored-expiry precision) +- `cargo test -p codex-proxy-runtime --test token_refresh timer -- --nocapture` +- `cargo test -p codex-proxy-runtime --test token_refresh -- --nocapture` +- `cargo fmt --all` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` + (failed once on `large_enum_variant` for `TokenRefreshOutcome::Refreshed(Account)`) +- `cargo fmt --all && cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `cargo test -p codex-proxy-runtime --test token_refresh` +- `cargo test -p codex-proxy-runtime --test tasks` +- `cargo test -p codex-proxy-runtime --test quota_refresh` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo fmt --all -- --check` +- `git diff --check` +- `codegraph sync . && codegraph status .` + (synced 2 changed files; index up to date with 271 files, 3,213 nodes, and + 8,855 edges) + +Remaining after this slice: + +- Token refresh now covers the old per-account timer map, immediate due refresh, + scheduled callback due-time semantics, next-refresh scheduling, process-local + in-flight guard, SQLite lease coordination, retry, delayed recovery, and + permanent-failure confirmation behavior. +- True live chunk-by-chunk stream proxying, live stream usage audit, and + WebSocket stream parity still need migration. + +### 2026-06-19 Slice 50: Responses Stream Done Termination + +Status: implemented and targeted verification passed. + +Changes: + +- All Responses `stream: true` bodies emitted by the server now append the + OpenAI compatible `data: [DONE]` SSE terminator when the body does not + already end with one. +- This covers buffered successful upstream SSE bodies and stream errors + produced by the proxy before a live upstream stream exists, such as no active + account, exhausted fallback, invalid upstream response, and other + `ResponseDispatchError` mappings. +- The change does not make the buffered HTTP SSE path a true live + chunk-by-chunk proxy, but it closes a client-visible stream lifecycle gap: + clients can now terminate cleanly after both buffered success and generated + error streams. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-server --test openai_responses_routes generated_stream_errors -- --nocapture` + failed because the route body ended after `event: response.failed` and did + not include `data: [DONE]\n\n`. +- After the first fix, generated stream errors passed. +- A second RED test, + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_proxy_sse_and_record_usage -- --nocapture`, + proved successful buffered upstream SSE bodies also lacked `[DONE]`. +- Moving the terminator normalization into `event_stream_response(...)` made + both generated error streams and successful buffered streams terminate + consistently without duplicating an existing terminator. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_responses_routes generated_stream_errors -- --nocapture` + (RED: missing `[DONE]` terminator) +- `cargo test -p codex-proxy-server --test openai_responses_routes generated_stream_errors -- --nocapture` +- `cargo test -p codex-proxy-server --test openai_responses_routes` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_proxy_sse_and_record_usage -- --nocapture` + (RED: successful buffered stream missing `[DONE]` terminator) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_proxy_sse_and_record_usage -- --nocapture` +- `cargo test -p codex-proxy-server --test openai_responses_routes generated_stream_errors -- --nocapture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo fmt --all` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo fmt --all -- --check` +- `git diff --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `codegraph sync . && codegraph status .` + (synced 2 changed files; index up to date with 271 files, 3,215 nodes, and + 8,879 edges) + +Remaining after this slice: + +- Responses `stream: true` success and generated error bodies now terminate + cleanly. +- True live chunk-by-chunk stream proxying, live stream usage audit, and + WebSocket stream parity still need migration. + +### 2026-06-19 Slice 51: WebSocket Transport Decision And Audit Snapshots + +Status: implemented and targeted verification passed. + +Changes: + +- Restored the old three-state Responses transport decision in + `core::serving::responses`: + - default Responses requests without history are `WebSocketPreferred`; + - requests with `previous_response_id` are `WebSocketRequired`; + - `force_http_sse` forces `HttpSse` and allows fallback. +- Added `http_sse_fallback_allowed(...)` so the no-fallback history path is a + first-class core policy instead of an implicit runtime/server detail. +- Added core coverage proving `use_websocket` and `force_http_sse` remain + internal routing flags and are not serialized into upstream Codex JSON. +- Restored pure WebSocket payload audit snapshots for Responses + `response.create` payloads, including ordered top-level key reporting and + redaction for instructions, input, previous response ID, prompt cache key, + client metadata, and tool definitions. +- Expanded opening audit snapshots to include a request line and redacted + ordered headers while preserving the existing `header_order` compatibility + field. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-core --test protocol codex_responses_transport -- --nocapture` + failed to compile because `CodexTransport`, `transport_for_request(...)`, and + `http_sse_fallback_allowed(...)` did not exist. +- After adding the transport policy, the four transport/serialization tests + passed and the full core protocol test suite passed. +- Before payload audit implementation, + `cargo test -p codex-proxy-core --test protocol codex_websocket_payload -- --nocapture` + failed to compile because `websocket_payload_audit_snapshot(...)` did not + exist. +- After implementing the snapshot builder/redaction, the payload audit test + passed. +- Before opening audit expansion, + `cargo test -p codex-proxy-adapters --test codex websocket_connection_opening_audit -- --nocapture` + failed to compile because `OpeningAuditSnapshot` had no `request_line` or + `headers` fields. +- After expanding the snapshot type and adapter builder, the opening audit test + and existing WebSocket adapter tests passed. + +Verification run so far: + +- `cargo test -p codex-proxy-core --test protocol codex_responses_transport -- --nocapture` + (RED: missing transport policy API) +- `cargo test -p codex-proxy-core --test protocol codex_responses_transport -- --nocapture` +- `cargo test -p codex-proxy-core --test protocol` +- `cargo test -p codex-proxy-core --test protocol codex_websocket_payload -- --nocapture` + (RED: missing payload audit API) +- `cargo test -p codex-proxy-core --test protocol codex_websocket_payload -- --nocapture` +- `cargo test -p codex-proxy-core --test protocol` +- `cargo test -p codex-proxy-adapters --test codex websocket_connection_opening_audit -- --nocapture` + (RED: missing opening audit fields) +- `cargo test -p codex-proxy-adapters --test codex websocket_connection_opening_audit -- --nocapture` +- `cargo test -p codex-proxy-adapters --test codex websocket_` +- `cargo test -p codex-proxy-core --test protocol` +- `cargo fmt --all` +- `cargo test -p codex-proxy-core --test protocol` +- `cargo test -p codex-proxy-adapters --test codex` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo fmt --all -- --check` +- `git diff --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `codegraph sync . && codegraph status .` + (synced 6 changed files; index up to date with 271 files, 3,242 nodes, and + 8,977 edges) + +Remaining after this slice: + +- WebSocket transport selection and audit snapshot primitives are no longer + empty shells. +- Real `tokio-tungstenite` connection/opening bytes, permessage-deflate, + ping/pong, pooled live sockets, live stream proxying, live stream audit, and + WebSocket serving parity still need migration. + +### 2026-06-19 Slice 52: WebSocket Core Codec Diagnostics + +Status: implemented and targeted verification passed. + +Changes: + +- Migrated more of the old WebSocket codec's pure event semantics into + `crates/core/src/protocol/codex/websocket.rs`: + - public WebSocket event to SSE frame conversion for public JSON events; + - `response.metadata` turn-state extraction from case-insensitive + `x-codex-turn-state` headers, including array-valued headers; + - terminal event detection for `response.completed`, `response.failed`, and + `error`; + - upstream WebSocket error-frame classification to HTTP-style status codes + for rate limit, quota, auth, forbidden, bad request, overload, and + WebSocket connection-limit codes; + - retry-after extraction from wrapped WebSocket `error.headers` values; + - `response.completed` shape validation against the old typed usage/end-turn + expectations; + - basic official stream-event shape diagnostics; + - missing-required-field diagnostics for `response.created`, + `response.output_text.delta`, `response.output_item.*`, and message output + items. +- Kept this slice in `core` as pure JSON/protocol logic. It does not introduce + socket IO, TLS, deflate, pooling, filesystem audit writes, or adapter + dependencies. +- Added crate-local protocol tests for the migrated semantics instead of + relying on the deleted root `tests/codex_gateway/websocket.rs` coverage. + +TDD evidence: + +- Before the first implementation batch, + `cargo test -p codex-proxy-core --test protocol codex_websocket_ -- --nocapture` + failed to compile because the new metadata, terminal-event, error + classification, wrapped retry-after, and completed-response parse APIs did + not exist. +- After the first implementation batch, the same test command found a real + implementation bug: the `response.completed` usage field was not mapped onto + the private deserialization struct, so the invalid-usage test failed. After + fixing the serde rename, the filtered test passed 8/8. +- Before the second diagnostics batch, the same filtered command failed to + compile because the event-shape and output-item diagnostic APIs did not exist. +- After implementing those diagnostics, the filtered command passed 12/12. + +Verification run so far: + +- `cargo test -p codex-proxy-core --test protocol codex_websocket_ -- --nocapture` + (RED: missing first codec API batch) +- `cargo test -p codex-proxy-core --test protocol codex_websocket_ -- --nocapture` + (GREEN after serde rename fix: 8/8 filtered tests passed) +- `cargo test -p codex-proxy-core --test protocol codex_websocket_ -- --nocapture` + (RED: missing second diagnostics API batch) +- `cargo test -p codex-proxy-core --test protocol codex_websocket_ -- --nocapture` + (GREEN: 12/12 filtered tests passed) +- `cargo fmt --all` +- `cargo test -p codex-proxy-core --test protocol` + (36/36 tests passed) +- `codegraph sync .` + (synced 2 changed files) +- `codegraph status .` + (index up to date with 271 files, 3,289 nodes, and 9,059 edges) + +Remaining after this slice: + +- WebSocket codec helpers are materially less skeletal, but they are not yet + wired into a live WebSocket response path in the new crate split. +- The remaining old codec diagnostics still need migration, especially + incomplete-response handling, custom/function/tool/local-shell/web-search/ + image-generation/compaction item validators, and reasoning summary index + checks. +- Real `tokio-tungstenite` connection/opening bytes, permessage-deflate, + ping/pong, pooled live sockets, live stream proxying, live stream audit, and + WebSocket serving parity still need migration. + +### 2026-06-19 Slice 53: WebSocket Core Item Validators + +Status: implemented and verification passed for the current batch. + +Changes: + +- Migrated the next old WebSocket codec diagnostics batch into + `crates/core/src/protocol/codex/websocket.rs`: + - `response.incomplete` reason extraction; + - missing `response` detection for `response.completed`; + - official required-field checks for custom-tool, reasoning-summary, and + reasoning-text delta events; + - output-item metadata validation; + - agent-message, reasoning, function-call, function-call-output, + custom-tool-call, custom-tool-call-output, tool-search-call, + tool-search-output, local-shell-call, web-search-call, + image-generation-call, compaction, and reasoning-summary-part validators. +- Added protocol tests covering each migrated validator category with concrete + old-style bad-frame examples. +- Kept all logic in `core` as pure `serde_json::Value` inspection, so adapter + and runtime wiring can reuse it without bringing IO or transport dependencies + into the protocol crate. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-core --test protocol codex_websocket_ -- --nocapture` + failed to compile because the new validator APIs did not exist. +- After implementation, the same filtered command passed 19/19 WebSocket tests. + +Verification run so far: + +- `cargo test -p codex-proxy-core --test protocol codex_websocket_ -- --nocapture` + (RED: missing second validator batch) +- `cargo test -p codex-proxy-core --test protocol codex_websocket_ -- --nocapture` + (GREEN: 19/19 filtered tests passed) +- `cargo fmt --all` +- `cargo test -p codex-proxy-core --test protocol` + (43/43 tests passed) +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` + (22/22 tests passed) +- `cargo fmt --all -- --check` +- `git diff --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `codegraph sync .` + (synced 2 changed files) +- `codegraph status .` + (index up to date with 271 files, 3,324 nodes, and 9,144 edges) + +Remaining after this slice: + +- Most of the old codec's pure bad-frame diagnostics are now represented in + `core`, but they are not yet wired into live WebSocket response forwarding. +- The old WebSocket request/response loop, real `tokio-tungstenite` + connection/opening bytes, permessage-deflate, ping/pong handling, pooled live + sockets, audit artifact IO, live stream proxying, and WebSocket serving parity + still need migration. + +### 2026-06-19 Slice 54: WebSocket Opening Request Descriptor + +Status: implemented and verification passed for the current batch. + +Changes: + +- Added adapter-level Responses WebSocket endpoint construction: + - `https://...` backend bases map to `wss://.../codex/responses`; + - `http://...` backend bases map to `ws://.../codex/responses`; + - existing backend base paths are preserved. +- Added `CodexWebSocketConnection::responses(...)` to build the standard + opening descriptor around business headers: + `Host`, `Connection`, `Upgrade`, `Sec-WebSocket-Version`, + `Sec-WebSocket-Key`, caller-provided business headers, then + `sec-websocket-extensions`. +- Added `CodexWebSocketRequest` and + `CodexWebSocketConnection::responses_create_request(...)` so adapter code can + hold the opening descriptor and first `response.create` text frame together. +- Exposed `websocket_response_create_payload(...)` and + `websocket_response_create_payload_text(...)` from `core` so the adapter does + not duplicate payload construction. +- Added adapter tests for endpoint conversion, opening header ordering, + redaction through the existing opening audit snapshot, and response-create + payload text generation. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-adapters --test codex websocket_ -- --nocapture` + failed to compile because `responses_websocket_endpoint(...)` and + `CodexWebSocketConnection::responses(...)` did not exist. +- After implementing the opening descriptor, the filtered adapter WebSocket + tests passed 5/5. +- Before adding the payload descriptor implementation, the same filtered command + failed to compile because + `CodexWebSocketConnection::responses_create_request(...)` did not exist. +- After exposing core payload text construction and adding the request + descriptor, the filtered adapter WebSocket tests passed 6/6. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test codex websocket_ -- --nocapture` + (RED: missing endpoint/opening descriptor API) +- `cargo test -p codex-proxy-adapters --test codex websocket_ -- --nocapture` + (GREEN: 5/5 filtered tests passed) +- `cargo test -p codex-proxy-adapters --test codex websocket_ -- --nocapture` + (RED: missing response-create request descriptor API) +- `cargo test -p codex-proxy-adapters --test codex websocket_ -- --nocapture` + (GREEN: 6/6 filtered tests passed) +- `cargo fmt --all` +- `cargo test -p codex-proxy-core --test protocol` + (43/43 tests passed) +- `cargo test -p codex-proxy-adapters --test codex` + (23/23 tests passed) +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` + (22/22 tests passed) +- `cargo fmt --all -- --check` +- `git diff --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `codegraph sync .` + (synced 3 changed files) +- `codegraph status .` + (index up to date with 271 files, 3,337 nodes, and 9,176 edges) + +Remaining after this slice: + +- The adapter can now describe the opening request and first text frame, but it + still does not open a network WebSocket, drive ping/pong, decompress frames, + pool live sockets, or forward events into the runtime/server response path. + +### 2026-06-19 Slice 55: Minimal Live WebSocket Exchange + +Status: implemented and verification passed for the current batch. + +Changes: + +- Added `futures` and `tokio-tungstenite` to the adapters crate using existing + workspace dependency versions. +- Added `execute_response_create_request(...)` for a prepared Responses + WebSocket request: + - builds a tungstenite opening request from the ordered descriptor; + - opens the WebSocket; + - sends the first `response.create` text frame; + - collects public text events into SSE frames until a terminal event; + - extracts completed usage from the collected SSE body; + - captures metadata turn-state frames without forwarding them into SSE. +- Added `CodexWebSocketExchange` and `CodexWebSocketExchangeError`. +- Restored the old response-failed error behavior for this one-shot adapter + path: + - classified WebSocket `response.failed` / `error` frames through the core + classifier; + - surfaced HTTP-style status codes; + - extracted retry-after from wrapped headers or official rate-limit error + messages. +- Added adapter tests with a local WebSocket server for: + - successful `response.completed` collection into SSE and usage extraction; + - `response.failed` rate-limit classification into an upstream error. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request -- --nocapture` + failed to compile because `futures`, `tokio_tungstenite`, and + `execute_response_create_request(...)` were missing. +- After adding the live one-shot exchange, the successful exchange test passed. +- Before adding upstream-error handling, + `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request_should_surface_response_failed -- --nocapture` + failed to compile because `CodexWebSocketExchangeError::Upstream` did not + exist. +- After wiring core error classification and retry-after extraction, the + response-failed test passed. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request -- --nocapture` + (RED: missing live dependency/function) +- `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request -- --nocapture` + (GREEN: success path passed) +- `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request_should_surface_response_failed -- --nocapture` + (RED: missing upstream error variant) +- `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request_should_surface_response_failed -- --nocapture` + (GREEN: response-failed error path passed) +- `cargo fmt --all` +- `cargo test -p codex-proxy-core --test protocol` + (43/43 tests passed) +- `cargo test -p codex-proxy-adapters --test codex` + (25/25 tests passed) +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` + (22/22 tests passed) +- `cargo fmt --all -- --check` +- `git diff --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `codegraph sync .` + (synced 2 changed files) +- `codegraph status .` + (index up to date with 271 files, 3,356 nodes, and 9,224 edges) + +Remaining after this slice: + +- A minimal adapter-level live WebSocket exchange exists, but it is not wired + into `CodexBackendClient`, runtime dispatch, or server routes. +- The current live adapter path is one-shot and buffered; it does not yet expose + a streaming body, pooled live sockets, ping/pong recovery, permessage-deflate + handling, audit artifact IO, or runtime fallback integration. + +### 2026-06-19 Slice 56: Codex Backend WebSocket Required Path + +Status: implemented and verification passed for the current batch. + +Changes: + +- Wired `CodexBackendClient::create_response(...)` through the core Responses + transport policy: + - `HttpSse` keeps the existing HTTP/SSE POST path; + - `WebSocketRequired` uses the adapter live WebSocket exchange; + - `WebSocketPreferred` attempts WebSocket and falls back to HTTP/SSE when the + core policy allows fallback. +- Mapped WebSocket `Upstream` errors back into the existing + `CodexClientError::Upstream` shape so runtime fallback/recovery code can keep + matching status code and retry-after semantics. +- Added WebSocket request encoding and generic WebSocket transport errors to + `CodexClientError`. +- Added a client-level adapter test proving a `CodexResponsesRequest` with + `previous_response_id` opens a local WebSocket server, sends + `response.create` with `previous_response_id`, collects `response.completed` + into SSE, and extracts usage. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-adapters --test codex codex_backend_client_should_use_websocket_when_previous_response_id_is_present -- --nocapture` + failed because the client sent an HTTP POST to the local WebSocket server. The + server rejected the request with `Protocol(WrongHttpMethod)` and the client + returned a reqwest incomplete-message error. +- After wiring the transport policy and WebSocket path, the same test passed. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test codex codex_backend_client_should_use_websocket_when_previous_response_id_is_present -- --nocapture` + (RED: client still used HTTP POST) +- `cargo test -p codex-proxy-adapters --test codex codex_backend_client_should_use_websocket_when_previous_response_id_is_present -- --nocapture` + (GREEN) +- `cargo test -p codex-proxy-adapters --test codex` + (26/26 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream previous_response -- --nocapture` + (3/3 filtered tests passed) +- `cargo test -p codex-proxy-core --test protocol` + (43/43 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (54/54 tests passed) +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` + (22/22 tests passed) +- `cargo fmt --all` +- `cargo fmt --all -- --check` +- `git diff --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `codegraph sync .` + (synced 2 changed files) +- `codegraph status .` + (index up to date with 271 files, 3,365 nodes, and 9,276 edges) + +Remaining after this slice: + +- Non-streaming `previous_response_id` requests can now take the required + WebSocket path in the adapter client, but the runtime/server surface is still + buffered and not a live streaming proxy. +- WebSocketPreferred fallback is wired, but full old parity still needs + streaming body forwarding, validator integration in the live loop, + permessage-deflate, ping/pong recovery, pooled live sockets, audit artifact + IO, and the broader old WebSocket behavior-test migration. + +### 2026-06-19 Slice 57: WebSocket Live Invalid-Event Filtering + +Status: implemented and verification passed for the current batch. + +Changes: + +- Wired the core WebSocket bad-frame validators into the live event-to-SSE path. +- `websocket_event_to_sse_frame(...)` now skips invalid public stream events + before encoding an SSE frame, matching the old transport behavior that + ignored malformed/non-forwardable WebSocket events instead of forwarding them + to clients. +- The skip decision reuses the migrated core diagnostics for metadata events, + missing `response` objects, missing delta payloads, output item required + fields, reasoning summary indexes, and the old agent/function/tool/local + shell/web/image/compaction validator categories. +- Added an adapter live WebSocket regression test proving an invalid + `response.output_text.delta` frame is dropped while a following + `response.completed` terminal frame still completes the exchange. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request_should_skip_invalid_stream_events -- --nocapture` + failed because the invalid delta frame was forwarded into the buffered SSE + body. +- After wiring the core skip decision into event-to-SSE conversion, the same + test passed. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request_should_skip_invalid_stream_events -- --nocapture` + (RED: invalid delta frame was forwarded) +- `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request_should_skip_invalid_stream_events -- --nocapture` + (GREEN) +- `cargo fmt --all` +- `cargo test -p codex-proxy-core --test protocol` + (43/43 tests passed) +- `cargo test -p codex-proxy-adapters --test codex` + (27/27 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (54/54 tests passed) +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` + (22/22 tests passed) +- `cargo fmt --all -- --check` +- `git diff --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `codegraph sync .` + (synced 2 changed files) +- `codegraph status .` + (index up to date with 271 files, 3,367 nodes, and 9,295 edges) + +Remaining after this slice: + +- The live WebSocket path now avoids forwarding the migrated invalid-event + categories, but it remains buffered and one-shot rather than a live streaming + body proxy. +- Full old parity still needs streaming body forwarding, permessage-deflate, + ping/pong recovery, pooled live sockets, audit artifact IO, and broader old + WebSocket behavior test migration. + +### 2026-06-19 Slice 58: WebSocket Live Control Frames And Handshake Metadata + +Status: implemented and verification passed for the current batch. + +Changes: + +- Extended the minimal live WebSocket exchange to preserve handshake response + metadata: + - `x-codex-turn-state`; + - `set-cookie`; + - retry/rate-limit style headers. +- Mapped those handshake values through `CodexBackendClient::create_response(...)` + so required WebSocket requests expose the same metadata shape as HTTP/SSE + responses. +- Added live handling for `response.incomplete`, returning a typed + `IncompleteResponse { reason }` error. +- Added live validation for malformed `response.completed` payloads using the + core completed-response parser. +- Changed terminal handling so a skipped invalid terminal frame, such as + `response.completed` without `response`, does not end the exchange + successfully. +- Matched the old behavior for unclassified success-status `error` frames by + ignoring them until a real terminal frame or close-before-terminal error. +- Added adapter tests for handshake metadata, `response.incomplete`, malformed + `response.completed`, missing-response completed frames, success-status + `error` frames, and client-level WebSocket metadata propagation. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request_should_ -- --nocapture` + failed to compile because `CodexWebSocketExchange` did not expose + `set_cookie_headers` / `rate_limit_headers`, and + `CodexWebSocketExchangeError` did not expose `IncompleteResponse` / + `InvalidCompletedResponse`. +- After implementing the live-loop and metadata changes, the same filtered + command passed 8/8 tests. +- The first full clippy run found the new `accept_hdr_async` test callback + matched the existing large-error lint pattern. The test was annotated with a + local `#[expect(clippy::result_large_err)]` and clippy then passed. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request_should_ -- --nocapture` + (RED: missing exchange fields and error variants) +- `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request_should_ -- --nocapture` + (GREEN: 8/8 filtered tests passed) +- `cargo fmt --all` +- `cargo test -p codex-proxy-adapters --test codex` + (32/32 tests passed) +- `cargo test -p codex-proxy-core --test protocol` + (43/43 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream previous_response -- --nocapture` + (3/3 filtered tests passed) +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` + (22/22 tests passed) +- `cargo fmt --all -- --check` +- `git diff --check` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (54/54 tests passed) +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` + (first run failed on the test callback large-error lint; rerun passed after a + local `#[expect]`) +- `cargo test -p codex-proxy-adapters --test codex codex_backend_client_should_use_websocket_when_previous_response_id_is_present -- --nocapture` + (client-level WebSocket metadata propagation passed) +- `cargo test -p codex-proxy-adapters --test codex` + (32/32 tests passed after the client-level assertion update) +- `cargo fmt --all -- --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- `codegraph sync . && codegraph status .` + (synced 5 changed files; index up to date with 272 files, 3,656 nodes, and + 10,555 edges before this documentation-only update) +- `codegraph sync . && codegraph status .` + (synced 9 changed files; index up to date with 272 files, 3,652 nodes, and + 10,554 edges before this documentation-only update) +- `codegraph sync . && codegraph status .` + (index up to date with 271 files, 3,378 nodes, and 9,334 edges) + +Remaining after this slice: + +- The one-shot live exchange now restores more old control-frame and metadata + behavior, but it is still buffered rather than a live streaming body proxy. +- Full old parity still needs streaming body forwarding, permessage-deflate, + ping/pong recovery, pooled live sockets, audit artifact IO, and broader old + WebSocket behavior test migration. + +### 2026-06-19 Slice 59: WebSocket Ping Coverage And Audit Artifact IO + +Status: implemented and verification passed for the current batch. + +Changes: + +- Migrated old server-ping coverage into `crates/adapters/tests/codex.rs`: + the local WebSocket server sends a ping before `response.completed`, waits for + the client pong, then completes the response. The test passed with the current + `tokio-tungstenite` read loop, so no extra ping production code was needed + for this one-shot path. +- Restored structured WebSocket audit artifact data in core: + - `WebSocketAuditArtifact`; + - `WebSocketAuditErrorSnapshot`; + - structured `WebSocketParityDiff` / `WebSocketParityDifference`; + - `websocket_audit_artifact_from_attempt(...)`; + - `websocket_parity_diff(...)`. +- Added explicit adapter IO for audit artifacts through + `write_websocket_audit_artifact_for_dir(...)`, with no file writes unless a + caller supplies a non-empty directory. +- Added environment-gated client wiring through `CODEX_PROXY_WS_AUDIT_DIR` so + WebSocket attempts can emit the redacted opening/payload artifact when + explicitly enabled. +- Kept artifact construction pure in `core`; filesystem writes stay in the + adapter layer. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-core --test protocol codex_websocket_audit -- --nocapture` + failed to compile because the artifact/diff types and builders did not exist. +- Before implementation, + `cargo test -p codex-proxy-adapters --test codex websocket_audit_artifact_should_require_explicit_directory -- --nocapture` + failed to compile because the core artifact types and adapter write function + did not exist. +- After implementation, the core audit artifact, core parity diff, and adapter + explicit-write tests passed. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test codex websocket_execute_response_create_request_should_reply_to_server_ping_before_terminal -- --nocapture` + (server ping coverage passed) +- `cargo test -p codex-proxy-core --test protocol codex_websocket_audit -- --nocapture` + (RED: missing artifact/diff APIs) +- `cargo test -p codex-proxy-adapters --test codex websocket_audit_artifact_should_require_explicit_directory -- --nocapture` + (RED: missing artifact types/write function) +- `cargo test -p codex-proxy-core --test protocol codex_websocket_audit -- --nocapture` + (GREEN: artifact test passed) +- `cargo test -p codex-proxy-core --test protocol codex_websocket_parity -- --nocapture` + (GREEN: parity diff test passed) +- `cargo test -p codex-proxy-adapters --test codex websocket_audit_artifact_should_require_explicit_directory -- --nocapture` + (GREEN) +- `cargo test -p codex-proxy-core --test protocol codex_websocket_ -- --nocapture` + (21/21 filtered WebSocket tests passed) +- `cargo test -p codex-proxy-adapters --test codex codex_backend_client_should_use_websocket_when_previous_response_id_is_present -- --nocapture` + (client WebSocket path passed with audit wiring present) +- `cargo fmt --all` +- `cargo test -p codex-proxy-core --test protocol` + (45/45 tests passed) +- `cargo test -p codex-proxy-adapters --test codex` + (34/34 tests passed) +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` + (22/22 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (54/54 tests passed) +- `cargo fmt --all -- --check` +- `git diff --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `codegraph sync . && codegraph status .` + (index up to date with 271 files, 3,404 nodes, and 9,433 edges) + +Remaining after this slice: + +- The one-shot live WebSocket path now has server-ping response coverage and + environment-gated audit artifact IO. +- It is still buffered rather than a live streaming body proxy. +- Full old parity still needs streaming body forwarding, live + permessage-deflate integration, pooled socket reuse/keepalive, and broader + old WebSocket behavior test migration. + +### 2026-06-19 Slice 60: WebSocket Deflate Frame Rewriter + +Status: implemented and verification passed for the current batch. + +Changes: + +- Added `flate2` to `codex-proxy-adapters`. +- Expanded `crates/adapters/src/codex/websocket/deflate.rs` beyond extension + detection: + - raw permessage-deflate payload inflation; + - WebSocket server frame parsing; + - compressed RSV1 data-frame rewriting into a normal uncompressed server + frame; + - uncompressed/partial/non-data frame no-op handling. +- Added adapter tests proving compressed server text frames are rewritten with + RSV1 cleared and original JSON payload restored, while uncompressed frames are + left unchanged. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-adapters --test codex websocket_deflate -- --nocapture` + failed to compile because `flate2` was not an adapters dependency and + `rewrite_permessage_deflate_server_frame(...)` did not exist. +- The first implementation attempt failed the compressed-frame test because a + single `Decompress::decompress_vec` call produced an empty payload for the + test fixture. The root cause was incomplete decoder driving, so the inflater + was changed to `DeflateDecoder::read_to_end(...)`. +- The test fixture then used `DeflateEncoder` to produce a real raw deflate + payload, and the filtered deflate tests passed 2/2. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test codex websocket_deflate -- --nocapture` + (RED: missing dependency/API) +- `cargo test -p codex-proxy-adapters --test codex websocket_deflate -- --nocapture` + (intermediate failure: empty inflated payload; fixed by full decoder read and + a real raw-deflate test fixture) +- `cargo test -p codex-proxy-adapters --test codex websocket_deflate -- --nocapture` + (GREEN: 2/2 filtered tests passed) +- `cargo fmt --all` +- `cargo test -p codex-proxy-adapters --test codex` + (36/36 tests passed) +- `cargo test -p codex-proxy-core --test protocol` + (45/45 tests passed) +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` + (22/22 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (54/54 tests passed) +- `cargo fmt --all -- --check` +- `git diff --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `find crates tests src -path '*/target' -prune -o -type f -empty -print | sort` + (no empty files reported) +- `rg -n "fn main\(\) \{\}|not wired yet|后续承载|todo!\(|unimplemented!\(|TODO\b|panic!\(\"TODO|stub|placeholder" crates src tests docs/architecture-audit.md` + (matches were limited to this audit document and the architecture + placeholder guard/test names) +- `codegraph sync . && codegraph status .` + (index up to date with 271 files, 3,422 nodes, and 9,482 edges) + +Remaining after this slice: + +- Deflate frame-level rewrite logic is no longer a one-line shell, but it is + not yet inserted into the live `tokio-tungstenite` connection path. Because + `tokio-tungstenite` rejects RSV1 frames before yielding messages, live + integration still needs a TLS-aware stream insertion point or a custom + connector. +- The WebSocket path is still buffered rather than a live streaming body proxy. +- Full old parity still needs streaming body forwarding, live + permessage-deflate integration, pooled socket reuse/keepalive, and broader + old WebSocket behavior test migration. + +### 2026-06-19 Slice 61: Live HTTP SSE Body Forwarding + +Status: implemented and targeted verification passed for this batch. + +Changes: + +- Added a live HTTP SSE response path to `CodexBackendClient`: + - `CodexBackendSseStream`; + - `CodexBackendStreamingResponse`; + - `CodexBackendClient::create_response_stream(...)`. +- Kept non-2xx upstream responses on the existing typed `CodexClientError::Upstream` + path, including capped error body reads and retry-after extraction, so HTTP + `429`/`402`/`401`/`403`/`5xx` fallback remains before any downstream bytes are + sent. +- Changed `ResponseDispatchService::stream(...)` from a buffered `String` result + to `ResponseDispatchStream`, a runtime-owned byte stream that does not expose + axum types. +- Added first-event SSE prefetch in runtime before the live stream is returned. + This preserves existing fallback/recovery when the first upstream event is a + recoverable `response.failed` while allowing normal delta streams to be + forwarded before upstream completion. +- Moved stream usage/session-affinity finalization to stream end: + - token usage is extracted from the observed SSE body and recorded after the + upstream stream ends; + - session affinity is recorded from `response.completed`; + - the runtime account-pool slot is held until the stream ends instead of being + released immediately after response headers; + - `[DONE]` is appended for live streams when the upstream body omits it. +- Wired the server Responses handler to return `Body::from_stream(...)` for + successful streaming Responses while keeping generated SSE error responses on + the existing string path. +- Added a server integration regression test with a raw chunked upstream that + writes the first SSE event, keeps the connection open, and only later sends + `response.completed`. The test proves the proxy returns the HTTP response and + first SSE chunk before upstream completion. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_forward_first_chunk_before_upstream_completes -- --nocapture` + failed as expected with: + `stream response should be returned before upstream completes: Elapsed(())`. +- After implementing the live stream path, the same test passed. +- The existing `responses_stream_should...` batch stayed green, including SSE + failure fallback, 429/402/401/403 fallback, 5xx retry, history recovery, usage + persistence, and generated stream error termination coverage. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_forward_first_chunk_before_upstream_completes -- --nocapture` + (RED: timeout because the old implementation waited for upstream completion) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_forward_first_chunk_before_upstream_completes -- --nocapture` + (GREEN: 1/1 filtered test passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should -- --nocapture` + (16/16 filtered stream tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (55/55 tests passed) +- `cargo fmt --all` +- `cargo check --workspace --all-targets` +- `codegraph sync . && codegraph status .` + (synced 4 changed files; index up to date with 271 files, 3,456 nodes, and + 9,738 edges) + +Remaining after this slice: + +- HTTP SSE `stream: true` now has a live chunk-by-chunk success path plus + first-event fallback preservation, but once normal downstream bytes have been + sent, later upstream `response.failed` events cannot be retried on another + account without violating HTTP streaming semantics. Late-failure stream audit + behavior still needs explicit policy and tests. +- WebSocket Responses are still buffered in the one-shot exchange path. +- Full WebSocket parity still needs live WebSocket body forwarding, live + permessage-deflate integration, pooled socket reuse/keepalive, and broader old + behavior-test migration. + +### 2026-06-19 Slice 62: WebSocket Pool Reuse And Error Discard + +Status: implemented and adapter verification passed for this batch. + +Changes: + +- Replaced the WebSocket pool policy-only struct with a real async pool state: + - `CodexWebSocketPoolKey` keyed by normalized base URL, account id, and + conversation id; + - idle/busy slots guarded by `tokio::sync::Mutex`; + - max-per-account admission checks; + - max-age recycling for idle connections; + - `acquire`, `put`, `discard`, and explicit `gc_sweep`. +- Added pooled connection metadata and a concrete `CodexWsStream` type for + `tokio-tungstenite` WebSocket streams. +- Added `CodexBackendClient::with_websocket_pool(...)` and optional pool key + derivation from request/context data. Pooling is only used when both an + upstream account id and a conversation key are available; otherwise the + existing one-shot WebSocket path is preserved. +- Split the WebSocket execution path so fresh and reused connections share the + same response collector: + - successful terminal responses return the socket to the pool; + - upstream error frames, transport errors, and invalid terminal behavior + discard the pool slot instead of returning a bad socket; + - bypass remains available when a key is already busy or account capacity is + exhausted. +- Migrated adapter tests proving: + - two same-account/same-conversation Responses requests reuse one accepted + upstream WebSocket connection; + - a pooled upstream error discards the connection, so the next request opens a + fresh upstream WebSocket and succeeds. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-adapters --test codex pooled_websocket -- --nocapture` + failed because `CodexBackendClient::with_websocket_pool(...)` did not exist. +- The first implementation attempt exposed a test setup false positive: the + requests did not pass `account_id`, so the pool was correctly bypassed. The + tests were corrected to pass the same account id, proving actual pool usage. +- After implementation and test correction, the filtered pooled WebSocket tests + passed 2/2. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test codex pooled_websocket -- --nocapture` + (RED: missing `with_websocket_pool` API) +- `cargo test -p codex-proxy-adapters --test codex pooled_websocket -- --nocapture` + (intermediate failure: test omitted account id and bypassed the pool; fixed) +- `cargo test -p codex-proxy-adapters --test codex pooled_websocket -- --nocapture` + (GREEN: 2/2 filtered tests passed) +- `cargo test -p codex-proxy-adapters --test codex` + (38/38 tests passed) +- `cargo fmt --all` +- `cargo check --workspace --all-targets` +- `cargo test -p codex-proxy-core --test protocol` + (45/45 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (55/55 tests passed) +- `cargo test --test architecture` + (22/22 tests passed) +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` + (first run failed on a `needless_collect` in pool GC and too-many-arguments in + the live SSE finalizer helpers; both were refactored, and the rerun passed) +- `cargo test -p codex-proxy-adapters --test codex` + (38/38 tests passed after the clippy refactor) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (55/55 tests passed after the clippy refactor) +- `cargo test --test architecture` + (22/22 tests passed after the clippy refactor) +- `cargo fmt --all -- --check` +- `git diff --check` +- `codegraph sync . && codegraph status .` + (final sync after lint cleanup synced 2 changed files; index up to date with + 271 files, 3,493 nodes, and 9,871 edges) + +Remaining after this slice: + +- WebSocket pooling now reuses successful same-conversation sockets and discards + errored sockets, but it does not yet have the old background maintenance task, + ping probe/timeout keepalive, liveness timeout, or shutdown/evict-account API. +- WebSocket Responses are still a buffered one-shot response collector from the + server/runtime perspective; live WebSocket body forwarding still needs a + streaming API shape. +- Live permessage-deflate integration is still not inserted into the + `tokio-tungstenite` stream path. + +### 2026-06-19 Slice 63: WebSocket Pool Lifecycle And Runtime Wiring + +Status: implemented and verified for this batch. + +Changes: + +- Restored the adapter pool lifecycle surface from the old WebSocket pool: + - added `CodexWebSocketPoolConfig` with `enabled`, `max_age`, + `max_per_account`, `maintenance_interval`, `ping_interval`, + `ping_timeout`, and `liveness_timeout`; + - preserved the existing `CodexWebSocketPool::new(max_per_account, max_age)` + test-friendly constructor without background maintenance; + - added `with_config(...)`, `with_default_max_age()`, and `with_limits(...)`; + - added `evict_account(...)`, `shutdown()`, and + `maintain_idle_connections()`; + - kept `gc_sweep()` as the public maintenance entry point while expanding it + beyond age-only cleanup. +- Added real idle lifecycle state: + - pool state now tracks `shutting_down`; + - slots include `Checking` so a socket under keepalive probe cannot be reused + concurrently; + - pooled connections track `last_activity_at` and `last_ping_at`. +- Restored idle maintenance behavior: + - max-age and liveness-timeout cleanup close idle sockets; + - ping keepalive sends a WebSocket `Ping`, waits up to `ping_timeout` for an + upstream response, keeps the socket on success, and closes/discards it on + timeout, close, or transport error; + - optional background maintenance starts when a Tokio runtime is available and + `maintenance_interval` is configured. +- Wired platform `ws_pool` config into runtime Codex client construction: + - `codex_proxy_runtime::upstream::codex_backend_client(...)` now accepts + `&WebSocketPoolConfig`; + - `Services::with_installation_id(...)` passes `config.ws_pool`; + - enabled configs inject a shared `CodexWebSocketPool` into + `CodexBackendClient`. + +TDD evidence: + +- Before implementation, + `cargo test -p codex-proxy-adapters --test codex websocket_pool -- --nocapture` + still only ran the old policy test, proving the lifecycle tests were not yet + covered by the existing filter. +- After adding lifecycle tests, the correct RED command + `cargo test -p codex-proxy-adapters --test codex websocket_pool -- --nocapture` + failed at compile time because `CodexWebSocketPoolConfig` and + `CodexWebSocketPool::with_config(...)` did not exist. +- After the initial adapter implementation, + `cargo test -p codex-proxy-adapters --test codex pooled_websocket -- --nocapture` + passed the manual keepalive/evict/shutdown/liveness tests but exposed a + background-maintenance test design issue: a 1 ms ping interval let the + background task repeatedly move the socket into `Checking`. The test was + narrowed to one expected background probe by using a long post-probe + `ping_interval`, then passed. +- Runtime wiring RED: + `cargo test -p codex-proxy-runtime --test upstream codex_backend_client_should_apply_configured_websocket_pool -- --nocapture` + failed because `codex_backend_client(...)` accepted only `base_url` and + `fingerprint`, with no `ws_pool` config argument. +- Runtime wiring GREEN: the same runtime upstream test passed after mapping + platform `WebSocketPoolConfig` to adapter `CodexWebSocketPoolConfig` and + injecting the pool. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test codex websocket_pool -- --nocapture` + (RED: missing pool config/constructor API) +- `cargo test -p codex-proxy-adapters --test codex pooled_websocket -- --nocapture` + (GREEN: 7/7 filtered pooled WebSocket tests passed) +- `cargo test -p codex-proxy-runtime --test upstream codex_backend_client_should_apply_configured_websocket_pool -- --nocapture` + (RED: runtime backend client constructor did not accept `ws_pool` config) +- `cargo test -p codex-proxy-runtime --test upstream -- --nocapture` + (GREEN: 1/1 runtime upstream wiring test passed) +- `cargo test -p codex-proxy-adapters --test codex` + (43/43 tests passed) +- `cargo check --workspace --all-targets` + (first run showed one warning caused by a temporarily removed test attribute; + fixed by restoring `#[tokio::test]`) +- `cargo check --workspace --all-targets` + (rerun passed without warnings) +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (55/55 tests passed) +- `cargo test --test architecture` + (22/22 tests passed) +- `codegraph sync .` + (synced 6 changed files; 1 added and 5 modified; 496 nodes in 671 ms) +- `codegraph status .` + (index up to date with 272 files, 3,535 nodes, and 10,058 edges) + +Remaining after this slice: + +- Adapter-level WebSocket pool reuse, error discard, keepalive maintenance, + liveness cleanup, shutdown/evict-account behavior, and runtime config + injection are restored. +- WebSocket Responses are still buffered through the one-shot exchange from the + server/runtime perspective; live WebSocket body forwarding still needs a + streaming API shape. +- Live permessage-deflate rewriting exists as tested frame helpers, but it is + still not inserted into the `tokio-tungstenite` stream path. +- Broader old WebSocket behavior-test migration remains incomplete, including + server/runtime dispatch parity for live WebSocket routes. + +### 2026-06-19 Slice 64: Server Responses WebSocketRequired Path + +Status: implemented and verified for this batch. + +Changes: + +- Removed the unconditional `force_http_sse = true` from OpenAI Responses + translation for non-streaming requests with `previous_response_id`. + Non-streaming history-continuation requests now reach + `CodexTransport::WebSocketRequired` instead of being forced through HTTP SSE. +- Kept current live streaming behavior intentionally conservative: + `stream: true` requests with `previous_response_id` still force HTTP SSE until + live WebSocket streaming is implemented. +- After history recovery strips stale `previous_response_id` state, the runtime + now also clears `use_websocket` and sets `force_http_sse = true` so the retry + matches a normal no-history OpenAI Responses request. +- Migrated server/runtime tests that previously used HTTP-only mocks for + non-streaming `previous_response_id`: + - session-affinity account preference now uses a captured WebSocket opening + and asserts `authorization`, `chatgpt-account-id`, and payload + `previous_response_id`; + - previous-response-not-found, unanswered-function-call, and invalid encrypted + reasoning recovery now use a first WebSocket `response.failed` frame followed + by an HTTP SSE retry after history is stripped; + - added an end-to-end server test proving two same-history non-streaming + `/v1/responses` requests use WebSocket and reuse the configured runtime + WebSocket pool. + +TDD evidence: + +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_with_previous_response_id_should_use_websocket_and_configured_pool -- --nocapture` + failed with `Protocol(WrongHttpMethod)`, proving the server path still sent an + HTTP POST to a WebSocket upstream for `previous_response_id`. +- GREEN: + the same test passed after changing the translator to leave + `force_http_sse=false` for non-streaming history-continuation requests. +- Full server test RED after the first fix exposed five old HTTP mock tests that + still assumed non-streaming `previous_response_id` used HTTP. Those tests were + migrated to WebSocket/hybrid upstream fixtures, and the full server test then + passed. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_with_previous_response_id_should_use_websocket_and_configured_pool -- --nocapture` + (RED: WebSocket upstream received HTTP POST / wrong method) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_with_previous_response_id_should_use_websocket_and_configured_pool -- --nocapture` + (GREEN: targeted server WebSocketRequired path passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (intermediate RED: five previous-response tests still used HTTP-only mocks) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_strip_history_after -- --nocapture` + (4/4 migrated non-streaming history recovery tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream previous_response -- --nocapture` + (4/4 previous-response filtered tests passed after fixing the affinity fixture + timestamp) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (56/56 tests passed) +- `cargo test -p codex-proxy-core --test protocol` + (46/46 tests passed) +- `cargo test -p codex-proxy-adapters --test codex` + (43/43 tests passed) +- `cargo test -p codex-proxy-runtime --test upstream -- --nocapture` + (1/1 test passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` + (first run flagged `accept_hdr_async` test callback large-error types; fixed + with local `#[expect(clippy::result_large_err)]` annotations and reran green) +- `cargo test --test architecture` + (22/22 tests passed) +- `git diff --check` +- `codegraph sync .` + (synced 4 changed files; 515 nodes in 449 ms) +- `codegraph status .` + (index up to date with 272 files, 3,550 nodes, and 10,120 edges) + +Remaining after this slice: + +- Non-streaming OpenAI Responses with `previous_response_id` now exercise the + server/runtime WebSocketRequired path and configured WebSocket pool. +- Streaming `previous_response_id` requests and live permessage-deflate were + still open at the end of Slice 64; both are addressed in Slice 65 below. +- Broader old WebSocket behavior-test migration remains incomplete beyond the + migrated previous-response server/runtime coverage above. + +### 2026-06-19 Slice 65: Live WebSocket Streaming And Deflate Integration + +Status: implemented and verified for this batch. + +Changes: + +- Enabled OpenAI Responses translation to keep `force_http_sse=false` for all + requests with `previous_response_id`, including `stream: true`, so streaming + history-continuation requests now reach `CodexTransport::WebSocketRequired`. +- Added live WebSocket SSE byte-stream forwarding in the adapter path used by + `CodexBackendClient::create_response_stream(...)`. The stream converts public + WebSocket events to SSE bytes before terminal completion, returns pooled + sockets after successful terminal events, and discards sockets on transport or + upstream errors. +- Replaced the live WebSocket opening path with the migrated original-handshake + flow: the adapter writes the ordered upgrade request, reads the opening + response and any preloaded frame bytes, validates `Sec-WebSocket-Accept`, wraps + the stream in a `PerMessageDeflateStream`, and then hands it to + `tokio-tungstenite`. +- Inserted tested live permessage-deflate support into both aggregated and + streaming WebSocket paths. Server RSV1 text/binary frames are inflated before + tungstenite parses them, while the existing pure frame rewriter tests remain + in place. +- Restored streaming history recovery for WebSocket-first failures that surface + while prefetching the first SSE chunk. `previous_response_not_found` and + `invalid_encrypted_content` now strip history and retry over HTTP SSE before + any downstream bytes are returned. +- Migrated server streaming history-recovery tests from HTTP-only mocks to a + first WebSocket `response.failed` frame followed by HTTP SSE recovery, and + added coverage that a streaming `previous_response_id` response forwards a + WebSocket chunk before upstream completion. +- Added live adapter coverage for negotiated permessage-deflate by using a raw + local WebSocket upstream that replies with + `Sec-WebSocket-Extensions: permessage-deflate` and an RSV1 compressed + `response.completed` frame. + +TDD evidence: + +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_with_previous_response_id_should_forward_websocket_chunks_before_completion -- --nocapture` + failed with `Protocol(WrongHttpMethod)`, proving streaming + `previous_response_id` still went to HTTP instead of WebSocket. +- GREEN: + the same test passed after enabling WebSocketRequired routing for streaming + `previous_response_id` and adding adapter live WebSocket byte streaming. +- RED: + `cargo test -p codex-proxy-adapters --test codex codex_backend_client_should_decode_live_permessage_deflate_websocket_frame -- --nocapture` + failed with `Protocol(NonZeroReservedBits)`, proving live RSV1 frames still + reached tungstenite without decompression. +- GREEN: + the same test passed after inserting `PerMessageDeflateStream` into the live + connection path. +- Full server test then exposed two streaming history-recovery tests that still + assumed HTTP-first behavior. They were migrated to WS-first recovery fixtures, + and the runtime prefetch error branch was fixed to strip history and retry + before returning downstream bytes. + +Verification run so far: + +- `cargo test -p codex-proxy-core --test protocol openai_streaming_response_with_previous_response_should_require_websocket -- --nocapture` + (1/1 filtered test passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_with_previous_response_id_should_forward_websocket_chunks_before_completion -- --nocapture` + (RED then GREEN; final filtered test passed) +- `cargo test -p codex-proxy-adapters --test codex codex_backend_client_should_decode_live_permessage_deflate_websocket_frame -- --nocapture` + (RED then GREEN; final filtered test passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_strip_history_after -- --nocapture` + (2/2 streaming WS-first history recovery tests passed after the prefetch + recovery fix) +- `cargo test -p codex-proxy-adapters --test codex websocket_ -- --nocapture` + (26/26 filtered WebSocket tests passed) +- `cargo test -p codex-proxy-adapters --test codex` + (44/44 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (57/57 tests passed) +- `cargo test -p codex-proxy-runtime --test upstream` + (1/1 test passed) +- `cargo test -p codex-proxy-core --test protocol` + (46/46 tests passed) +- `cargo check --workspace --all-targets` +- `cargo fmt --all -- --check` +- `cargo test --test architecture` + (22/22 tests passed) +- `git diff --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` + (first run flagged one new `accept_hdr_async` test helper large-error + callback; fixed with a local + `#[expect(clippy::result_large_err)]` annotation and reran green) +- `codegraph sync . && codegraph status .` + (synced 9 changed files; index up to date with 272 files, 3,595 nodes, and + 10,382 edges before the documentation-only update) + +Remaining after this slice: + +- Live WebSocket streaming body forwarding, streaming `previous_response_id` + WebSocket dispatch, and live permessage-deflate integration are now covered by + adapter/server/core tests. +- Remaining old WebSocket behavior-test migration still needs a focused pass for + parity/audit edge cases not represented by the restored live, pool, deflate, + ping, history-recovery, and previous-response coverage. +- Late-failure stream audit/reporting policy after downstream bytes have already + been sent remains a serving-layer gap. + +### 2026-06-19 Slice 66: Live Stream Late-Failure SSE Reporting + +Status: implemented and verified for this batch. + +Changes: + +- Restored the old serving-layer client reporting policy for live stream + failures after at least one downstream SSE chunk has already been sent. The + runtime no longer forwards late upstream body read errors as Hyper body + errors; it appends a synthetic OpenAI-compatible + `response.failed` event with `code: "stream_disconnected"` and then appends + `data: [DONE]`. +- Covered the clean premature EOF case separately. If the upstream stream ends + without `response.completed`, `response.failed`, or `error`, the runtime now + emits the same `stream_disconnected` failure event before `[DONE]`. +- Preserved terminal-event semantics: if a live stream has already received + `response.completed`, `response.failed`, or `error`, the runtime only appends + `[DONE]` when needed and does not synthesize a second failure. +- Added partial-SSE safety before synthetic failures. If upstream closes while + a frame is incomplete, the runtime first emits the missing blank line so the + synthetic `response.failed` event is parsed as a separate SSE frame. +- Kept retry behavior conservative: once bytes have been sent downstream, the + serving layer does not attempt account fallback or history replay, because + doing so would violate HTTP streaming semantics. + +TDD evidence: + +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_emit_failed_event -- --nocapture` + failed for both new tests. The abrupt-close case surfaced a body + `UnexpectedEof` error to the downstream reader, and the clean-close case + lacked `event: response.failed`. +- GREEN: + the same filtered command passed after runtime stream finalization converted + late read errors and missing-terminal EOF into SSE + `response.failed(stream_disconnected)` frames followed by `[DONE]`. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_emit_failed_event -- --nocapture` + (RED then GREEN; final filtered tests passed 2/2) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should -- --nocapture` + (18/18 filtered stream tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream -- --nocapture` + (59/59 tests passed) +- `cargo test -p codex-proxy-runtime --test upstream -- --nocapture` + (1/1 test passed) +- `cargo test -p codex-proxy-core --test protocol -- --nocapture` + (46/46 tests passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture -- --nocapture` + (22/22 tests passed) +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` + +Remaining after this slice: + +- Client-visible late-failure SSE reporting is restored for live stream body + errors and clean premature EOF after downstream bytes have been sent. +- Full persisted stream audit parity remains incomplete. The new crate split + still lacks the old `StreamAudit` / `WebSocketStreamAudit` event-log style + artifacts for all late-success and late-failure stream outcomes. +- Remaining old WebSocket behavior-test migration still needs a focused pass for + parity/audit edge cases not represented by the restored live, pool, deflate, + ping, history-recovery, previous-response, and late-failure coverage. + +### 2026-06-19 Slice 67: Live Stream Event-Log Audit Restoration + +Status: implemented and verified for this batch. + +Changes: + +- Added a runtime write path for admin event logs through + `AdminLogService::record(...)`, reusing the existing core `EventLogService` + policy: when logging is enabled, all events are recorded; when disabled, + error-level events are still retained. +- Injected the shared `AdminLogService` into `ResponseDispatchService` and + carried request audit context into `LiveResponseStreamContext`: request id, + account id, model, route, stream metadata, status code, and latency. +- Restored persisted `v1.response` event-log records for completed live + Responses streams. Successful terminal streams now write an info event with + status `200`, `stream: true`, `completed: true`, and `responseId`. +- Restored persisted `v1.response` event-log records for late failed live + Responses streams. Synthetic `stream_disconnected` failures, upstream + `response.failed` events, and missing-terminal streams now write error/warn + events with failure metadata and the mapped status code. +- Updated the admin logs error mapper for the new append failure variant so + management routes remain exhaustive. + +TDD evidence: + +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_record_event_log -- --nocapture` + failed for both new tests with `expected a v1.response event log`, proving live + stream completion and late-disconnect paths still did not persist audit + records. +- GREEN: + the same filtered command passed after wiring `AdminLogService` into the live + stream finalizer and writing completed/failed stream events. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_record_event_log -- --nocapture` + (RED then GREEN; final filtered tests passed 2/2) +- `cargo test -p codex-proxy-server --test admin_logs_routes -- --nocapture` + (2/2 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should -- --nocapture` + (20/20 filtered stream tests passed) +- `cargo test -p codex-proxy-runtime --test upstream -- --nocapture` + (1/1 test passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream -- --nocapture` + (61/61 tests passed) +- `cargo test --test architecture -- --nocapture` + (22/22 tests passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `codegraph sync . && codegraph status .` + (synced 3 changed files; index up to date with 272 files, 3,623 nodes, and + 10,460 edges before the documentation-only update) + +Remaining after this slice: + +- Persisted live stream event-log audit is restored for completed streams and + late-failure streams at the runtime/server boundary. +- Full old `StreamAudit` / `WebSocketStreamAudit` parity is still broader than + this slice. Body capture policy, rate-limit/header snapshots, WebSocket pool + update snapshots, and every old parity edge-case artifact still need a focused + pass before calling stream audit migration complete. +- Remaining old WebSocket behavior-test migration still needs a focused pass for + parity/audit edge cases not represented by the restored live, pool, deflate, + ping, history-recovery, previous-response, late-failure, and event-log + coverage. + +### 2026-06-19 Slice 68: Stream Audit Metadata And Log Policy Batch + +Status: implemented and verified for this batch. + +Changes: + +- Restored the old admin log state mutation surface for migrated routes: + `PATCH /api/admin/logs/state` now updates `enabled`, `capacity`, and + `captureBody`, rejects zero capacity, and trims persisted `event_logs` to the + configured newest-N capacity without adding a schema migration. +- Restored the old log body-capture policy in `AdminLogService::record(...)`. + Events can attach `body`, `rawBody`, `requestBody`, `responseBody`, and + `upstreamBody`; when `captureBody` is false these keys are removed before + persistence, and when true they are retained. +- Extended live Responses stream event logs with old `StreamAudit` metadata: + completed and failed streams now include normalized `usage`, + `rateLimitHeaders`, and capture-policy-controlled `requestBody` / + `responseBody`. +- Added an actual upstream transport marker to streaming adapter responses. + Runtime stream audit records `transport: "websocket"` only when the adapter + actually used WebSocket, so WebSocketPreferred HTTP fallback cannot be logged + as a WebSocket stream by request-policy inference. +- Restored live WebSocket stream rate-limit update snapshots. Adapter live + forwarding now captures internal `codex.rate_limits` events, converts them + back into the old `x-codex-*` header-pair shape, shares those updates with the + runtime stream finalizer, and merges them into persisted + `rateLimitHeaders`. +- Preserved the existing sensitive-cookie stance: `set-cookie` response headers + are still not added to event metadata. The restored header snapshot is scoped + to rate-limit and Codex review/quota headers. + +TDD evidence: + +- RED: + `cargo test -p codex-proxy-server --test admin_logs_routes admin_logs_should_update_state_and_trim_to_capacity -- --nocapture` + failed with HTTP `405`, proving the migrated admin log state update route was + missing. +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_record_event_log_after_completed_stream -- --nocapture` + failed because stream metadata lacked `usage` and rate-limit header snapshot + fields. +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_preserve_body_metadata_when_capture_body_enabled -- --nocapture` + failed because capture-body-enabled logs still had no `requestBody` / + `responseBody`. +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_record_event_log_after_late_disconnect -- --nocapture` + failed because late-failure stream logs lacked `rateLimitHeaders`. +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_with_previous_response_id_should_record_websocket_audit_metadata -- --nocapture` + first failed on missing `transport: "websocket"`, then exposed missing live + WebSocket rate-limit update propagation. +- GREEN: + the same targeted tests passed after restoring admin log state mutation, + body-capture policy, actual stream transport propagation, HTTP stream + metadata enrichment, and WebSocket `codex.rate_limits` update sharing. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test admin_logs_routes -- --nocapture` + (3/3 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream -- --nocapture` + (63/63 tests passed) +- `cargo test -p codex-proxy-adapters --test codex websocket -- --nocapture` + (27/27 filtered WebSocket tests passed) +- `cargo test -p codex-proxy-core --test protocol websocket -- --nocapture` + (25/25 filtered WebSocket/protocol tests passed) +- `cargo test -p codex-proxy-runtime --test upstream -- --nocapture` + (1/1 test passed) +- `cargo test -p codex-proxy-runtime --test tasks -- --nocapture` + (9/9 tests passed) +- `cargo test --test architecture -- --nocapture` + (22/22 tests passed) +- `cargo check --workspace --all-targets` +- `cargo fmt --all -- --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` + +Remaining after this slice: + +- Stream audit metadata parity is restored for body capture policy, usage, + rate-limit/header snapshots, actual WebSocket transport marking, and live + WebSocket rate-limit update snapshots. +- Remaining old WebSocket behavior-test migration still needs a focused pass for + parity/audit edge cases not represented by the restored live, pool, deflate, + ping, history-recovery, previous-response, late-failure, event-log, and + rate-limit-update coverage. +- Full persisted stream audit parity still needs any old pool-specific artifact + details that are separate from event-log metadata. + +### 2026-06-19 Slice 69: Serving Toy Helper Cleanup Guard + +Status: implemented and verified for this batch. + +Changes: + +- Added an architecture guard that fails if the serving modules keep the + previously audited toy helpers: + `fallback::should_fallback(...)`, `recovery::is_recoverable(...)`, and + `responses::prefers_websocket(...)`. +- The same guard verifies that `quota::quota_reached(...)`, which remains a + useful pure threshold helper, has a non-test production caller. +- Replaced `fallback::should_fallback(...)` with explicit pure status helpers: + `status_code_is_rate_limited`, `status_code_is_quota_exhausted`, and + `status_code_is_transient_upstream`. +- Replaced `recovery::is_recoverable(...)` with the narrower + `status_code_allows_same_account_retry(...)`, matching the runtime's current + same-account retry behavior instead of implying broader recovery semantics. +- Removed `responses::prefers_websocket(...)`; production transport decisions + already use `transport_for_request(...)` and `http_sse_fallback_allowed(...)`. +- Wired runtime dispatch classification and admin quota-warning threshold + matching through the core serving helpers, so these modules are no longer just + architecture-shaped files with uncalled convenience functions. + +TDD evidence: + +- RED: + `cargo test --test architecture serving_modules_should_not_keep_uncalled_toy_helpers -- --nocapture` + failed and listed all four audited helpers: `should_fallback`, + `is_recoverable`, `prefers_websocket`, and uncalled `quota_reached`. +- GREEN: + the same architecture test passed after removing/renaming the toy helpers and + wiring runtime production callers through the remaining pure policy helpers. + +Verification run so far: + +- `cargo test --test architecture serving_modules_should_not_keep_uncalled_toy_helpers -- --nocapture` + (RED then GREEN; final filtered test passed) +- `cargo test -p codex-proxy-server --test admin_accounts_routes quota_warnings -- --nocapture` + (3/3 filtered quota-warning tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream -- --nocapture` + (63/63 tests passed) +- `cargo test --test architecture -- --nocapture` + (23/23 tests passed) +- `cargo test -p codex-proxy-core --test protocol -- --nocapture` + (46/46 tests passed) +- `cargo check --workspace --all-targets` +- `cargo fmt --all -- --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` + +Remaining after this slice: + +- The specific audited uncalled serving toy helpers are removed or wired to + production callers, and architecture tests now catch their return. +- Broader serving dispatch parity is still not complete: old fallback/recovery + edge-case behavior should continue to be migrated through runtime/server tests + rather than inferred from helper presence. + +### 2026-06-19 Slice 70: Admin Settings Patch And Test Migration Guards + +Status: implemented and verified for this batch. + +Changes: + +- Migrated the old admin settings PATCH behavior into the crate split: + - `core::admin::settings` now owns the retained settings patch model, + validation, allowed rotation strategies, quota-warning threshold validation, + and pure patch application. + - `platform::config::AppConfig::write_settings_overlay(...)` now writes the + retained runtime settings subset to a local YAML overlay without putting + filesystem IO in `core`. + - `RuntimeSettingsService` now stores a same-process current config snapshot, + applies core settings patches, writes `local.yaml`, and exposes an injected + local-config path for tests. + - `PATCH /api/admin/settings` is restored, admin-session gated, rejects + unknown/invalid fields, persists retained fields to local YAML, and updates + the same-process `GET /api/admin/settings` view. +- Added crate-local tests migrated from the old `tests/admin/settings_route.rs` + coverage for PATCH auth, retained-field persistence, reload through + `AppConfig::load_from_dir(...)`, same-process GET visibility, and + unknown/invalid field rejection. +- Added `tests/architecture/test_migration.rs` with guards that: + - keep root `tests/` limited to architecture tests plus shared fixtures; + - require high-risk baseline suites to have crate-local migration targets; + - fail when exported core service modules have no production callers. +- Wired previously thin core service modules into production paths: + - `AccountService` now supplies account quota/cloudflare availability helpers + used by `AccountPool`; + - `UsageService` now converts normalized Codex token usage into + `AccountUsageDelta` for runtime persistence; + - `SettingsService` is now used by runtime settings updates. + +TDD evidence: + +- RED: + `cargo test -p codex-proxy-core --test admin_settings -- --nocapture` + failed before implementation because the retained settings patch types and + full settings fields did not exist. +- RED: + `cargo test -p codex-proxy-server --test admin_settings_routes admin_settings_patch -- --nocapture` + failed before implementation because `AppState` had no local-config-path + constructor and the server test lacked the YAML writer dependency; the route + was also still only registered as `GET`. +- GREEN: + both targets passed after implementing the core/runtime/platform/server + settings path. + +Verification run so far: + +- `cargo test -p codex-proxy-core --test admin_settings -- --nocapture` + (2/2 tests passed) +- `cargo test -p codex-proxy-server --test admin_settings_routes admin_settings_patch -- --nocapture` + (3/3 filtered PATCH tests passed) +- `cargo test --test architecture test_migration -- --nocapture` + (3/3 filtered architecture migration tests passed) +- `cargo test -p codex-proxy-server --test admin_settings_routes -- --nocapture` + (6/6 tests passed) +- `cargo test -p codex-proxy-core --test account_pool -- --nocapture` + (25/25 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream -- --nocapture` + (63/63 tests passed) +- `cargo test -p codex-proxy-platform --test config -- --nocapture` + (3/3 tests passed) +- `cargo test --test architecture -- --nocapture` + (26/26 tests passed) +- `cargo fmt --all` +- `cargo check --workspace --all-targets` + +Remaining after this slice: + +- The old admin settings PATCH/local YAML behavior is restored for the retained + settings surface. +- Root behavior test files are now guarded against returning under `tests/`; + high-risk old suites now have crate-local migration-target evidence, but this + does not by itself prove every individual old assertion has a one-to-one + migrated equivalent. +- Generic empty/placeholder detection is stronger for exported core service + modules, but broader non-empty semantic completeness still depends on + behavior tests and focused old-code parity passes. + +### 2026-06-19 Slice 71: Admin Client-Key Assertion Coverage Migration + +Status: crate-local behavior coverage migrated and verified for this batch. + +Changes: + +- Migrated additional old `tests/admin/client_keys_route.rs` assertions into + `crates/server/tests/admin_client_keys_routes.rs`. +- Added crate-local coverage for: + - label update, label clearing, overlong-label rejection, and missing-key + `404`; + - batch delete with found/missing IDs, authorization invalidation, and empty + ID rejection; + - export metadata shape without plaintext, hash, or pepper material; + - import of exported metadata with forced key rotation, one-time plaintext + return, old plaintext rejection, and new plaintext authorization. + +TDD/coverage evidence: + +- The migrated tests passed on the first run, so this slice did not require + production-code changes. That means the behavior had already been restored in + earlier migration work, but the old assertions were not yet present in the + crate-local suite. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test admin_client_keys_routes -- --nocapture` + (7/7 tests passed) +- `cargo fmt --all` +- `cargo fmt --all -- --check` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- `codegraph sync . && codegraph status .` + (synced 1 changed file; index up to date with 274 files, 3,737 nodes, and + 10,716 edges before this documentation-only update) + +Remaining after this slice: + +- The highest-risk old client-key route assertions are now represented in the + server crate-local suite. +- Assertion-level migration is still incomplete for other old suites called out + in this audit, especially the largest WebSocket and serving suites. + +### 2026-06-19 Slice 72: Admin Account Manual Create, Refresh-Only, And CLI Import + +Status: implemented and verified for this batch. + +Changes: + +- Migrated a larger batch of old `tests/admin/accounts/import_export.rs` + assertions into `crates/server/tests/admin_accounts_routes.rs`. +- Strengthened manual account create coverage for: + - ignoring caller-supplied metadata/status and deriving email/account/user/plan + from JWT claims; + - encrypting both access and refresh tokens at rest; + - immediately syncing created/updated active accounts into the runtime account + pool without requiring process restart; + - rejecting missing, invalid, expired, and account-claim-less tokens. +- Restored refresh-token-only manual create: + - `AdminAccountService` now accepts a `TokenRefresher` port; + - production service assembly injects the OpenAI OAuth client as the default + refresher; + - tests can inject a static refresher through an `AppState` constructor; + - refresh-only create stores a rotated refresh token when returned; + - refresh-only create preserves the input refresh token for new accounts when + the exchange omits rotation; + - refresh-only update preserves the existing stored refresh token when the + exchange omits rotation. +- Restored same-account manual update behavior: + - same ChatGPT account/user claims update the existing row rather than + inserting a duplicate; + - access token, email, plan, status, and pool snapshot are updated; + - existing refresh token is preserved when the new request does not provide a + replacement. +- Restored `POST /api/admin/accounts/import-cli`: + - server reads `codexHome/auth.json`; + - runtime parses `access_token` / `refresh_token` from the CLI auth payload and + reuses the same manual-create path; + - caller metadata is ignored in favor of JWT claims; + - imported tokens are encrypted and not returned in the response; + - imported active accounts are synced into the runtime account pool. +- Updated `docs/architecture.md` and the directory-shape whitelist for the new + `crates/server/src/admin_api/accounts/import_cli.rs` handler. +- Added runtime account-pool synchronization for successful admin account + create/import/delete paths. Label/status updates now attempt pool sync + best-effort so metadata writes do not fail when the existing stored token is + intentionally corrupt or undecryptable. +- Added `AccountStore::get_pool_account(...)` and a SQLite implementation that + reads a single pool snapshot by ID. This avoids full-table pool scans during + admin write synchronization and prevents unrelated corrupt token rows from + breaking otherwise valid create/import sync. + +TDD/debugging evidence: + +- RED: + `cargo test -p codex-proxy-server --test admin_accounts_routes -- --nocapture` + failed before implementation because `AppState` had no injected token + refresher constructor. The missing constructor blocked refresh-only migration + tests exactly where the runtime had no refresh-token-only create path. +- GREEN attempt: + the same command initially reached 22/23 tests and exposed a regression in + `admin_accounts_lifecycle_should_update_and_delete_accounts`: label/status + updates returned `500` because the new pool sync tried to decrypt a deliberately + corrupt token row. +- Root-cause fix: + sync now reads the target pool account by ID, and label/status updates treat + pool sync as best-effort. Manual create/import still fail visibly if their own + valid account cannot be synced. +- GREEN: + `cargo test -p codex-proxy-server --test admin_accounts_routes -- --nocapture` + passed with 23/23 tests. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test admin_accounts_routes -- --nocapture` + (23/23 tests passed) +- `cargo test -p codex-proxy-adapters --test account_repository -- --nocapture` + (8/8 tests passed) +- `cargo test -p codex-proxy-runtime --test account_pool_restore -- --nocapture` + (1/1 test passed) +- `cargo test --test architecture test_migration -- --nocapture` + (3/3 filtered architecture migration tests passed) +- `cargo check --workspace --all-targets` +- `cargo fmt --all` +- `cargo fmt --all -- --check` +- `cargo test --test architecture -- --nocapture` + (26/26 architecture tests passed after adding `import_cli.rs` to the + architecture whitelist) +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- `codegraph sync . && codegraph status .` + (synced 9 changed files; index up to date with 275 files, 3,780 nodes, and + 10,931 edges) + +Current migration counters after this slice: + +- Approximately 358 `#[test]` / `#[tokio::test]` occurrences. +- 15 root Rust test files, all architecture tests. +- 35 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- The highest-risk old manual account create/refresh-only/import-cli assertions + from `tests/admin/accounts/import_export.rs` are now represented in the server + crate-local suite. +- Other admin account old assertions are still not fully migrated, especially + OAuth-route edge cases and any import/export cases not covered by this batch. +- The largest remaining unmigrated suites are still WebSocket/serving behavior + suites called out below; this slice should not be read as full migration + completion. + +### 2026-06-19 Slice 73: Admin Account OAuth Route Migration + +Status: implemented and verified for this batch. + +Changes: + +- Migrated a broad batch of old `tests/admin/accounts/oauth.rs` behavior into + `crates/server/tests/admin_accounts_routes.rs`. +- Restored `GET /api/admin/auth/status`: + - requires an admin session cookie; + - reads account metadata without decrypting tokens; + - reports `authenticated`, a sanitized active user, and pool counts by status; + - does not expose access token, refresh token, or token aliases. +- Restored `POST /api/admin/auth/logout`: + - deletes all accounts and associated `account_usage`, `account_cookies`, and + `account_refresh_leases` rows explicitly; + - clears the runtime account pool so old accounts cannot still be scheduled. +- Restored device-code OAuth routes: + - `POST /api/admin/auth/device-login` calls the injected OAuth client and + returns device code, user code, verification URI, expiry, and poll interval; + - `GET /api/admin/auth/device-poll/{device_code}` maps + `authorization_pending` / `slow_down` to a non-importing pending response; + - successful polling imports the returned token pair through the same + admin-account create path used in Slice 72, preserving encryption, claims + derivation, and runtime pool sync while not returning secrets. +- Restored PKCE relay/callback flow: + - `POST /api/admin/auth/code-relay` parses a callback URL, exchanges + code/state through the injected OAuth client, imports the account, and + returns a sanitized success envelope; + - invalid callback URLs return `400`; + - `GET /auth/callback` exchanges code/state, imports the account, and + redirects back to the original admin host with `303 See Other`; + - unsupported legacy callback paths remain unregistered. +- `AdminOAuthService` now owns an injected `OAuthClient` port in addition to the + PKCE session store. +- Runtime service assembly now supports both production OpenAI OAuth clients and + test-injected OAuth clients. The test constructor still builds a full + `AppState`; only the OAuth upstream port is substituted. + +TDD evidence: + +- RED: + `cargo test -p codex-proxy-server --test admin_accounts_routes admin_auth -- --nocapture` + failed before implementation because `AppState` had no + `with_pool_secret_api_key_hasher_and_oauth_client(...)` constructor. This + showed the runtime had no testable OAuth-client port for the old device/code + relay behaviors. +- GREEN: + the same filtered command passed with 11/11 OAuth tests after implementing the + runtime OAuth client port and server routes. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_auth -- --nocapture` + (11/11 filtered OAuth tests passed) +- `cargo test -p codex-proxy-server --test admin_accounts_routes -- --nocapture` + (33/33 tests passed) +- `cargo test -p codex-proxy-adapters --test account_repository -- --nocapture` + (8/8 tests passed) +- `cargo test -p codex-proxy-runtime --test account_pool_restore -- --nocapture` + (1/1 test passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture -- --nocapture` + (26/26 architecture tests passed) +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- `codegraph sync . && codegraph status .` + (synced 7 changed files; index up to date with 275 files, 3,841 nodes, and + 11,275 edges) + +Current migration counters after this slice: + +- Approximately 368 `#[test]` / `#[tokio::test]` occurrences. +- 15 root Rust test files, all architecture tests. +- 35 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- The highest-risk old admin account OAuth route assertions are now represented + in the server crate-local suite. +- Some admin account import/export edge cases outside the Slice 72 and Slice 73 + batches still needed assertion-level comparison with the old baseline. + **Closed for the audited import/export edge cases in Slice 74.** +- The largest remaining unmigrated suites are still WebSocket/serving behavior + suites called out below; this slice should not be read as full migration + completion. + +### 2026-06-19 Slice 74: Admin Import/Export Strictness And Route Assertion Sweep + +Status: implemented and verified for this larger batch. + +Changes: + +- Migrated the remaining audited old `tests/admin/accounts/import_export.rs` + assertions into `crates/server/tests/admin_accounts_routes.rs`: + - `POST /api/admin/accounts/import` now has crate-local coverage for missing + admin session cookies; + - imported tokens are verified encrypted at rest and absent from list + responses; + - non-native external export shapes are rejected without inserting accounts; + - native account objects with unknown fields are rejected instead of silently + importing; + - native containers with unknown fields are rejected instead of silently + importing; + - account export filters by requested IDs and returns only native account token + material for the selected accounts; + - unsupported `format=external` and `format=full` exports are rejected. +- Tightened runtime native account import parsing: + - `parse_account_import_payload(...)` now returns `Result` instead of silently + filtering every shape; + - native top-level envelopes, native account containers, account arrays, and + single account objects use explicit key allow-lists; + - `sourceFormat` is accepted only when it is empty or `native`; + - current native export metadata fields (`addedAt` / `updatedAt`, and snake + case aliases) are accepted so native exports remain re-importable; + - legacy/external fields such as `type`, `legacy`, `credentials`, + `legacyField`, or `legacyContainer` now map to + `No importable accounts found`. +- Folded additional old `tests/admin/client_keys_route.rs` assertions into the + existing crate-local client-key suite: + - unknown client API keys are rejected before any key is created; + - re-enabled keys authorize `/v1/models` again; + - invalid key status values return `400`; + - deleted keys disappear from the list, no longer authorize, and repeated + deletion returns `404`. +- Expanded `crates/server/tests/admin_logs_routes.rs` from 3 to 9 tests using + old `tests/admin/logs_route.rs` behavior: + - log list requires an admin session cookie; + - cursor pagination includes `requestId`, page limit, first item, and + `nextCursor`; + - unsupported level filters return `400`; + - detail lookup returns `404` for missing events; + - log-state PATCH requires an admin session cookie; + - zero log capacity is rejected; + - clearing logs is verified by reading the empty list afterward. +- Expanded `crates/server/tests/admin_models_routes.rs` with old + `tests/admin/models_route.rs` assertions: + - admin model refresh now verifies configured fingerprint headers reach the + mocked Codex model endpoint; + - refreshed backend models are visible through `/v1/models/catalog` using a + stored client API key. + +TDD/coverage evidence: + +- RED: + `cargo test -p codex-proxy-server --test admin_accounts_routes admin_accounts_import -- --nocapture` + failed with 4/7 passing before implementation. The failures showed unknown + native account/container fields were accepted as `200`, and the non-native + shape error message was still `no importable accounts found` instead of the + old `No importable accounts found` surface. +- GREEN: + the same filtered import command passed with 7/7 after strict native import + parsing and the message restoration. +- Coverage-only migrations: + the added client-key, log, and model assertions passed after test migration, + so they document behavior already restored in earlier slices rather than + requiring new production changes. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_accounts_import -- --nocapture` + (RED first: 4/7 passed, 3 failed; GREEN after implementation: 7/7 passed) +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_accounts_export -- --nocapture` + (4/4 filtered export tests passed) +- `cargo test -p codex-proxy-server --test admin_accounts_routes -- --nocapture` + (41/41 tests passed) +- `cargo test -p codex-proxy-server --test admin_client_keys_routes -- --nocapture` + (7/7 tests passed) +- `cargo test -p codex-proxy-server --test admin_logs_routes -- --nocapture` + (9/9 tests passed) +- `cargo test -p codex-proxy-server --test admin_models_routes -- --nocapture` + (2/2 tests passed) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --tests -- --nocapture` + (136 server tests passed) +- `cargo test -p codex-proxy-runtime --tests -- --nocapture` + (30 runtime tests passed) +- `cargo test --test architecture -- --nocapture` + (26/26 architecture tests passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `codegraph sync . && codegraph status .` + (synced 5 changed files; index up to date with 275 files, 3,863 nodes, and + 11,376 edges) + +Current migration counters after this slice: + +- Approximately 382 `#[test]` / `#[tokio::test]` occurrences. +- 15 root Rust test files, all architecture tests. +- 35 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- The audited remaining old admin account import/export assertions are now + represented in the server crate-local suite. +- The client-key route suite now carries the old missing assertions in the + current crate-local file. +- The logs route suite still does not reproduce every historical fixture line + one-for-one, but the old auth, cursor, filter, state, detail, and clear route + behaviors are covered at the crate boundary. +- The admin models route suite now verifies the old fingerprint-header and + catalog-read assertions. + +### 2026-06-19 Slice 75: Admin Session Login, Usage Stats, And API Contract Migration + +Status: implemented and verified for this larger batch. + +Changes: + +- Migrated admin session persistence behavior into + `crates/adapters/tests/admin_sessions.rs` and + `crates/adapters/src/sqlite/admin_sessions.rs`: + - the SQLite store can create a default admin user exactly once; + - the first admin user can be loaded with its password hash; + - sessions are created with a user ID and expiry timestamp; + - existing session validation and cleanup behavior remain in the adapter. +- Restored runtime admin login/session behavior in + `crates/runtime/src/services.rs`: + - `AdminSessionService` now owns default-admin bootstrap, credential + verification, session creation, and TTL assignment; + - admin password hashing/verification is delegated to `platform::identity`; + - username validation is delegated to the pure `core::admin::auth` service. +- Added `POST /api/admin/login` in `crates/server/src/admin_api/session.rs` and + routed it from `crates/server/src/admin_api/router.rs`: + - valid credentials return the standard admin envelope plus `expiresAt`; + - successful login sets the `cpr_admin_session` HTTP-only cookie; + - invalid credentials return the old `40102` admin error surface and do not + set a cookie. +- Restored server startup default-admin bootstrap in + `crates/server/src/main.rs`, guarded by + `tests/architecture/server_startup.rs`. +- Migrated remaining audited usage-stats assertions into + `crates/server/tests/admin_accounts_routes.rs`: + - account usage stats require an admin session cookie; + - cursor pagination preserves account usage ordering and `nextCursor`. +- Added API response contract tests in + `crates/server/tests/admin_api_contract.rs`: + - request IDs stay camelCase; + - page envelopes expose camelCase pagination; + - HTTP status stays outside successful admin envelope bodies; + - error bodies keep `data: null`; + - `AdminError` serializes through the current response shape. +- Added a platform password boundary test in + `crates/platform/tests/admin_password.rs` so admin password hashes do not + accept client API key strings. + +TDD/coverage evidence: + +- RED: + `cargo test -p codex-proxy-adapters --test admin_sessions -- --nocapture` + failed before implementation because the SQLite admin-session store did not + expose default-admin or session-creation persistence methods. +- RED: + `cargo test -p codex-proxy-server --test admin_session_routes -- --nocapture` + failed before implementation with `404` for `/api/admin/login`. +- RED: + `cargo test --test architecture server_entrypoint_should_ensure_default_admin_before_serving -- --nocapture` + failed before startup wiring because `main.rs` did not bootstrap the default + admin user from config. +- GREEN: + after implementation the adapter, server login, architecture startup, + platform password, usage-stats, and API contract test filters all passed. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test admin_sessions -- --nocapture` + (2/2 tests passed after RED) +- `cargo test -p codex-proxy-server --test admin_session_routes -- --nocapture` + (2/2 tests passed after RED) +- `cargo test --test architecture server_entrypoint_should_ensure_default_admin_before_serving -- --nocapture` + (1/1 filtered architecture test passed after RED) +- `cargo test -p codex-proxy-platform --test admin_password -- --nocapture` + (1/1 test passed) +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_usage_stats -- --nocapture` + (3/3 filtered usage-stat tests passed) +- `cargo test -p codex-proxy-server --test admin_api_contract -- --nocapture` + (5/5 tests passed) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --tests -- --nocapture` + (145 server tests passed) +- `cargo test -p codex-proxy-adapters --tests -- --nocapture` + (68 adapter tests passed) +- `cargo test -p codex-proxy-platform --tests -- --nocapture` + (21 platform tests passed) +- `cargo test -p codex-proxy-runtime --tests -- --nocapture` + (30 runtime tests passed) +- `cargo test --test architecture -- --nocapture` + (27/27 architecture tests passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- `codegraph sync .` +- `codegraph status .` + (index up to date with 279 files, 3,919 nodes, and 11,545 edges) + +Current migration counters after this slice: + +- Approximately 395 `#[test]` / `#[tokio::test]` occurrences. +- 15 root Rust test files, all architecture tests. +- 39 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- The audited admin-session login/startup bootstrap behavior is represented in + crate-local server, runtime, adapter, platform, and architecture coverage. +- Usage-stats route authorization and cursor pagination are now represented in + the server crate-local suite. +- Old account lifecycle, cookie/quota, and broader serving/WebSocket suites + still need larger migrated coverage batches. +- The largest remaining unmigrated/high-risk area is still the old WebSocket and + serving behavior surface called out below; this slice should not be read as + full migration completion. + +### 2026-06-19 Slice 76: Admin Account Batch Status, Refresh, And Reset Usage Migration + +Status: implemented and verified for this larger batch. + +Changes: + +- Migrated old admin account lifecycle/cookie-quota behavior into + `crates/server/tests/admin_accounts_routes.rs`: + - `POST /api/admin/accounts/batch-status` updates found accounts and reports + missing IDs plus invalid statuses; + - `POST /api/admin/accounts/{account_id}/refresh` refreshes tokens, syncs the + runtime pool, and does not return secret material; + - refresh invalid-grant failures mark the account `expired`; + - `POST /api/admin/accounts/{account_id}/reset-usage` clears persisted local + counters and syncs the runtime pool. +- Added the persisted reset path in `SqliteAccountStore::reset_usage()` so + cumulative/window counters and `last_used_at` are cleared without changing the + retained quota window reset metadata. +- Added `AdminAccountService::batch_update_status()`, + `AdminAccountService::refresh_account()`, and + `AdminAccountService::reset_usage()` in `crates/runtime/src/services.rs`. +- Added server request/response handlers in + `crates/server/src/admin_api/accounts/lifecycle.rs` and routed them from + `crates/server/src/admin_api/router.rs`. +- Added `crates/core/tests/account_pool.rs` coverage for in-memory usage reset + semantics. + +TDD/coverage evidence: + +- RED: + `cargo test -p codex-proxy-server --test admin_accounts_routes admin_accounts_batch_status -- --nocapture` + failed before implementation with HTTP `405` instead of `200`. +- RED: + `cargo test -p codex-proxy-server --test admin_accounts_routes admin_account_refresh_should -- --nocapture` + failed before implementation with HTTP `404` instead of `200`. +- RED: + `cargo test -p codex-proxy-server --test admin_accounts_routes admin_account_reset_usage -- --nocapture` + failed before implementation with HTTP `404` instead of `200`. +- GREEN: + after implementation the batch-status, refresh, reset-usage, and core + account-pool reset filters all passed. + +Verification run so far: + +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_accounts_batch_status -- --nocapture` + (1/1 filtered test passed after RED) +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_account_refresh_should -- --nocapture` + (2/2 filtered tests passed after RED) +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_account_reset_usage -- --nocapture` + (1/1 filtered test passed after RED) +- `cargo test -p codex-proxy-core --test account_pool reset_usage_should_clear_runtime_counters_and_preserve_window_reset -- --nocapture` + (1/1 filtered test passed) +- `cargo test -p codex-proxy-server --test admin_accounts_routes -- --nocapture` + (47/47 tests passed) +- `cargo test -p codex-proxy-core --test account_pool -- --nocapture` + (26/26 tests passed) +- `cargo fmt --all` +- `codegraph sync . && codegraph status .` + (index up to date with 279 files, 3,948 nodes, and 11,879 edges) +- `cargo test -p codex-proxy-server --tests -- --nocapture` + (149 server tests passed) +- `cargo test -p codex-proxy-core --tests -- --nocapture` + (100 core tests passed) +- `cargo test -p codex-proxy-adapters --tests -- --nocapture` + (68 adapter tests passed) +- `cargo test -p codex-proxy-runtime --tests -- --nocapture` + (30 runtime tests passed) +- `cargo test --test architecture -- --nocapture` + (27/27 architecture tests passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` + +Current migration counters after this slice: + +- Approximately 400 `#[test]` / `#[tokio::test]` occurrences. +- 15 root Rust test files, all architecture tests. +- 39 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- The audited batch-status, manual refresh, invalid refresh-token expiry, and + reset-usage behavior from old admin account lifecycle/cookie-quota coverage is + now represented in crate-local server/core suites. +- Remaining old admin cookie/quota behavior still needs a focused pass for the + health-check OAuth-refresh semantics and exact request-field validation. +- The largest remaining unmigrated/high-risk area is still the old serving and + WebSocket behavior surface. + +### 2026-06-19 Slice 77: Admin Health-Check OAuth Refresh Migration + +Status: implemented and verified for the remaining old admin cookie/quota +health-check behavior. + +Changes: + +- Migrated the remaining `tests/admin/accounts/cookies_quota.rs` + health-check assertions into `crates/server/tests/admin_accounts_routes.rs`: + - health-check now skips active accounts that have no refresh token instead + of probing Codex usage and marking them alive; + - health-check uses the configured `TokenRefresher` to determine alive/dead + account status and does not touch the Codex usage backend; + - invalid refresh-token failures persist `expired` status; + - snake_case `stagger_ms` stays rejected by the camelCase API contract. +- Reworked `AdminAccountService::health_check_accounts()` to restore the old + refresh-token probe path with `buffer_unordered(concurrency)` and the existing + stagger delay. +- Added test-local `HealthCheckTokenRefresher` coverage so the server route + proves refresh calls, response sanitization, and persisted status changes. +- Re-ran empty-directory and placeholder scans after the implementation: + - empty directory scan returned no paths outside ignored locations; + - placeholder scan only matched architecture guard tests themselves; + - migration/compatibility keyword scan matched docs, architecture tests, + product compatibility wording, and old-format rejection tests, not new + migration shims in runtime/server source. + +TDD/coverage evidence: + +- RED: + `cargo test -p codex-proxy-server --test admin_accounts_routes health_check -- --nocapture` + failed 2/3 before implementation: + - `admin_account_health_check_should_skip_account_without_refresh_token` + reported `alive = 1` instead of `0`; + - `admin_accounts_health_check_should_refresh_oauth_without_touching_codex_backend` + reported `alive = 0` instead of `1`. +- The same RED run showed + `admin_accounts_health_check_should_reject_unsupported_stagger_ms_field` + already passed, confirming the route already rejected unknown snake_case + request fields. +- GREEN: + after implementation the health-check filter passed 3/3. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test admin_accounts_routes health_check -- --nocapture` + (3/3 filtered tests passed after RED) +- `cargo test -p codex-proxy-server --test admin_accounts_routes -- --nocapture` + (49/49 tests passed) +- `cargo test -p codex-proxy-runtime --tests -- --nocapture` + (30 runtime tests passed) +- `cargo test -p codex-proxy-server --tests -- --nocapture` + (151 server tests passed) +- `cargo test --test architecture -- --nocapture` + (27/27 architecture tests passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- `codegraph sync . && codegraph status .` + (synced 2 changed files; index up to date with 279 files, 3,955 nodes, and + 11,930 edges) +- `find . -path './.git' -prune -o -path './target' -prune -o -path './.codegraph' -prune -o -path './web/node_modules' -prune -o -type d -empty -print` + (no output) +- `rg "not wired yet|后续承载|todo!|unimplemented!|stub|placeholder" -n crates src tests --glob '!target/**'` + (only architecture placeholder guard tests matched) +- `rg "迁移|兼容|compat|legacy|shim|deprecated" -n crates src tests docs/architecture-audit.md docs/architecture.md --glob '!target/**'` + (matches were docs, architecture tests, product compatibility comments, and + old-format rejection tests) + +Current migration counters after this slice: + +- Approximately 402 `#[test]` / `#[tokio::test]` occurrences. +- 15 root Rust test files, all architecture tests. +- 39 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- The old admin `cookies_quota.rs` behavior surface audited in this pass is now + represented in crate-local tests. +- The largest remaining unmigrated/high-risk area is still the old serving and + WebSocket behavior surface, especially the historical HTTP SSE/WebSocket + suites. + +### 2026-06-19 Slice 78: Responses Default Streaming Migration + +Status: implemented and verified for this focused HTTP SSE behavior. + +Changes: + +- Migrated the old `tests/codex_serving/responses_http_sse.rs` + `v1_responses_should_default_to_streaming_when_stream_is_omitted` behavior + into current crate-local coverage: + - `OpenAiResponsesRequest` now deserializes omitted `stream` as `true`; + - explicit `"stream": false` remains supported and continues to drive the + non-streaming dispatch path; + - `/v1/responses` with omitted `stream` returns an SSE + `response.failed` plus `data: [DONE]` when no upstream account is + available, matching the streaming error surface. +- Added protocol coverage in `crates/core/tests/protocol.rs`. +- Added route coverage in `crates/server/tests/openai_responses_routes.rs`. + +TDD/coverage evidence: + +- RED: + `cargo test -p codex-proxy-core --test protocol openai_response_request_should_default_missing_stream_to_true -- --nocapture` + failed because the translated Codex request had `stream == false`. +- RED: + `cargo test -p codex-proxy-server --test openai_responses_routes responses_route_should_default_omitted_stream_to_sse -- --nocapture` + failed with HTTP `503` instead of streaming HTTP `200`. +- GREEN: + after changing the request default, both filtered tests passed. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-core --test protocol openai_response_request_should_default_missing_stream_to_true -- --nocapture` + (1/1 filtered test passed after RED) +- `cargo test -p codex-proxy-server --test openai_responses_routes responses_route_should_default_omitted_stream_to_sse -- --nocapture` + (1/1 filtered test passed after RED) +- `cargo test -p codex-proxy-core --test protocol -- --nocapture` + (47/47 tests passed) +- `cargo test -p codex-proxy-server --test openai_responses_routes -- --nocapture` + (5/5 tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream -- --nocapture` + (63/63 tests passed) +- `cargo test --test architecture -- --nocapture` + (27/27 architecture tests passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- `cargo test -p codex-proxy-server --tests -- --nocapture` + (152 server tests passed) +- `codegraph sync . && codegraph status .` + (synced 3 changed files; index up to date with 279 files, 3,958 nodes, and + 11,927 edges) + +Current migration counters after this slice: + +- Approximately 404 `#[test]` / `#[tokio::test]` occurrences. +- 15 root Rust test files, all architecture tests. +- 39 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- One old HTTP SSE default-streaming behavior is restored. +- The old Responses HTTP SSE suite still has broader gaps around review/compact + routes, tuple schema reconversion, passive rate-limit caching, cookie path + scoping, and disconnect behavior. +- The old Responses WebSocket and gateway WebSocket suites remain the largest + high-risk migration areas. + +### 2026-06-19 Slice 79: Responses HTTP SSE Review/Compact/Tuple Migration + +Status: implemented and verified as one continuous migration batch from the old +`tests/codex_serving/responses_http_sse.rs` suite. + +Changes: + +- Migrated the old review route behavior: + `/v1/responses/review` now forces `x-openai-subagent=review` into the Codex + upstream request metadata and header path instead of behaving like a plain + `/v1/responses` call. +- Migrated the old compact route behavior: + `/v1/responses/compact` now builds a `CodexCompactRequest`, dispatches to + `/codex/responses/compact`, keeps JSON response semantics, strips + Responses-only fields such as `stream`, `store`, and `prompt_cache_key`, and + uses the same account fallback classifications for rate limit/quota/auth/model + errors. +- Expanded `OpenAiResponsesRequest` back toward the old request surface so + `instructions`, `reasoning`, `tools`, `tool_choice`, + `parallel_tool_calls`, `text.format`, `include`, `client_metadata`, + context fields, and transport hints are parsed and translated. +- Restored tuple schema conversion for Responses HTTP SSE: + - tuple-shaped JSON schema is converted before the Codex upstream request; + - non-streaming completed Responses output is reconverted for the client; + - live streaming SSE `response.output_text.delta`, + `response.output_item.done`, and `response.completed` events are + reconverted event-by-event when a tuple schema is present, without buffering + the whole stream. +- Added crate-local migration coverage in: + - `crates/core/tests/protocol.rs`; + - `crates/server/tests/openai_chat_upstream.rs`. + +TDD/coverage evidence: + +- RED: + `cargo test -p codex-proxy-core --test protocol openai_` failed at compile + time because `translate_response_to_compact` did not exist. +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_review_route_should_force_review_subagent_upstream` + failed with HTTP `502` instead of `200`, proving the review route was not + forcing the upstream subagent. +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_compact` + failed because the compact route still used the normal Responses/SSE handler + and returned a non-JSON error body. +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream tuple` + showed the upstream schema conversion and non-stream response reconversion + passing, but `responses_stream_should_reconvert_tuple_schema_output_for_client` + failed because live SSE was still forwarding the Codex tuple object form + unchanged. +- GREEN: + after implementing compact dispatch, review forcing, request translation, and + tuple SSE event reconversion, the focused review/compact/tuple tests passed. + +Verification run so far: + +- `cargo fmt --all` +- `cargo test -p codex-proxy-core --test protocol openai_` + (7/7 filtered tests passed after RED) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_review_route_should_force_review_subagent_upstream` + (1/1 filtered test passed after RED) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_compact` + (3/3 filtered tests passed after RED) +- `cargo test -p codex-proxy-server --test openai_chat_upstream tuple` + (3/3 filtered tests passed after RED) +- `cargo test -p codex-proxy-server --tests` + (159 server tests passed) +- `cargo test -p codex-proxy-core --test protocol` + (50/50 protocol tests passed) +- `cargo test --test architecture` + (27/27 architecture tests passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- Empty directory scan excluding `.git`, `target`, and `web/node_modules` found + no empty source directories. A raw scan only found + `web/node_modules/.vite-temp`, which is generated dependency-cache state and + was not removed. +- Placeholder-marker scan only matched the architecture guard tests that define + the forbidden markers. +- `codegraph sync . && codegraph status .` + (synced 5 changed files; index up to date with 279 files, 4,011 nodes, and + 12,197 edges) + +Current migration counters after this slice: + +- Approximately 414 `#[test]` / `#[tokio::test]` occurrences. +- 59 Rust files contain tests. +- 15 root Rust test files, all architecture tests. +- 39 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- The review, compact, and tuple-schema portion of the old Responses HTTP SSE + suite is now represented in crate-local tests. +- The next HTTP SSE audit pass should stay on the old suite order and + cross-check passive rate-limit caching, cookie path scoping, disconnect/live + stream audit behavior, and any remaining old fixture/golden gaps before + moving away from this historical suite. +- The old Responses WebSocket and gateway WebSocket suites remain high-risk + migration areas, but they should be handled by following their original test + surfaces rather than by unrelated module-by-module cleanup. + +### 2026-06-19 Slice 80: Responses HTTP SSE Non-Stream Side-Effects Migration + +Status: implemented and verified as the next migration batch from the old +`tests/codex_serving/responses_http_sse.rs` suite, following the old +non-stream success/side-effect path instead of doing directory-driven cleanup. + +Changes: + +- Migrated old HTTP SSE non-stream side effects into crate-local server tests: + - imported account dispatch with `authorization` and `chatgpt-account-id` + headers; + - non-stream account usage persistence and `v1.response` event-log metadata; + - logging-disabled skip behavior; + - successful and failed image-generation usage counters; + - upstream `Set-Cookie` capture for `cf_clearance`; + - cookie replay scoped by `/codex/responses` path; + - passive rate-limit header caching into `accounts.quota_json`, + cooldown/verify flags, and `account_usage` window metadata; + - empty completed response retry exhaustion with three recorded empty + attempts. +- Added adapter-local cookie migration coverage in + `crates/adapters/tests/cookies.rs` for captured `cf_clearance`, path/domain + replay, and ignoring non-capturable `Set-Cookie` headers. +- Restored `SqliteCookieStore::capture_set_cookie()` with the old allowlist and + Domain/Path/Max-Age/Expires parsing semantics. +- Extended the account usage port so runtime can record empty responses and + image request success/failure counters without bypassing `core` boundaries. +- Added account store port methods for quota JSON read/write and rate-limit + window sync, with SQLite implementations ported from the old repository + behavior. +- Updated `ResponseDispatchService::complete` to capture cookies, sync passive + rate-limit state, record enhanced Responses usage, record non-stream event + logs through `AdminLogService`, and retry empty completed responses twice + before returning the upstream-response error. +- Updated runtime cookie replay to request cookies for `/codex/responses` and + `/codex/responses/compact`, preserving old path-scoped cookie ordering. + +TDD/coverage evidence: + +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_use_imported_account_record_usage_cookie_and_event_log` + failed to compile because `SqliteCookieStore::capture_set_cookie()` did not + exist, proving the old cookie capture API had not been migrated. +- RED: + after adding the adapter API, + `responses_should_use_imported_account_record_usage_cookie_and_event_log` + failed because the persisted cookie header was `None` instead of + `Some("cf_clearance=new")`, proving the non-stream runtime path still ignored + upstream `Set-Cookie` side effects. +- GREEN: + after wiring adapter/runtime/core behavior, the imported-account/cookie/log + test and the image usage, failed image attempt, cookie path, passive + rate-limit, empty-response, and logging-disabled tests all passed. + +Verification run so far: + +- `cargo test -p codex-proxy-adapters --test cookies` + (2/2 cookie tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_use_imported_account_record_usage_cookie_and_event_log` + (1/1 filtered test passed after RED) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_record_image_generation` + (1/1 filtered success-image test passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_record_failed_image_generation_attempt_when_tool_has_no_output` + (1/1 filtered failed-image test passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_scope_upstream_cookie_by_codex_response_path` + (1/1 filtered cookie-path test passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_passively_cache_rate_limit_headers` + (1/1 filtered passive rate-limit test passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_record_empty_response_attempts` + (1/1 filtered empty-response test passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_skip_event_log_when_logging_disabled` + (1/1 filtered logging-disabled test passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_` + (34/34 filtered Responses tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (77/77 upstream tests passed after final clippy cleanup) +- `cargo test -p codex-proxy-adapters --tests` + (70 adapter tests passed, including the new cookie tests) +- `cargo test -p codex-proxy-core --tests` + (104 core tests passed) +- `cargo test -p codex-proxy-runtime --tests` + (30 runtime tests passed) +- `cargo test -p codex-proxy-server --tests` + (166 server tests passed) +- `cargo test --test architecture` + (27/27 architecture tests passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- Empty directory scan excluding `.git`, `target`, `.codegraph`, and + `web/node_modules` found no empty source directories. +- Placeholder-marker scan only matched the architecture guard tests that define + the forbidden markers. +- `codegraph sync . && codegraph status .` + (synced 4 changed code files before this doc update; index up to date with + 280 files, 4,055 nodes, and 12,310 edges) + +Current migration counters after this slice: + +- Approximately 423 `#[test]` / `#[tokio::test]` occurrences. +- 60 Rust files contain tests. +- 15 root Rust test files, all architecture tests. +- 40 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- The old HTTP SSE non-stream side-effect block around imported accounts, + image usage, cookie path scoping, passive rate-limit caching, empty-response + retries, and logging-disabled behavior is now represented in crate-local + tests. +- The next HTTP SSE pass should continue down the same old suite and verify + parity fields/context headers, reasoning include defaults, sanitized + reasoning/compaction input, non-stream text reconstruction, done-item fallback, + passthrough stream logging, and downstream disconnect behavior. +- The old Responses WebSocket and gateway WebSocket suites remain high-risk, + but should be handled only when the migration line reaches those suites. + +### 2026-06-19 Slice 81: Responses HTTP SSE Parity/Reasoning/Stream Disconnect Migration + +Status: implemented and verified as the next contiguous batch from the old +`tests/codex_serving/responses_http_sse.rs` suite. This slice stayed on the +HTTP SSE migration line and did not split work by module. + +Changes: + +- Added a core pure `apply_response_model_options(...)` policy and runtime + wiring so Responses requests now apply: + - parsed model id; + - reasoning effort from the request body, model suffix, or model config; + - default reasoning summary `auto`; + - default `include: ["reasoning.encrypted_content"]` only when the client did + not send a non-empty include list; + - service tier from request body, model suffix, or config; + - upstream service-tier normalization from `fast` to `priority`. +- Added `ModelService::config()` so runtime dispatch can apply the same model + option rules without duplicating model configuration. +- Restored Responses prompt-cache/conversation identity behavior in runtime: + - `complete(...)` and `stream(...)` call `ensure_prompt_cache_key(...)` after + previous-response affinity lookup; + - Codex upstream context now uses `build_conversation_identity(...)` to pass + account-scoped `session_id` and `x-codex-window-id`. +- Restored `/v1/responses` header-context fallbacks for + `x-codex-turn-state`, `x-codex-turn-metadata`, `x-codex-beta-features`, + `x-responsesapi-include-timing-metrics`, `version`, + `x-codex-window-id`, and `x-codex-parent-thread-id`, while preserving body + fields when present. +- Restored normal-route `x-openai-subagent` propagation for valid subagent + headers; the review route still forces `review`. +- Fixed WebSocket pool reuse for previous-response requests after prompt-cache + restoration by keying the adapter pool on `previous_response_id` when the + client did not explicitly provide a prompt cache key. +- Restored downstream disconnect propagation for live HTTP SSE streams: + dropping the downstream body sends a cancellation signal, releases the account + promptly, and drops the upstream SSE body/socket. +- Added crate-local tests for parity fields/context headers, include + preservation, reasoning/compaction sanitization, text-delta reconstruction, + done-item fallback, stream disconnect, and pure model-option defaults. + +TDD/coverage evidence: + +- RED: + `cargo test -p codex-proxy-core --test protocol codex_responses_model_options_should_apply_suffix_defaults_and_include_reasoning` + failed to compile because `apply_response_model_options(...)` did not exist. +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_forward_parity_fields_context_headers_and_account_scoped_identity` + failed with upstream `service_tier` `null` instead of `"priority"`. +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should_close_http_sse_upstream_when_client_disconnects` + failed because the upstream socket stayed open after the downstream body was + dropped. +- Regression found during GREEN verification: + `responses_with_previous_response_id_should_use_websocket_and_configured_pool` + accepted two upstream WebSocket connections instead of one after prompt-cache + keys were restored; the adapter pool key now preserves old previous-response + reuse semantics. +- GREEN: + the focused core model-option tests, new server parity/reasoning/sanitization/ + reconstruction/done-item/disconnect tests, and the existing previous-response + WebSocket pool test passed after implementation. + +Verification run for this slice: + +- `cargo test -p codex-proxy-core --test protocol codex_responses_model_options` + (2/2 filtered model-option tests passed) +- `cargo test -p codex-proxy-core --test protocol` + (52/52 protocol tests passed) +- `cargo test -p codex-proxy-core --tests` + (all core tests passed) +- `cargo test -p codex-proxy-runtime --tests` + (all runtime tests passed) +- `cargo test -p codex-proxy-adapters --tests` + (70/70 adapter tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (83/83 upstream tests passed) +- `cargo test -p codex-proxy-server --tests` + (all server tests passed) +- `cargo test --test architecture` + (27/27 architecture tests passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- Empty directory scan excluding `.git`, `target`, `.codegraph`, and + `web/node_modules` found no empty source directories. +- Placeholder-marker scan matched only architecture guard tests and expected API + compatibility/alias text; no source placeholder markers were found outside + those guard/compatibility contexts. +- `codegraph sync . && codegraph status .` + (synced 7 changed files before this doc update; index up to date with + 280 files, 4,082 nodes, and 12,430 edges) + +Current migration counters after this slice: + +- Approximately 431 `#[test]` / `#[tokio::test]` occurrences. +- 60 Rust files contain tests. +- 15 root Rust test files, all architecture tests. +- 40 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- The old HTTP SSE tail around parity fields, reasoning include defaults, + sanitized reasoning/compaction input, non-stream text reconstruction, + done-item fallback, stream usage/logging coverage, and downstream disconnect + behavior is now substantially represented in crate-local tests. +- The next migration pass should first audit the earlier old HTTP SSE tests for + any assertion-level gaps that were only partially covered, then move to the + old Responses WebSocket and gateway WebSocket suites in their original + behavior order. + +### 2026-06-19 Slice 82: Responses HTTP SSE Early-Block/Stagger Migration + +Status: implemented and verified as the assertion-level pass over the earlier +block of the old `tests/codex_serving/responses_http_sse.rs` suite. This stayed +on the same historical HTTP SSE behavior line. + +Changes: + +- Added crate-local server tests for old early-block HTTP SSE behavior: + invalid JSON rejection without upstream traffic, non-object JSON rejection + without upstream traffic, no-accounts Responses error shape, explicit + `use_websocket: false` HTTP SSE transport, and same-account request + staggering before sending upstream. +- Restored the non-streaming `/v1/responses` no-available-accounts response + body to the old Responses-specific shape with top-level `type: "error"` and + `error.code: "no_available_accounts"`, matching the already-migrated compact + route. +- Restored runtime request-interval staggering: + `RuntimeAccountPoolService` now stores `request_interval_ms`, uses the + `previous_slot_at` metadata returned by the core account pool, and waits in + runtime before Chat/Responses upstream dispatch. The async sleep remains in + `runtime`; `core` still only owns pure scheduling metadata. + +TDD/coverage evidence: + +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_return_no_available_accounts_error_when_no_accounts_are_available` + failed because the non-streaming Responses path returned the generic + OpenAI-style error body, leaving top-level `type` as `null` instead of + `"error"`. +- GREEN: + after mapping `NoActiveAccount` / `AccountStore` through + `responses_no_available_accounts_response()`, the no-accounts test passed. +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_stagger_same_account_requests_before_sending_upstream` + failed because the second upstream request was sent after about 1.4ms instead + of waiting for the configured request interval. +- GREEN: + after adding runtime request-interval waiting from `previous_slot_at`, the + same stagger test passed. + +Verification run for this slice: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_reject` + (2/2 invalid-json/non-object tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_honor_explicit_http_sse_transport` + (1/1 filtered explicit HTTP SSE test passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_` + (44/44 filtered Responses tests passed) +- `cargo test -p codex-proxy-runtime --tests` + (all runtime tests passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (88/88 upstream tests passed) +- `cargo test -p codex-proxy-server --tests` + (all server tests passed) +- `cargo test --test architecture` + (27/27 architecture tests passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `git diff --check` +- Empty directory scan excluding `.git`, `target`, `.codegraph`, and + `web/node_modules` found no empty source directories. +- Placeholder-marker scan matched only architecture audit text and architecture + guard tests; no source placeholder markers were found outside those contexts. +- `codegraph sync . && codegraph status .` + (synced 3 changed files after implementation; index up to date with + 280 files, 4,089 nodes, and 12,464 edges) + +Current migration counters after this slice: + +- 436 `#[test]` / `#[tokio::test]` occurrences. +- 60 Rust files contain tests. +- 15 root Rust test files, all architecture tests. +- 40 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- The old Responses HTTP SSE suite has crate-local coverage for the early + rejection/no-account/explicit-transport/stagger block, the review/compact/ + tuple block, the non-stream side-effect block, and the parity/reasoning/ + disconnect tail. Further work should only revisit this suite for exact + assertion-level gaps found by direct old/new comparison. +- The next migration stage should move to the old Responses WebSocket and + gateway WebSocket suites in their original behavior order, not to unrelated + directory cleanup. + +### 2026-06-19 Slice 83: Responses WebSocket Default/Pool/Metadata Bulk Migration + +Status: implemented and verified as the first large contiguous block from the +old `tests/codex_serving/responses_websocket.rs` suite. + +Changes: + +- Migrated old Responses WebSocket assertions into crate-local server tests for: + default `/v1/responses` WebSocket upstream while still serving downstream SSE, + ignored camelCase `useWebSocket`, first WebSocket stream frame reaching the + downstream client before terminal completion, synthesized + `response.failed(stream_disconnected)` after early upstream WebSocket close, + recorded-conversation WebSocket pool reuse, disabled-pool non-reuse while + keeping the recorded conversation key, and WebSocket `response.metadata` + turn-state persistence for continuation requests. +- Restored the old default transport: Responses requests without + `use_websocket: false` now remain `WebSocketPreferred`; only snake_case + `use_websocket: false` forces HTTP SSE. +- Adjusted HTTP-SSE-specific server tests to opt into `use_websocket: false` + explicitly, so those tests keep covering HTTP behavior after the default + WebSocket route was restored. +- Fixed WebSocket pool key selection so continuation requests use the restored + conversation/prompt-cache key before falling back to `previous_response_id`. + This lets a response recorded in session affinity reuse the original pooled + WebSocket. +- Propagated live WebSocket `response.metadata` turn-state updates from the + adapter stream into runtime stream finalization so completed streaming + responses record the latest turn state in session affinity. + +TDD/coverage evidence: + +- RED: + `cargo test -p codex-proxy-server --test openai_chat_upstream responses_websocket -- --nocapture` + first failed because the new tests referenced a missing local helper. After + adding the helper, the same filter failed for the intended behavior gaps: + default/no-history Responses requests were still sent as HTTP POST instead of + WebSocket, disabled/default pool cases did not receive WebSocket payloads, and + WebSocket metadata turn state was not present on the continuation request. +- GREEN: + after restoring WebSocketPreferred defaults and adding live turn-state + propagation, four of the five WebSocket-filtered tests passed. The remaining + recorded-conversation pool test still opened a second socket because the + adapter pool key preferred `previous_response_id` over the restored + conversation key. +- GREEN: + after changing pool-key priority to + `prompt_cache_key -> client_conversation_id -> previous_response_id`, the + filtered WebSocket block passed 5/5. +- Regression cleanup: + the full server upstream suite initially exposed HTTP-SSE fixture pollution + after the default transport change. Those HTTP-specific tests were made + explicit with `use_websocket: false`; `previous_response_id` WebSocket tests + were kept on the WebSocketRequired path. + +Verification run for this slice: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_websocket -- --nocapture` + (RED then GREEN; final run 5/5 passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_use_websocket_upstream_by_default_while_serving_sse -- --nocapture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_ignore_camel_case_use_websocket_field -- --nocapture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` + (95/95 passed) +- `cargo test -p codex-proxy-core --test protocol websocket -- --nocapture` + (25/25 filtered WebSocket/protocol tests passed) +- `cargo test -p codex-proxy-core --test protocol` + (52/52 passed) +- `cargo test -p codex-proxy-adapters --test codex websocket_ -- --nocapture` + (26/26 filtered WebSocket adapter tests passed) +- `cargo test -p codex-proxy-adapters --test codex pooled_websocket -- --nocapture` + (7/7 filtered pooled WebSocket tests passed) +- `cargo test -p codex-proxy-runtime --tests` + (all runtime tests passed) +- `cargo fmt --all` +- `cargo fmt --all -- --check` +- `codegraph sync . && codegraph status .` + (synced 5 changed files; index up to date with 280 files, 4,100 nodes, and + 12,495 edges) + +Current migration counters after this slice: + +- 443 `#[test]` / `#[tokio::test]` occurrences. +- 60 Rust files contain tests. +- 15 root Rust test files, all architecture tests. +- 40 crate-local Rust test files under `crates/*/tests`. + +Remaining after this slice: + +- The first seven old Responses WebSocket behaviors are now represented in + crate-local server coverage. +- The next Responses WebSocket block should migrate the implicit resume, + reasoning replay, function-call replay exclusion, SQLite-restored affinity, + cross-window exclusion, replay eviction, restored-full-history retry, + admin-status pool eviction, and fallback/error-account routing cases from the + rest of old `tests/codex_serving/responses_websocket.rs`. +- The old gateway WebSocket suite remains the largest unmigrated pure + adapter/protocol behavior block after the Responses WebSocket suite. + +### 2026-06-19 Slice 84: Responses WebSocket Implicit Resume / Reasoning Replay Bulk Migration + +Status: implemented and verified as the next contiguous block from the old +`tests/codex_serving/responses_websocket.rs` suite. + +Changes: + +- Added pure serving modules for WebSocket implicit resume and reasoning replay: + `crates/core/src/serving/implicit_resume.rs` and + `crates/core/src/serving/reasoning_replay.rs`. +- Extended Responses protocol metadata extraction so completed WebSocket + responses can record replayable reasoning items, function-call IDs, + instructions hash, conversation identity, variant hash, and usage input + tokens. +- Extended session affinity so runtime can match the latest response by + conversation plus request variant, reject cross-window/cross-variant implicit + resume, and distinguish unmatched function-call-output continuations. +- Updated Responses dispatch to: + - apply implicit resume when the client omits `previous_response_id`; + - prepend cached reasoning replay items only when the continuation matches; + - restore the full original request history when implicit resume recovery is + rejected by upstream `previous_response_not_found`; + - evict reasoning replay cache entries after invalid encrypted content; + - record affinity/replay metadata after non-streaming and streaming WebSocket + completions. +- Updated admin account status changes to evict pooled WebSockets for both the + internal account id and the upstream `chatgpt-account-id`, matching the pool + key used by live sockets. + +Migrated tests: + +- `responses_websocket_should_implicitly_resume_full_history_with_reasoning_replay` +- `responses_websocket_should_not_implicitly_resume_unmatched_function_call_output` +- `responses_websocket_should_implicitly_resume_after_sqlite_affinity_restore` +- `responses_websocket_pool_should_be_evicted_after_admin_account_status_cycle` +- `responses_websocket_should_not_implicitly_resume_self_contained_function_call_replay` +- `responses_websocket_should_not_implicitly_resume_across_codex_windows` +- `responses_websocket_should_evict_reasoning_replay_after_invalid_encrypted_content` +- `responses_websocket_should_restore_full_history_when_implicit_resume_previous_response_is_missing` + +TDD/debug evidence: + +- The new WebSocket implicit-resume tests initially exposed that the translator + was forwarding the downstream non-streaming `stream: false` flag into the + upstream Codex request. The old service used the OpenAI `stream` flag only for + downstream response shape and still sent Codex Responses requests with the + upstream streaming default. Removing that assignment restored old behavior. +- The admin status-cycle test initially showed pooled sockets were not evicted + because the pool key uses upstream `chatgpt-account-id`, while the admin route + receives the internal account id. Evicting both identifiers closed the gap. + +Verification run for this slice: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream implicitly_resume -- --nocapture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream implicit_resume -- --nocapture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream reasoning_replay -- --nocapture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream affinity_restore -- --nocapture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream status_cycle -- --nocapture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream -- --nocapture` + (103/103 passed) +- `cargo test -p codex-proxy-core --tests` +- `cargo test -p codex-proxy-adapters --test codex codex_backend_client_should_close_idle_pooled_websocket_when_account_is_evicted -- --nocapture` +- `cargo test -p codex-proxy-runtime --tests` +- `cargo test -p codex-proxy-adapters --tests` +- `cargo test -p codex-proxy-server --test admin_accounts_routes -- --nocapture` + (49/49 passed) +- `cargo fmt --all` +- `cargo fmt --all -- --check` +- `cargo check --workspace --tests` +- `codegraph sync . && codegraph status .` + (index up to date with 282 files, 4,198 nodes, and 12,743 edges) + +Current migration counters after this slice: + +- 454 `#[test]` / `#[tokio::test]` occurrences. +- 62 Rust files contain tests. +- 14 root Rust test files, all architecture tests. +- 40 crate-local Rust test files under `crates/*/tests`. +- Empty source-file scan found no empty files under `tests`, `crates`, or `src`. +- Empty-directory scan found no empty directories outside ignored locations + after excluding `.git`, `target`, `.codegraph`, and `web/node_modules`. + +Remaining after this slice: + +- The old Responses WebSocket default/pool/metadata and implicit + resume/reasoning-replay blocks are now represented in crate-local server + coverage. +- The remaining Responses WebSocket block should continue with fallback and + account-status routing cases from the old suite: previous-response routing, + WebSocket 429 fallback, 401/402/model/path-block no-fallback stream errors, + quota-classified `response.failed`, and any pool-specific audit artifact + edges still not represented. +- The old gateway WebSocket suite remains a high-risk unmigrated adapter/protocol + behavior block after the Responses WebSocket suite. + +### 2026-06-19 Slice 85: Responses WebSocket Fallback / Status Routing Bulk Migration + +Status: implemented and verified for the remaining audited fallback/status +block from old `tests/codex_serving/responses_websocket.rs`. + +Changes: + +- Migrated old Responses WebSocket fallback/status assertions into + `crates/server/tests/openai_chat_upstream.rs` for: + - previous-response-id routing to the recorded account; + - WebSocket 429 fallback for streaming and non-streaming continuations; + - fallback exhaustion after mixed 429 -> 401 errors; + - WebSocket 402 quota exhaustion with no fallback; + - model-unsupported and Cloudflare path-block stream failures with no + fallback; + - WebSocket `response.failed` quota classification and fallback retry; + - rate-limit exhaustion across two fallback accounts. +- Updated WebSocket opening failures in `crates/adapters` so non-101 upgrade + responses preserve upstream HTTP status, retry-after, and body for runtime + fallback classification instead of collapsing into an opaque transport error. +- Kept `WebSocketPreferred` HTTP SSE fallback only for non-upstream transport + failures. Business/status errors from the WebSocket upgrade now reach runtime + fallback logic. +- Kept history recovery retries on WebSocket instead of forcing HTTP SSE after + stripping `previous_response_id`. Older crate-local history recovery tests + were updated to assert WebSocket -> WebSocket retry, matching the old + Responses WebSocket suite. +- Added last-exhausted-account classification for Responses dispatch fallback + exhaustion so mixed failures such as 429 followed by 401 return the final + authentication failure instead of a stale rate-limit aggregate. +- Preserved current aggregate error wording while restoring the old WebSocket + stream fragments required by the baseline for model-unsupported and + Cloudflare path-block no-fallback cases. +- Updated `docs/architecture.md` and the architecture source-tree whitelist for + `core/src/serving/implicit_resume.rs` and + `core/src/serving/reasoning_replay.rs`, which were introduced in Slice 84. + +Migrated tests added in this slice: + +- `responses_websocket_should_route_previous_response_id_to_recorded_account` +- `responses_websocket_non_stream_previous_response_not_found_should_strip_history_and_retry_same_account` +- `responses_websocket_stream_previous_response_not_found_should_strip_history_and_retry_same_account` +- `responses_websocket_non_stream_unanswered_function_call_should_strip_history_and_retry_same_account` +- `responses_websocket_previous_response_id_should_retry_fallback_account_after_429` +- `responses_websocket_non_stream_previous_response_id_should_retry_fallback_account_after_429` +- `responses_websocket_without_history_should_mark_expired_after_fallback_401` +- `responses_websocket_without_history_should_return_rate_limit_stream_error_when_fallback_accounts_exhausted` +- `responses_websocket_response_failed_quota_should_retry_fallback_account` +- `responses_websocket_without_history_should_return_quota_stream_error_when_402_has_no_fallback` +- `responses_websocket_without_history_should_return_model_unsupported_stream_error_when_no_fallback` +- `responses_websocket_with_history_should_return_path_block_stream_error_when_no_fallback` + +TDD/debug evidence: + +- Initial focused run of `responses_websocket` failed for WebSocket upgrade + status handling, history recovery transport selection, pool-reuse test + coverage, and no-fallback stream error fragments. The new tests now exercise + those paths directly. +- The mixed 429 -> 401 case identified a real runtime aggregation bug: fixed + priority returned rate-limit when the last exhausted account failed auth. The + runtime now records the last exhaustion class for Responses complete/stream + dispatch. +- The `strip_history` regression run exposed stale crate-local tests that still + expected WebSocket recovery to fall back to HTTP SSE. Those helpers now model + WebSocket failure followed by WebSocket success, matching the migrated old + Responses WebSocket behavior. +- `cargo test --test architecture` exposed that the Slice 84 core serving files + were not reflected in the architecture whitelist; the architecture document + and test whitelist now match the actual tree. + +Verification run for this slice: + +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_websocket -- --nocapture` + (25/25 passed) +- `cargo test -p codex-proxy-server --test openai_chat_upstream strip_history -- --nocapture` + (9/9 passed) +- `cargo test -p codex-proxy-adapters --test codex websocket_ -- --nocapture` + (26/26 passed) +- `cargo test -p codex-proxy-adapters --tests` +- `cargo test -p codex-proxy-runtime --tests` +- `cargo test -p codex-proxy-server --test openai_chat_upstream -- --nocapture` + (115/115 passed) +- `cargo check --workspace --tests` +- `cargo fmt --all -- --check` +- `cargo test --test architecture` + (27/27 passed) +- `cargo test --workspace --tests` +- `codegraph sync .` +- `codegraph status .` + (index up to date with 282 files, 4,228 nodes, and 12,874 edges) +- Empty file scan under `tests`, `crates`, and `src` found no empty Rust/TOML/MD + files. +- Empty-directory scan found no empty directories outside ignored `.git`, + `target`, `.codegraph`, and `web/node_modules` paths. + +Current migration counters after this slice: + +- 466 `#[test]` / `#[tokio::test]` occurrences. +- 62 Rust files contain tests. +- 15 root Rust test files, all architecture tests. +- 40 crate-local Rust test files under `crates/*/tests`. +- `crates/server/tests/openai_chat_upstream.rs` now has 115 tests. + +Remaining after this slice: + +- The audited old Responses WebSocket fallback/status block is now represented + in crate-local server coverage alongside the default/pool/metadata and + implicit-resume/reasoning-replay blocks. +- The old gateway WebSocket suite remains a high-risk adapter/protocol behavior + block for the next migration stage, especially pure codec diagnostics, + opening/audit artifacts, transport edge cases, and pool parity. +- Continue auditing old pre-migration behavior against + `f07442a28b0c186bd22b598a53cc4856ab4b2445`; passing current workspace tests + still does not by itself prove all old behavior suites have been migrated. + +### 2026-06-19 Slice 86: Gateway WebSocket Adapter/Protocol Bulk Migration + +Status: implemented and verified for the next audited block from old +`tests/codex_gateway/websocket.rs` and `tests/codex_gateway/websocket/pool.rs`. + +Changes: + +- Migrated old gateway WebSocket pool parity into + `crates/adapters/tests/codex.rs` for: + - bypassing a busy pooled key with one-shot connections; + - bypassing new conversation keys after the per-account pool cap; + - evicting idle connections when maintenance ping times out; + - garbage-collecting expired idle connections. +- Migrated old gateway WebSocket adapter/protocol boundary coverage for: + - opening failure status/body/retry-after preservation; + - binary upstream frames surfacing as explicit WebSocket errors instead of + being silently ignored until close; + - wrapped `error` frames preserving explicit status and `retry-after`; + - `websocket_connection_limit_reached` mapping to retryable 503; + - internal `codex.rate_limits` events updating captured rate-limit headers + without forwarding to downstream SSE; + - `response.metadata` turn-state updates without downstream forwarding; + - live stream partial output followed by close-before-terminal error. +- Added `CodexWebSocketExchangeError::UnexpectedBinaryEvent` and wired both + buffered and live-stream WebSocket paths to emit it for upstream binary + events before a terminal response. +- Kept Ping/Pong and Close handling separate: ping/pong/control frames continue + through tungstenite behavior, close before terminal remains + `ClosedBeforeTerminal`, and binary application data is now a distinct + protocol error. + +Migrated tests added in this slice: + +- `websocket_pool_should_bypass_busy_key_with_one_shot_connections` +- `websocket_pool_should_bypass_new_keys_after_account_cap` +- `websocket_pool_should_evict_idle_connection_when_ping_times_out` +- `websocket_pool_should_gc_expired_idle_connections` +- `websocket_execute_response_create_request_should_preserve_opening_error_status_body_and_retry_after` +- `websocket_execute_response_create_request_should_reject_binary_event` +- `websocket_execute_response_create_request_should_surface_wrapped_error_status_and_retry_after` +- `websocket_execute_response_create_request_should_surface_connection_limit_as_503` +- `websocket_execute_response_create_request_should_capture_internal_metadata_and_rate_limit_events` +- `codex_backend_client_stream_should_reject_binary_websocket_event` +- `codex_backend_client_stream_should_error_when_websocket_closes_before_terminal` + +Verification run for this slice so far: + +- `cargo fmt -p codex-proxy-adapters` +- `cargo test -p codex-proxy-adapters --test codex websocket_ -- --nocapture` + (37/37 passed) +- `cargo test -p codex-proxy-adapters --tests` + (`crates/adapters/tests/codex.rs`: 55/55 passed) +- `codegraph sync . && codegraph status .` + (index up to date with 282 files, 4,243 nodes, and 12,942 edges) +- `cargo fmt --all -- --check` +- `cargo check --workspace --tests` +- `cargo test --test architecture` + (27/27 passed) +- `cargo test --workspace --tests` + (workspace test suite passed; notable migrated suites: + `crates/adapters/tests/codex.rs` 55/55, + `crates/server/tests/openai_chat_upstream.rs` 115/115) + +Current migration counters after this slice: + +- 477 `#[test]` / `#[tokio::test]` occurrences. +- 62 Rust files contain tests. +- 15 root Rust test files, all architecture tests. +- 40 crate-local Rust test files under `crates/*/tests`. +- `crates/adapters/tests/codex.rs` now has 55 tests, including 37 WebSocket + filtered tests. + +Remaining after this slice: + +- The old gateway WebSocket transport/pool/status/internal-event block is now + represented in crate-local adapter/core coverage instead of root behavior + tests. +- Remaining gateway WebSocket risk is narrower: byte-for-byte capture harness + parity, old audit artifact edge cases not represented by current redacted + artifacts/event logs, and any assertion-level official-client event-shape + cases not already covered by the pure core validators. +- Continue auditing old pre-migration behavior against + `f07442a28b0c186bd22b598a53cc4856ab4b2445`; this slice does not close + serving fallback/recovery, admin import/export, or the remaining + non-WebSocket removed suites. + +### 2026-06-19 Slice 87: Gateway WebSocket Core Event Classification Migration + +Status: implemented and verified for the pure protocol/assertion block from old +`tests/codex_gateway/websocket.rs`. + +Changes: + +- Migrated old gateway WebSocket core assertions into + `crates/core/tests/protocol.rs` for: + - special upstream error code classification: quota, auth, banned, previous + response missing, overload, and usage-not-included; + - success-status `error` frames and unmapped `error` frames being ignored + instead of treated as terminal upstream failures; + - malformed `response.completed` shapes with missing id or incomplete usage; + - optional `null` fields matching missing-field behavior for message, + custom-tool, and context-compaction events; + - optional field type mismatches for tool search, custom tool output, image + generation, and context compaction items; + - nested shape mismatches for local-shell and web-search actions; + - invalid JSON and official event shape mismatches being skipped before SSE + forwarding. +- No production code changes were needed in this slice; the migrated assertions + validated that the current pure core validators already cover these old + gateway behaviors. + +Migrated tests added in this slice: + +- `codex_websocket_error_frame_should_classify_old_gateway_special_codes` +- `codex_websocket_error_frame_should_ignore_success_status_and_unmapped_error` +- `codex_websocket_response_completed_parse_error_should_reject_missing_id_and_incomplete_usage` +- `codex_websocket_optional_null_fields_should_match_missing_field_behavior` +- `codex_websocket_output_item_optional_fields_should_reject_old_gateway_edge_cases` +- `codex_websocket_local_shell_and_web_search_should_reject_nested_shape_mismatches` +- `codex_websocket_event_to_sse_frame_should_skip_invalid_json_and_shape_mismatches` + +Verification run for this slice so far: + +- `cargo fmt -p codex-proxy-core` +- `cargo test -p codex-proxy-core --test protocol -- --nocapture` + (59/59 passed) +- `codegraph sync . && codegraph status .` + (index up to date with 282 files, 4,250 nodes, and 12,957 edges) +- `cargo fmt --all -- --check` +- `cargo check --workspace --tests` +- `cargo test --test architecture` + (27/27 passed) +- `cargo test --workspace --tests` + (workspace test suite passed; notable migrated suites: + `crates/core/tests/protocol.rs` 59/59, + `crates/adapters/tests/codex.rs` 55/55, + `crates/server/tests/openai_chat_upstream.rs` 115/115) +- Empty Rust/TOML/MD file scan under `tests`, `crates`, and `src` found no + empty files. +- Empty-directory scan found no empty directories outside ignored `.git`, + `target`, `.codegraph`, and `web/node_modules` paths. +- Placeholder marker scan only found the architecture guard/test strings that + enforce placeholder detection. + +Current migration counters after this slice: + +- 484 `#[test]` / `#[tokio::test]` occurrences. +- 62 Rust files contain tests. +- 15 root Rust test files, all architecture tests. +- 40 crate-local Rust test files under `crates/*/tests`. +- `crates/core/tests/protocol.rs` now has 59 tests. + +Remaining after this slice: + +- The old gateway WebSocket pure event-shape/status-classification block is now + represented in `core` coverage. +- Remaining gateway WebSocket risk is primarily the old capture harness and + byte-for-byte audit artifact parity, not the core event validator or adapter + transport/status/pool behavior already migrated in Slices 86 and 87. + +### 2026-06-19 Slice 88: Gateway WebSocket Security Chain Header/Payload Migration + +Status: implemented and verified for the old gateway WebSocket +security-chain/header-body parity block. + +Changes: + +- Added adapter-level live WebSocket coverage proving the real + `CodexBackendClient` path forwards security-chain context through the + WebSocket opening and first `response.create` payload: + - `session_id` derives `x-client-request-id`, `prompt_cache_key`, + `session-id`, and `thread-id`; + - client metadata preserves string-only user metadata and drops non-string + values; + - installation id, window id, turn metadata, and parent thread id are added + to `client_metadata`; + - `x-openai-subagent` is promoted from metadata into the opening headers. +- Fixed the WebSocket request header builder so HTTP/SSE-only headers are not + forwarded into WebSocket openings: + - removes `content-type`; + - removes `accept`; + - replaces HTTP `session_id` with WebSocket `session-id` and `thread-id`. + +Migrated test added in this slice: + +- `codex_backend_client_websocket_should_forward_security_chain_headers_and_payload_fields` + +Verification run for this slice so far: + +- `cargo fmt -p codex-proxy-adapters` +- `cargo test -p codex-proxy-adapters --test codex websocket_ -- --nocapture` + (38/38 passed) +- `cargo test -p codex-proxy-adapters --tests` + (`crates/adapters/tests/codex.rs`: 56/56 passed) +- `codegraph sync . && codegraph status .` + (index up to date with 282 files, 4,252 nodes, and 12,947 edges) +- `cargo fmt --all -- --check` +- `cargo check --workspace --tests` +- `cargo test --test architecture` + (27/27 passed) +- `cargo test --workspace --tests` + (workspace test suite passed; notable migrated suites: + `crates/adapters/tests/codex.rs` 56/56, + `crates/core/tests/protocol.rs` 59/59, + `crates/server/tests/openai_chat_upstream.rs` 115/115) +- Empty Rust/TOML/MD file scan under `tests`, `crates`, and `src` found no + empty files. +- Empty-directory scan found no empty directories outside ignored `.git`, + `target`, `.codegraph`, and `web/node_modules` paths. +- Placeholder marker scan only found the architecture guard/test strings that + enforce placeholder detection. + +Current migration counters after this slice: + +- 485 `#[test]` / `#[tokio::test]` occurrences. +- 62 Rust files contain tests. +- 15 root Rust test files, all architecture tests. +- 40 crate-local Rust test files under `crates/*/tests`. +- `crates/adapters/tests/codex.rs` now has 56 tests, including 38 WebSocket + filtered tests. + +Remaining after this slice: + +- The old gateway WebSocket security-chain header/body block is now represented + in adapter coverage and production WebSocket headers no longer carry HTTP/SSE + body headers. +- Remaining gateway WebSocket risk is now focused on byte-for-byte capture + harness/audit artifact parity and any still-unmapped assertion-level cases + from the historical root suite. + +### 2026-06-19 Slice 89: Gateway WebSocket Capture Harness / Idle Timeout Migration + +Status: implemented and verified for the old gateway WebSocket capture harness, +artifact redaction, ordered first-frame serialization, and silent-upstream +timeout block. + +Changes: + +- Migrated the old gateway WebSocket capture harness assertions into + crate-local core/adapter coverage: + - `CodexWebSocketConnection::opening_request_text()` now exposes the actual + WebSocket opening request text used by the adapter for capture/audit parity. + - Adapter tests assert the raw opening request line, standard WebSocket + headers, business-header order, session/thread headers, turn metadata, and + permessage-deflate offer order. + - Core and adapter tests assert the first `response.create` WebSocket frame + is serialized in the old client field order rather than `serde_json::Map` + key order. + - The audit artifact write test now builds real redacted opening/payload + snapshots and asserts access tokens, account IDs, cookies, user prompt text, + prompt cache keys, and thread metadata are not written. +- Restored the old silent-upstream behavior: + WebSocket receive loops now fail with an explicit 20-second + `ReceiveIdleTimeout` when the upstream accepts the opening and first frame but + sends no events. The timeout is shared by buffered and live streaming + WebSocket response paths. + +TDD/coverage evidence: + +- RED: + `cargo test -p codex-proxy-adapters --test codex capture -- --nocapture` + initially failed because `CodexWebSocketConnection::opening_request_text()` + did not exist. +- RED: + after adding the raw opening API, the same capture filter failed because + `websocket_response_create_payload_text(...)` serialized the first frame in + sorted `serde_json::Map` key order, with `client_metadata` before `type`. +- RED: + `cargo test -p codex-proxy-core --test protocol codex_websocket_response_create_payload_text_should_preserve_old_field_order -- --nocapture` + failed for the same sorted field order. +- RED: + `cargo test -p codex-proxy-adapters --test codex silent -- --nocapture` + failed to compile because the adapter had no `ReceiveIdleTimeout` error + variant. +- GREEN: + the core ordered serializer, adapter raw-opening capture tests, artifact + redaction test, and silent-upstream timeout test now pass. + +Verification run for this slice so far: + +- `cargo test -p codex-proxy-core --test protocol codex_websocket_response_create_payload_text_should_preserve_old_field_order -- --nocapture` +- `cargo test -p codex-proxy-adapters --test codex capture -- --nocapture` +- `cargo test -p codex-proxy-adapters --test codex silent -- --nocapture` +- `cargo test -p codex-proxy-adapters --test codex websocket_audit_artifact_should_require_explicit_directory -- --nocapture` +- `cargo test -p codex-proxy-core --test protocol websocket -- --nocapture` + (33/33 filtered WebSocket/core protocol tests passed) +- `cargo test -p codex-proxy-adapters --test codex websocket_ -- --nocapture` + (41/41 filtered WebSocket adapter tests passed) +- `cargo test -p codex-proxy-adapters --tests` + (`crates/adapters/tests/codex.rs`: 59/59 passed) +- `cargo test -p codex-proxy-core --test protocol -- --nocapture` + (`crates/core/tests/protocol.rs`: 60/60 passed) +- `cargo fmt --all -- --check` +- `cargo check --workspace --tests` +- `cargo test --test architecture` + (27/27 architecture tests passed) +- `cargo test --workspace --tests` + (workspace test suite passed; notable migrated suites: + `crates/adapters/tests/codex.rs` 59/59, + `crates/core/tests/protocol.rs` 60/60, + `crates/server/tests/openai_chat_upstream.rs` 115/115) +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- Empty Rust/TOML/MD file scan under `tests`, `crates`, and `src` found no + empty files. +- Empty-directory scan found no empty directories outside ignored `.git`, + `target`, `.codegraph`, and `web/node_modules` paths. +- Placeholder marker scan only found architecture guard/test strings, + documented API/product compatibility wording, and model alias terminology; no + source placeholder or migration shim introduced by this slice. + +Current migration counters after this slice: + +- 488 `#[test]` / `#[tokio::test]` occurrences. +- 62 Rust files contain tests. +- 15 root Rust test files, all architecture tests. +- 40 crate-local Rust test files under `crates/*/tests`. +- `crates/adapters/tests/codex.rs` now has 59 tests. + +Remaining after this slice: + +- The old gateway WebSocket byte-for-byte capture harness is now represented by + raw opening text, first-frame order, redacted audit artifact, opening-status, + deflate, ping, binary-frame, idle-timeout, wrapped-error, connection-limit, + metadata, rate-limit, pool, previous-response, history-recovery, and + security-chain coverage. +- Remaining gateway WebSocket risk is now limited to assertion-level parity + gaps found by any later direct old/new comparison; no separate high-risk + capture harness block is currently known. + +## Verification After Current Implementation Pass + +Commands run after Slice 12: + +- `cargo fmt --all -- --check` +- `cargo test -p codex-proxy-runtime --test tasks` +- `cargo test --test architecture` +- `cargo check --workspace --all-targets` +- `cargo test -p codex-proxy-runtime --test token_refresh` +- `cargo test -p codex-proxy-runtime --test quota_refresh` +- `cargo test -p codex-proxy-adapters --test account_repository account_repository_should_update_quota_json_and_fetched_at` +- `cargo test -p codex-proxy-adapters --test codex websocket_` +- `cargo test --workspace --all-targets` +- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` +- `cargo test -p codex-proxy-adapters --test codex` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo test -p codex-proxy-core --test models` +- `cargo test -p codex-proxy-core --test session_affinity` +- `cargo test -p codex-proxy-adapters --test session_affinity` +- `cargo test -p codex-proxy-runtime --test session_affinity` +- `cargo check -p codex-proxy-server --all-targets` +- `cargo test -p codex-proxy-runtime --test account_pool_restore` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_prefer_session_affinity_account_for_previous_response` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_record_session_affinity_for_completed_response` +- `cargo check --workspace --all-targets` +- `cargo fmt --all -- --check` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-runtime --test tasks session_affinity_cleanup_task_should_delete_only_expired_affinities` +- `cargo test -p codex-proxy-runtime --test tasks start_background_tasks_should_register_migrated_runtime_tasks` +- `cargo test -p codex-proxy-runtime --test tasks` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_account_quota_should_fetch_usage_store_quota_and_not_return_secrets` +- `cargo test -p codex-proxy-server --test admin_accounts_routes` +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test admin_accounts_routes admin_account_quota_should_return_bad_gateway_when_usage_fetch_fails` +- `cargo test -p codex-proxy-server --test admin_accounts_routes` +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_fallback_to_next_account_after_rate_limit` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_fallback_to_next_account_after_rate_limit` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo fmt --all -- --check` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_fallback_to_next_account_after_rate_limit` (red check before Slice 18 implementation; failed as expected on `quota_limit_reached`) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_fallback_to_next_account_after_rate_limit` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream strip_history` (red check before Slice 37 implementation; failed 3/3 as expected because no history recovery retry was sent) +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream strip_history` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream strip_history` +- `cargo test -p codex-proxy-core --test protocol openai_response_request_should_translate_to_codex_request` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream model_unsupported` (red check before Slice 36 implementation; failed 4/4 as expected on missing model-unsupported fallback/no-fallback behavior) +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream model_unsupported` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream model_unsupported` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream cloudflare` (red check before Slice 35 implementation; failed 7/7 as expected on missing Cloudflare fallback/cooldown/cookie/path-block behavior) +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream cloudflare` +- `cargo fmt --all` +- `cargo test -p codex-proxy-core cloudflare` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` (first post-implementation run failed on the new architecture whitelist entry, then the whitelist/docs were updated) +- `cargo test --test architecture` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream cloudflare` +- `cargo test -p codex-proxy-core cloudflare` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_retry_same_account_after_5xx_before_fallback` (red check before Slice 19 implementation; failed as expected with HTTP `502`) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_retry_same_account_after_5xx_before_fallback` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_record_request_count_when_5xx_retries_are_exhausted` (red check before Slice 20 implementation; failed as expected with `request_count = 0`) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_record_request_count_when_5xx_retries_are_exhausted` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo test -p codex-proxy-runtime --test account_pool_restore` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_fallback_to_next_account_after_rate_limit` (red check before Slice 21 implementation; failed as expected with secondary token usage `(1, 0, 0)`) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_fallback_to_next_account_after_rate_limit` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo test -p codex-proxy-adapters --test account_repository` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_dispatch_to_codex_and_return_openai_response` (red check before Slice 22 implementation; failed as expected with Chat usage `(1, 0, 0)`) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_dispatch_to_codex_and_return_openai_response` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_mark_quota_exhausted_after_402_and_fallback` (red check before Slice 23 implementation; failed as expected with HTTP `502`) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_mark_quota_exhausted_after_402_and_fallback` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo test -p codex-proxy-adapters --test account_repository` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_classify_sse_quota_failure_and_fallback` (red check before Slice 24 implementation; failed as expected with HTTP `502`) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_classify_sse_quota_failure_and_fallback` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_mark_quota_exhausted_after_402_and_fallback` (red check before Slice 25 implementation; failed as expected with HTTP `502`) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_mark_quota_exhausted_after_402_and_fallback` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_return_quota_error_when_402_fallback_is_exhausted` (red check before Slice 26 implementation; failed as expected with HTTP `503`) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_return_quota_error_when_402_fallback_is_exhausted` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_return_quota_error_when_402_fallback_is_exhausted` (red check before Slice 27 implementation; failed as expected with HTTP `503`) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_return_quota_error_when_402_fallback_is_exhausted` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_return_rate_limit_error_when_429_fallback_is_exhausted` (red check before Slice 28 implementation; failed as expected with HTTP `503`) +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream chat_completions_should_return_rate_limit_error_when_429_fallback_is_exhausted` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_return_rate_limit_error_when_429_fallback_is_exhausted` (red check before Slice 29 implementation; failed as expected with HTTP `503`) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_should_return_rate_limit_error_when_429_fallback_is_exhausted` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should` (red check before Slice 30 implementation; after tightening the 5xx assertion, all four stream tests failed as expected on missing `text/event-stream`) +- `cargo test -p codex-proxy-server --test openai_chat_upstream responses_stream_should` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream 401` (red check before Slice 31 implementation; failed as expected with HTTP `502` for Chat/Responses fallback and missing stream fallback usage) +- `cargo test -p codex-proxy-server --test openai_chat_upstream sse_failed_event` (red check before Slice 31 implementation; failed as expected with HTTP `502`) +- `cargo test -p codex-proxy-server --test openai_chat_upstream 401` +- `cargo test -p codex-proxy-server --test openai_chat_upstream sse_failed_event` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream auth_error_when_401` (red check before Slice 32 implementation; failed as expected with HTTP `503`/no-active stream error) +- `cargo test -p codex-proxy-server --test openai_chat_upstream auth_error_when_401` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream deactivated` (red check before Slice 33 implementation; failed as expected with persisted status `expired`) +- `cargo test -p codex-proxy-server --test openai_chat_upstream deactivated` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` +- `cargo test -p codex-proxy-server --test openai_chat_upstream after_403` (red check before Slice 34 implementation; failed as expected with HTTP `502` and missing fallback stream response id) +- `cargo test -p codex-proxy-server --test openai_chat_upstream after_403` +- `cargo fmt --all` +- `cargo test -p codex-proxy-server --test openai_chat_upstream` +- `cargo check --workspace --all-targets` +- `cargo test --test architecture` +- `codegraph sync .` + +Result: + +- Commands listed before Slice 7 completed successfully in the previous pass. +- The Slice 7 through Slice 12 targeted commands completed successfully in the + current pass. +- `codegraph sync .` was rerun after the Slice 7 implementation and synced 12 + changed files. +- `codegraph sync .` was rerun after the Slice 8 implementation, syncing 10 + changed files, then rerun after formatting and synced 1 changed file. +- `codegraph sync .` was rerun after the Slice 9 implementation, syncing 5 + changed files, then rerun after formatting and synced 1 changed file. +- `codegraph sync .` was rerun after the Slice 10 implementation, syncing 2 + changed files, then rerun after formatting and synced 1 changed file. +- `codegraph sync .` was rerun after the Slice 11 and Slice 12 implementation + work, then rerun after fixing the `previous_response_id` protocol test and + synced 2 changed files. +- `codegraph sync .` was rerun after the Slice 13 session-affinity cleanup task + work and synced 5 changed files, then rerun after architecture + spec/whitelist documentation updates and synced 2 changed files. +- `codegraph sync .` was rerun after the Slice 14 explicit admin quota route + migration and synced 5 changed files. +- `codegraph sync .` was rerun after the Slice 15 quota upstream failure shape + migration and synced 2 changed files. +- `codegraph sync .` was rerun after the Slice 16 Responses rate-limit account + fallback migration and synced 2 changed files. +- `codegraph sync .` was rerun after the Slice 17 Chat rate-limit account + fallback migration and synced 2 changed files. +- `codegraph sync .` was rerun after the Slice 18 persisted quota cooldown + transition and synced 5 changed files. +- `codegraph sync .` was rerun after the Slice 19 Responses 5xx same-account + retry migration and synced 2 changed files. +- `codegraph sync .` was rerun after the Slice 20 selected-account request-count + persistence and synced 5 changed files. +- `codegraph sync .` was rerun after the Slice 21 Responses token usage delta + persistence and synced 6 changed files. +- `codegraph sync .` was rerun after the Slice 22 Chat token usage delta + persistence and synced 2 changed files. +- `codegraph sync .` was rerun after the Slice 23 Responses HTTP 402 quota + exhausted fallback and synced 5 changed files. +- `codegraph sync .` was rerun after the Slice 24 Responses SSE quota failure + fallback and synced 2 changed files. +- `codegraph sync .` was rerun after the Slice 25 Chat HTTP 402 quota fallback + and synced 2 changed files. +- `codegraph sync .` was rerun after the Slice 26 Chat HTTP 402 no-fallback + quota error aggregation and synced 3 changed files. +- `codegraph sync .` was rerun after the Slice 27 Responses HTTP 402 + no-fallback quota error aggregation and synced 3 changed files. +- `codegraph sync .` was rerun after the Slice 28 Chat HTTP 429 no-fallback + rate-limit error aggregation and synced 3 changed files. +- `codegraph sync .` was rerun after the Slice 29 Responses HTTP 429 + no-fallback rate-limit error aggregation and synced 3 changed files. +- `codegraph sync .` was rerun after the Slice 30 buffered Responses + `stream: true` fallback and usage work and synced 3 changed files. +- `codegraph sync .` was rerun after the Slice 31 auth failure account + recovery work and synced 2 changed files. +- `codegraph sync .` was rerun after the Slice 32 auth no-fallback expired + aggregation work and synced 4 changed files. +- `codegraph sync .` was rerun after the Slice 33 deactivated auth status + classification work and synced 2 changed files. +- `codegraph sync .` was rerun after the Slice 34 HTTP 403 banned fallback + work and synced 2 changed files. +- `codegraph sync .` was rerun after adding the Slice 35 Cloudflare tests and + synced 1 changed file, then rerun after implementation/formatting and synced + 10 changed files. +- `codegraph sync .` was rerun after adding the Slice 36 model-unsupported + tests and synced 1 changed file, then rerun after implementation/formatting + and synced 2 changed files. +- `codegraph sync .` was rerun after adding the Slice 37 history recovery tests + and synced 1 changed file, then rerun after implementation/formatting and + synced 4 changed files. +- `codegraph sync .` was rerun after Slice 39 token refresh lease coordination + implementation and synced 7 changed files. +- `codegraph sync .` was rerun after Slice 40 token refresh refreshing recovery + implementation and synced 3 changed files. +- `codegraph sync .` was rerun after Slice 41 token refresh retry confirmation + implementation and synced 2 changed files. +- `codegraph sync .` was rerun after Slice 42 token refresh in-flight guard + implementation and synced 2 changed files. +- `codegraph sync .` was rerun after Slice 43 token refresh delayed recovery + implementation and synced 2 changed files. +- `codegraph sync .` was rerun after Slice 44 Responses stream no-fallback + coverage and synced 1 changed file. +- `codegraph sync .` was rerun after Slice 45 admin quota edge-case coverage + and synced 1 changed file. +- `codegraph sync .` was rerun before Slice 49 token refresh timer work and + synced the 2 changed WIP files already present at the start of the pass. +- `codegraph sync . && codegraph status .` was rerun after Slice 49 token + refresh timer implementation and reported the index up to date with 271 files, + 3,213 nodes, and 8,855 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 50 Responses + stream done termination and reported the index up to date with 271 files, + 3,215 nodes, and 8,879 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 51 WebSocket + transport decision and audit snapshots and reported the index up to date with + 271 files, 3,242 nodes, and 8,977 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 52 WebSocket + core codec diagnostics and reported the index up to date with 271 files, + 3,289 nodes, and 9,059 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 53 WebSocket + core item validators and reported the index up to date with 271 files, 3,324 + nodes, and 9,144 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 54 WebSocket + opening request descriptor and reported the index up to date with 271 files, + 3,337 nodes, and 9,176 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 55 minimal live + WebSocket exchange and reported the index up to date with 271 files, 3,356 + nodes, and 9,224 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 56 Codex + backend WebSocket required path and reported the index up to date with 271 + files, 3,365 nodes, and 9,276 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 57 WebSocket + live invalid-event filtering and reported the index up to date with 271 + files, 3,367 nodes, and 9,295 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 58 WebSocket + live control-frame and handshake metadata work and reported the index up to + date with 271 files, 3,378 nodes, and 9,334 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 59 WebSocket + ping coverage and audit artifact IO and reported the index up to date with + 271 files, 3,404 nodes, and 9,433 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 60 WebSocket + deflate frame rewriter work and reported the index up to date with 271 files, + 3,422 nodes, and 9,482 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 61 live HTTP + SSE body forwarding and reported the index up to date with 271 files, 3,456 + nodes, and 9,738 edges. +- `codegraph sync . && codegraph status .` was rerun after Slice 62 WebSocket + pool reuse/error-discard work and reported the index up to date with 271 files, + 3,493 nodes, and 9,871 edges after lint cleanup. +- `codegraph sync . && codegraph status .` was rerun during Slice 63 WebSocket + pool lifecycle/runtime wiring work and reported the index up to date with 272 + files, 3,535 nodes, and 10,058 edges before the documentation-only update. +- `codegraph sync . && codegraph status .` was rerun during Slice 64 server + WebSocketRequired work and reported the index up to date with 272 files, 3,550 + nodes, and 10,120 edges before the documentation-only update. +- `codegraph sync . && codegraph status .` was rerun after Slice 65 live + WebSocket streaming/deflate work and reported the index up to date with 272 + files, 3,595 nodes, and 10,382 edges before the documentation-only update. +- `codegraph sync . && codegraph status .` was rerun after Slice 66 live stream + late-failure SSE reporting work and reported the index up to date with 272 + files, 3,609 nodes, and 10,418 edges before the documentation-only update. +- `codegraph sync . && codegraph status .` was rerun after Slice 67 live stream + event-log audit work and reported the index up to date with 272 files, 3,623 + nodes, and 10,460 edges before the documentation-only update. +- `codegraph sync . && codegraph status .` was rerun after Slice 68 stream audit + metadata/log-policy work and reported the index up to date with 272 files, + 3,652 nodes, and 10,554 edges before the documentation-only update. +- `codegraph sync . && codegraph status .` was rerun after Slice 69 serving toy + helper cleanup and reported the index up to date with 272 files, 3,656 nodes, + and 10,555 edges before the documentation-only update. +- `codegraph sync . && codegraph status .` was rerun after Slice 70 admin + settings/test-migration guard work and reported the index up to date with 274 + files, 3,733 nodes, and 10,666 edges before this documentation-only update. +- `codegraph sync . && codegraph status .` was rerun after Slice 71 admin + client-key assertion coverage migration and reported the index up to date with + 274 files, 3,737 nodes, and 10,716 edges before this documentation-only + update. + +Historical high-risk gap log before the current full-audit pass: + +- Full WebSocket parity still needs a focused pass for byte-for-byte capture + harness/audit artifact edge cases not represented by the restored opening, + live streaming, deflate, pool, ping, binary-frame rejection, opening-status + preservation, wrapped-error status/retry-after handling, connection-limit + mapping, history-recovery, previous-response, event-log, rate-limit-update, + metadata, special error-code classification, optional/null event field + handling, invalid event-shape skipping, security-chain header/body forwarding, + HTTP-header exclusion from WebSocket openings, default/pool/metadata route + behavior, and implicit resume/reasoning-replay coverage. Live WebSocket body + forwarding, streaming `previous_response_id` dispatch, live deflate stream + integration, actual WebSocket stream audit transport marking, internal + `codex.rate_limits` update snapshots, old gateway WebSocket + pool/status/internal-event parity, security-chain parity, pure + event-shape/status-classification parity, and the latest Responses WebSocket + implicit-resume/replay cases are now implemented and covered by tests. +- Serving dispatch/fallback/recovery remains materially reduced, though + non-streaming Chat and Responses now retry the next account after an HTTP 429 + and persist the quota cooldown transition for that fallback path, and + non-streaming Responses now retry the same account for transient upstream + HTTP 5xx responses. HTTP SSE `stream: true` now has a live chunk-by-chunk + success path, first-event SSE failure fallback preservation, completed usage + persistence, HTTP 429/402 fallback, same-account 5xx retry, and + client-visible late-failure `response.failed(stream_disconnected)` reporting. + Persisted live stream event-log audit is restored for completed and + late-failure streams, including body-capture policy, usage, rate-limit + snapshots, and actual WebSocket transport metadata. Full old stream audit + artifacts remain incomplete only for any pool-specific artifact details that + are separate from event-log metadata. +- Selected-account request-count persistence is restored, and persisted token + usage deltas are restored for non-streaming Responses and Chat completions; + buffered HTTP SSE completed usage is restored for `stream: true`; live + streaming usage audit remains incomplete. +- Non-streaming Responses HTTP 402 and SSE quota failure fallback are restored, + and buffered HTTP SSE stream HTTP 402 fallback plus no-fallback aggregation + coverage are restored; buffered Responses stream success/error bodies now + append `[DONE]`; true live chunk-by-chunk stream behavior remains incomplete. +- Multi-account and no-fallback Chat/Responses HTTP 402 quota behavior is + restored for non-streaming paths, and buffered HTTP SSE stream HTTP 402 + fallback is restored; live streaming/SSE no-fallback aggregation remains + incomplete. +- Chat and Responses HTTP 429 no-fallback aggregation is restored for + non-streaming paths; buffered HTTP SSE stream HTTP 429 multi-account fallback + and no-fallback aggregation coverage are restored. +- HTTP 401 and token-invalid SSE account recovery are restored for the covered + Chat/Responses paths, including no-fallback expired aggregation and + deactivated/banned status classification; ordinary HTTP 403 banned fallback + is restored for covered Chat/Responses paths; Cloudflare challenge/path-block + recovery is restored for the covered Chat/Responses HTTP and buffered SSE + paths, but live stream/WebSocket Cloudflare parity remains incomplete. +- Responses model-unsupported fallback is restored for covered non-streaming and + buffered SSE paths, including exhausted fallback error mapping; Chat HTTP + model-unsupported fallback is restored for the covered non-streaming path. +- Responses request-history recovery is restored for covered SSE failure paths + in non-streaming and buffered SSE modes; invalid encrypted reasoning replay + recovery is restored for covered HTTP error, non-streaming SSE failure, + buffered SSE failure, and WebSocket implicit-resume replay paths. Responses + WebSocket implicit resume, reasoning replay, cross-window rejection, + function-call replay exclusion, SQLite-restored affinity, replay eviction, + and admin status-cycle pool eviction are covered by crate-local server tests. +- Token refresh now uses the SQLite refresh lease table during scans, persists + `refreshing` before scheduled refresh calls, recovers persisted `refreshing` + accounts after restart, restores `active` after retry exhaustion, retries + transient transport failures, and confirms permanent failures twice before + persisting terminal statuses. Process-local in-flight protection is also + restored for concurrent scans, retry exhaustion now schedules a process-local + delayed recovery window, and the old per-account timer map/next-refresh + scheduling behavior is restored for background scheduling. +- Quota refresh now covers the old active quota-locked scan, per-account minimum + refresh window, usage fetch/persist path, and inter-account request spacing. +- Fingerprint update startup/update coverage is migrated at adapter and runtime + wrapper levels. +- Admin explicit quota route happy path, upstream failure shape, inactive-account + rejection, persistence-failure response, quota-warning auth, and + invalid/below-threshold quota-warning edge coverage are restored. +- Admin settings `PATCH /api/admin/settings` retained-field update behavior is + restored for the current architecture: pure validation in `core`, local YAML + overlay writing in `platform`, mutable same-process config view in `runtime`, + and HTTP parsing/error mapping in `server`. +- Assets static serving is no longer an empty scaffold: the assets crate now + serves SPA/index and `/assets/*` files with cache/security headers, and the + server router uses it as the final fallback behind API routes. +- Root behavior tests are now guarded against returning under `tests/`, and the + audited high-risk old suites have crate-local migration-target evidence. + Remaining risk is assertion-level parity inside those broad suites, not the + directory placement itself. + +## What Is In Good Shape + +- Root tracked Rust source has no remaining `src/` facade. +- Architecture boundary tests import member crates directly: + - `codex_proxy_core` + - `codex_proxy_runtime` + - `codex_proxy_server` +- Tracked legacy source directories such as `src/codex`, `src/admin`, + `src/runtime`, and `src/platform` are removed. +- Current top-level Rust tests are architecture tests only. Shared fixtures + remain under `tests/fixtures`, which is acceptable as shared test data. +- The current workspace has crate-local tests under `crates/*/tests`. +- `assets` now owns real SPA/static-file routing and header helpers, and server + uses it as the final fallback after OpenAI/admin routes. + +## Historical Blocking Findings + +The findings in this section were the original blockers found early in the +migration. They are retained to show the audit trail, but the current verdict +and the 2026-06-19 full-audit pass above supersede any item that has since been +closed. + +### 1. Server Entrypoint Startup Was Incomplete + +Audit finding before Slice 2: + +- `crates/server/src/main.rs` is `fn main() {}`. + +Baseline: + +- `src/main.rs` loaded configuration, initialized tracing, built state, restored + accounts, started background tasks, served Axum, handled `ctrl_c`, and shut + tasks down. + +Impact: + +- The server binary does not implement the real startup path. +- This alone prevents claiming the runtime/server migration is complete. + +Progress: + +- Slice 2 replaced the empty main with config loading, logging initialization, + SQLite connection, secret/API key material loading, runtime state + construction, Axum serving, and `ctrl_c` graceful shutdown. +- Slice 3 wired the server entrypoint to start the runtime background task + coordinator and shut it down after Axum graceful shutdown returns. + +Remaining: + +- Startup now restores persisted accounts into the runtime pool and restores + session-affinity state. Chat/Responses dispatch uses the restored account + pool; Responses dispatch also uses restored session affinity for + preferred-account selection and records completed response affinities. +- Baseline parity remains incomplete in remaining serving fallback/recovery + branches, quota/account-status WebSocket routing, stream audit artifact edge + cases, and broader behavior-test parity. Responses WebSocket implicit + resume/reasoning replay is now covered in Slice 84. + +### 2. Xtask Commands Were Not Wired + +Audit finding before Slice 2: + +- `crates/xtask/src/build_web.rs` prints + `web build orchestration is not wired yet`. +- `crates/xtask/src/release.rs` prints + `release packaging is not wired yet`. +- `crates/xtask/src/check_architecture.rs` only prints + `cargo test --test architecture`; it does not run it. + +Impact: + +- `xtask` exists as a scaffold, not as the automation crate described in + `docs/architecture.md`. + +Progress: + +- Slice 2 wired: + - `check-architecture` -> `cargo test --test architecture`; + - `build-web` -> `pnpm install --frozen-lockfile` and `pnpm build`; + - `release` -> Rust format/test/clippy gates plus web build. + +Remaining: + +- Keep `release` in the final verification gate after the remaining migration + slices stabilize. + +### 3. Runtime Tasks Were Mostly Empty + +Current after Slice 4: + +- `crates/runtime/src/tasks/token_refresh.rs` contains a real refresh scan, + persistence path, and periodic task loop. +- `crates/runtime/src/tasks/quota_refresh.rs` contains a real quota-locked + account scan, Codex usage fetch, normalized quota persistence path, and + periodic task loop. + +Baseline examples: + +- `src/codex/tasks/token_refresh.rs`: 882 lines. +- `src/codex/tasks/quota_refresh.rs`: 117 lines. +- `src/codex/tasks/cookie_cleanup.rs`: 66 lines. +- `src/admin/tasks/session_cleanup.rs`: 59 lines. + +Impact: + +- Background task behavior has not been fully migrated. +- The current task coordinator cannot be treated as equivalent to the previous + scheduler set. + +Progress: + +- Slice 3 reintroduced runtime task structs for cookie cleanup and admin session + cleanup, plus SQLite cleanup methods and runtime tests. +- Slice 3 also wired the top-level coordinator to start `cookie_cleanup`, + `session_cleanup`, and `model_refresh`, and wired server shutdown to stop the + coordinator. +- Slice 4 implemented token/quota task loops and wired both into the top-level + coordinator. +- Slice 5 wired fingerprint update into the top-level coordinator. +- Slice 13 wired expired session-affinity cleanup into the top-level + coordinator. +- Slice 40 restored token-refresh `refreshing` persistence before refresher + calls, restart/crash recovery from persisted `refreshing`, and active + restoration after transport failure. +- Slice 41 restored token-refresh retry attempts, transient retry success, and + two-hit permanent failure confirmation. +- Slice 42 restored token-refresh process-local in-flight tracking for + duplicate concurrent scans. +- Slice 43 restored token-refresh delayed recovery after retry exhaustion. +- Slice 45 migrated explicit admin quota route edge-case coverage for inactive + accounts, persistence failure, quota-warning auth, and invalid/below-threshold + snapshots. +- Slice 46 migrated runtime-level fingerprint update startup coverage for + initial appcast apply, no-update no-op, and `repository: None` semantics. +- Slice 47 restored quota refresh inter-account request spacing with the old + 3-second default. +- Slice 49 restored token-refresh per-account timer scheduling, immediate due + refresh, planned-trigger-time callbacks, and next-refresh scheduling after a + successful scheduled refresh. + +Remaining: + +- No token-refresh-task-specific gap is currently known after restoring lease + coordination, refreshing recovery, retry/permanent-failure handling, + in-flight protection, delayed recovery, and per-account timer parity. +- No further quota-refresh-task-specific gap was found in this pass after + restoring request spacing; broader quota helper/admin behavior should continue + to be covered through serving/admin tests. + +### 4. WebSocket Implementation Was Reduced To Shells + +Current after Slice 86: + +- `crates/adapters/src/codex/websocket/connect.rs`: now stores endpoint and + ordered opening headers, and can emit an opening audit snapshot with request + line plus redacted sensitive headers. It also builds Responses WebSocket + endpoints, standard opening header descriptors, and a prepared request that + pairs the opening descriptor with the first `response.create` text frame. It + now has a minimal one-shot live exchange that opens the WebSocket, sends the + prepared text frame, buffers public events into SSE until terminal, extracts + usage, skips invalid public stream events through the migrated core + validators, captures handshake turn-state/cookies/rate-limit headers, handles + `response.incomplete` and malformed `response.completed` frames, ignores + success-status `error` frames, responds to server ping frames through the + tungstenite read loop, rejects upstream binary application frames as explicit + protocol errors, preserves non-101 opening status/body/retry-after, captures + internal metadata and `codex.rate_limits` events without downstream forwarding, + forwards security-chain headers/body metadata without leaking HTTP/SSE-only + `content-type`/`accept` headers into the WebSocket opening, emits + environment-gated redacted audit artifacts, and surfaces classified upstream + error frames including wrapped errors and connection-limit frames. +- `crates/adapters/src/codex/websocket/pool.rs`: now owns an async idle/busy + connection pool keyed by base URL, account id, and conversation id. It supports + max-age/max-per-account policy, acquire/put/discard, explicit idle GC, reuse of + successful sockets, one-shot bypass for busy keys or over-cap new keys, + maintenance ping/liveness cleanup, and discard of errored sockets. +- `crates/adapters/src/codex/websocket/deflate.rs`: now detects + permessage-deflate responses and contains tested raw payload inflation plus + compressed server data-frame rewriting. It is not yet inserted into the live + connection path. +- `crates/adapters/src/codex/websocket/opening.rs`: simple audit snapshot helper. + It also provides explicit and environment-gated WebSocket audit artifact + writes. +- `crates/core/src/protocol/codex/websocket.rs`: basic frame wrapper + encode/decode plus payload/opening audit snapshot types and Responses + payload redaction; pure event-to-SSE conversion; metadata turn-state + extraction; terminal-event detection; error-frame classification; wrapped + retry-after parsing; `response.completed` shape validation; and migrated + stream/output-item diagnostics for the old core bad-frame categories. It now + also owns pure WebSocket audit artifact and parity-diff structures plus status + classification for old gateway wrapped errors and connection-limit frames. +- `crates/core/src/serving/responses.rs`: restores the old Responses transport + decision policy for `HttpSse`, `WebSocketPreferred`, and + `WebSocketRequired`. +- `crates/adapters/src/codex/client.rs`: `create_response(...)` now honors the + core transport policy and uses the live WebSocket exchange for required + WebSocket requests, while preserving HTTP/SSE fallback when allowed. + `create_response_stream(...)` now exposes a live SSE byte stream for the + OpenAI `stream: true` path and propagates WebSocket stream protocol errors + through the stream item. + +Baseline: + +- `src/codex/gateway/transport/websocket/codec.rs`: 1,147 lines. +- `src/codex/gateway/transport/websocket/mod.rs`: 896 lines. +- `src/codex/gateway/transport/websocket/pool.rs`: 496 lines. +- `src/codex/gateway/transport/websocket/opening.rs`: 495 lines. +- `src/codex/gateway/transport/websocket/deflate.rs`: 273 lines. +- `src/codex/gateway/transport/websocket/audit.rs`: 229 lines. + +Impact: + +- WebSocket transport, pooling, opening handshake, live forwarding integration, + and most adapter/protocol parity behavior are now migrated into crate-local + coverage. +- Remaining risk is concentrated in old byte-for-byte capture harness and audit + artifact edge cases rather than the previously empty/shell adapter behavior. + +Progress: + +- Slice 6 removed the empty adapter structs and added basic tests/guards for the + new connection and pool policy types. +- Slice 51 restored transport decision and opening/payload audit snapshot + primitives. +- Slice 52 restored a first core-only batch of WebSocket codec diagnostics and + error classification. +- Slice 53 restored the remaining large batch of old pure output-item and + stream bad-frame validators into core. +- Slice 54 restored adapter-level endpoint conversion, standard opening header + descriptor construction, and response-create text-frame preparation. +- Slice 55 restored a minimal adapter-level live WebSocket one-shot exchange and + response-failed error classification. +- Slice 56 wired `CodexBackendClient::create_response(...)` to use that live + WebSocket path for required Responses requests. +- Slice 57 wired the migrated core bad-frame validators into live WebSocket + event forwarding so invalid stream events are dropped before SSE encoding. +- Slice 58 restored live handshake metadata propagation, `response.incomplete` + error mapping, malformed `response.completed` validation, invalid terminal + skip semantics, and success-status `error` frame ignoring. +- Slice 59 migrated server-ping response coverage and restored core/adapters + WebSocket audit artifact construction, diffing, explicit writes, and + environment-gated client wiring. +- Slice 60 restored tested frame-level permessage-deflate inflation/rewrite + helpers in the adapter crate. +- Slice 61 restored live HTTP SSE body forwarding for normal `stream: true` + success paths while preserving first-event `response.failed` fallback/recovery + before downstream bytes are sent. +- Slice 62 restored adapter-level WebSocket pool reuse for successful + same-account/same-conversation requests and discard after upstream/transport + errors. +- Slice 63 restored adapter-level WebSocket pool keepalive maintenance, + liveness timeout cleanup, shutdown/evict-account APIs, and runtime + `ws_pool` config injection into `CodexBackendClient`. +- Slice 64 restored non-streaming server/runtime `previous_response_id` + dispatch through WebSocketRequired, including session-affinity header coverage, + history-recovery retry behavior, and configured pool reuse. +- Slice 65 restored live WebSocket streaming body forwarding, streaming + `previous_response_id` WebSocketRequired dispatch, live permessage-deflate + insertion before tungstenite parsing, and WS-first streaming history recovery + before downstream bytes are returned. +- Slice 66 restored serving-layer late-failure SSE reporting for live streams + after downstream bytes have already been sent. Late body read errors and + premature EOF without a terminal event now append + `response.failed(stream_disconnected)` and `[DONE]` instead of surfacing a + transport error to downstream clients. +- Slice 67 restored persisted `v1.response` event-log audit records for live + stream completion and late-failure outcomes, including request id, account id, + model, route, status, latency, and stream/failure metadata. +- Slice 68 restored stream audit metadata and log policy parity for body capture + cleaning/preservation, admin log state updates and capacity trimming, + completed/failed stream usage metadata, rate-limit header snapshots, actual + WebSocket transport marking, and live WebSocket `codex.rate_limits` update + propagation into persisted event logs. +- Slice 69 removed or replaced the audited serving toy helpers and added an + architecture guard so `should_fallback`, `is_recoverable`, + `prefers_websocket`, or an uncalled `quota_reached` cannot silently return. +- Slice 86 restored the next old gateway WebSocket adapter/protocol batch: + pool busy/over-cap/maintenance GC behavior, opening failure status/body + preservation, binary-frame errors, wrapped-error status/retry-after + propagation, connection-limit 503 mapping, internal metadata/rate-limit event + capture, and stream close-before-terminal errors. +- Slice 87 restored the old gateway WebSocket pure protocol assertion batch: + special error-code classification, success-status/unmapped error skipping, + malformed `response.completed` validation, optional `null` handling, + optional-field type rejection, nested local-shell/web-search validation, and + invalid JSON/event-shape skipping before SSE forwarding. +- Slice 88 restored the old gateway WebSocket security-chain header/payload + block and fixed production WebSocket header construction so HTTP/SSE-only + `content-type`, `accept`, and `session_id` are not sent in the opening. + WebSocket openings now send `session-id` and `thread-id` when session context + is available. + +Remaining: + +- Remaining old WebSocket behavior-test migration still needs a focused pass for + byte-for-byte capture harness and audit artifact edge cases not represented by + the restored live, pool, deflate, ping, binary-frame, opening-status, + wrapped-error, connection-limit, history-recovery, previous-response, + late-failure, event-log, metadata, security-chain, rate-limit-update, and pure + event-shape validator coverage. Full persisted stream audit parity with the + old `StreamAudit` / `WebSocketStreamAudit` artifacts remains open only for + artifact details that are separate from event-log metadata and redacted + WebSocket audit snapshots. + +### 5. Serving Dispatch Logic Required Bulk Migration + +Current symptoms: + +- The previously audited uncalled toy helpers in + `crates/core/src/serving/fallback.rs`, + `crates/core/src/serving/recovery.rs`, and + `crates/core/src/serving/responses.rs` were removed or replaced in Slice 69. + Runtime production code now calls the remaining pure serving policy helpers + for rate-limit/quota/transient status classification, same-account retry + status classification, and quota-warning threshold matching. +- Architecture tests now fail if those exact toy helper names return or if + `quota_reached(...)` loses its non-test production caller. +- `crates/runtime/src/services.rs` dispatch paths now acquire accounts from the + restored runtime account pool and Responses dispatch uses restored + session-affinity state. HTTP SSE `stream: true` now forwards normal success + bodies live, records usage/session affinity at stream end, and converts late + stream disconnects into SSE `response.failed` events. It also writes + persisted `v1.response` event logs for completed and late-failed live streams, + but full old stream-audit artifacts remain materially reduced. +- Responses WebSocket implicit resume and reasoning replay now run through + runtime production paths, record replay/affinity metadata after completions, + restore full history on implicit-resume rejection, and evict replay state after + invalid encrypted content. + +Baseline examples: + +- `src/codex/serving/dispatch/mod.rs`: 2,063 lines. +- `src/codex/serving/responses.rs`: 1,313 lines. +- `src/codex/serving/dispatch/fallback.rs`: 632 lines. +- `src/codex/serving/dispatch/affinity.rs`: 698 lines. +- `src/codex/serving/dispatch/stream.rs`: 638 lines. +- `src/codex/serving/dispatch/reasoning_replay.rs`: 288 lines. + +Impact: + +- Account fallback, quota state transitions, stream audit artifacts, and the + remaining recovery transitions are not fully proven to be preserved. Affinity, + Responses implicit resume, reasoning replay, and the audited Responses + WebSocket fallback/status-routing block now have dedicated migrated server + coverage, but the old gateway WebSocket adapter/protocol suite remains + high-risk. +- The old behavior should not be considered migrated until these paths are + reimplemented in the intended `core`/`runtime`/`adapters` split and covered by + tests. + +### 6. Core Filesystem IO Violation + +Audit finding before Slice 1: + +- `crates/core/src/gateway/installation.rs` imports `std::fs`, `Path`, + `PathBuf`, and uses `dirs::home_dir()`. +- `crates/core/Cargo.toml` depends on `dirs`. + +Architecture rule: + +- `core` must not own filesystem paths, environment reads, or concrete IO. + +Impact: + +- This is a direct architecture violation. +- Installation ID persistence and `~/.codex/installation_id` compatibility + should move to `platform` or `adapters`, with `core` retaining only pure + identity rules if needed. + +Progress: + +- Fixed in Slice 1 by moving file resolution/persistence to + `platform::storage` and keeping only UUID generation/parsing in `core`. +- Added an architecture test guard so this class of `core` IO dependency + regresses visibly. + +### 7. Test Migration Was Incomplete + +Baseline test count: + +- Approximately 471 `#[test]` / `#[tokio::test]` occurrences. +- 66 top-level Rust test files. + +Current test count: + +- 495 `#[test]` / `#[tokio::test]` occurrences. +- 65 Rust files contain tests. +- 15 root Rust test files, all architecture tests. +- 43 crate-local Rust test files under `crates/*/tests`. + +Large removed baseline tests include: + +- `tests/codex_gateway/websocket.rs`: 4,983 lines, 79 tests. +- `tests/codex_serving/responses_websocket.rs`: 2,863 lines, 28 tests. +- `tests/codex_serving/responses_http_sse.rs`: 1,766 lines, 28 tests. +- `tests/codex_serving/upstream_fallback.rs`: 2,068 lines, 27 tests. +- `tests/codex_serving/chat_completions.rs`: 910 lines, 12 tests. +- `tests/admin/accounts/import_export.rs`: 863 lines, 15 tests. +- `tests/admin/accounts/oauth.rs`: 678 lines, 16 tests. +- `tests/admin/client_keys_route.rs`: 637 lines, 7 tests. + +Impact: + +- This was a valid blocker before the full mapping audit. +- The current architecture guard now maps all 46 old root behavior files to + crate-local test files. Remaining risk is assertion-level parity inside those + broad suites, not a missing crate-local migration target. + +### 8. Architecture Tests Did Not Catch Empty Implementations + +Current architecture tests verify: + +- Directory whitelist. +- Root package has no facade target; boundary tests import member crates directly. +- Workspace dependency direction. +- Some forbidden imports and legacy path references. +- Obvious placeholder markers: empty `main`, `not wired yet`, and `后续承载`. +- Server startup owns background task lifecycle and injects the runtime + installation ID. +- Specific serving toy helper regressions closed in Slice 69: + `should_fallback`, `is_recoverable`, `prefers_websocket`, and an uncalled + `quota_reached`. + +They do not currently verify: + +- Full assertion-level parity for every old behavior test. +- Generic semantic completeness beyond the explicit placeholder/empty-shell + markers and full baseline test-suite mapping. + +Impact: + +- This was a valid concern before the full source scan and test-migration + mapping guard. Current scans found no empty source files/directories or empty + function/module shells in `crates src tests`, and architecture tests now + reject root behavior tests outside the approved architecture/fixture shape. + +## Tests And Fixtures Status + +Current root `tests/` after local empty-directory cleanup: + +- `tests/architecture/*`: root-level architecture tests. +- `tests/fixtures/*`: shared fixtures. + +This shape is reasonable if the policy is: + +- crate behavior tests live under `crates//tests`; +- root `tests/` keeps only cross-workspace architecture tests and shared + fixtures. + +However, many old behavior tests still need to be migrated into the appropriate +crate-level test suites. + +## Migration Compatibility / Legacy Surface + +Tracked legacy source modules are removed from `src/`, which is aligned with the +architecture goal of not keeping old root modules. + +Remaining concerns: + +- Current Rust source scan does not find `后续承载` or `not wired yet` outside + the architecture placeholder guard itself. +- Some older documents still reference pre-migration paths. Those are not source + code, but they should be marked historical or updated after the migration is + actually complete. +- `serde(alias = ...)` and OpenAI-compatible wording in current source are API + compatibility features, not migration compatibility shims. +- Codex Desktop installation ID lookup is product interoperability with the + upstream client identity model, not a migration compatibility shim. + +## Recommended Completion Plan + +1. Restore the real binary startup path in `crates/server/src/main.rs` and keep + bootstrap/state assembly inside `runtime`. **Startup config/logging/storage, + background-task ownership, installation ID injection, account-pool restore, + and session-affinity restore are done in Slices 2, 3, 7, 8, and 9.** +2. Move installation ID persistence out of `core`; add an architecture test that + forbids `std::fs`, `std::env`, `dirs::`, `Path`, and `PathBuf` in `core/src`. + **Done in Slice 1; startup/request-context injection done in Slice 7.** +3. Re-migrate serving dispatch behavior into the intended split: + - pure decisions in `core/src/serving`; + - IO through `core` ports; + - adapter calls in `adapters`; + - wiring in `runtime`; + - HTTP mapping in `server`. +4. Re-migrate WebSocket transport and pooling: + - protocol frame types and pure parsing in `core`; + - `tokio-tungstenite`, TLS opening, pooling, deflate, and audit artifact IO in + `adapters`. +5. Rebuild background tasks in `runtime/src/tasks`. +6. Either implement `assets` and `xtask` according to `docs/architecture.md`, or + change the architecture spec before treating them as complete. **`xtask` + wiring is done in Slice 2; `assets` static serving and server fallback wiring + are done in Slice 48.** +7. Migrate deleted behavior tests into `crates/*/tests`, prioritizing: + - WebSocket codec/transport/pool tests; + - Responses HTTP SSE and WebSocket tests; + - upstream fallback/recovery tests; + - admin account OAuth/import/export/cookie/quota tests; + - startup/background task tests. +8. Add architecture tests that fail on: + - empty required files; + - `not wired yet` strings; **done for Rust sources in Slice 2** + - empty `main`; **done for Rust sources in Slice 2** + - uncalled placeholder functions in required architecture modules; **done for + the audited serving toy helpers in Slice 69; broader generic detection is + still open** + - root `tests/` behavior suites outside architecture tests. + +## Evidence Commands + +Commands used during this audit include: + +```bash +codegraph sync . +codegraph status . +git diff --shortstat f07442a28b0c186bd22b598a53cc4856ab4b2445..HEAD -- src tests crates +git diff --name-status f07442a28b0c186bd22b598a53cc4856ab4b2445..HEAD -- src tests crates +git diff --numstat f07442a28b0c186bd22b598a53cc4856ab4b2445..HEAD -- +git grep -n "#\\[\\(tokio::test\\|test\\)\\]" f07442a28b0c186bd22b598a53cc4856ab4b2445 -- tests src +rg "#\\[(tokio::test|test)\\]" -n tests crates src --glob '!target/**' +rg "not wired yet|后续承载|TODO|FIXME|todo!|unimplemented!|stub|placeholder|迁移|兼容|compat|legacy|alias|shim|deprecated" -n crates src tests +rg "std::fs|\\bfs::|std::env|env::|dirs::|std::path|PathBuf|Path::" -n crates/core/src crates/core/Cargo.toml +find tests -type d -empty -print -delete +find . -path './.git' -prune -o -path './target' -prune -o -path './.codegraph' -prune -o -type d -empty -print ``` -src/ -├── codex/ # OpenAI/Codex 对接层（核心业务逻辑） -├── service/ # 业务服务层（Chat, Responses, Diagnostics） -├── http/ # HTTP 路由层（/v1/* 和 /api/admin/*） -├── auth/ # 认证模块（API keys, sessions） -├── scheduler/ # 后台调度器（refresh, quota, model） -├── storage/ # 数据持久化（SQLite） -├── logs/ # 日志和事件追踪 -├── config/ # 配置管理 -└── utils/ # 工具函数（crypto, pagination） -``` - ---- - -## ✅ 架构优点 - -### 1. 清晰的关注点分离 - -**三层认证边界设计：** -``` -Admin Auth (密码 + Session) → /api/admin/* -Client API Keys (cpr_ 前缀) → /v1/* -Upstream Tokens (内部加密) → ChatGPT Backend -``` - -✅ **评价：** 完美！三种凭证系统互不混淆，避免了常见的认证边界模糊问题。 - -### 2. 错误处理模式规范 - -```rust -// 库代码：typed thiserror -pub enum AccountServiceError { ... } - -// main.rs：anyhow::Result -// 测试代码：anyhow::Result -``` - -**审计结果：** -- ✅ 0 个 `unwrap()` / `expect()` 在生产代码中 -- ✅ 0 个 `panic!` / `unimplemented!()` 在生产代码中 -- ✅ 0 个 `TODO` / `FIXME` 注释 -- ✅ 0 个 `unsafe` 代码块 - -✅ **评价：** 错误处理非常规范，符合 Rust 最佳实践。 - -### 3. 服务层设计（Service Layer） - -**服务列表：** -``` -AccountService - 账户 CRUD、OAuth、配额 -ModelService - 模型目录、后端刷新 -LogService - 事件日志、分页 -UsageService - 使用统计 -AdminAuthService - 管理员认证 -ChatService - Chat Completions 转换 -ResponsesService - Responses 转换 -DiagnosticsService - 诊断和健康检查 -``` - -✅ **评价：** 服务职责清晰，每个服务聚焦单一领域。 - -### 4. 数据加密策略 - -**加密实现：** -- 🔐 AES-256-GCM 加密存储（tokens, refresh tokens, cookies） -- 🔑 Argon2id 密码哈希（admin passwords） -- 🔒 HMAC-SHA256 + pepper（client API keys） -- 🚫 从不解密 tokens 用于列表视图 - -✅ **评价：** 安全措施完善，加密算法选择合理。 - -### 5. 并发和同步 - -**模式：** -- `Arc` - 不可变共享状态 -- `Arc>` - 账户池（唯一的可变共享状态） -- `tokio::sync::Mutex` - 避免了 `std::sync::Mutex` 跨 await 问题 - -✅ **评价：** 并发设计合理，避免了常见的死锁陷阱。 - -### 6. 后台调度器生命周期 - -**已实现的调度器：** -1. RefreshScheduler - JWT 令牌刷新 -2. SessionCleanupScheduler - 会话清理 -3. QuotaRefresher - 配额解锁 -4. ModelRefresher - 模型刷新 - -**关闭机制：** -```rust -tokio::select! { - _ = ticker.tick() => { /* work */ } - _ = shutdown_rx.recv() => { break; } -} -``` - -✅ **评价：** 优雅关闭机制完善，使用 `tokio::select!` 正确处理信号。 - -### 7. 测试覆盖 - -**测试策略：** -- 单元测试：内联 `#[cfg(test)] mod tests` -- 集成测试：42 个独立测试文件 -- Mock 策略：`:memory:` SQLite, `wiremock` HTTP, `tempfile` 文件系统 - -✅ **评价：** 测试覆盖率高（142 个测试），分层合理。 - ---- - -## ⚠️ 潜在问题和改进建议 - -### 1. 🟡 Clone 使用频率较高 - -**统计：** -- `codex/` 模块：92 次 `.clone()` -- 全项目：276 次 string allocations - -**问题：** -```rust -// 示例：可能过度克隆 -pub fn some_method(&self) -> SomeType { - self.config.clone() // AppConfig 可能很大 -} -``` - -**建议：** -- 对于配置类数据，考虑使用 `Arc` 而不是直接 `Config` -- 审查热路径（request handling）中的 clone 调用 -- 考虑使用 `&str` 而不是 `String` 作为参数类型 - -**优先级：** 🟡 中等（性能优化，非阻塞问题） - -### 2. 🟡 AccountService 方法分散 - -**当前结构：** -``` -src/codex/accounts/service/ -├── mod.rs # 主结构 -├── cookies.rs # Cookie 操作 -├── health.rs # 健康检查 -├── import.rs # 导入逻辑 -├── lifecycle.rs # CRUD 操作 -├── quota.rs # 配额管理 -├── refresh.rs # 令牌刷新 -└── runtime_pool.rs # 运行时池同步 -``` - -**问题：** -- 7 个 `impl AccountService` 块分散在不同文件 -- 难以一眼看出所有可用方法 -- 可能导致方法重复或不一致 - -**建议：** -- 考虑使用 trait 将功能分组： - ```rust - trait AccountLifecycle { ... } - trait AccountQuota { ... } - trait AccountAuth { ... } - ``` -- 或者在 `mod.rs` 中保留所有公共方法签名的注释索引 - -**优先级：** 🟡 中等（代码可维护性） - -### 3. 🟢 指纹匹配不完整 - -**缺失的头部：** -- ❌ `sec-ch-ua` (Chromium Client Hints) -- ❌ `accept-language` -- ❌ `accept-encoding` -- ❌ `x-codex-installation-id` -- ❌ 头部顺序精确匹配 - -**风险：** -- 可能被 Cloudflare 高级指纹识别检测 -- 可能触发 403 封禁 - -**建议：** -1. **短期**：监控生产环境 Cloudflare 403 错误率 -2. **中期**：如果出现封禁，优先添加 `x-codex-installation-id` -3. **长期**：实现完整的 Desktop 指纹镜像 - -**优先级：** 🟢 低（等待实际问题反馈） - -### 4. 🟡 错误类型可以更精细 - -**当前：** -```rust -pub enum AccountServiceError { - RepositoryUnavailable, // 太宽泛 - List, // 没有上下文 - ... -} -``` - -**建议：** -```rust -pub enum AccountServiceError { - RepositoryNotConfigured, - DatabaseQueryFailed(sqlx::Error), - ListAccountsFailed { reason: String }, - ... -} -``` - -**优点：** -- 更好的错误追踪 -- 更有用的日志信息 -- 更容易调试 - -**优先级：** 🟡 中等（可观测性改进） - -### 5. 🟢 配置热重载未完全实现 - -**当前状态：** -- `/api/admin/settings PATCH` 写入 `local.yaml` -- 更新 `AppState` 中的 config -- ⚠️ 但已构造的服务（如调度器）仍使用旧配置 - -**文档说明：** -> "Already constructed runtime services still use construction-time config until restart" - -**建议：** -- 要么实现完整的热重载（复杂） -- 要么在文档中明确标注哪些配置需要重启 - -**优先级：** 🟢 低（已有文档说明） - -### 6. 🟡 依赖版本管理 - -**当前状态：** -- ✅ TLS 相关依赖已固定（`reqwest = 0.12.28`, `rustls = 0.23.36`） -- ✅ 有明确的升级策略（需要指纹验证） -- ⚠️ `serde_yaml = 0.9.34+deprecated` - 标记为已弃用 - -**建议：** -- 监控 `serde_yaml` 的替代方案 -- 考虑迁移到 `serde_yml` 或其他维护的 YAML 库 - -**优先级：** 🟡 中等（技术债务） - -### 7. 🟢 日志体量可能较大 - -**当前配置：** -```yaml -logging: - capacity: 2000 # SQLite 内存容量 - retention_days: 14 -``` - -**潜在问题：** -- 高流量场景下，SQLite 写入可能成为瓶颈 -- 14 天 * 大量请求 = 大量磁盘空间 - -**建议：** -- 考虑异步批量写入 -- 添加日志采样（例如只记录 1% 的成功请求） -- 考虑使用专用日志系统（如 Loki） - -**优先级：** 🟢 低（等待生产环境验证） - ---- - -## 🎯 安全审计 - -### ✅ 通过的安全检查 - -1. ✅ **SQL 注入防护**：使用 `sqlx` 参数化查询 -2. ✅ **密码存储**：Argon2id（行业标准） -3. ✅ **Token 加密**：AES-256-GCM -4. ✅ **API Key 哈希**：HMAC-SHA256 + pepper -5. ✅ **Session 安全**：HttpOnly cookies -6. ✅ **无 unsafe 代码**：`unsafe_code = "forbid"` -7. ✅ **输入验证**：JWT claims 验证，状态检查 - -### ⚠️ 需要关注的点 - -1. **Refresh Token 一次性使用** - - 当前实现：有保存逻辑 - - ⚠️ 需要确认：网络失败时不会重试已消费的 RT - -2. **Rate Limiting** - - 当前：依赖上游 ChatGPT 的限流 - - 建议：考虑添加本地 rate limiting（防止恶意客户端） - -3. **Admin Session TTL** - - 默认：1440 分钟（24 小时） - - 建议：考虑更短的 TTL 或 sliding expiration - ---- - -## 📊 性能评估 - -### 预期瓶颈 - -1. **AccountPool 互斥锁** - - `Arc>` - 所有请求共享 - - 高并发下可能成为争用点 - - 建议：考虑 sharding 或 lock-free 结构 - -2. **SQLite 写入** - - 事件日志、使用统计都写 SQLite - - WAL 模式已启用（好！） - - 建议：监控写入延迟 - -3. **String 分配** - - 276 次字符串分配 - - 在热路径中可能影响性能 - - 建议：使用 profiler（如 `cargo flamegraph`）找出热点 - -### 优化建议 - -```rust -// 当前 -pub fn get_model(&self, model: String) -> ... { } - -// 优化后 -pub fn get_model(&self, model: &str) -> ... { } -``` - ---- - -## 🏆 总体评分 - -| 维度 | 评分 | 说明 | -|------|------|------| -| **架构设计** | ⭐⭐⭐⭐⭐ | 分层清晰，职责明确 | -| **代码质量** | ⭐⭐⭐⭐⭐ | 无 unwrap/panic，错误处理规范 | -| **安全性** | ⭐⭐⭐⭐☆ | 加密完善，需关注 RT 一次性使用 | -| **测试覆盖** | ⭐⭐⭐⭐☆ | 142 测试，覆盖主要路径 | -| **文档完整性** | ⭐⭐⭐⭐⭐ | CLAUDE.md + 详细注释 | -| **可维护性** | ⭐⭐⭐⭐☆ | 服务分层好，但方法分散 | -| **性能** | ⭐⭐⭐⭐☆ | 设计合理，有优化空间 | - -**总评：⭐⭐⭐⭐☆ (4.6/5.0)** - ---- - -## 🚀 优先改进清单 - -### 立即处理（本周） -1. ✅ 无（当前代码质量已经很高） - -### 短期改进（本月） -1. 🟡 审查热路径中的 `.clone()` 调用 -2. 🟡 迁移 `serde_yaml` 到维护的替代品 -3. 🟡 为 `AccountService` 添加方法索引文档 - -### 长期改进（本季度） -1. 🟢 监控生产环境指纹匹配情况 -2. 🟢 考虑实现 AccountPool sharding（如果并发成为瓶颈） -3. 🟢 添加更细粒度的错误类型 - ---- - -## ✅ 结论 - -**codex-proxy-rs 是一个高质量的 Rust 项目**，具有： -- ✅ 清晰的架构分层 -- ✅ 规范的错误处理 -- ✅ 完善的安全措施 -- ✅ 良好的测试覆盖 -- ✅ 零严重缺陷 - -**可以安全地部署到生产环境**，同时保持对上述中等优先级改进的关注。 - -**推荐：** 先部署小规模试点，监控性能和 Cloudflare 行为，然后根据实际数据进行针对性优化。 diff --git a/docs/architecture.md b/docs/architecture.md index 49aa0419e..98369409b 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -1,241 +1,291 @@ -# Architecture +# 架构规范 -This document is the canonical Rust workspace architecture for -`codex-proxy-rs`. The Rust implementation must converge to this directory -shape and naming exactly. A migration is not complete while any extra Rust -source file, extra Rust source directory, or forbidden transitional name remains. +本文档定义 `codex-proxy-rs` 的正式 Rust workspace 架构。每个 Rust 源文件、目录、 +依赖方向和命名约定在最终完成时必须与本文档完全一致。 -This document intentionally governs Rust workspace source layout. The `web/` -frontend has its own package layout and is not part of the Rust source -white-list below. +不允许过渡期别名。不允许兼容性重导出。迁移必须直接收敛到此布局。 -## Architectural Style +本文档仅约束 Rust workspace 源码布局。`web/` 前端和 `docs/` 目录有各自的约定， +不属于 Rust 源码白名单。 -The application is a modular monolith with hexagonal boundaries. +## 目录 -- `core` owns domain model, use cases, protocol semantics, policy, and ports. -- `adapters` owns concrete implementations of `core` ports. -- `runtime` owns composition, application state, and background task wiring. -- `server` owns the Axum HTTP boundary. -- `platform` owns shared infrastructure primitives. -- `assets` owns compiled frontend asset serving. -- `xtask` owns local automation. +1. [架构风格](#架构风格) +2. [核心规则](#核心规则) +3. [根 Workspace 布局](#根-workspace-布局) +4. [依赖方向](#依赖方向) +5. [Crate 职责](#crate-职责) +6. [精确 Rust 源码形状](#精确-rust-源码形状) +7. [Cargo 集成测试形状](#cargo-集成测试形状) +8. [端口（Port）放置](#端口port放置) +9. [领域模型归属](#领域模型归属) +10. [WebSocket 架构——硬拆分](#websocket-架构硬拆分) +11. [迁移计划](#迁移计划) +12. [架构测试](#架构测试) +13. [完成标准](#完成标准) -## Non-Negotiable Rules +## 架构风格 -- Rust source directories and files must match this document exactly at - completion. -- Do not keep transitional module names after migration. -- Do not hide structure drift behind compatibility re-exports. -- `core` must not depend on any project crate in the final architecture. -- `core` must not own concrete SQLx, Reqwest, Axum, filesystem path, TLS, or - upstream IO implementations. -- `server` must not decide business policy. -- `runtime` must not contain raw SQL, HTTP handlers, or domain rules. -- `adapters` must not define domain policy. -- Architecture tests must enforce directory shape, dependency direction, - forbidden names, and forbidden imports. +本应用是一个重组为多 crate workspace 的模块化单体，具有严格的六边形边界。 +每个 crate 有且仅有一个明确定义的职责。 -## Root Workspace Layout +- **`core`** 拥有领域模型、用例、协议语义、策略和端口（port trait）。 + 它是最内层——不依赖本 workspace 中的任何其他 crate。 +- **`adapters`** 拥有 `core` 端口的具体实现——SQLx 仓储、Reqwest HTTP 客户端、 + tokio-tungstenite WebSocket 客户端、OAuth 客户端、指纹更新客户端。 +- **`runtime`** 拥有应用组装——`AppState`、依赖注入、后台任务启停、配置映射。 +- **`server`** 拥有 Axum HTTP 边界——路由、处理器、中间件、响应信封、 + SSE 帧封装、错误到 HTTP 的映射。 +- **`platform`** 拥有共享基础设施原语——配置加载、加密、身份哈希、 + SQLite 连接初始化、schema 文件、文件系统路径、日志原语、JSON 和分页帮手。 +- **`assets`** 拥有编译后的前端静态资源服务——SPA 回退、静态缓存、安全头。 +- **`xtask`** 拥有本地自动化——前端构建、架构检查、发布命令。 + +## 核心规则 + +### 依赖规则（不可协商） + +1. `core` 不得依赖任何其他 workspace crate。 +2. `core` 不得依赖 `axum`、`sqlx`、`reqwest`、`tokio-tungstenite`、 + `tungstenite`、`rustls`、`tokio-rustls` 或 `hyper`。 +3. `core` 不得拥有文件系统路径、环境变量读取或具体网络 IO。 +4. `platform` 不得依赖 `core`、`server`、`runtime`、`adapters` 或 `assets`。 +5. `adapters` 不得依赖 `runtime` 或 `server`。 +6. `runtime` 不得依赖 `server`。 +7. `server` 不得包含 SQL 查询、具体上游 Codex IO、账号选择策略或后台任务编排。 +8. `assets` 不得依赖任何 workspace crate。 + +### 结构规则 + +1. [精确 Rust 源码形状](#精确-rust-源码形状) 中列出的每个文件和目录都必须存在。 +2. `crates/*/src/` 下不得存在白名单外的任何 Rust 源文件或目录，根目录不得保留 + `src/`。 +3. 每个 `Cargo.toml` 的 `[dependencies]` 条目必须至少被该 crate `src/` 中的一个 + `use` 语句引用。 +4. 根 package 不提供 library 或 binary target；类似 `crates/server/facade.rs` 的文件被禁止。 +5. 端口 trait 存放在各自领域模块的 `ports.rs` 文件中。 +6. 仓储实现存放在 `adapters/src/sqlite/` 中。 +7. 每条 `TODO` 注释必须引用一个 GitHub issue：`// TODO(#123): ...` + +### Rust 最佳实践（不可协商） + +以下规则适用于每个 crate，由架构测试和 CI 强制执行。 + +- **所有权**：参数优先使用 `&T` 而非 `.clone()`。参数使用 `&str`， + 返回自有数据使用 `String`。`Copy` 类型（≤ 24 字节）按值传递。 +- **错误处理**：可失败操作返回 `Result`。测试外禁止 + `unwrap()` 或 `expect()`。库错误类型使用 `thiserror`。用 `?` 传播。 +- **Lint**：`cargo clippy --all-targets --all-features --locked -- -D warnings` + 必须通过。使用 `#[expect(clippy::lint)]` 并附上理由注释。 +- **测试**：每个测试一个断言。测试名称描述行为： + `某操作_在某条件下_应产生某结果()`。生成输出的快照测试使用 `cargo insta`。 + 公共 API 示例使用 doc test。 +- **文档**：`//` 解释**为什么**（设计理由、安全不变量）。`///` 解释**什么** + （公共 API 契约）。在 `core`、`adapters` 和 `platform` 上启用 + `#![deny(missing_docs)]`。 +- **泛型**：性能关键路径优先使用静态分发（泛型）。仅在异构集合需要时使用 + `dyn Trait`。在 API 边界 Box，而非内部。 +- **性能**：绝不在循环中 clone。优先使用迭代器而非手动索引。避免中间的 + `.collect()` 调用。用 `--release` 运行基准测试。 + +## 根 Workspace 布局 ```text codex-proxy-rs/ -├── Cargo.toml +├── Cargo.toml # [workspace] 含 members = ["crates/*"] ├── Cargo.lock ├── README.md ├── AGENTS.md -├── src/ -│ └── lib.rs ├── crates/ │ ├── core/ +│ │ ├── Cargo.toml +│ │ └── src/ │ ├── adapters/ +│ │ ├── Cargo.toml +│ │ └── src/ │ ├── runtime/ +│ │ ├── Cargo.toml +│ │ └── src/ │ ├── server/ +│ │ ├── Cargo.toml +│ │ └── src/ │ ├── platform/ +│ │ ├── Cargo.toml +│ │ └── src/ │ ├── assets/ +│ │ ├── Cargo.toml +│ │ └── src/ │ └── xtask/ +│ ├── Cargo.toml +│ └── src/ ├── tests/ ├── web/ └── docs/ ``` -Root `src/lib.rs` is the only root facade. The root `Cargo.toml` must use the -default library path or explicitly point to `src/lib.rs`. Files such as -`crates/server/facade.rs` are forbidden. +**根 `Cargo.toml`** 必须声明 `[workspace]`，其中 `members = ["crates/*"]`， +并定义 workspace 级别的依赖、lint 设置和 profile 覆盖。根 package 不提供 +library 或 binary target；根 `tests/` 只承载跨 workspace 架构测试。 -## Dependency Direction - -Final dependency graph: +## 依赖方向 ```text -root facade -> server - -server -> runtime + core + assets + platform -runtime -> core + adapters + platform -adapters -> core + platform -core -> no project crate dependencies -assets -> no project crate dependencies -xtask -> workspace tooling only -platform -> no domain crate dependencies +server ──► runtime + core + assets + platform +runtime ──► core + adapters + platform +adapters ──► core + platform +core ──► （无 workspace crate 依赖） +assets ──► （无 workspace crate 依赖） +xtask ──► （仅 workspace 工具链） +platform ──► （无 workspace crate 依赖） ``` -Forbidden project-crate dependencies: +**被禁止的项目 crate 依赖**（由架构测试强制执行）： -- `core` must not depend on `platform`, `server`, `runtime`, `adapters`, or - `assets`. -- `platform` must not depend on `core`, `server`, `runtime`, `adapters`, or - `assets`. -- `adapters` must not depend on `runtime` or `server`. -- `runtime` must not depend on `server`. -- `assets` must not depend on project crates. +| 源 Crate | 被禁止的依赖 | +|---------|------------| +| `core` | `platform`、`server`、`runtime`、`adapters`、`assets` | +| `platform` | `core`、`server`、`runtime`、`adapters`、`assets` | +| `adapters` | `runtime`、`server` | +| `runtime` | `server` | +| `assets` | 所有 workspace crate | -`runtime` maps `platform` configuration into `core` configuration DTOs before -constructing core services. This keeps `core` independent from config loading, -filesystem paths, environment variables, and storage bootstrap concerns. +`runtime` 在构造 `core` 服务前将 `platform::config` 类型映射为 `core` 的配置 DTO。 +这使 `core` 独立于配置文件路径、YAML 解析和环境变量。 -## Crate Responsibilities +## Crate 职责 ### `crates/core` -Owns domain concepts, invariants, use case orchestration, protocol semantics, -and ports. +拥有领域概念、不变量、用例编排、协议语义和端口。此 crate 零 IO 依赖。 -Allowed: +**允许：** +- 账号池和调度策略（`accounts/pool.rs`） +- 管理员业务用例（`admin/`） +- 认证/会话领域行为（`auth/`） +- 模型目录规则（`models/`） +- 用量和事件策略（`usage/`、`events/`） +- OpenAI/Codex 协议数据结构和纯转换（`protocol/`） +- WebSocket 消息编解码（单个帧的序列化、反序列化、验证——**非**连接建立） +- 请求分发、回退、重试、配额检查、亲和性、恢复和用量跟踪策略（`serving/`） +- 面向上游的端口 trait（`gateway/ports.rs`） +- 类型化领域和应用错误（`error.rs` 或各模块的 `errors.rs`） -- account pool and scheduling policy -- admin business use cases -- authentication/session domain behavior -- model catalog rules -- usage and event policy -- OpenAI/Codex protocol data structures and pure conversion -- request dispatch, fallback, retry, quota, affinity, recovery, and usage policy -- port traits -- typed domain and application errors - -Forbidden: - -- Axum extractors, routers, responses, or `IntoResponse` -- SQLx pools, queries, rows, transactions, or migrations -- concrete `reqwest::Client` -- concrete WebSocket/TLS connection setup -- concrete filesystem path decisions -- platform config loader/types -- public re-exports of project crates +**禁止：** +- `axum` 提取器、路由、响应、`IntoResponse` 实现 +- `sqlx` 连接池、查询、行、事务、迁移 +- 具体的 `reqwest::Client` 构造或使用 +- 具体的 WebSocket/TLS 连接建立（`tokio_tungstenite`、`tokio-rustls`、`rustls`） +- 文件系统路径构造或环境变量读取 +- 平台配置加载器类型（`platform::config::*`） ### `crates/adapters` -Owns concrete implementations of ports defined by `core`. +拥有 `core` 端口的具体实现。每个实现结构体以其使用的技术命名。 -Allowed: +**允许：** +- `SqliteAccountStore` —— 账号存储端口的 SQLx 实现 +- `SqliteAdminSessionStore` —— 管理员会话端口的 SQLx 实现 +- `SqliteClientKeyStore` —— 客户端密钥端口的 SQLx 实现 +- `SqliteEventLogStore` —— 事件日志端口的 SQLx 实现 +- `SqliteModelSnapshotStore` —— 模型快照端口的 SQLx 实现 +- `SqliteSessionAffinityStore` —— 会话亲和性的 SQLx 实现 +- `SqliteCookieStore` —— Cookie 端口的 SQLx 实现 +- `SqliteUsageStore` —— 用量跟踪端口的 SQLx 实现 +- `ReqwestCodexClient` —— 上游 Codex HTTP/SSE 的 Reqwest 实现 +- `ReqwestFingerprintClient` —— 指纹更新的 Reqwest 实现 +- `OpenAiOAuthClient` —— OpenAI OAuth 流程的 Reqwest 实现 +- `CodexWebSocketConnection` —— WS 连接建立和消息收发的 tokio-tungstenite 实现 +- `CodexWebSocketPool` —— WS 连接池的 tokio-tungstenite 实现 +- `ReqwestTlsConfig` —— 自定义 CA 配置的 rustls 实现 +- 适配器特定错误从框架错误到 `core` 错误类型的翻译 -- SQLx implementations of storage ports -- encrypted SQLite stores -- Reqwest HTTP clients -- concrete OpenAI OAuth client -- concrete Codex upstream HTTP/SSE/WebSocket clients -- concrete fingerprint update client -- adapter-specific translation from framework errors into `core` errors - -Forbidden: - -- business policy -- Axum route handlers -- runtime task orchestration +**禁止：** +- 业务策略决策 +- Axum 路由处理器 +- 运行时任务编排 +- 领域不变量 ### `crates/runtime` -Owns application composition. +拥有应用组装。这是 "胶水" 层。 -Allowed: +**允许：** +- `AppState` 结构体构造 +- 服务构造（用 `adapters` 实现实例化 `core` 服务） +- 仓储和适配器构造 +- 向 `server` 处理器注入依赖 +- 后台任务启动和协调关闭 +- 启动恢复流程（从数据库加载账号到池中） +- `platform::config` → `core` 配置 DTO 映射 -- `AppState` -- service construction -- repository construction -- adapter construction -- dependency injection -- background task startup and shutdown -- startup restore flows -- platform-to-core config mapping - -Forbidden: - -- HTTP request/response mapping -- raw SQL query strings -- concrete upstream protocol side effects -- account selection, fallback, quota, model, or recovery policy +**禁止：** +- HTTP 请求/响应映射 +- 原始 SQL 查询字符串 +- 具体上游协议的副作用 +- 账号选择、回退、配额、模型或恢复策略决策 ### `crates/server` -Owns the HTTP boundary. +拥有 HTTP 边界。所有 Axum 类型都在这里。 -Allowed: +**允许：** +- Axum 路由和路由处理器 +- 提取器和中间件（`request_id`、tracing、CORS、auth） +- 管理员响应信封（`AdminResponse`、`AdminEnvelope`） +- OpenAI 兼容的 HTTP 响应体和错误形状 +- HTTP 状态码和头映射 +- 客户端响应的 SSE 帧封装 +- 路由本地的 OpenAI API 错误/事件帧封装 -- Axum routers and route handlers -- extractors and middleware -- request IDs, tracing, CORS, and auth extraction -- admin response envelopes -- OpenAI-compatible HTTP response bodies -- HTTP status and header mapping -- SSE framing for client responses - -Forbidden: - -- SQL queries -- concrete upstream Codex IO -- account selection, fallback, quota, model, or recovery policy -- background task orchestration +**禁止：** +- SQL 查询 +- 具体上游 Codex IO +- 账号选择、回退、配额、模型或恢复策略 +- 后台任务编排 ### `crates/platform` -Owns cross-cutting infrastructure primitives. +拥有无领域知识的跨领域基础设施原语。 -Allowed: +**允许：** +- 配置加载（`config/loader.rs`）和配置 schema 类型（`config/types.rs`） +- 加密原语：AES-GCM 密钥箱（`crypto/secret_box.rs`）、HMAC-SHA256 哈希（`crypto/hash.rs`） +- 身份哈希：管理员密码（Argon2）、客户端 API 密钥（HMAC） +- SQLite 连接初始化（`storage/sqlite.rs`）和 schema 文件（`storage/schema.sql`） +- 文件系统路径帮手（`storage/paths.rs`） +- 日志原语（`logging/`）—— tracing subscriber 初始化、轮转 +- JSON 帮手（`json/`）—— cursor 分页、信封帮手 -- config loading and config schema types -- crypto primitives -- identity hashing primitives -- SQLite connection setup and schema files -- filesystem path primitives -- logging primitives -- JSON and pagination helpers - -Forbidden: - -- Codex/OpenAI business use cases -- admin business use cases -- Axum handlers -- implementations of `core` ports +**禁止：** +- Codex/OpenAI 业务用例 +- 管理员业务用例 +- Axum 处理器 +- `core` 端口的实现 +- 领域模型类型 ### `crates/assets` -Owns static frontend asset serving. +拥有前端静态资源服务。 -Allowed: - -- root `index.html` -- `/assets/*` -- SPA fallback after API routes -- static cache and security headers +**允许：** +- 根 `/` → `index.html` 响应 +- `/assets/*` 静态文件服务 +- SPA 回退（API 路由后的任何未匹配路径 → `index.html`） +- 静态缓存和安全头 ### `crates/xtask` -Owns local automation. +拥有本地自动化。它是二进制 crate，不是 library。 -Allowed: +**允许：** +- 前端构建编排（`build_web.rs`） +- 架构检查（`check_architecture.rs`） +- 发布/打包命令（`release.rs`） -- frontend build orchestration -- architecture checks -- release/package commands -- developer maintenance commands +## 精确 Rust 源码形状 -## Exact Rust Source Shape - -### Root Facade - -```text -src/ -└── lib.rs -``` +下面列出的每个文件和目录在最终完成时必须存在。`crates/*/src/` 下 +不得存在其他 Rust 源文件或目录，根目录不得保留 `src/`。 ### Core @@ -247,81 +297,91 @@ crates/core/ ├── error.rs ├── admin/ │ ├── mod.rs - │ ├── ports.rs - │ ├── auth.rs - │ ├── client_keys.rs - │ └── settings.rs + │ ├── ports.rs # AdminSessionStore、ClientKeyStore trait + │ ├── auth.rs # 管理员密码验证逻辑 + │ ├── client_keys.rs # API 密钥服务（仅业务逻辑） + │ └── settings.rs # 设置领域逻辑 ├── accounts/ │ ├── mod.rs - │ ├── model.rs - │ ├── ports.rs - │ ├── service.rs - │ ├── pool.rs - │ ├── lifecycle.rs - │ ├── cookies.rs - │ ├── jwt.rs - │ └── usage.rs + │ ├── model.rs # Account、AccountStatus、NewAccount + │ ├── ports.rs # AccountStore、CookieStore、UsageStore trait + │ ├── service.rs # 账号管理用例 + │ ├── pool.rs # AccountPool 调度 + │ ├── lifecycle.rs # 状态转换、过期规则 + │ ├── cloudflare.rs # Cloudflare path-block 纯状态追踪 + │ ├── cookies.rs # Cookie 捕获/重放策略 + │ ├── jwt.rs # JWT 过期检查（纯逻辑，无网络） + │ └── usage.rs # 用量累积策略 ├── auth/ │ ├── mod.rs - │ ├── ports.rs - │ ├── oauth.rs - │ └── session.rs + │ ├── ports.rs # OAuthClient trait + │ ├── oauth.rs # OAuth 流程编排（无 HTTP） + │ └── session.rs # 管理员会话领域逻辑 ├── models/ │ ├── mod.rs - │ ├── model.rs - │ ├── ports.rs - │ ├── catalog.rs - │ └── service.rs + │ ├── model.rs # Model、ModelPlan、ModelCatalog + │ ├── ports.rs # ModelSnapshotStore trait + │ ├── catalog.rs # 模型目录规则 + │ └── service.rs # 模型服务用例 ├── events/ │ ├── mod.rs - │ ├── model.rs - │ ├── ports.rs - │ └── service.rs + │ ├── model.rs # EventLog、EventLevel + │ ├── ports.rs # EventLogStore trait + │ └── service.rs # 日志策略（启用、容量、保留） ├── usage/ │ ├── mod.rs - │ ├── model.rs - │ ├── ports.rs - │ └── service.rs + │ ├── model.rs # UsageSnapshot、UsageWindow + │ ├── ports.rs # UsageStore trait + │ └── service.rs # 用量聚合策略 ├── protocol/ │ ├── mod.rs │ ├── openai/ │ │ ├── mod.rs - │ │ ├── chat.rs - │ │ ├── responses.rs - │ │ ├── models.rs - │ │ ├── errors.rs - │ │ └── schema.rs + │ │ ├── chat.rs # ChatCompletionRequest/Response + │ │ ├── responses.rs # Response API 类型 + │ │ ├── models.rs # Model list 类型 + │ │ └── errors.rs # OpenAI 错误类型 │ └── codex/ │ ├── mod.rs - │ ├── chat.rs - │ ├── responses.rs - │ ├── events.rs - │ ├── sse.rs - │ ├── websocket.rs - │ └── schema.rs + │ ├── chat.rs # Codex 特定 chat 扩展 + │ ├── responses.rs # CodexResponsesRequest + │ ├── events.rs # Codex SSE 事件类型 + │ ├── sse.rs # SSE 事件解析（纯函数，无 IO） + │ ├── websocket.rs # WS 帧类型、payload 编解码 + │ └── schema.rs # JSON schema 验证规则 ├── gateway/ │ ├── mod.rs - │ ├── ports.rs - │ ├── fingerprint.rs - │ ├── conversation.rs - │ └── installation.rs + │ ├── ports.rs # CodexUpstreamClient、FingerprintClient trait + │ ├── fingerprint.rs # 指纹模型和版本规则 + │ ├── conversation.rs # 会话标识派生 + │ └── installation.rs # 安装 ID 生成规则 └── serving/ ├── mod.rs - ├── chat.rs - ├── responses.rs - ├── errors.rs - ├── routing.rs - ├── fallback.rs - ├── affinity.rs - ├── quota.rs - ├── stream.rs - ├── recovery.rs - └── usage.rs + ├── chat.rs # Chat 补全编排 + ├── responses.rs # Response 创建编排 + ├── errors.rs # 服务错误类型 + ├── routing.rs # 账号选择路由策略 + ├── fallback.rs # 回退策略规则 + ├── affinity.rs # 会话亲和性策略 + ├── quota.rs # 配额检查策略 + ├── implicit_resume.rs # Responses 隐式续接策略 + ├── reasoning_replay.rs # Responses reasoning replay 缓存策略 + ├── stream.rs # 流生命周期策略 + ├── recovery.rs # 错误恢复规则 + └── usage.rs # 用量跟踪策略 ``` -`protocol` contains only protocol data structures and pure conversion. -`gateway` contains upstream-facing ports and transport-neutral request context. -Concrete HTTP, SSE, WebSocket, TLS, and cookie-jar behavior lives in adapters. +**`core` 的关键设计决策：** + +- `protocol/codex/websocket.rs` 包含 WS 帧类型定义和纯编解码 + （序列化/反序列化单个消息）。它**不**打开连接、不执行 TLS 握手、不调用 + `tokio_tungstenite`。 +- `protocol/codex/sse.rs` 包含 SSE 事件解析，作为 `&[u8]` → `Vec` 的纯函数。 + 无网络 IO。 +- `serving/` 模块包含策略决策（"用哪个账号"、"是否重试"），但将所有 IO + 委托给 `gateway/ports.rs` 中定义的 trait。 +- `core/src/` 中没有任何文件以 `_repository.rs` 结尾。 + 存储实现位于 `adapters` 中。 ### Adapters @@ -332,38 +392,43 @@ crates/adapters/ ├── lib.rs ├── sqlite/ │ ├── mod.rs - │ ├── accounts.rs - │ ├── account_tokens.rs - │ ├── account_usage.rs - │ ├── refresh_leases.rs - │ ├── cookies.rs - │ ├── events.rs - │ ├── models.rs - │ ├── session_affinity.rs - │ ├── admin_sessions.rs - │ └── client_keys.rs - ├── http/ - │ ├── mod.rs - │ ├── reqwest_client.rs - │ └── headers.rs + │ ├── accounts.rs # SqliteAccountStore + │ ├── account_tokens.rs # token 在 DB 中的加解密/轮换 + │ ├── account_usage.rs # SqliteUsageStore + │ ├── refresh_leases.rs # 分布式刷新锁 + │ ├── cookies.rs # SqliteCookieStore + │ ├── events.rs # SqliteEventLogStore + │ ├── models.rs # SqliteModelSnapshotStore + │ ├── session_affinity.rs # SqliteSessionAffinityStore + │ ├── admin_sessions.rs # SqliteAdminSessionStore + │ └── client_keys.rs # SqliteClientKeyStore ├── codex/ │ ├── mod.rs - │ ├── client.rs - │ ├── sse.rs - │ ├── models.rs - │ ├── fingerprint.rs + │ ├── client.rs # ReqwestCodexClient（HTTP + SSE） + │ ├── models.rs # 通过 Reqwest 获取模型列表 + │ ├── fingerprint.rs # ReqwestFingerprintClient │ └── websocket/ │ ├── mod.rs - │ ├── client.rs - │ ├── pool.rs - │ ├── codec.rs - │ ├── deflate.rs - │ └── opening.rs + │ ├── connect.rs # CodexWebSocketConnection — tokio-tungstenite + │ ├── pool.rs # CodexWebSocketPool + │ ├── deflate.rs # permessage-deflate 处理 + │ └── opening.rs # WS 打开握手（TLS + Upgrade） └── oauth/ ├── mod.rs - └── openai.rs + └── openai.rs # OpenAiOAuthClient（Reqwest） ``` +**`adapters` 的关键设计决策：** + +- `sqlite/` 模块按**存储什么**命名，而非怎么存储。`SqliteAccountStore` 实现 + `core::accounts::ports::AccountStore`。 +- `codex/websocket/connect.rs` 包装 `tokio_tungstenite::connect_async`， + 返回 `core::protocol::codex::websocket::WsConnection`（trait 或 opaque handle）。 + `core` 永远不会看到 `tokio_tungstenite` 类型。 +- `codex/websocket/opening.rs` 处理原始 TCP/TLS 连接、自定义头顺序和 WebSocket 升级。 + 它产生 `OpeningAuditSnapshot` 用于一致性测试。 +- `oauth/openai.rs` 使用 Reqwest 实现 `core::auth::ports::OAuthClient`。 + ### Runtime ```text @@ -371,25 +436,32 @@ crates/runtime/ ├── Cargo.toml └── src/ ├── lib.rs - ├── bootstrap.rs - ├── state.rs - ├── services.rs - ├── repositories.rs - ├── upstream.rs - ├── config.rs + ├── bootstrap.rs # 启动：加载配置、初始化 DB、恢复池 + ├── state.rs # AppState 构造 + ├── services.rs # 服务工厂函数 + ├── repositories.rs # 仓储工厂函数 + ├── upstream.rs # 适配器工厂函数 + ├── config.rs # platform::config → core 配置 DTO 映射 └── tasks/ ├── mod.rs - ├── coordinator.rs - ├── token_refresh.rs - ├── quota_refresh.rs - ├── model_refresh.rs - ├── cookie_cleanup.rs - ├── session_cleanup.rs - └── fingerprint_update.rs + ├── coordinator.rs # 后台任务生成/关闭 + ├── token_refresh.rs # 定期 token 刷新任务接线 + ├── quota_refresh.rs # 定期配额刷新任务接线 + ├── model_refresh.rs # 定期模型刷新任务接线 + ├── cookie_cleanup.rs # 定期 cookie 清理任务接线 + ├── session_cleanup.rs # 定期会话清理任务接线 + ├── session_affinity_cleanup.rs # 定期会话亲和性清理任务接线 + └── fingerprint_update.rs # 定期指纹更新任务接线 ``` -`runtime/src/config.rs` maps `platform` configuration into `core` configuration -DTOs. +**`runtime` 的关键设计决策：** + +- 任务文件仅包含接线：它们从 `core` 接收 `Arc`，从配置中获取间隔， + 并启动 `tokio::spawn` 循环。它们不包含领域策略。 +- `runtime/src/config.rs` 将 `platform::config::AppConfig` 翻译为 + `core::admin::AdminConfig`、`core::accounts::PoolConfig` 等。 +- `runtime/src/bootstrap.rs` 编排启动序列，但将所有具体 IO + 委托给 `adapters` 工厂。 ### Server @@ -397,46 +469,47 @@ DTOs. crates/server/ ├── Cargo.toml └── src/ - ├── main.rs + ├── main.rs # #[tokio::main] 入口 ├── lib.rs - ├── router.rs + ├── router.rs # 顶层 Axum Router ├── error/ │ ├── mod.rs - │ ├── admin.rs - │ └── openai.rs + │ ├── admin.rs # 管理员错误 → HTTP 响应 + │ └── openai.rs # 领域错误 → OpenAI 错误形状 ├── middleware/ │ ├── mod.rs - │ ├── request_id.rs - │ ├── trace.rs - │ ├── auth.rs - │ └── cors.rs + │ ├── request_id.rs # X-Request-Id 头 + 扩展 + │ ├── trace.rs # HTTP tracing 层 + │ ├── auth.rs # API 密钥 / 会话提取 + │ └── cors.rs # CORS 配置 ├── admin_api/ │ ├── mod.rs - │ ├── router.rs - │ ├── response.rs - │ ├── session.rs - │ ├── settings.rs - │ ├── diagnostics.rs - │ ├── models.rs - │ ├── usage.rs + │ ├── router.rs # /api/admin/* 路由 + │ ├── response.rs # AdminEnvelope、AdminResponse + │ ├── session.rs # 登录/登出处理器 + │ ├── settings.rs # GET/PATCH /api/admin/settings + │ ├── diagnostics.rs # GET /api/admin/diagnostics + │ ├── models.rs # POST /api/admin/refresh-models + │ ├── usage.rs # GET /api/admin/usage-stats │ ├── accounts/ │ │ ├── mod.rs │ │ ├── list.rs │ │ ├── create.rs │ │ ├── import.rs + │ │ ├── import_cli.rs # Codex CLI auth.json 导入 │ │ ├── export.rs - │ │ ├── lifecycle.rs - │ │ ├── quota.rs - │ │ ├── cookies.rs - │ │ ├── oauth.rs - │ │ └── health.rs + │ │ ├── lifecycle.rs # status/label/delete/batch + │ │ ├── quota.rs # 配额警告 + │ │ ├── cookies.rs # cookie 管理 + │ │ ├── oauth.rs # OAuth 流程处理器 + │ │ └── health.rs # 健康检查 │ ├── client_keys/ │ │ ├── mod.rs │ │ ├── list.rs │ │ ├── create.rs │ │ ├── import.rs │ │ ├── export.rs - │ │ └── lifecycle.rs + │ │ └── lifecycle.rs # label/status/delete/batch │ └── logs/ │ ├── mod.rs │ ├── query.rs @@ -444,19 +517,24 @@ crates/server/ │ └── state.rs └── openai_api/ ├── mod.rs - ├── router.rs - ├── auth.rs - ├── chat.rs - ├── responses.rs - ├── models.rs - ├── diagnostics.rs - ├── error.rs - └── sse.rs + ├── router.rs # /v1/* 路由 + ├── auth.rs # Bearer token 提取 + ├── chat.rs # POST /v1/chat/completions + ├── responses.rs # POST /v1/responses + ├── models.rs # GET /v1/models + ├── diagnostics.rs # GET /debug/* + ├── error.rs # OpenAI 错误 → HTTP 响应 + └── sse.rs # 客户端的 SSE 事件帧封装 ``` -`server/src/error/openai.rs` maps application errors into OpenAI-compatible HTTP -responses. `server/src/openai_api/error.rs` and `server/src/openai_api/sse.rs` -own route-local OpenAI API error/event framing. +**`server` 的关键设计决策：** + +- 路由处理器是薄的：提取认证、解析请求体、调用 `core::serving::*` 用例、 + 将响应映射为 HTTP。处理器中没有业务逻辑。 +- `server/src/openai_api/sse.rs` 将 `core::protocol::codex::events::SseEvent` + 帧封装为 HTTP 响应的 `data: {...}\n\n`。它拥有线格式，不拥有事件语义。 +- `server/src/error/openai.rs` 将 `core::serving::errors::ServingError` + 映射为 OpenAI 兼容的 `{"error": {...}}` JSON 响应。 ### Platform @@ -467,29 +545,40 @@ crates/platform/ ├── lib.rs ├── config/ │ ├── mod.rs - │ ├── loader.rs - │ └── types.rs + │ ├── loader.rs # YAML 配置加载（含 local.yaml 覆盖） + │ └── types.rs # AppConfig、ServerConfig 等 ├── crypto/ │ ├── mod.rs - │ ├── secret_box.rs - │ └── hash.rs + │ ├── secret_box.rs # AES-256-GCM 加解密 + │ └── hash.rs # HMAC-SHA256 + 常数时间比较 ├── identity/ │ ├── mod.rs - │ ├── admin_password.rs - │ └── client_key.rs + │ ├── admin_password.rs # Argon2 密码哈希/验证 + │ └── client_key.rs # API 密钥生成 + HMAC 哈希 ├── storage/ │ ├── mod.rs - │ ├── sqlite.rs - │ ├── schema.sql - │ └── paths.rs + │ ├── sqlite.rs # SqlitePool 构造 + WAL pragma + │ ├── schema.sql # 所有表的 DDL + │ └── paths.rs # data/ 目录路径帮手 ├── logging/ │ ├── mod.rs - │ └── rotation.rs + │ └── rotation.rs # tracing-subscriber 初始化 + 文件轮转 └── json/ ├── mod.rs - └── pagination.rs + └── pagination.rs # Cursor 编解码 + Page ``` +**`platform` 的关键设计决策：** + +- `config/types.rs` 定义配置 schema 类型，供 `config/loader.rs` 使用。 + 这些是纯数据，不是领域类型。 +- `crypto/secret_box.rs` 是 `aes_gcm` 上的薄包装，输出 base64 编码密文。 + 它不知道账号或 token。 +- `identity/client_key.rs` 包含基于 HMAC 的 API 密钥哈希器，但**不**包含端口 trait—— + 那在 `core::admin::ports` 中。 +- `storage/schema.sql` 是 DDL 的唯一真实来源。它在连接时被 + `storage/sqlite.rs` 通过 `include_str!` 引入。 + ### Assets ```text @@ -497,8 +586,8 @@ crates/assets/ ├── Cargo.toml └── src/ ├── lib.rs - ├── router.rs - └── headers.rs + ├── router.rs # / 和 /assets/* 路由 + └── headers.rs # Cache-Control、CSP、安全头 ``` ### Xtask @@ -507,173 +596,296 @@ crates/assets/ crates/xtask/ ├── Cargo.toml └── src/ - ├── main.rs - ├── build_web.rs - ├── check_architecture.rs - └── release.rs + ├── main.rs # CLI 分发器 + ├── build_web.rs # 构建 web/ 前端 + ├── check_architecture.rs # 运行架构测试 + └── release.rs # 发布打包 ``` -## Cargo Integration Test Shape +## 端口（Port）放置 -Cargo integration test entry files must remain at `tests/*.rs`; submodules live -under matching directories. +端口是 `core` 拥有的 trait。每个端口 trait 在其领域模块中存放在名为 `ports.rs` 的文件中。 -```text -tests/ -├── admin.rs -├── openai_api.rs -├── accounts.rs -├── codex_upstream.rs -├── runtime.rs -├── architecture.rs -├── admin/ -├── openai_api/ -├── accounts/ -├── codex_upstream/ -├── runtime/ -├── architecture/ -│ ├── mod.rs -│ ├── dependency_direction.rs -│ ├── forbidden_imports.rs -│ └── directory_shape.rs -└── support/ - ├── mod.rs - ├── app.rs - ├── fixtures.rs - └── wiremock.rs +| 端口 Trait | 位置 | 实现（在 `adapters` 中） | +|---|---|---| +| `AccountStore` | `core/src/accounts/ports.rs` | `SqliteAccountStore` | +| `CookieStore` | `core/src/accounts/ports.rs` | `SqliteCookieStore` | +| `UsageStore` | `core/src/accounts/ports.rs` | `SqliteUsageStore` | +| `AdminSessionStore` | `core/src/admin/ports.rs` | `SqliteAdminSessionStore` | +| `ClientKeyStore` | `core/src/admin/ports.rs` | `SqliteClientKeyStore` | +| `OAuthClient` | `core/src/auth/ports.rs` | `OpenAiOAuthClient` | +| `ModelSnapshotStore` | `core/src/models/ports.rs` | `SqliteModelSnapshotStore` | +| `EventLogStore` | `core/src/events/ports.rs` | `SqliteEventLogStore` | +| `SessionAffinityStore` | `core/src/serving/affinity.rs` | `SqliteSessionAffinityStore` | +| `CodexUpstreamClient` | `core/src/gateway/ports.rs` | `ReqwestCodexClient` | +| `CodexWebSocketTransport` | `core/src/gateway/ports.rs` | `CodexWebSocketConnection` | +| `FingerprintClient` | `core/src/gateway/ports.rs` | `ReqwestFingerprintClient` | + +每个 `ports.rs` 文件仅包含 trait 定义和相关类型别名 +（例如 `AccountStoreResult = Result`）。 + +## 领域模型归属 + +领域模型是纯数据结构体，不包含 IO。所有模型类型归 `core` 所有。 + +| 模型 | 位置 | 备注 | +|------|------|------| +| `Account` | `core/src/accounts/model.rs` | 运行时账号表示 | +| `AccountStatus` | `core/src/accounts/model.rs` | 枚举：Active、Expired、Disabled、Banned | +| `StoredAccount` | `adapters/src/sqlite/accounts.rs` | DB 行类型——适配器内部 | +| `EventLog` | `core/src/events/model.rs` | 结构化事件 | +| `EventLevel` | `core/src/events/model.rs` | Info、Warn、Error | +| `Model`、`ModelPlan` | `core/src/models/model.rs` | 模型目录条目 | +| `UsageWindow` | `core/src/usage/model.rs` | 聚合用量窗口 | +| `Fingerprint` | `core/src/gateway/fingerprint.rs` | 设备指纹 | +| `ChatCompletionRequest` | `core/src/protocol/openai/chat.rs` | OpenAI 请求/响应类型 | +| `CodexResponsesRequest` | `core/src/protocol/codex/responses.rs` | Codex 请求类型 | +| `SseEvent` | `core/src/protocol/codex/sse.rs` | 解析后的 SSE 事件 | +| `CodexWebSocketResponse` | `core/src/protocol/codex/websocket.rs` | WS 响应聚合 | +| `OpeningAuditSnapshot` | `core/src/protocol/codex/websocket.rs` | WS 握手审计快照 | +| `PayloadAuditSnapshot` | `core/src/protocol/codex/websocket.rs` | WS payload 审计快照 | +| `WsParityDiff` | `core/src/protocol/codex/websocket.rs` | 一致性差异报告 | + +**规则**：`StoredAccount` 和类似的 DB 行类型是 `adapters` 内部类型。它们 +**不**作为 `core` 端口 trait 的输入或输出暴露。端口 trait 使用方法签名 +中的领域模型类型（`Account`、`EventLog` 等）。适配器在内部将 DB 行 +转换为领域类型。 + +## WebSocket 架构——硬拆分 + +WebSocket 代码是架构上最具挑战性的拆分，因为当前 +`src/codex/gateway/transport/websocket/` 目录混淆了三个不同的关注点： + +### 1. 纯协议编解码 → `core` + +处理 WS 消息格式、验证和单个帧序列化/反序列化的文件。这些无 IO 依赖。 + +- **`core/src/protocol/codex/websocket.rs`**：帧类型定义、 + `CodexWebSocketResponse`、`CodexWebSocketStreamResponse`、`WsMessage`。 + 纯数据类型。 +- `CodexResponsesRequest` ↔ WS 文本帧 JSON 的序列化/反序列化逻辑。纯转换。 + +### 2. 连接管理 → `adapters` + +处理实际 TCP/TLS 连接、WebSocket 握手和消息收发的文件。 +这些依赖 `tokio-tungstenite`、`tokio-rustls`、`rustls`。 + +- **`adapters/src/codex/websocket/connect.rs`**：包装 + `tokio_tungstenite::connect_async`，返回实现 `core` 中定义的 + `WsTransport` trait 的 opaque handle。 +- **`adapters/src/codex/websocket/pool.rs`**：连接池管理、保活、驱逐。 +- **`adapters/src/codex/websocket/opening.rs`**：原始 TLS 连接、 + 自定义头顺序、WebSocket 升级握手、accept key 验证。 + 产生 `OpeningAuditSnapshot`（一个 `core` 类型）用于一致性对比。 +- **`adapters/src/codex/websocket/deflate.rs`**：permessage-deflate + 扩展协商和帧压缩/解压。 + +### 3. 审计快照 → `core` + +审计快照类型是纯数据，用于比较代理的行为与官方 Codex Desktop 客户端。 +它们属于 `core`，因为代表协议级别的一致性关注点。 + +- **`core/src/protocol/codex/websocket.rs`**（或单独的 `audit.rs`）： + `OpeningAuditSnapshot`、`PayloadAuditSnapshot`、`WsParityDiff`。 + +### Trait 桥接 + +WebSocket 的 `core` 和 `adapters` 之间的桥接是 `core/src/gateway/ports.rs` 中的一个 trait： + +```rust +// core/src/gateway/ports.rs + +#[async_trait] +pub trait CodexWebSocketTransport: Send + Sync { + async fn connect( + &self, + request: &CodexResponsesRequest, + context: &RequestContext, + ) -> Result; + + async fn connect_with_pool( + &self, + request: &CodexResponsesRequest, + context: &RequestContext, + ) -> Result; +} ``` -Additional fixture files are allowed under `tests/fixtures/`. +`adapters` 用 `tokio-tungstenite` 实现此 trait。`core` 调用该 trait 而无需知道具体实现。 -## Docs Shape +## 迁移计划 -```text -docs/ -├── architecture.md -├── adr/ -├── api/ -└── operations/ -``` +迁移分八个阶段进行。每个阶段必须产出**可编译、测试通过的 workspace**， +然后才能开始下一阶段。不允许多个阶段交错进行，不允许过渡期别名。 -Historical migration plans may remain under `docs/superpowers/` while active -development is ongoing. They are not part of the final architecture white-list. +### 阶段 1：Workspace 脚手架 -## Web Shape +1. 在根 `Cargo.toml` 中添加 `[workspace]`，`members = ["crates/*"]`。 +2. 创建 `crates/` 目录及所有七个子目录，每个含 stub `Cargo.toml`。 +3. 在根 `Cargo.toml` 中定义 workspace 级别依赖、lint 设置和 profile 覆盖。 +4. 验证 `cargo check --workspace` 通过（所有 crate 为空）。 -The `web/` directory is outside this Rust workspace architecture white-list. -It is governed by frontend build tooling and must not be checked by Rust -directory-shape tests. `web/node_modules/` and `web/dist/` are build artifacts -from the frontend toolchain, not Rust architecture inputs. +**验收**：`cargo check --workspace` 编译通过。workspace 成员数量架构测试通过。 -## Required Names +### 阶段 2：Platform 提取 -- Root facade: `src/lib.rs` -- HTTP admin boundary: `server/src/admin_api` -- HTTP OpenAI-compatible boundary: `server/src/openai_api` -- Core port files: `ports.rs` -- Concrete SQLite adapters: `adapters/src/sqlite` -- Concrete OAuth adapter: `adapters/src/oauth/openai.rs` -- Concrete Codex upstream adapters: `adapters/src/codex` -- Concrete Codex WebSocket adapter: `adapters/src/codex/websocket` -- Shared middleware: `server/src/middleware` -- HTTP response errors: `server/src/error` -- OpenAI API route-local SSE framing: `server/src/openai_api/sse.rs` -- Business errors: `core/src/error.rs` or domain-local `errors.rs` +1. 从当前 `src/platform/`、`src/config/` 和 `src/utils/` 创建 `crates/platform/`。 +2. 移动文件：`platform/` 直接映射。`config/` → `platform/src/config/`。 + `utils/` → `platform/src/json/` 和 `platform/src/`（合并小型帮手）。 +3. 更新 `platform` 的 `Cargo.toml` 依赖：仅提取 `serde`、`serde_yaml`、 + `config`、`aes-gcm`、`argon2`、`hmac`、`sha2`、`base64`、`rand`、 + `sqlx`、`tracing`、`tracing-subscriber`、`tracing-appender`、`chrono`。 +4. 删除旧的 `src/platform/`、`src/config/`、`src/utils/`。 +5. 更新所有 `use crate::platform::` → `use platform::`。 +6. 更新所有 `use crate::config::` → `use platform::config::`。 +7. 更新所有 `use crate::utils::` → `use platform::json::`。 -## Forbidden Names +**验收**：`cargo test -p platform` 通过。所有现有测试通过。 -- `crates/admin` -- `crates/server/facade.rs` -- `server/src/admin.rs` -- `server/src/codex_http` -- `server/src/openai_api/stream.rs` -- `core/src/**/repository.rs` -- `core/src/**/repository/` -- `adapters/src/*_repository.rs` -- `runtime/src/*_repository.rs` -- `core/src/codex/gateway/transport` -- `core/src/codex/gateway/oauth/client.rs` -- `core/src/codex/serving/http` -- `platform/src/platform` -- `platform/src/utils` +### 阶段 3：Core 提取——协议 + 模型 -## Port Placement +1. 创建 `crates/core/`，含空的 `src/`。 +2. 将 `src/codex/gateway/protocol/` 和 `src/codex/serving/http/` + 中的协议类型和请求/响应类型移动到 `core/src/protocol/`。 +3. 将领域模型从 `src/codex/accounts/model.rs`、`src/codex/models/`、 + `src/codex/events/model.rs` 移动到 `core/src/accounts/model.rs`、 + `core/src/models/` 等。 +4. 在 `core/src/*/ports.rs` 中定义端口 trait——从当前内联 trait 定义 + 和仓储签名中提取。 +5. Core 的 `Cargo.toml` 依赖：仅 `serde`、`serde_json`、`chrono`、 + `thiserror`、`async-trait`、`uuid`。**零 IO crate**。 -Ports are traits owned by `core`. +**验收**：`cargo check -p core` 通过。Core 零 IO crate 依赖。 -- account storage ports live in `core/src/accounts/ports.rs` -- admin session/client-key ports live in `core/src/admin/ports.rs` -- OAuth and token refresh ports live in `core/src/auth/ports.rs` -- model snapshot ports live in `core/src/models/ports.rs` -- event log ports live in `core/src/events/ports.rs` -- usage ports live in `core/src/usage/ports.rs` -- upstream Codex and fingerprint ports live in `core/src/gateway/ports.rs` +### 阶段 4：Adapters 提取 -Concrete implementations live in `adapters`. +1. 创建 `crates/adapters/`，含 `src/sqlite/`、`src/codex/`、`src/oauth/`。 +2. 将所有 `*_repository.rs` 文件和 `repository/` 目录移动到 + `adapters/src/sqlite/`。重命名去掉 `_repository` 后缀 + （例如 `accounts.rs` 而非 `account_repository.rs`）。 +3. 移动 `src/codex/gateway/transport/http_client.rs` → + `adapters/src/codex/client.rs`。 +4. 移动 `src/codex/gateway/transport/endpoints.rs` → + `adapters/src/codex/endpoints.rs`。 +5. 移动 `src/codex/gateway/transport/custom_ca.rs` → + `adapters/src/codex/tls.rs`。 +6. 移动 `src/codex/gateway/oauth/client.rs` → + `adapters/src/oauth/openai.rs`。 +7. 移动 `src/codex/gateway/fingerprint/update_checker.rs` → + `adapters/src/codex/fingerprint.rs`。 +8. 拆分 `src/codex/gateway/transport/websocket/`： + - `codec.rs` 协议逻辑 → `core/src/protocol/codex/websocket.rs` + - `mod.rs`、`pool.rs`、`opening.rs`、`deflate.rs` → + `adapters/src/codex/websocket/` +9. 在适配器结构体上实现 `core` 端口 trait。 -- `SqliteAccountStore` lives in `adapters/src/sqlite/accounts.rs` -- `SqliteAdminSessionStore` lives in `adapters/src/sqlite/admin_sessions.rs` -- `OpenAiOAuthClient` lives in `adapters/src/oauth/openai.rs` -- `ReqwestCodexClient` lives in `adapters/src/codex/client.rs` -- `CodexWebSocketPool` implementation lives in `adapters/src/codex/websocket/pool.rs` +**验收**：`cargo test -p adapters` 通过。架构测试确认 `core` 依赖中无 +`axum`、`sqlx`、`reqwest` 或 `tokio-tungstenite`。 -## Migration Rules +### 阶段 5：Core 提取——业务逻辑 -Every migration step must move the codebase closer to this document. +1. 将剩余 `src/codex/accounts/` 业务逻辑（pool、lifecycle、cookies、 + jwt、usage policy）移动到 `core/src/accounts/`。 +2. 将 `src/codex/serving/dispatch/`（affinity、fallback、routing、 + recovery、stream、usage、implicit_resume、reasoning_replay）移动到 + `core/src/serving/`。 +3. 将 `src/admin/session/service.rs`、`src/admin/client_keys/service.rs`、 + `src/admin/settings.rs` 移动到 `core/src/admin/`。 +4. 将 `src/codex/events/service.rs` 移动到 `core/src/events/service.rs`。 +5. 所有移动的文件必须使用端口 trait，而非具体实现。 + 替换 `AccountRepository` 为 `dyn AccountStore` 等。 -Required order: +**验收**：`cargo test -p core` 通过。所有 core 业务逻辑测试通过。 -1. Add architecture tests for directory shape and forbidden names. -2. Move the root facade to `src/lib.rs`. -3. Rename server HTTP boundaries to `admin_api` and `openai_api`. -4. Move admin business code into `core/src/admin`. -5. Move all core repository traits into `ports.rs`. -6. Move all SQLx implementations into `adapters/src/sqlite`. -7. Move all Reqwest/SSE/WebSocket/OAuth implementations into `adapters`. -8. Collapse runtime into composition-only files. -9. Collapse platform into config/crypto/identity/storage/logging/json primitives. -10. Update facade exports after internal structure is correct. -11. Run full verification only after the structure is fully aligned. +### 阶段 6：Server 提取 -Do not use transitional aliases as a substitute for migration. A transitional -alias is acceptable only if it is removed in the same migration series before -claiming completion. +1. 移动 `src/admin/api/` → `crates/server/src/admin_api/`。 +2. 移动 `src/codex/serving/http/` → `crates/server/src/openai_api/`。 +3. 移动 `src/runtime/router.rs` → `crates/server/src/router.rs`。 +4. 移动 `src/platform/http/` 中间件 → `crates/server/src/middleware/`。 +5. 移动 `src/web/` → `crates/assets/`。 +6. 移动 `src/main.rs` → `crates/server/src/main.rs`。 +7. 所有 Axum 类型（`Router`、`IntoResponse`、提取器）必须仅存在于 + `server` 和 `assets` 中。 -## Architecture Tests +**验收**：`cargo test -p server` 通过。架构测试确认 `core`、`adapters`、 +`runtime`、`platform` 中无 Axum 类型。 -Architecture tests must enforce: +### 阶段 7：Runtime 清理 -- every required Rust source file exists -- every required Rust source directory exists -- forbidden Rust source files and directories do not exist -- no extra Rust source files exist under `src/` and `crates/*/src/` -- no extra Rust source directories exist under `src/` and `crates/*/src/` -- `Cargo.toml` workspace members match this document -- root library path is `src/lib.rs` -- `core` has no project-crate dependencies -- `core` has no `axum`, `sqlx`, or concrete `reqwest::Client` -- `server` has no SQLx queries -- `runtime` has no SQLx query strings and no Axum handlers -- `adapters` owns concrete SQLx/Reqwest/OAuth/Codex IO -- `core` ports are declared in `ports.rs` -- root facade does not hide invalid internal directory names +1. `runtime/src/` 仅保留 `bootstrap.rs`、`state.rs`、`services.rs`、 + `repositories.rs`、`upstream.rs`、`config.rs` 和 `tasks/`。 +2. 移除 `runtime/` 中残留的任何 HTTP、SQL 或领域策略。 +3. `runtime` 通过将 `core` 服务与 `adapters` 实现接线来构造 `AppState`。 -The required architecture test files are: +**验收**：`cargo test -p runtime` 通过。架构测试确认 `runtime` 中无 +SQL 查询字符串或 Axum 处理器。 -```text -tests/architecture/dependency_direction.rs -tests/architecture/forbidden_imports.rs -tests/architecture/directory_shape.rs -``` +### 阶段 8：架构强制执行 -## Completion Criteria +1. 编写 `tests/architecture/directory_shape.rs` —— 验证 + [精确 Rust 源码形状](#精确-rust-源码形状) 中的每个文件和目录存在， + 且无额外文件/目录。 +2. 编写 `tests/architecture/dependency_direction.rs` —— 通过 + `cargo metadata` + 解析 `Cargo.toml` 依赖验证 crate 级别依赖规则。 +3. 编写 `tests/architecture/forbidden_imports.rs` —— grep 查找被禁止的 + crate 使用（`core` 中的 `use axum::`、`runtime` 中的 `use sqlx::` 等）。 +4. 运行完整 CI：`cargo fmt --check`、`cargo check --workspace --all-targets`、 + `cargo test --workspace`、`cargo clippy --workspace --all-targets --all-features --locked -- -D warnings`。 -The migration is complete only when all of the following are true: +**验收**：所有架构测试通过。所有 CI 检查通过。 -- The Rust workspace source tree matches this document exactly. -- Forbidden names are absent from Rust source directories. -- Architecture tests enforce this document. -- `cargo fmt --check` passes. -- `cargo check --workspace --all-targets` passes. -- `cargo test --workspace` passes. -- `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` passes. +## 架构测试 + +架构测试是验证 workspace 符合本文档的集成测试。 +它们作为 `cargo test --workspace` 的一部分运行。 + +### `tests/architecture/directory_shape.rs` + +验证： +- [精确 Rust 源码形状](#精确-rust-源码形状) 中的每个必需文件存在。 +- 每个必需目录存在。 +- `crates/*/src/` 下无额外 Rust 源文件或目录，根目录无 `src/`。 +- `Cargo.toml` workspace 成员匹配 `["crates/core", "crates/adapters", + "crates/runtime", "crates/server", "crates/platform", "crates/assets", + "crates/xtask"]`。 +- 根 package 不提供 library 或 binary target。 + +### `tests/architecture/dependency_direction.rs` + +验证： +- `core/Cargo.toml` 无 workspace crate 依赖。 +- `platform/Cargo.toml` 无 workspace crate 依赖。 +- `adapters/Cargo.toml` 不依赖 `runtime` 或 `server`。 +- `runtime/Cargo.toml` 不依赖 `server`。 +- `assets/Cargo.toml` 无 workspace crate 依赖。 +- `xtask/Cargo.toml` 不依赖 `assets`。 + +### `tests/architecture/forbidden_imports.rs` + +验证： +- `core`、`adapters`、`runtime`、`platform` 中无 `use axum::`。 +- `core`、`server`、`runtime` 中无 `use sqlx::`。 +- `core`、`server`、`runtime` 中无 `use reqwest::`。 +- `core`、`server`、`runtime` 中无 `use tokio_tungstenite::` 或 + `use tungstenite::`。 +- `core` 中无 `use rustls::` 或 `use tokio_rustls::`。 +- `core/src/` 下无名为 `*_repository.rs` 的文件。 +- `core/src/` 下无名为 `repository/` 的目录。 +- `core/src/` 下无名为 `transport/` 的目录。 +- `crates/server/src/` 下无名为 `facade.rs` 的文件。 + +## 完成标准 + +当**所有**以下条件为真时，迁移完成： + +1. Rust workspace 源码树与 [精确 Rust 源码形状](#精确-rust-源码形状) 完全匹配。 +2. 被禁止的名称（架构测试中列出）不存在。 +3. 所有架构测试通过（`cargo test --test architecture`）。 +4. `cargo fmt --check --workspace` 通过。 +5. `cargo check --workspace --all-targets` 通过。 +6. `cargo test --workspace` 通过（所有 473+ 现有测试 + 新的架构测试）。 +7. `cargo clippy --workspace --all-targets --all-features --locked -- -D warnings` 通过。 +8. `cargo doc --workspace --no-deps` 无错误生成。 +9. `core`、`adapters` 和 `platform` 上启用了 `#![deny(missing_docs)]`。 +10. 非测试代码中无不安全的 `unwrap()`、`expect()`、`panic!` 或 `todo!`。 diff --git a/src/admin/api/accounts/create.rs b/src/admin/api/accounts/create.rs deleted file mode 100644 index 383b90b41..000000000 --- a/src/admin/api/accounts/create.rs +++ /dev/null @@ -1,42 +0,0 @@ -use axum::{ - extract::State, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, Json, -}; -use serde::Deserialize; - -use crate::{platform::http::request_id::RequestId, runtime::state::AppState}; - -use super::super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; -use super::{admin_account_data_from_stored, validated_account_import_error}; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct CreateAccountRequest { - pub token: Option, - pub refresh_token: Option, -} - -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), - )) -} diff --git a/src/admin/api/accounts/delete.rs b/src/admin/api/accounts/delete.rs deleted file mode 100644 index 03f4d2b82..000000000 --- a/src/admin/api/accounts/delete.rs +++ /dev/null @@ -1,98 +0,0 @@ -use axum::{ - extract::{Path, State}, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, Json, -}; -use serde::{Deserialize, Serialize}; - -use crate::{platform::http::request_id::RequestId, runtime::state::AppState}; - -use super::super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; -use super::account_service_error; - -#[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, -} - -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, - ), - )) -} diff --git a/src/admin/api/accounts/export.rs b/src/admin/api/accounts/export.rs deleted file mode 100644 index f9de8e99c..000000000 --- a/src/admin/api/accounts/export.rs +++ /dev/null @@ -1,85 +0,0 @@ -use axum::{ - extract::{Query, State}, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, -}; -use secrecy::ExposeSecret; -use serde::Deserialize; -use serde_json::{json, Value}; - -use crate::{ - codex::accounts::repository::StoredAccount, platform::http::request_id::RequestId, - runtime::state::AppState, -}; - -use super::super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; -use super::{account_export_ids, account_service_error, account_status_value}; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountExportQuery { - pub ids: Option, - pub format: Option, -} - -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(); - validate_account_export_format(query.format.as_deref()) - .map_err(|message| 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 = native_account_export(accounts); - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(data, request_id), - )) -} - -fn validate_account_export_format(value: Option<&str>) -> Result<(), &'static str> { - match value.unwrap_or("native") { - "" | "native" => Ok(()), - _ => 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(), - }) -} diff --git a/src/admin/api/accounts/import.rs b/src/admin/api/accounts/import.rs deleted file mode 100644 index af0a8f68b..000000000 --- a/src/admin/api/accounts/import.rs +++ /dev/null @@ -1,280 +0,0 @@ -use axum::{ - body::Bytes, - extract::State, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, Json, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; - -use crate::{ - codex::{ - accounts::service::{AccountImportEntry, StoreImportAccountError}, - gateway::oauth::{default_codex_home, read_cli_auth_from_home}, - }, - platform::http::request_id::RequestId, - runtime::state::AppState, - utils::json::first_string, -}; - -use super::super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; -use super::validated_account_import_error; - -const NATIVE_IMPORT_CONTAINER_KEYS: [&str; 1] = ["accounts"]; -const NATIVE_IMPORT_ACCOUNT_KEYS: [&str; 10] = [ - "id", - "email", - "accountId", - "userId", - "label", - "planType", - "token", - "refreshToken", - "accessTokenExpiresAt", - "status", -]; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum AccountImportFormat { - Native, - CodexCli, -} - -impl AccountImportFormat { - fn as_str(self) -> &'static str { - match self { - Self::Native => "native", - 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, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct ImportCliAuthRequest { - pub codex_home: Option, -} - -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(repository_unavailable(&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, - ), - )) -} - -fn repository_unavailable(request_id: &str) -> AdminError { - AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - "Account repository is not initialized", - 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, - ), - } -} - -fn parse_account_import_payload(payload: &Value) -> ParsedAccountImportPayload { - ParsedAccountImportPayload { - accounts: parse_native_account_payload(payload), - source_format: AccountImportFormat::Native, - } -} - -fn parse_native_account_payload(payload: &Value) -> Vec { - if !has_only_allowed_keys(payload, &NATIVE_IMPORT_CONTAINER_KEYS) { - return Vec::new(); - } - let Some(accounts) = payload.get("accounts").and_then(Value::as_array) else { - return Vec::new(); - }; - accounts - .iter() - .filter_map(account_entry_from_value) - .collect() -} - -fn account_entry_from_value(value: &Value) -> Option { - if !has_only_allowed_keys(value, &NATIVE_IMPORT_ACCOUNT_KEYS) { - return None; - } - - let entry = AccountImportEntry { - id: first_string(value, &[&["id"]]), - email: first_string(value, &[&["email"]]), - account_id: first_string(value, &[&["accountId"]]), - user_id: first_string(value, &[&["userId"]]), - label: label_from_value(value), - plan_type: first_string(value, &[&["planType"]]), - token: first_string(value, &[&["token"]]), - refresh_token: first_string(value, &[&["refreshToken"]]), - access_token_expires_at: first_string(value, &[&["accessTokenExpiresAt"]]), - status: first_string(value, &[&["status"]]), - }; - (entry.token.is_some() || entry.refresh_token.is_some()).then_some(entry) -} - -fn has_only_allowed_keys(value: &Value, allowed: &[&str]) -> bool { - value - .as_object() - .is_some_and(|object| object.keys().all(|key| allowed.contains(&key.as_str()))) -} - -fn label_from_value(value: &Value) -> Option { - normalized_label(first_string(value, &[&["label"]])) -} - -fn normalized_label(value: Option) -> Option { - value.map(|label| label.chars().take(64).collect()) -} - -fn empty_to_none(value: Option) -> Option { - value - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) -} diff --git a/src/admin/api/accounts/lifecycle.rs b/src/admin/api/accounts/lifecycle.rs deleted file mode 100644 index 49b0878d2..000000000 --- a/src/admin/api/accounts/lifecycle.rs +++ /dev/null @@ -1,291 +0,0 @@ -use axum::{ - extract::{Path, State}, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, Json, -}; -use serde::{Deserialize, Serialize}; - -use crate::{ - codex::accounts::{model::AccountStatus, service::RefreshAccountError}, - platform::http::request_id::RequestId, - runtime::state::AppState, -}; - -use super::super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; -use super::{account_probe_data_from_service, account_service_error, account_status_value}; - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct ResetAccountUsageData { - pub id: String, - pub reset: bool, -} - -#[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, 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, -} - -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 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 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, - ), - )) -} - -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::LeaseAcquire => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - "Failed to acquire refresh lease", - request_id, - ), - RefreshAccountError::TokenRefresherUnavailable => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - "Token refresher is not initialized", - request_id, - ), - RefreshAccountError::StoreStatus => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - "Failed to store account status", - request_id, - ), - RefreshAccountError::StoreRefreshed => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - "Failed to store refreshed account", - request_id, - ), - } -} - -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}")), - } -} diff --git a/src/admin/api/accounts/list.rs b/src/admin/api/accounts/list.rs deleted file mode 100644 index 488111e90..000000000 --- a/src/admin/api/accounts/list.rs +++ /dev/null @@ -1,51 +0,0 @@ -use axum::{ - extract::{Query, State}, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, -}; -use serde::Deserialize; - -use crate::{ - platform::http::request_id::RequestId, - runtime::state::AppState, - utils::pagination::{clamp_limit, Page}, -}; - -use super::super::{require_admin_session, AdminError, AdminPageEnvelope, AdminResponse}; -use super::{account_service_error, AdminAccountData}; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountsQuery { - pub cursor: Option, - pub limit: Option, -} - -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), - )) -} diff --git a/src/admin/api/accounts/mod.rs b/src/admin/api/accounts/mod.rs deleted file mode 100644 index ba3aafbf7..000000000 --- a/src/admin/api/accounts/mod.rs +++ /dev/null @@ -1,319 +0,0 @@ -use axum::http::StatusCode; -use serde::Serialize; - -use crate::{ - codex::accounts::service::{AccountServiceError, ValidatedAccountImportError}, - codex::accounts::{ - model::AccountStatus, - repository::{StoredAccount, StoredAccountMetadata}, - }, -}; - -use super::AdminError; - -pub mod cookies; -pub mod create; -pub mod delete; -pub mod export; -pub mod health; -pub mod import; -pub mod lifecycle; -pub mod list; -pub mod oauth; -pub mod quota; - -pub use cookies::{ - delete_account_cookies, get_account_cookies, set_account_cookies, AccountCookiesData, - DeleteAccountCookiesData, SetAccountCookiesRequest, -}; -pub use create::{create_account, CreateAccountRequest}; -pub use delete::{ - batch_delete_accounts, delete_account, BatchDeleteAccountsData, BatchDeleteAccountsRequest, - DeleteAccountData, -}; -pub use export::{export_accounts, AccountExportQuery}; -pub use health::{health_check_accounts, AccountProbeData, HealthCheckData, HealthCheckRequest}; -pub use import::{import_accounts, import_cli_auth, AccountImportData, ImportCliAuthRequest}; -pub use lifecycle::{ - batch_update_account_status, refresh_account, reset_account_usage, update_account_label, - update_account_status, BatchUpdateAccountStatusData, BatchUpdateAccountStatusRequest, - ResetAccountUsageData, UpdateAccountLabelData, UpdateAccountLabelRequest, - UpdateAccountStatusData, UpdateAccountStatusRequest, -}; -pub use list::{accounts, AccountsQuery}; -pub use oauth::{ - auth_callback, auth_code_relay, auth_device_login, auth_device_poll, auth_login_start, - AdminAuthCallbackQuery, AdminAuthCodeRelayData, AdminAuthCodeRelayRequest, - AdminAuthDeviceLoginData, AdminAuthDevicePollData, AdminAuthLoginStartData, -}; -pub use quota::{AccountQuotaData, AccountQuotaWarningsData}; - -pub(super) use health::account_probe_data_from_service; -pub(super) use quota::{account_quota, quota_warnings}; - -// 管理端账号接口按能力拆分；这里仅保留共享转换和对路由层稳定的 re-export。 -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, - ), - AccountServiceError::QuotaWarnings => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - "Failed to load account quota warnings", - request_id, - ), - } -} - -#[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, - ), - } -} - -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() -} - -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/admin/api/accounts/oauth.rs b/src/admin/api/accounts/oauth.rs deleted file mode 100644 index bd79cd1f1..000000000 --- a/src/admin/api/accounts/oauth.rs +++ /dev/null @@ -1,338 +0,0 @@ -use axum::{ - extract::{Path, Query, State}, - http::{ - header::{HeaderValue, HOST, LOCATION}, - HeaderMap, StatusCode, - }, - response::{IntoResponse, Response}, - Extension, Json, -}; -use reqwest::Url; -use serde::{Deserialize, Serialize}; - -use crate::{ - admin::session::service::{AdminAuthOAuthError, AdminAuthPkceExchange}, - codex::gateway::oauth::{DeviceCode, OAuthError}, - config::AppConfig, - platform::http::request_id::RequestId, - runtime::state::AppState, -}; - -use super::super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; -use super::validated_account_import_error; - -// OAuth 成功后会导入或更新账号资产，因此 handler 归在 accounts 域；密码登录仍留在 auth。 -#[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 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) -} - -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/admin/api/accounts/quota.rs b/src/admin/api/accounts/quota.rs deleted file mode 100644 index 674a88455..000000000 --- a/src/admin/api/accounts/quota.rs +++ /dev/null @@ -1,195 +0,0 @@ -use axum::{ - extract::{Path, State}, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, -}; -use chrono::SecondsFormat; -use serde::Serialize; -use serde_json::{json, Value}; - -use crate::{ - codex::accounts::service::{ - AccountQuotaError, AccountQuotaWarning as ServiceAccountQuotaWarning, - AccountQuotaWarnings as ServiceAccountQuotaWarnings, - }, - platform::http::request_id::RequestId, - runtime::state::AppState, -}; - -use super::super::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; -use super::{account_service_error, account_status_value}; - -pub(crate) 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(), - } - } -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountQuotaData { - pub quota: Value, - pub raw: Value, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountQuotaWarningsData { - pub warnings: Vec, - pub updated_at: Option, -} - -impl From for AccountQuotaWarningsData { - fn from(warnings: ServiceAccountQuotaWarnings) -> Self { - Self { - warnings: warnings.warnings.into_iter().map(Into::into).collect(), - updated_at: warnings - .updated_at - .map(|value| value.to_rfc3339_opts(SecondsFormat::Secs, true)), - } - } -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AccountQuotaWarningData { - pub account_id: String, - pub email: Option, - pub window: String, - pub level: String, - pub used_percent: f64, - pub reset_at: Option, -} - -impl From for AccountQuotaWarningData { - fn from(warning: ServiceAccountQuotaWarning) -> Self { - Self { - account_id: warning.account_id, - email: warning.email, - window: warning.window.as_str().to_string(), - level: warning.level.as_str().to_string(), - used_percent: warning.used_percent, - reset_at: warning.reset_at, - } - } -} - -pub(crate) 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(crate) async fn quota_warnings( - 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 warnings = state - .services - .accounts - .quota_warnings() - .await - .map_err(|error| account_service_error(error, &request_id))?; - - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountQuotaWarningsData::from(warnings), request_id), - )) -} - -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(), - }, - } -} diff --git a/src/admin/api/client_keys/create.rs b/src/admin/api/client_keys/create.rs deleted file mode 100644 index 6769c0739..000000000 --- a/src/admin/api/client_keys/create.rs +++ /dev/null @@ -1,28 +0,0 @@ -use axum::{ - extract::State, http::HeaderMap, http::StatusCode, response::IntoResponse, Extension, Json, -}; - -use crate::{platform::http::request_id::RequestId, runtime::state::AppState}; - -use super::{api_key_service_error, CreateApiKeyRequest, CreatedClientApiKeyData}; -use crate::admin::api::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; - -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)), - } -} diff --git a/src/admin/api/client_keys/export.rs b/src/admin/api/client_keys/export.rs deleted file mode 100644 index 5a0bf0b8e..000000000 --- a/src/admin/api/client_keys/export.rs +++ /dev/null @@ -1,44 +0,0 @@ -use axum::{ - extract::{Query, State}, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, -}; - -use crate::{platform::http::request_id::RequestId, runtime::state::AppState}; - -use super::{ - api_key_service_error, ApiKeyExportQuery, ClientApiKeyExportData, ClientApiKeyExportEntry, -}; -use crate::admin::api::{ - account_export_ids, require_admin_session, AdminEnvelope, AdminError, AdminResponse, -}; - -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), - )) -} diff --git a/src/admin/api/client_keys/import.rs b/src/admin/api/client_keys/import.rs deleted file mode 100644 index 04081b56e..000000000 --- a/src/admin/api/client_keys/import.rs +++ /dev/null @@ -1,40 +0,0 @@ -use axum::{ - extract::State, http::HeaderMap, http::StatusCode, response::IntoResponse, Extension, Json, -}; -use serde_json::Value; - -use crate::{platform::http::request_id::RequestId, runtime::state::AppState}; - -use super::{api_key_service_error, ClientApiKeyImportData, ImportedClientApiKeyData}; -use crate::admin::api::{require_admin_session, AdminEnvelope, AdminError, AdminResponse}; - -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, - ), - )) -} diff --git a/src/admin/api/logs/mod.rs b/src/admin/api/logs/mod.rs deleted file mode 100644 index 399699eb8..000000000 --- a/src/admin/api/logs/mod.rs +++ /dev/null @@ -1,60 +0,0 @@ -use axum::http::StatusCode; - -use crate::codex::events::service::LogServiceError; - -use super::AdminError; - -pub mod detail; -pub mod query; -pub mod state; - -pub use detail::log_detail; -pub use query::logs; -pub use state::{clear_logs, logs_state, update_logs_state}; - -pub(super) 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, - ), - LogServiceError::Get => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - "Failed to load event log", - request_id, - ), - LogServiceError::Count => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - "Failed to count event logs", - request_id, - ), - LogServiceError::Clear => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - "Failed to clear event logs", - request_id, - ), - LogServiceError::Write => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - "Failed to write event log", - request_id, - ), - LogServiceError::InvalidCapacity => AdminError::new( - StatusCode::UNPROCESSABLE_ENTITY, - 42201, - "logsCapacity must be greater than 0", - request_id, - ), - } -} diff --git a/src/admin/api/logs/query.rs b/src/admin/api/logs/query.rs deleted file mode 100644 index 6ba3cb07c..000000000 --- a/src/admin/api/logs/query.rs +++ /dev/null @@ -1,71 +0,0 @@ -use axum::{ - extract::{Query, State}, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, -}; -use serde::Deserialize; - -use crate::{ - codex::events::service::LogListFilter, platform::http::request_id::RequestId, - runtime::state::AppState, utils::pagination::clamp_limit, -}; - -use super::log_service_error; -use crate::admin::api::{require_admin_session, AdminError, AdminPageEnvelope, AdminResponse}; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct LogsQuery { - pub cursor: Option, - pub limit: Option, - pub kind: Option, - pub level: Option, - pub request_id: Option, - pub account_id: Option, - pub route: Option, - pub model: Option, - pub status_code: Option, - pub search: 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)); - let cursor = query.cursor.clone(); - match state.services.logs.list(cursor, limit, query.into()).await { - Ok(page) => Ok(AdminResponse::new( - StatusCode::OK, - AdminPageEnvelope::ok(page, limit, request_id), - )), - Err(error) => Err(log_service_error(error, request_id)), - } -} - -impl From for LogListFilter { - fn from(query: LogsQuery) -> Self { - Self { - kind: non_empty(query.kind), - level: non_empty(query.level), - request_id: non_empty(query.request_id), - account_id: non_empty(query.account_id), - route: non_empty(query.route), - model: non_empty(query.model), - status_code: query.status_code, - search: non_empty(query.search), - } - } -} - -fn non_empty(value: Option) -> Option { - value - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) -} diff --git a/src/admin/api/mod.rs b/src/admin/api/mod.rs deleted file mode 100644 index be0f22d6b..000000000 --- a/src/admin/api/mod.rs +++ /dev/null @@ -1,20 +0,0 @@ -pub mod accounts; -pub mod client_keys; -pub mod diagnostics; -pub mod logs; -pub mod models; -pub mod response; -pub mod router; -pub mod session; -pub mod settings; -pub mod usage; - -use self::{ - accounts::{account_export_ids, account_status_value}, - session::require_admin_session, -}; - -pub use self::{ - response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse, PageMeta}, - router::router, -}; diff --git a/src/admin/api/models.rs b/src/admin/api/models.rs deleted file mode 100644 index 33b33c88d..000000000 --- a/src/admin/api/models.rs +++ /dev/null @@ -1,114 +0,0 @@ -use axum::{ - extract::State, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, -}; -use serde::Serialize; - -use crate::{ - codex::models::service::{ModelRefreshResult, ModelServiceError}, - platform::http::request_id::RequestId, - runtime::state::AppState, -}; - -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/admin/api/session.rs b/src/admin/api/session.rs deleted file mode 100644 index a2b1d27be..000000000 --- a/src/admin/api/session.rs +++ /dev/null @@ -1,338 +0,0 @@ -use axum::{ - extract::State, - http::{ - header::{HeaderValue, SET_COOKIE}, - HeaderMap, StatusCode, - }, - response::{IntoResponse, Response}, - Extension, Json, -}; -use serde::{Deserialize, Serialize}; - -use crate::{ - admin::session::service::{ - AdminAuthPoolSummary as ServiceAdminAuthPoolSummary, AdminAuthServiceError, - AdminAuthStatus as ServiceAdminAuthStatus, AdminAuthUser as ServiceAdminAuthUser, - AdminSessionValidationError, - }, - platform::http::{auth::admin_session_id, request_id::RequestId}, - runtime::state::AppState, -}; - -use super::{account_status_value, 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, -} - -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 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, - ), - )) -} diff --git a/src/admin/client_keys/mod.rs b/src/admin/client_keys/mod.rs deleted file mode 100644 index 1f278a4d5..000000000 --- a/src/admin/client_keys/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod service; diff --git a/src/admin/client_keys/service.rs b/src/admin/client_keys/service.rs deleted file mode 100644 index 688b473fb..000000000 --- a/src/admin/client_keys/service.rs +++ /dev/null @@ -1,334 +0,0 @@ -use serde_json::Value; - -use crate::{ - platform::identity::{ - client_key::ApiKeyHasher, - client_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/admin/mod.rs b/src/admin/mod.rs deleted file mode 100644 index 1c84f2fc2..000000000 --- a/src/admin/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub mod api; -pub mod client_keys; -pub mod session; -pub mod settings; -pub mod tasks; diff --git a/src/admin/session/mod.rs b/src/admin/session/mod.rs deleted file mode 100644 index 5acc82448..000000000 --- a/src/admin/session/mod.rs +++ /dev/null @@ -1,2 +0,0 @@ -pub mod repository; -pub mod service; diff --git a/src/admin/session/service.rs b/src/admin/session/service.rs deleted file mode 100644 index 3ee9f101e..000000000 --- a/src/admin/session/service.rs +++ /dev/null @@ -1,401 +0,0 @@ -use std::sync::Arc; - -use chrono::{DateTime, Duration, Utc}; -use tokio::sync::Mutex; -use uuid::Uuid; - -use crate::{ - admin::session::repository::AdminAuthRepository, - codex::accounts::service::{AccountService, ValidatedAccountImportError}, - codex::accounts::{ - model::AccountStatus, - pool::AccountPool, - repository::{AccountRepository, StoredAccountMetadata}, - }, - codex::gateway::oauth::{ - DeviceCode, OAuthClient, OAuthConfig, OAuthError, PkceSession, PkceSessionStore, - }, - config::AppConfig, - platform::identity::admin_session::verify_admin_password, -}; - -#[derive(Clone)] -pub struct AdminAuthService { - config: Arc, - auth_repository: Option, - account_repository: Option, - account_pool: Arc>, - oauth_config: OAuthConfig, - oauth_client: Option>, - oauth_sessions: Arc>, - accounts: AccountService, -} - -pub struct AdminAuthDependencies { - pub auth_repository: Option, - pub account_repository: Option, - pub account_pool: Arc>, - pub oauth_config: OAuthConfig, - pub oauth_client: Option>, - pub oauth_sessions: Arc>, - pub 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, -} - -impl AdminAuthService { - pub fn new(config: Arc, dependencies: AdminAuthDependencies) -> Self { - let AdminAuthDependencies { - auth_repository, - account_repository, - account_pool, - oauth_config, - oauth_client, - oauth_sessions, - accounts, - } = dependencies; - Self { - config, - auth_repository, - account_repository, - account_pool, - oauth_config, - oauth_client, - oauth_sessions, - accounts, - } - } - - pub async fn login(&self, password: &str) -> Result { - let auth_repository = self.auth_repository()?; - let admin = auth_repository - .load_first_admin() - .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()); - auth_repository - .create_session(&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 auth_repository = self - .auth_repository() - .map_err(|_| AdminSessionValidationError::DatabaseUnavailable)?; - auth_repository - .validate_session(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, &self.oauth_config); - 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 auth_repository(&self) -> Result<&AdminAuthRepository, AdminAuthServiceError> { - self.auth_repository - .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)) - } - } - } -} - -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/admin/settings.rs b/src/admin/settings.rs deleted file mode 100644 index 2c902dc45..000000000 --- a/src/admin/settings.rs +++ /dev/null @@ -1,290 +0,0 @@ -use std::{ - collections::BTreeMap, - path::{Path, PathBuf}, - sync::Arc, -}; - -use thiserror::Error; -use tokio::sync::RwLock; - -use crate::config::{ - AppConfig, AuthConfig, LoggingConfig, ModelConfig, QuotaConfig, QuotaWarningThresholds, - UsageStatsConfig, -}; - -const ROTATION_STRATEGIES: [&str; 3] = ["least_used", "round_robin", "sticky"]; - -#[derive(Clone)] -pub struct SettingsService { - current: Arc>>, - local_config_path: Arc, -} - -#[derive(Debug, Default)] -pub struct SettingsUpdate { - pub default_model: Option, - pub default_reasoning_effort: Option, - pub service_tier: Option, - pub model_aliases: Option>, - pub refresh_enabled: Option, - pub refresh_margin_seconds: Option, - pub refresh_concurrency: Option, - pub max_concurrent_per_account: Option, - pub request_interval_ms: Option, - pub rotation_strategy: Option, - pub tier_priority: Option>, - pub quota_refresh_interval_minutes: Option, - pub quota_warning_thresholds: Option, - pub quota_skip_exhausted: Option, - pub logs_enabled: Option, - pub logs_capacity: Option, - pub logs_capture_body: Option, - pub usage_history_retention_days: Option, -} - -#[derive(Debug, Error)] -pub enum SettingsServiceError { - #[error("invalid setting `{field}`: {message}")] - InvalidField { field: String, message: String }, - #[error("serialize settings overlay")] - Serialize(#[source] serde_yml::Error), - #[error("create settings directory")] - CreateDirectory(#[source] std::io::Error), - #[error("write settings overlay")] - Write(#[source] std::io::Error), -} - -#[derive(Debug, serde::Serialize)] -struct SettingsConfigOverlay { - model: ModelConfig, - auth: AuthConfig, - quota: QuotaConfig, - usage_stats: UsageStatsConfig, - logging: LoggingConfig, -} - -impl SettingsService { - pub fn new(config: Arc, local_config_path: impl Into) -> Self { - Self { - current: Arc::new(RwLock::new(config)), - local_config_path: Arc::new(local_config_path.into()), - } - } - - pub async fn current(&self) -> Arc { - self.current.read().await.clone() - } - - pub fn local_config_path(&self) -> &Path { - self.local_config_path.as_ref() - } - - pub async fn update( - &self, - update: SettingsUpdate, - ) -> Result, SettingsServiceError> { - let mut current = self.current.write().await; - let mut next = (**current).clone(); - apply_update(&mut next, update)?; - let next = Arc::new(next); - self.write_overlay(&next).await?; - *current = next.clone(); - Ok(next) - } - - async fn write_overlay(&self, config: &AppConfig) -> Result<(), SettingsServiceError> { - if let Some(parent) = self - .local_config_path - .parent() - .filter(|path| !path.as_os_str().is_empty()) - { - tokio::fs::create_dir_all(parent) - .await - .map_err(SettingsServiceError::CreateDirectory)?; - } - - let overlay = SettingsConfigOverlay::from(config); - let content = serde_yml::to_string(&overlay).map_err(SettingsServiceError::Serialize)?; - tokio::fs::write(self.local_config_path.as_ref(), content) - .await - .map_err(SettingsServiceError::Write) - } -} - -impl From<&AppConfig> for SettingsConfigOverlay { - fn from(config: &AppConfig) -> Self { - Self { - model: config.model.clone(), - auth: config.auth.clone(), - quota: config.quota.clone(), - usage_stats: config.usage_stats.clone(), - logging: config.logging.clone(), - } - } -} - -fn apply_update( - config: &mut AppConfig, - update: SettingsUpdate, -) -> Result<(), SettingsServiceError> { - if let Some(default_model) = update.default_model { - config.model.default_model = non_empty_string("defaultModel", default_model)?; - } - if let Some(default_reasoning_effort) = update.default_reasoning_effort { - config.model.default_reasoning_effort = Some(non_empty_string( - "defaultReasoningEffort", - default_reasoning_effort, - )?); - } - if let Some(service_tier) = update.service_tier { - config.model.service_tier = Some(non_empty_string("serviceTier", service_tier)?); - } - if let Some(model_aliases) = update.model_aliases { - config.model.aliases = validate_aliases(model_aliases)?; - } - if let Some(refresh_enabled) = update.refresh_enabled { - config.auth.refresh_enabled = refresh_enabled; - } - if let Some(refresh_margin_seconds) = update.refresh_margin_seconds { - config.auth.refresh_margin_seconds = refresh_margin_seconds; - } - if let Some(refresh_concurrency) = update.refresh_concurrency { - config.auth.refresh_concurrency = positive_u32("refreshConcurrency", refresh_concurrency)?; - } - if let Some(max_concurrent_per_account) = update.max_concurrent_per_account { - config.auth.max_concurrent_per_account = - positive_usize("maxConcurrentPerAccount", max_concurrent_per_account)?; - } - if let Some(request_interval_ms) = update.request_interval_ms { - config.auth.request_interval_ms = request_interval_ms; - } - if let Some(rotation_strategy) = update.rotation_strategy { - config.auth.rotation_strategy = validate_rotation_strategy(rotation_strategy)?; - } - if let Some(tier_priority) = update.tier_priority { - config.auth.tier_priority = validate_tier_priority(tier_priority)?; - } - if let Some(quota_refresh_interval_minutes) = update.quota_refresh_interval_minutes { - config.quota.refresh_interval_minutes = positive_u64( - "quotaRefreshIntervalMinutes", - quota_refresh_interval_minutes, - )?; - } - if let Some(quota_warning_thresholds) = update.quota_warning_thresholds { - validate_thresholds("a_warning_thresholds)?; - config.quota.warning_thresholds = quota_warning_thresholds; - } - if let Some(quota_skip_exhausted) = update.quota_skip_exhausted { - config.quota.skip_exhausted = quota_skip_exhausted; - } - if let Some(logs_enabled) = update.logs_enabled { - config.logging.enabled = logs_enabled; - } - if let Some(logs_capacity) = update.logs_capacity { - config.logging.capacity = positive_u32("logsCapacity", logs_capacity)?; - } - if let Some(logs_capture_body) = update.logs_capture_body { - config.logging.capture_body = logs_capture_body; - } - if let Some(usage_history_retention_days) = update.usage_history_retention_days { - config.usage_stats.history_retention_days = Some(positive_u64( - "usageHistoryRetentionDays", - usage_history_retention_days, - )?); - } - - Ok(()) -} - -fn validate_aliases( - aliases: BTreeMap, -) -> Result, SettingsServiceError> { - let mut validated = BTreeMap::new(); - for (alias, target) in aliases { - let alias = non_empty_string("modelAliases", alias)?; - let target = non_empty_string("modelAliases", target)?; - validated.insert(alias, target); - } - Ok(validated) -} - -fn validate_tier_priority(tiers: Vec) -> Result, SettingsServiceError> { - let mut validated = Vec::with_capacity(tiers.len()); - for tier in tiers { - validated.push(non_empty_string("tierPriority", tier)?); - } - Ok(validated) -} - -fn validate_rotation_strategy(strategy: String) -> Result { - let strategy = non_empty_string("rotationStrategy", strategy)?; - if ROTATION_STRATEGIES.contains(&strategy.as_str()) { - Ok(strategy) - } else { - Err(invalid_field( - "rotationStrategy", - "must be one of least_used, round_robin, sticky", - )) - } -} - -fn validate_thresholds(thresholds: &QuotaWarningThresholds) -> Result<(), SettingsServiceError> { - if thresholds.primary.is_empty() && thresholds.secondary.is_empty() { - return Err(invalid_field( - "quotaWarningThresholds", - "must include at least one threshold", - )); - } - if thresholds - .primary - .iter() - .chain(thresholds.secondary.iter()) - .any(|threshold| *threshold > 100) - { - return Err(invalid_field( - "quotaWarningThresholds", - "thresholds must be between 0 and 100", - )); - } - Ok(()) -} - -fn non_empty_string(field: &str, value: String) -> Result { - let value = value.trim().to_string(); - if value.is_empty() { - Err(invalid_field(field, "must not be empty")) - } else { - Ok(value) - } -} - -fn positive_u32(field: &str, value: u32) -> Result { - if value == 0 { - Err(invalid_field(field, "must be greater than 0")) - } else { - Ok(value) - } -} - -fn positive_u64(field: &str, value: u64) -> Result { - if value == 0 { - Err(invalid_field(field, "must be greater than 0")) - } else { - Ok(value) - } -} - -fn positive_usize(field: &str, value: usize) -> Result { - if value == 0 { - Err(invalid_field(field, "must be greater than 0")) - } else { - Ok(value) - } -} - -fn invalid_field(field: &str, message: impl Into) -> SettingsServiceError { - SettingsServiceError::InvalidField { - field: field.to_string(), - message: message.into(), - } -} diff --git a/src/admin/tasks/mod.rs b/src/admin/tasks/mod.rs deleted file mode 100644 index b4501c361..000000000 --- a/src/admin/tasks/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod session_cleanup; diff --git a/src/admin/tasks/session_cleanup.rs b/src/admin/tasks/session_cleanup.rs deleted file mode 100644 index e43ced5a2..000000000 --- a/src/admin/tasks/session_cleanup.rs +++ /dev/null @@ -1,59 +0,0 @@ -use chrono::Utc; -use tokio::time::{interval, Duration}; -use tracing::{error, info}; - -use crate::admin::session::repository::AdminAuthRepository; -use crate::runtime::tasks::types::SchedulerHandle; - -/// 会话清理调度器 - 定期删除过期的管理员会话 -pub struct SessionCleanupScheduler { - repository: AdminAuthRepository, - interval_secs: u64, -} - -impl SessionCleanupScheduler { - pub fn new(repository: AdminAuthRepository, interval_secs: u64) -> Self { - Self { - repository, - interval_secs, - } - } - - /// 启动会话清理调度器 - pub fn start(self) -> SchedulerHandle { - let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); - - tokio::spawn(async move { - let mut ticker = interval(Duration::from_secs(self.interval_secs)); - info!(interval_secs = self.interval_secs, "会话清理调度器已启动"); - - loop { - tokio::select! { - _ = ticker.tick() => { - match self.cleanup_expired_sessions().await { - Ok(count) if count > 0 => { - info!(count, "已清理过期会话"); - } - Ok(_) => { - // 没有过期会话，不输出日志 - } - Err(e) => { - error!(error = %e, "清理过期会话失败"); - } - } - } - _ = shutdown_rx.recv() => { - info!("会话清理调度器已关闭"); - break; - } - } - } - }); - - SchedulerHandle::new(shutdown_tx) - } - - async fn cleanup_expired_sessions(&self) -> Result { - self.repository.cleanup_expired_sessions(Utc::now()).await - } -} diff --git a/src/codex/accounts/cloudflare_challenge.rs b/src/codex/accounts/cloudflare_challenge.rs deleted file mode 100644 index bb8292214..000000000 --- a/src/codex/accounts/cloudflare_challenge.rs +++ /dev/null @@ -1,143 +0,0 @@ -use std::collections::HashMap; -use std::sync::Arc; - -use chrono::{DateTime, Duration, Utc}; -use tokio::sync::RwLock; - -/// Cloudflare Path Block Tracker -/// -/// 追踪每个账户连续遇到的 Cloudflare path-block 404 错误。 -/// 当一个账户连续遇到 3 次 path-block 时，自动将其标记为 Disabled。 -/// -/// Path-block 的特征： -/// - HTTP 404 状态码 -/// - 空响应体或 Cloudflare 页面 -/// - 通常发生在 /backend-api/responses 路径 -const PATH_BLOCK_THRESHOLD: u32 = 3; -const PATH_BLOCK_STALE_AFTER: Duration = Duration::hours(1); - -#[derive(Debug, Clone, Copy)] -struct PathBlockState { - count: u32, - last_at: DateTime, -} - -#[derive(Debug, Clone)] -pub struct CfPathBlockTracker { - /// 每个账户的连续 path-block 计数 - counts: Arc>>, -} - -impl Default for CfPathBlockTracker { - fn default() -> Self { - Self::new() - } -} - -impl CfPathBlockTracker { - pub fn new() -> Self { - Self { - counts: Arc::new(RwLock::new(HashMap::new())), - } - } - - /// 记录一次 path-block 404 - /// - /// 返回当前连续计数，调用者应该在达到阈值时采取行动 - pub async fn record_path_block(&self, account_id: &str) -> u32 { - let mut counts = self.counts.write().await; - let now = Utc::now(); - let count = counts - .get(account_id) - .filter(|state| now.signed_duration_since(state.last_at) <= PATH_BLOCK_STALE_AFTER) - .map(|state| state.count + 1) - .unwrap_or(1); - counts.insert( - account_id.to_string(), - PathBlockState { - count, - last_at: now, - }, - ); - count - } - - /// 重置账户的 path-block 计数（成功请求后调用） - pub async fn reset(&self, account_id: &str) { - let mut counts = self.counts.write().await; - counts.remove(account_id); - } - - /// 获取账户当前的 path-block 计数 - pub async fn get_count(&self, account_id: &str) -> u32 { - let counts = self.counts.read().await; - counts - .get(account_id) - .filter(|state| { - Utc::now().signed_duration_since(state.last_at) <= PATH_BLOCK_STALE_AFTER - }) - .map(|state| state.count) - .unwrap_or(0) - } - - /// 检查是否应该禁用账户 - pub async fn should_disable(&self, account_id: &str) -> bool { - self.get_count(account_id).await >= PATH_BLOCK_THRESHOLD - } - - /// 清除所有计数（用于测试或管理操作） - pub async fn clear_all(&self) { - let mut counts = self.counts.write().await; - counts.clear(); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_path_block_counter() { - let tracker = CfPathBlockTracker::new(); - - // 第一次 - let count = tracker.record_path_block("acct_1").await; - assert_eq!(count, 1); - assert!(!tracker.should_disable("acct_1").await); - - // 第二次 - let count = tracker.record_path_block("acct_1").await; - assert_eq!(count, 2); - assert!(!tracker.should_disable("acct_1").await); - - // 第三次 - 应该触发禁用 - let count = tracker.record_path_block("acct_1").await; - assert_eq!(count, 3); - assert!(tracker.should_disable("acct_1").await); - } - - #[tokio::test] - async fn test_reset_clears_count() { - let tracker = CfPathBlockTracker::new(); - - tracker.record_path_block("acct_1").await; - tracker.record_path_block("acct_1").await; - assert_eq!(tracker.get_count("acct_1").await, 2); - - // 重置后计数应该清零 - tracker.reset("acct_1").await; - assert_eq!(tracker.get_count("acct_1").await, 0); - } - - #[tokio::test] - async fn test_multiple_accounts() { - let tracker = CfPathBlockTracker::new(); - - tracker.record_path_block("acct_1").await; - tracker.record_path_block("acct_2").await; - tracker.record_path_block("acct_2").await; - - assert_eq!(tracker.get_count("acct_1").await, 1); - assert_eq!(tracker.get_count("acct_2").await, 2); - } -} diff --git a/src/codex/accounts/cookies/jar.rs b/src/codex/accounts/cookies/jar.rs deleted file mode 100644 index 6d6d03068..000000000 --- a/src/codex/accounts/cookies/jar.rs +++ /dev/null @@ -1,67 +0,0 @@ -use std::collections::HashMap; - -const CAPTURABLE_COOKIES: &[&str] = &["cf_clearance"]; - -#[derive(Debug, Clone)] -struct StoredCookie { - domain: String, - name: String, - value: String, -} - -#[derive(Debug, Default, Clone)] -pub struct CookieJar { - by_account: HashMap>, -} - -impl CookieJar { - pub fn capture_set_cookie(&mut self, account_id: &str, raw: &str) { - let mut parts = raw.split(';').map(str::trim); - let Some(name_value) = parts.next() else { - return; - }; - let Some((name, value)) = name_value.split_once('=') else { - return; - }; - - if !CAPTURABLE_COOKIES.contains(&name) { - return; - } - - let mut domain = "chatgpt.com".to_string(); - for part in parts { - if let Some(value) = part.strip_prefix("Domain=") { - domain = value.trim_start_matches('.').to_string(); - } - } - - let account = self.by_account.entry(account_id.to_string()).or_default(); - account.retain(|cookie| !(cookie.domain == domain && cookie.name == name)); - account.push(StoredCookie { - domain, - name: name.to_string(), - value: value.to_string(), - }); - } - - pub fn cookie_header(&self, account_id: &str, domain: &str) -> Option { - let cookies = self.by_account.get(account_id)?; - let pairs = cookies - .iter() - .filter(|cookie| domain_matches(domain, &cookie.domain)) - .map(|cookie| format!("{}={}", cookie.name, cookie.value)) - .collect::>(); - if pairs.is_empty() { - None - } else { - Some(pairs.join("; ")) - } - } -} - -fn domain_matches(request_domain: &str, cookie_domain: &str) -> bool { - request_domain == cookie_domain - || request_domain - .strip_suffix(cookie_domain) - .is_some_and(|prefix| prefix.ends_with('.')) -} diff --git a/src/codex/accounts/cookies/mod.rs b/src/codex/accounts/cookies/mod.rs deleted file mode 100644 index 1a846c964..000000000 --- a/src/codex/accounts/cookies/mod.rs +++ /dev/null @@ -1,2 +0,0 @@ -pub mod jar; -pub mod repository; diff --git a/src/codex/accounts/lifecycle.rs b/src/codex/accounts/lifecycle.rs deleted file mode 100644 index 3a30cbfa2..000000000 --- a/src/codex/accounts/lifecycle.rs +++ /dev/null @@ -1,9 +0,0 @@ -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum AccountLifecycleEvent { - Added, - Refreshed, - Expired, - QuotaExhausted, - Banned, - Disabled, -} diff --git a/src/codex/accounts/mod.rs b/src/codex/accounts/mod.rs deleted file mode 100644 index 8e6014fdb..000000000 --- a/src/codex/accounts/mod.rs +++ /dev/null @@ -1,9 +0,0 @@ -pub mod cloudflare_challenge; -pub mod cookies; -pub mod jwt; -pub mod lifecycle; -pub mod model; -pub mod pool; -pub mod repository; -pub mod service; -pub mod usage_snapshots; diff --git a/src/codex/accounts/repository/accounts.rs b/src/codex/accounts/repository/accounts.rs deleted file mode 100644 index 2f485e5b4..000000000 --- a/src/codex/accounts/repository/accounts.rs +++ /dev/null @@ -1,754 +0,0 @@ -use chrono::{DateTime, Utc}; -use secrecy::{ExposeSecret, SecretString}; -use sqlx::{Row, SqlitePool}; -use thiserror::Error; - -use crate::{ - codex::accounts::model::{Account, AccountStatus}, - platform::crypto::{CryptoError, SecretBox}, - utils::pagination::{decode_cursor, encode_cursor, Page}, -}; - -#[derive(Debug, Error)] -pub enum AccountRepositoryError { - #[error("database error: {0}")] - Database(#[from] sqlx::Error), - #[error("account secret encryption error: {0}")] - Crypto(#[from] CryptoError), - #[error("invalid account status: {0}")] - InvalidStatus(String), - #[error("invalid account timestamp: {0}")] - InvalidTimestamp(#[from] chrono::ParseError), - #[error("invalid pagination cursor")] - InvalidCursor, -} - -pub type AccountRepositoryResult = Result; - -#[derive(Debug)] -pub struct NewAccount { - pub id: String, - pub email: Option, - pub account_id: Option, - pub user_id: Option, - pub label: Option, - pub plan_type: Option, - pub access_token: SecretString, - pub refresh_token: Option, - pub access_token_expires_at: Option>, - pub status: AccountStatus, -} - -#[derive(Debug)] -pub struct TokenUpdate { - pub access_token: SecretString, - pub refresh_token: Option, - pub access_token_expires_at: Option>, -} - -#[derive(Debug)] -pub struct AccountClaimsUpdate { - pub email: Option, - pub account_id: Option, - pub user_id: Option, - pub plan_type: Option, - pub access_token: SecretString, - pub refresh_token: Option, - pub access_token_expires_at: Option>, - pub status: AccountStatus, -} - -#[derive(Debug, Clone)] -pub struct StoredAccount { - pub id: String, - pub email: Option, - pub account_id: Option, - pub user_id: Option, - pub label: Option, - pub plan_type: Option, - pub access_token: SecretString, - pub refresh_token: Option, - pub access_token_expires_at: Option>, - pub status: AccountStatus, - pub added_at: DateTime, - pub updated_at: DateTime, -} - -#[derive(Debug, Clone)] -pub struct StoredAccountMetadata { - pub id: String, - pub email: Option, - pub account_id: Option, - pub user_id: Option, - pub label: Option, - pub plan_type: Option, - pub access_token_expires_at: Option>, - pub status: AccountStatus, - pub added_at: DateTime, - pub updated_at: DateTime, -} - -#[derive(Debug, Clone, Copy, Default)] -pub struct UsageDelta { - pub input_tokens: i64, - pub output_tokens: i64, - pub cached_tokens: i64, - pub image_input_tokens: i64, - pub image_output_tokens: i64, - pub image_request_count: i64, - pub image_request_failed_count: i64, - pub empty_response_count: i64, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AccountUsageRecord { - pub account_id: String, - pub request_count: i64, - pub empty_response_count: i64, - pub input_tokens: i64, - pub output_tokens: i64, - pub cached_tokens: i64, - pub image_input_tokens: i64, - pub image_output_tokens: i64, - pub image_request_count: i64, - pub image_request_failed_count: i64, - pub window_request_count: i64, - pub window_input_tokens: i64, - pub window_output_tokens: i64, - pub window_cached_tokens: i64, - pub window_image_input_tokens: i64, - pub window_image_output_tokens: i64, - pub window_image_request_count: i64, - pub window_image_request_failed_count: i64, - pub window_started_at: Option>, - pub window_reset_at: Option>, - pub limit_window_seconds: Option, - pub last_used_at: Option>, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AccountUsageListRecord { - pub account_id: String, - pub email: Option, - pub label: Option, - pub plan_type: Option, - pub request_count: i64, - pub empty_response_count: i64, - pub input_tokens: i64, - pub output_tokens: i64, - pub cached_tokens: i64, - pub image_input_tokens: i64, - pub image_output_tokens: i64, - pub image_request_count: i64, - pub image_request_failed_count: i64, - pub last_used_at: Option>, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AccountQuotaSnapshot { - pub account_id: String, - pub email: Option, - pub quota_json: String, - pub quota_fetched_at: Option>, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct AccountUsageSummary { - pub account_count: i64, - pub request_count: i64, - pub empty_response_count: i64, - pub input_tokens: i64, - pub output_tokens: i64, - pub cached_tokens: i64, - pub image_input_tokens: i64, - pub image_output_tokens: i64, - pub image_request_count: i64, - pub image_request_failed_count: i64, -} - -#[derive(Clone)] -pub struct AccountRepository { - pub(super) pool: SqlitePool, - pub(super) secret_box: SecretBox, -} - -const INSERT_ACCOUNT_SQL: &str = r" -insert into accounts ( - id, - email, - account_id, - user_id, - label, - plan_type, - access_token_cipher, - refresh_token_cipher, - access_token_expires_at, - status, - added_at, - updated_at -) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"; - -const SELECT_STORED_ACCOUNT_BY_ID_SQL: &str = r" -select - id, - email, - account_id, - user_id, - label, - plan_type, - access_token_cipher, - refresh_token_cipher, - access_token_expires_at, - status, - added_at, - updated_at -from accounts -where id = ?"; - -const SELECT_STORED_ACCOUNT_BY_CHATGPT_IDENTITY_SQL: &str = r" -select - id, - email, - account_id, - user_id, - label, - plan_type, - access_token_cipher, - refresh_token_cipher, - access_token_expires_at, - status, - added_at, - updated_at -from accounts -where account_id = ? - and ((user_id is null and ? is null) or user_id = ?) -order by added_at asc -limit 1"; - -const LIST_STORED_ACCOUNTS_AFTER_CURSOR_SQL: &str = r" -select - id, - email, - account_id, - user_id, - label, - plan_type, - access_token_cipher, - refresh_token_cipher, - access_token_expires_at, - status, - added_at, - updated_at -from accounts -where added_at < ? - or (added_at = ? and id < ?) -order by added_at desc, id desc -limit ?"; - -const LIST_STORED_ACCOUNTS_SQL: &str = r" -select - id, - email, - account_id, - user_id, - label, - plan_type, - access_token_cipher, - refresh_token_cipher, - access_token_expires_at, - status, - added_at, - updated_at -from accounts -order by added_at desc, id desc -limit ?"; - -const LIST_ALL_STORED_ACCOUNTS_SQL: &str = r" -select - id, - email, - account_id, - user_id, - label, - plan_type, - access_token_cipher, - refresh_token_cipher, - access_token_expires_at, - status, - added_at, - updated_at -from accounts -order by added_at desc, id desc"; - -const LIST_METADATA_AFTER_CURSOR_SQL: &str = r" -select - id, - email, - account_id, - user_id, - label, - plan_type, - access_token_expires_at, - status, - added_at, - updated_at -from accounts -where added_at < ? - or (added_at = ? and id < ?) -order by added_at desc, id desc -limit ?"; - -const LIST_METADATA_SQL: &str = r" -select - id, - email, - account_id, - user_id, - label, - plan_type, - access_token_expires_at, - status, - added_at, - updated_at -from accounts -order by added_at desc, id desc -limit ?"; - -const LIST_ALL_METADATA_SQL: &str = r" -select - id, - email, - account_id, - user_id, - label, - plan_type, - access_token_expires_at, - status, - added_at, - updated_at -from accounts -order by added_at desc, id desc"; - -const LIST_POOL_ACCOUNTS_SQL: &str = r" -select - accounts.id, - email, - accounts.account_id, - user_id, - label, - plan_type, - access_token_cipher, - refresh_token_cipher, - access_token_expires_at, - status, - added_at, - updated_at, - quota_limit_reached, - quota_verify_required, - quota_cooldown_until, - cloudflare_cooldown_until, - coalesce(account_usage.request_count, 0) as usage_request_count, - coalesce(account_usage.empty_response_count, 0) as usage_empty_response_count, - coalesce(account_usage.image_input_tokens, 0) as usage_image_input_tokens, - coalesce(account_usage.image_output_tokens, 0) as usage_image_output_tokens, - coalesce(account_usage.image_request_count, 0) as usage_image_request_count, - coalesce(account_usage.image_request_failed_count, 0) as usage_image_request_failed_count, - coalesce(account_usage.window_request_count, 0) as usage_window_request_count, - coalesce(account_usage.window_input_tokens, 0) as usage_window_input_tokens, - coalesce(account_usage.window_output_tokens, 0) as usage_window_output_tokens, - coalesce(account_usage.window_cached_tokens, 0) as usage_window_cached_tokens, - coalesce(account_usage.window_image_input_tokens, 0) as usage_window_image_input_tokens, - coalesce(account_usage.window_image_output_tokens, 0) as usage_window_image_output_tokens, - coalesce(account_usage.window_image_request_count, 0) as usage_window_image_request_count, - coalesce(account_usage.window_image_request_failed_count, 0) as usage_window_image_request_failed_count, - account_usage.window_started_at as usage_window_started_at, - account_usage.window_reset_at as usage_window_reset_at, - account_usage.limit_window_seconds as usage_limit_window_seconds, - account_usage.last_used_at as usage_last_used_at -from accounts -left join account_usage on account_usage.account_id = accounts.id -order by added_at desc, accounts.id desc"; - -impl AccountRepository { - pub fn new(pool: SqlitePool, secret_box: SecretBox) -> Self { - Self { pool, secret_box } - } - - pub async fn insert(&self, account: NewAccount) -> AccountRepositoryResult<()> { - let now = Utc::now().to_rfc3339(); - let access_token_cipher = self.secret_box.encrypt(&account.access_token)?; - let refresh_token_cipher = account - .refresh_token - .as_ref() - .map(|token| self.secret_box.encrypt(token)) - .transpose()?; - sqlx::query(INSERT_ACCOUNT_SQL) - .bind(account.id) - .bind(account.email) - .bind(account.account_id) - .bind(account.user_id) - .bind(account.label) - .bind(account.plan_type) - .bind(access_token_cipher) - .bind(refresh_token_cipher) - .bind( - account - .access_token_expires_at - .map(|value| value.to_rfc3339()), - ) - .bind(status_to_db(account.status)) - .bind(&now) - .bind(now) - .execute(&self.pool) - .await?; - Ok(()) - } - - pub async fn get(&self, id: &str) -> AccountRepositoryResult> { - let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_ID_SQL) - .bind(id) - .fetch_optional(&self.pool) - .await?; - row.map(|row| self.account_from_row(&row)).transpose() - } - - pub async fn exists(&self, id: &str) -> AccountRepositoryResult { - let row = sqlx::query("select 1 from accounts where id = ?") - .bind(id) - .fetch_optional(&self.pool) - .await?; - Ok(row.is_some()) - } - - pub async fn find_by_chatgpt_identity( - &self, - account_id: &str, - user_id: Option<&str>, - ) -> AccountRepositoryResult> { - let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_CHATGPT_IDENTITY_SQL) - .bind(account_id) - .bind(user_id) - .bind(user_id) - .fetch_optional(&self.pool) - .await?; - row.map(|row| self.account_from_row(&row)).transpose() - } - - pub async fn list( - &self, - cursor: Option, - limit: u32, - ) -> AccountRepositoryResult> { - let fetch_limit = i64::from(limit) + 1; - let rows = if let Some(cursor) = cursor { - let (created_at, id) = - decode_cursor(&cursor).ok_or(AccountRepositoryError::InvalidCursor)?; - sqlx::query(LIST_STORED_ACCOUNTS_AFTER_CURSOR_SQL) - .bind(&created_at) - .bind(created_at) - .bind(id) - .bind(fetch_limit) - .fetch_all(&self.pool) - .await? - } else { - sqlx::query(LIST_STORED_ACCOUNTS_SQL) - .bind(fetch_limit) - .fetch_all(&self.pool) - .await? - }; - - let has_next = rows.len() > limit as usize; - let take_count = rows.len().min(limit as usize); - let mut items = Vec::with_capacity(take_count); - for row in rows.into_iter().take(take_count) { - items.push(self.account_from_row(&row)?); - } - let next_cursor = if has_next { - items - .last() - .map(|account| encode_cursor(&account.added_at.to_rfc3339(), &account.id)) - } else { - None - }; - Ok(Page { items, next_cursor }) - } - - pub async fn list_all(&self) -> AccountRepositoryResult> { - let rows = sqlx::query(LIST_ALL_STORED_ACCOUNTS_SQL) - .fetch_all(&self.pool) - .await?; - let mut accounts = Vec::with_capacity(rows.len()); - for row in rows { - accounts.push(self.account_from_row(&row)?); - } - Ok(accounts) - } - - pub async fn list_metadata( - &self, - cursor: Option, - limit: u32, - ) -> AccountRepositoryResult> { - let fetch_limit = i64::from(limit) + 1; - let rows = if let Some(cursor) = cursor { - let (created_at, id) = - decode_cursor(&cursor).ok_or(AccountRepositoryError::InvalidCursor)?; - sqlx::query(LIST_METADATA_AFTER_CURSOR_SQL) - .bind(&created_at) - .bind(created_at) - .bind(id) - .bind(fetch_limit) - .fetch_all(&self.pool) - .await? - } else { - sqlx::query(LIST_METADATA_SQL) - .bind(fetch_limit) - .fetch_all(&self.pool) - .await? - }; - - let has_next = rows.len() > limit as usize; - let take_count = rows.len().min(limit as usize); - let mut items = Vec::with_capacity(take_count); - for row in rows.into_iter().take(take_count) { - items.push(metadata_from_row(&row)?); - } - let next_cursor = if has_next { - items - .last() - .map(|account| encode_cursor(&account.added_at.to_rfc3339(), &account.id)) - } else { - None - }; - Ok(Page { items, next_cursor }) - } - - pub async fn list_all_metadata(&self) -> AccountRepositoryResult> { - let rows = sqlx::query(LIST_ALL_METADATA_SQL) - .fetch_all(&self.pool) - .await?; - let mut accounts = Vec::with_capacity(rows.len()); - for row in rows { - accounts.push(metadata_from_row(&row)?); - } - Ok(accounts) - } - - pub async fn list_pool_accounts(&self) -> AccountRepositoryResult> { - let rows = sqlx::query(LIST_POOL_ACCOUNTS_SQL) - .fetch_all(&self.pool) - .await?; - let mut accounts = Vec::with_capacity(rows.len()); - for row in rows { - let stored = self.account_from_row(&row)?; - let access_token = stored.access_token.expose_secret().to_string(); - let refresh_token = stored - .refresh_token - .as_ref() - .map(|token| token.expose_secret().to_string()); - accounts.push(Account { - id: stored.id, - email: stored.email, - account_id: stored.account_id, - user_id: stored.user_id, - label: stored.label, - plan_type: stored.plan_type, - access_token, - refresh_token, - access_token_expires_at: stored.access_token_expires_at, - status: stored.status, - quota_limit_reached: row.get::("quota_limit_reached") != 0, - quota_verify_required: row.get::("quota_verify_required") != 0, - quota_cooldown_until: parse_optional_rfc3339( - row.get::, _>("quota_cooldown_until"), - )?, - cloudflare_cooldown_until: parse_optional_rfc3339( - row.get::, _>("cloudflare_cooldown_until"), - )?, - request_count: row.get::("usage_request_count").max(0) as u64, - empty_response_count: row.get::("usage_empty_response_count").max(0) as u64, - image_input_tokens: nonnegative_i64_to_u64( - row.get::("usage_image_input_tokens"), - ), - image_output_tokens: nonnegative_i64_to_u64( - row.get::("usage_image_output_tokens"), - ), - image_request_count: nonnegative_i64_to_u64( - row.get::("usage_image_request_count"), - ), - image_request_failed_count: nonnegative_i64_to_u64( - row.get::("usage_image_request_failed_count"), - ), - window_request_count: nonnegative_i64_to_u64( - row.get::("usage_window_request_count"), - ), - window_input_tokens: nonnegative_i64_to_u64( - row.get::("usage_window_input_tokens"), - ), - window_output_tokens: nonnegative_i64_to_u64( - row.get::("usage_window_output_tokens"), - ), - window_cached_tokens: nonnegative_i64_to_u64( - row.get::("usage_window_cached_tokens"), - ), - window_image_input_tokens: nonnegative_i64_to_u64( - row.get::("usage_window_image_input_tokens"), - ), - window_image_output_tokens: nonnegative_i64_to_u64( - row.get::("usage_window_image_output_tokens"), - ), - window_image_request_count: nonnegative_i64_to_u64( - row.get::("usage_window_image_request_count"), - ), - window_image_request_failed_count: nonnegative_i64_to_u64( - row.get::("usage_window_image_request_failed_count"), - ), - window_started_at: parse_optional_rfc3339( - row.get::, _>("usage_window_started_at"), - )?, - window_reset_at: parse_optional_rfc3339( - row.get::, _>("usage_window_reset_at"), - )?, - limit_window_seconds: optional_positive_i64_to_u64( - row.get::, _>("usage_limit_window_seconds"), - ), - added_at: stored.added_at.to_rfc3339(), - last_used_at: row.get("usage_last_used_at"), - }); - } - Ok(accounts) - } - - pub async fn set_status( - &self, - id: &str, - status: AccountStatus, - ) -> AccountRepositoryResult { - let result = sqlx::query("update accounts set status = ?, updated_at = ? where id = ?") - .bind(status_to_db(status)) - .bind(Utc::now().to_rfc3339()) - .bind(id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - pub async fn set_label( - &self, - id: &str, - label: Option, - ) -> AccountRepositoryResult { - let result = sqlx::query("update accounts set label = ?, updated_at = ? where id = ?") - .bind(label) - .bind(Utc::now().to_rfc3339()) - .bind(id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - pub async fn delete(&self, id: &str) -> AccountRepositoryResult { - let result = sqlx::query("delete from accounts where id = ?") - .bind(id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - pub async fn delete_all(&self) -> AccountRepositoryResult { - let result = sqlx::query("delete from accounts") - .execute(&self.pool) - .await?; - Ok(result.rows_affected()) - } - - fn account_from_row( - &self, - row: &sqlx::sqlite::SqliteRow, - ) -> AccountRepositoryResult { - let access_token_cipher = row.get::("access_token_cipher"); - let access_token = self.secret_box.decrypt(&access_token_cipher)?; - let refresh_token = match row.get::, _>("refresh_token_cipher") { - Some(cipher) => Some(self.secret_box.decrypt(&cipher)?), - None => None, - }; - Ok(StoredAccount { - id: row.get("id"), - email: row.get("email"), - account_id: row.get("account_id"), - user_id: row.get("user_id"), - label: row.get("label"), - plan_type: row.get("plan_type"), - access_token, - refresh_token, - access_token_expires_at: parse_optional_rfc3339( - row.get::, _>("access_token_expires_at"), - )?, - status: status_from_db(&row.get::("status"))?, - added_at: parse_rfc3339(&row.get::("added_at"))?, - updated_at: parse_rfc3339(&row.get::("updated_at"))?, - }) - } -} - -pub(crate) fn status_to_db(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", - } -} - -fn status_from_db(value: &str) -> AccountRepositoryResult { - match value { - "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(AccountRepositoryError::InvalidStatus(other.to_string())), - } -} - -pub(crate) fn parse_optional_rfc3339( - value: Option, -) -> AccountRepositoryResult>> { - value.as_deref().map(parse_rfc3339).transpose() -} - -pub(crate) fn parse_rfc3339(value: &str) -> AccountRepositoryResult> { - Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc)) -} - -fn nonnegative_i64_to_u64(value: i64) -> u64 { - value.max(0) as u64 -} - -pub(crate) fn optional_positive_i64_to_u64(value: Option) -> Option { - value - .and_then(|value| u64::try_from(value).ok()) - .filter(|value| *value > 0) -} - -fn metadata_from_row( - row: &sqlx::sqlite::SqliteRow, -) -> AccountRepositoryResult { - Ok(StoredAccountMetadata { - id: row.get("id"), - email: row.get("email"), - account_id: row.get("account_id"), - user_id: row.get("user_id"), - label: row.get("label"), - plan_type: row.get("plan_type"), - access_token_expires_at: parse_optional_rfc3339( - row.get::, _>("access_token_expires_at"), - )?, - status: status_from_db(&row.get::("status"))?, - added_at: parse_rfc3339(&row.get::("added_at"))?, - updated_at: parse_rfc3339(&row.get::("updated_at"))?, - }) -} diff --git a/src/codex/accounts/repository/leases.rs b/src/codex/accounts/repository/leases.rs deleted file mode 100644 index e7b5ff9c9..000000000 --- a/src/codex/accounts/repository/leases.rs +++ /dev/null @@ -1,51 +0,0 @@ -use chrono::{DateTime, Utc}; - -use super::{AccountRepository, AccountRepositoryResult}; - -const TRY_ACQUIRE_REFRESH_LEASE_SQL: &str = r" -insert into account_refresh_leases ( - account_id, - owner, - expires_at, - updated_at -) values (?, ?, ?, ?) -on conflict(account_id) do update set - owner = excluded.owner, - expires_at = excluded.expires_at, - updated_at = excluded.updated_at -where account_refresh_leases.expires_at <= ? - or account_refresh_leases.owner = excluded.owner"; - -impl AccountRepository { - pub async fn try_acquire_refresh_lease( - &self, - account_id: &str, - owner: &str, - lease_until: DateTime, - ) -> AccountRepositoryResult { - let now = Utc::now().to_rfc3339(); - let result = sqlx::query(TRY_ACQUIRE_REFRESH_LEASE_SQL) - .bind(account_id) - .bind(owner) - .bind(lease_until.to_rfc3339()) - .bind(&now) - .bind(now) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - pub async fn release_refresh_lease( - &self, - account_id: &str, - owner: &str, - ) -> AccountRepositoryResult { - let result = - sqlx::query("delete from account_refresh_leases where account_id = ? and owner = ?") - .bind(account_id) - .bind(owner) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } -} diff --git a/src/codex/accounts/repository/mod.rs b/src/codex/accounts/repository/mod.rs deleted file mode 100644 index 4fb4a3374..000000000 --- a/src/codex/accounts/repository/mod.rs +++ /dev/null @@ -1,9 +0,0 @@ -mod accounts; -mod leases; -mod quotas; -mod tokens; -mod usage; - -pub use accounts::*; -pub(crate) use accounts::{optional_positive_i64_to_u64, parse_optional_rfc3339, status_to_db}; -pub use usage::AccountUsageRepository; diff --git a/src/codex/accounts/repository/quotas.rs b/src/codex/accounts/repository/quotas.rs deleted file mode 100644 index 136b48ef2..000000000 --- a/src/codex/accounts/repository/quotas.rs +++ /dev/null @@ -1,146 +0,0 @@ -use chrono::{DateTime, Utc}; -use sqlx::Row; - -use super::{ - parse_optional_rfc3339, AccountQuotaSnapshot, AccountRepository, AccountRepositoryResult, -}; - -const LIST_QUOTA_SNAPSHOTS_SQL: &str = r" -select - id, - email, - quota_json, - quota_fetched_at -from accounts -where quota_json is not null - and trim(quota_json) <> '' -order by coalesce(quota_fetched_at, '') desc, id desc"; - -const UPDATE_QUOTA_JSON_SQL: &str = r" -update accounts -set - quota_json = ?, - quota_fetched_at = ?, - updated_at = ? -where id = ?"; - -const SET_QUOTA_COOLDOWN_UNTIL_SQL: &str = r" -update accounts -set - quota_limit_reached = 1, - quota_cooldown_until = case - when quota_cooldown_until is not null and quota_cooldown_until > ? - then quota_cooldown_until - else ? - end, - updated_at = ? -where id = ?"; - -const SET_QUOTA_VERIFY_REQUIRED_SQL: &str = r" -update accounts -set - quota_verify_required = ?, - updated_at = ? -where id = ?"; - -const SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL: &str = r" -update accounts -set - cloudflare_cooldown_until = ?, - updated_at = ? -where id = ?"; - -impl AccountRepository { - pub async fn list_quota_snapshots(&self) -> AccountRepositoryResult> { - let rows = sqlx::query(LIST_QUOTA_SNAPSHOTS_SQL) - .fetch_all(&self.pool) - .await?; - let mut snapshots = Vec::with_capacity(rows.len()); - for row in rows { - snapshots.push(quota_snapshot_from_row(&row)?); - } - Ok(snapshots) - } - - pub async fn update_quota_json( - &self, - account_id: &str, - quota_json: &str, - ) -> AccountRepositoryResult { - let now = Utc::now().to_rfc3339(); - let result = sqlx::query(UPDATE_QUOTA_JSON_SQL) - .bind(quota_json) - .bind(&now) - .bind(now) - .bind(account_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - pub async fn get_quota_json( - &self, - account_id: &str, - ) -> AccountRepositoryResult> { - let row = sqlx::query("select quota_json from accounts where id = ?") - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - Ok(row.and_then(|row| row.get("quota_json"))) - } - - pub async fn set_quota_cooldown_until( - &self, - id: &str, - cooldown_until: DateTime, - ) -> AccountRepositoryResult { - let cooldown_until = cooldown_until.to_rfc3339(); - let result = sqlx::query(SET_QUOTA_COOLDOWN_UNTIL_SQL) - .bind(&cooldown_until) - .bind(cooldown_until) - .bind(Utc::now().to_rfc3339()) - .bind(id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - pub async fn set_quota_verify_required( - &self, - id: &str, - required: bool, - ) -> AccountRepositoryResult { - let result = sqlx::query(SET_QUOTA_VERIFY_REQUIRED_SQL) - .bind(if required { 1_i64 } else { 0 }) - .bind(Utc::now().to_rfc3339()) - .bind(id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - pub async fn set_cloudflare_cooldown_until( - &self, - id: &str, - cooldown_until: DateTime, - ) -> AccountRepositoryResult { - let result = sqlx::query(SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL) - .bind(cooldown_until.to_rfc3339()) - .bind(Utc::now().to_rfc3339()) - .bind(id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } -} - -fn quota_snapshot_from_row( - row: &sqlx::sqlite::SqliteRow, -) -> AccountRepositoryResult { - Ok(AccountQuotaSnapshot { - account_id: row.get("id"), - email: row.get("email"), - quota_json: row.get("quota_json"), - quota_fetched_at: parse_optional_rfc3339(row.get::, _>("quota_fetched_at"))?, - }) -} diff --git a/src/codex/accounts/repository/tokens.rs b/src/codex/accounts/repository/tokens.rs deleted file mode 100644 index a284ca625..000000000 --- a/src/codex/accounts/repository/tokens.rs +++ /dev/null @@ -1,203 +0,0 @@ -use chrono::{DateTime, Utc}; -use secrecy::SecretString; - -use super::{ - status_to_db, AccountClaimsUpdate, AccountRepository, AccountRepositoryResult, TokenUpdate, -}; - -const UPDATE_TOKENS_WITH_REFRESH_SQL: &str = r" -update accounts -set - access_token_cipher = ?, - refresh_token_cipher = ?, - access_token_expires_at = ?, - status = case - when status in ('disabled', 'banned') then status - else 'active' - end, - updated_at = ? -where id = ?"; - -const UPDATE_TOKENS_PRESERVING_REFRESH_SQL: &str = r" -update accounts -set - access_token_cipher = ?, - access_token_expires_at = ?, - status = case - when status in ('disabled', 'banned') then status - else 'active' - end, - updated_at = ? -where id = ?"; - -const UPDATE_ACCOUNT_CLAIMS_WITH_REFRESH_SQL: &str = r" -update accounts -set - email = ?, - account_id = ?, - user_id = ?, - plan_type = ?, - access_token_cipher = ?, - refresh_token_cipher = ?, - access_token_expires_at = ?, - status = ?, - updated_at = ? -where id = ?"; - -const UPDATE_ACCOUNT_CLAIMS_PRESERVING_REFRESH_SQL: &str = r" -update accounts -set - email = ?, - account_id = ?, - user_id = ?, - plan_type = ?, - access_token_cipher = ?, - access_token_expires_at = ?, - status = ?, - updated_at = ? -where id = ?"; - -#[derive(Debug)] -struct TokenWrite { - access_token_cipher: String, - refresh_token_cipher: Option, - access_token_expires_at: Option, - updated_at: String, -} - -impl AccountRepository { - fn prepare_token_write( - &self, - access_token: &SecretString, - refresh_token: Option<&SecretString>, - access_token_expires_at: Option>, - ) -> AccountRepositoryResult { - Ok(TokenWrite { - access_token_cipher: self.secret_box.encrypt(access_token)?, - refresh_token_cipher: refresh_token - .map(|token| self.secret_box.encrypt(token)) - .transpose()?, - access_token_expires_at: access_token_expires_at.map(|value| value.to_rfc3339()), - updated_at: Utc::now().to_rfc3339(), - }) - } - - pub async fn update_tokens( - &self, - id: &str, - update: TokenUpdate, - ) -> AccountRepositoryResult { - let TokenWrite { - access_token_cipher, - refresh_token_cipher, - access_token_expires_at, - updated_at, - } = self.prepare_token_write( - &update.access_token, - update.refresh_token.as_ref(), - update.access_token_expires_at, - )?; - - let result = if let Some(refresh_token_cipher) = refresh_token_cipher { - sqlx::query(UPDATE_TOKENS_WITH_REFRESH_SQL) - .bind(access_token_cipher) - .bind(refresh_token_cipher) - .bind(access_token_expires_at) - .bind(updated_at) - .bind(id) - .execute(&self.pool) - .await? - } else { - // 刷新接口可能不返回新的 refresh_token；此时必须保留旧 RT，避免账号失去后续刷新能力。 - sqlx::query(UPDATE_TOKENS_PRESERVING_REFRESH_SQL) - .bind(access_token_cipher) - .bind(access_token_expires_at) - .bind(updated_at) - .bind(id) - .execute(&self.pool) - .await? - }; - Ok(result.rows_affected() > 0) - } - - pub async fn update_from_claims( - &self, - id: &str, - update: AccountClaimsUpdate, - ) -> AccountRepositoryResult { - let TokenWrite { - access_token_cipher, - refresh_token_cipher, - access_token_expires_at, - updated_at, - } = self.prepare_token_write( - &update.access_token, - update.refresh_token.as_ref(), - update.access_token_expires_at, - )?; - let status = status_to_db(update.status); - - let result = if let Some(refresh_token_cipher) = refresh_token_cipher { - sqlx::query(UPDATE_ACCOUNT_CLAIMS_WITH_REFRESH_SQL) - .bind(update.email) - .bind(update.account_id) - .bind(update.user_id) - .bind(update.plan_type) - .bind(access_token_cipher) - .bind(refresh_token_cipher) - .bind(access_token_expires_at) - .bind(status) - .bind(updated_at) - .bind(id) - .execute(&self.pool) - .await? - } else { - // OpenAI 刷新/导入未给新 RT 时保留原值，避免把可继续刷新的账号写坏。 - sqlx::query(UPDATE_ACCOUNT_CLAIMS_PRESERVING_REFRESH_SQL) - .bind(update.email) - .bind(update.account_id) - .bind(update.user_id) - .bind(update.plan_type) - .bind(access_token_cipher) - .bind(access_token_expires_at) - .bind(status) - .bind(updated_at) - .bind(id) - .execute(&self.pool) - .await? - }; - Ok(result.rows_affected() > 0) - } -} - -#[cfg(test)] -mod tests { - use sqlx::SqlitePool; - - use super::*; - use crate::platform::crypto::SecretBox; - - #[tokio::test] - async fn prepare_token_write_should_encrypt_tokens_and_format_timestamps() { - let repository = AccountRepository::new( - SqlitePool::connect_lazy("sqlite::memory:").unwrap(), - SecretBox::new([55u8; 32]), - ); - let access_token = SecretString::new("access-secret".to_string().into()); - let refresh_token = SecretString::new("refresh-secret".to_string().into()); - let expires_at = super::super::accounts::parse_rfc3339("2026-06-14T00:00:00Z").unwrap(); - - let write = repository - .prepare_token_write(&access_token, Some(&refresh_token), Some(expires_at)) - .unwrap(); - - assert!(write.access_token_cipher.starts_with("v1:")); - assert!(!write.access_token_cipher.contains("access-secret")); - assert!(write.refresh_token_cipher.unwrap().starts_with("v1:")); - assert_eq!( - write.access_token_expires_at.as_deref(), - Some("2026-06-14T00:00:00+00:00") - ); - assert!(super::super::accounts::parse_rfc3339(&write.updated_at).is_ok()); - } -} diff --git a/src/codex/accounts/repository/usage.rs b/src/codex/accounts/repository/usage.rs deleted file mode 100644 index 27bc4acee..000000000 --- a/src/codex/accounts/repository/usage.rs +++ /dev/null @@ -1,491 +0,0 @@ -use chrono::{DateTime, Utc}; -use sqlx::{Row, SqlitePool}; - -use crate::utils::pagination::{decode_cursor, encode_cursor, Page}; - -use super::{ - optional_positive_i64_to_u64, parse_optional_rfc3339, AccountRepository, - AccountRepositoryError, AccountRepositoryResult, AccountUsageListRecord, AccountUsageRecord, - AccountUsageSummary, UsageDelta, -}; - -const RECORD_USAGE_SQL: &str = r" -insert into account_usage ( - account_id, - request_count, - empty_response_count, - input_tokens, - output_tokens, - cached_tokens, - image_input_tokens, - image_output_tokens, - image_request_count, - image_request_failed_count, - window_request_count, - window_input_tokens, - window_output_tokens, - window_cached_tokens, - window_image_input_tokens, - window_image_output_tokens, - window_image_request_count, - window_image_request_failed_count, - last_used_at -) values (?, 1, ?, ?, ?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?, ?, ?, ?, ?) -on conflict(account_id) do update set - request_count = request_count + 1, - empty_response_count = empty_response_count + excluded.empty_response_count, - input_tokens = input_tokens + excluded.input_tokens, - output_tokens = output_tokens + excluded.output_tokens, - cached_tokens = cached_tokens + excluded.cached_tokens, - image_input_tokens = image_input_tokens + excluded.image_input_tokens, - image_output_tokens = image_output_tokens + excluded.image_output_tokens, - image_request_count = image_request_count + excluded.image_request_count, - image_request_failed_count = image_request_failed_count + excluded.image_request_failed_count, - window_request_count = window_request_count + 1, - window_input_tokens = window_input_tokens + excluded.window_input_tokens, - window_output_tokens = window_output_tokens + excluded.window_output_tokens, - window_cached_tokens = window_cached_tokens + excluded.window_cached_tokens, - window_image_input_tokens = window_image_input_tokens + excluded.window_image_input_tokens, - window_image_output_tokens = window_image_output_tokens + excluded.window_image_output_tokens, - window_image_request_count = window_image_request_count + excluded.window_image_request_count, - window_image_request_failed_count = window_image_request_failed_count + excluded.window_image_request_failed_count, - window_started_at = case - when account_usage.window_started_at is null - and (account_usage.window_reset_at is not null or account_usage.limit_window_seconds is not null) - then excluded.last_used_at - else account_usage.window_started_at - end, - last_used_at = excluded.last_used_at"; - -const GET_USAGE_SQL: &str = r" -select - account_id, - request_count, - empty_response_count, - input_tokens, - output_tokens, - cached_tokens, - image_input_tokens, - image_output_tokens, - image_request_count, - image_request_failed_count, - window_request_count, - window_input_tokens, - window_output_tokens, - window_cached_tokens, - window_image_input_tokens, - window_image_output_tokens, - window_image_request_count, - window_image_request_failed_count, - window_started_at, - window_reset_at, - limit_window_seconds, - last_used_at -from account_usage -where account_id = ?"; - -const SELECT_RATE_LIMIT_WINDOW_SQL: &str = r" -select - window_reset_at, - limit_window_seconds -from account_usage -where account_id = ?"; - -const SYNC_RATE_LIMIT_WINDOW_RESET_SQL: &str = r" -insert into account_usage ( - account_id, - window_request_count, - window_input_tokens, - window_output_tokens, - window_cached_tokens, - window_image_input_tokens, - window_image_output_tokens, - window_image_request_count, - window_image_request_failed_count, - window_started_at, - window_reset_at, - limit_window_seconds -) values (?, 0, 0, 0, 0, 0, 0, 0, 0, ?, ?, ?) -on conflict(account_id) do update set - window_request_count = 0, - window_input_tokens = 0, - window_output_tokens = 0, - window_cached_tokens = 0, - window_image_input_tokens = 0, - window_image_output_tokens = 0, - window_image_request_count = 0, - window_image_request_failed_count = 0, - window_started_at = excluded.window_started_at, - window_reset_at = excluded.window_reset_at, - limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)"; - -const SYNC_RATE_LIMIT_WINDOW_SQL: &str = r" -insert into account_usage ( - account_id, - window_reset_at, - limit_window_seconds -) values (?, ?, ?) -on conflict(account_id) do update set - window_reset_at = excluded.window_reset_at, - limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)"; - -const RESET_ACCOUNT_USAGE_SQL: &str = r" -insert into account_usage ( - account_id, - request_count, - empty_response_count, - input_tokens, - output_tokens, - cached_tokens, - image_input_tokens, - image_output_tokens, - image_request_count, - image_request_failed_count, - window_request_count, - window_input_tokens, - window_output_tokens, - window_cached_tokens, - window_image_input_tokens, - window_image_output_tokens, - window_image_request_count, - window_image_request_failed_count, - window_started_at, - window_reset_at, - last_used_at -) values (?, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, null, null, null) -on conflict(account_id) do update set - request_count = 0, - empty_response_count = 0, - input_tokens = 0, - output_tokens = 0, - cached_tokens = 0, - image_input_tokens = 0, - image_output_tokens = 0, - image_request_count = 0, - image_request_failed_count = 0, - window_request_count = 0, - window_input_tokens = 0, - window_output_tokens = 0, - window_cached_tokens = 0, - window_image_input_tokens = 0, - window_image_output_tokens = 0, - window_image_request_count = 0, - window_image_request_failed_count = 0, - window_started_at = case - when account_usage.window_reset_at is not null then ? - else null - end, - window_reset_at = account_usage.window_reset_at, - last_used_at = null"; - -const LIST_ACCOUNT_USAGE_AFTER_CURSOR_SQL: &str = r" -select - account_usage.account_id, - accounts.email, - accounts.label, - accounts.plan_type, - account_usage.request_count, - account_usage.empty_response_count, - account_usage.input_tokens, - account_usage.output_tokens, - account_usage.cached_tokens, - account_usage.image_input_tokens, - account_usage.image_output_tokens, - account_usage.image_request_count, - account_usage.image_request_failed_count, - account_usage.last_used_at -from account_usage -left join accounts on accounts.id = account_usage.account_id -where coalesce(account_usage.last_used_at, '') < ? - or ( - coalesce(account_usage.last_used_at, '') = ? - and account_usage.account_id < ? - ) -order by coalesce(account_usage.last_used_at, '') desc, account_usage.account_id desc -limit ?"; - -const LIST_ACCOUNT_USAGE_SQL: &str = r" -select - account_usage.account_id, - accounts.email, - accounts.label, - accounts.plan_type, - account_usage.request_count, - account_usage.empty_response_count, - account_usage.input_tokens, - account_usage.output_tokens, - account_usage.cached_tokens, - account_usage.image_input_tokens, - account_usage.image_output_tokens, - account_usage.image_request_count, - account_usage.image_request_failed_count, - account_usage.last_used_at -from account_usage -left join accounts on accounts.id = account_usage.account_id -order by coalesce(account_usage.last_used_at, '') desc, account_usage.account_id desc -limit ?"; - -const ACCOUNT_USAGE_SUMMARY_SQL: &str = r" -select - count(*) as account_count, - coalesce(sum(request_count), 0) as request_count, - coalesce(sum(empty_response_count), 0) as empty_response_count, - coalesce(sum(input_tokens), 0) as input_tokens, - coalesce(sum(output_tokens), 0) as output_tokens, - coalesce(sum(cached_tokens), 0) as cached_tokens, - coalesce(sum(image_input_tokens), 0) as image_input_tokens, - coalesce(sum(image_output_tokens), 0) as image_output_tokens, - coalesce(sum(image_request_count), 0) as image_request_count, - coalesce(sum(image_request_failed_count), 0) as image_request_failed_count -from account_usage"; - -impl AccountRepository { - pub async fn record_usage( - &self, - account_id: &str, - usage: UsageDelta, - ) -> AccountRepositoryResult<()> { - let now = Utc::now().to_rfc3339(); - sqlx::query(RECORD_USAGE_SQL) - .bind(account_id) - .bind(usage.empty_response_count) - .bind(usage.input_tokens) - .bind(usage.output_tokens) - .bind(usage.cached_tokens) - .bind(usage.image_input_tokens) - .bind(usage.image_output_tokens) - .bind(usage.image_request_count) - .bind(usage.image_request_failed_count) - .bind(usage.input_tokens) - .bind(usage.output_tokens) - .bind(usage.cached_tokens) - .bind(usage.image_input_tokens) - .bind(usage.image_output_tokens) - .bind(usage.image_request_count) - .bind(usage.image_request_failed_count) - .bind(now) - .execute(&self.pool) - .await?; - Ok(()) - } - - pub async fn get_usage( - &self, - account_id: &str, - ) -> AccountRepositoryResult> { - let row = sqlx::query(GET_USAGE_SQL) - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - row.map(|row| usage_from_row(&row)).transpose() - } - - pub async fn sync_rate_limit_window( - &self, - account_id: &str, - reset_at: DateTime, - limit_window_seconds: Option, - ) -> AccountRepositoryResult<()> { - let row = sqlx::query(SELECT_RATE_LIMIT_WINDOW_SQL) - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - let existing_reset_at = row - .as_ref() - .map(|row| parse_optional_rfc3339(row.get::, _>("window_reset_at"))) - .transpose()? - .flatten(); - let existing_limit_window_seconds = row - .as_ref() - .and_then(|row| optional_positive_i64_to_u64(row.get("limit_window_seconds"))); - let limit_window_seconds_db = limit_window_seconds.map(u64_to_i64_saturating); - let reset_at_db = reset_at.to_rfc3339(); - - if should_reset_usage_window( - existing_reset_at, - existing_limit_window_seconds, - reset_at, - limit_window_seconds, - ) { - sqlx::query(SYNC_RATE_LIMIT_WINDOW_RESET_SQL) - .bind(account_id) - .bind(Utc::now().to_rfc3339()) - .bind(reset_at_db) - .bind(limit_window_seconds_db) - .execute(&self.pool) - .await?; - } else { - sqlx::query(SYNC_RATE_LIMIT_WINDOW_SQL) - .bind(account_id) - .bind(reset_at_db) - .bind(limit_window_seconds_db) - .execute(&self.pool) - .await?; - } - Ok(()) - } -} - -#[derive(Clone)] -pub struct AccountUsageRepository { - pool: SqlitePool, -} - -impl AccountUsageRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - pub async fn list( - &self, - cursor: Option, - limit: u32, - ) -> AccountRepositoryResult> { - let fetch_limit = i64::from(limit) + 1; - let rows = if let Some(cursor) = cursor { - let (last_used_at, account_id) = - decode_cursor(&cursor).ok_or(AccountRepositoryError::InvalidCursor)?; - sqlx::query(LIST_ACCOUNT_USAGE_AFTER_CURSOR_SQL) - .bind(&last_used_at) - .bind(last_used_at) - .bind(account_id) - .bind(fetch_limit) - .fetch_all(&self.pool) - .await? - } else { - sqlx::query(LIST_ACCOUNT_USAGE_SQL) - .bind(fetch_limit) - .fetch_all(&self.pool) - .await? - }; - - let has_next = rows.len() > limit as usize; - let take_count = rows.len().min(limit as usize); - let mut items = Vec::with_capacity(take_count); - for row in rows.into_iter().take(take_count) { - items.push(usage_list_from_row(&row)?); - } - let next_cursor = if has_next { - items.last().map(|usage| { - encode_cursor( - &usage - .last_used_at - .map(|value| value.to_rfc3339()) - .unwrap_or_default(), - &usage.account_id, - ) - }) - } else { - None - }; - Ok(Page { items, next_cursor }) - } - - pub async fn summary(&self) -> AccountRepositoryResult { - let row = sqlx::query(ACCOUNT_USAGE_SUMMARY_SQL) - .fetch_one(&self.pool) - .await?; - Ok(AccountUsageSummary { - account_count: row.get("account_count"), - request_count: row.get("request_count"), - empty_response_count: row.get("empty_response_count"), - input_tokens: row.get("input_tokens"), - output_tokens: row.get("output_tokens"), - cached_tokens: row.get("cached_tokens"), - image_input_tokens: row.get("image_input_tokens"), - image_output_tokens: row.get("image_output_tokens"), - image_request_count: row.get("image_request_count"), - image_request_failed_count: row.get("image_request_failed_count"), - }) - } - - pub async fn reset_account(&self, account_id: &str) -> AccountRepositoryResult<()> { - sqlx::query(RESET_ACCOUNT_USAGE_SQL) - .bind(account_id) - .bind(Utc::now().to_rfc3339()) - .execute(&self.pool) - .await?; - Ok(()) - } -} - -fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> AccountRepositoryResult { - Ok(AccountUsageRecord { - account_id: row.get("account_id"), - request_count: row.get("request_count"), - empty_response_count: row.get("empty_response_count"), - input_tokens: row.get("input_tokens"), - output_tokens: row.get("output_tokens"), - cached_tokens: row.get("cached_tokens"), - image_input_tokens: row.get("image_input_tokens"), - image_output_tokens: row.get("image_output_tokens"), - image_request_count: row.get("image_request_count"), - image_request_failed_count: row.get("image_request_failed_count"), - window_request_count: row.get("window_request_count"), - window_input_tokens: row.get("window_input_tokens"), - window_output_tokens: row.get("window_output_tokens"), - window_cached_tokens: row.get("window_cached_tokens"), - window_image_input_tokens: row.get("window_image_input_tokens"), - window_image_output_tokens: row.get("window_image_output_tokens"), - window_image_request_count: row.get("window_image_request_count"), - window_image_request_failed_count: row.get("window_image_request_failed_count"), - window_started_at: parse_optional_rfc3339( - row.get::, _>("window_started_at"), - )?, - window_reset_at: parse_optional_rfc3339(row.get::, _>("window_reset_at"))?, - limit_window_seconds: optional_positive_i64_to_u64(row.get("limit_window_seconds")), - last_used_at: parse_optional_rfc3339(row.get::, _>("last_used_at"))?, - }) -} - -fn usage_list_from_row( - row: &sqlx::sqlite::SqliteRow, -) -> AccountRepositoryResult { - Ok(AccountUsageListRecord { - account_id: row.get("account_id"), - email: row.get("email"), - label: row.get("label"), - plan_type: row.get("plan_type"), - request_count: row.get("request_count"), - empty_response_count: row.get("empty_response_count"), - input_tokens: row.get("input_tokens"), - output_tokens: row.get("output_tokens"), - cached_tokens: row.get("cached_tokens"), - image_input_tokens: row.get("image_input_tokens"), - image_output_tokens: row.get("image_output_tokens"), - image_request_count: row.get("image_request_count"), - image_request_failed_count: row.get("image_request_failed_count"), - last_used_at: parse_optional_rfc3339(row.get::, _>("last_used_at"))?, - }) -} - -fn should_reset_usage_window( - existing_reset_at: Option>, - existing_limit_window_seconds: Option, - new_reset_at: DateTime, - new_limit_window_seconds: Option, -) -> bool { - let Some(existing_reset_at) = existing_reset_at else { - return false; - }; - if existing_reset_at == new_reset_at { - return false; - } - let drift = existing_reset_at - .signed_duration_since(new_reset_at) - .num_seconds() - .unsigned_abs(); - let window_seconds = new_limit_window_seconds - .or(existing_limit_window_seconds) - .unwrap_or(0); - let threshold = if window_seconds > 0 { - window_seconds / 2 - } else { - 3_600 - }; - drift >= threshold -} - -fn u64_to_i64_saturating(value: u64) -> i64 { - value.min(i64::MAX as u64) as i64 -} diff --git a/src/codex/accounts/service/cookies.rs b/src/codex/accounts/service/cookies.rs deleted file mode 100644 index e9432ccee..000000000 --- a/src/codex/accounts/service/cookies.rs +++ /dev/null @@ -1,51 +0,0 @@ -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 deleted file mode 100644 index fceb2f0a8..000000000 --- a/src/codex/accounts/service/health.rs +++ /dev/null @@ -1,286 +0,0 @@ -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::gateway::transport::endpoints::primary_usage_request_path, - codex::gateway::transport::http_client::{ - build_reqwest_client, CodexBackendClient, CodexClientError, CodexRequestContext, - }, -}; - -use super::{ - cookies::request_domain, AccountProbeOutcome, AccountProbeResult, AccountService, - HealthCheckError, -}; - -impl AccountService { - pub async fn health_check_accounts( - &self, - ids: Option>, - concurrency: usize, - stagger_ms: u64, - ) -> 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(); - 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_oauth_refresh(account).await - } - }) - .buffer_unordered(concurrency) - .collect::>() - .await; - Ok(results) - } - - async fn probe_account_with_oauth_refresh(&self, account: StoredAccount) -> AccountProbeResult { - let started_at = std::time::Instant::now(); - let previous_status = account.status; - match self.refresh_account(&account.id).await { - Ok(mut result) => { - result.previous_status = previous_status; - result - } - Err(error) => AccountProbeResult { - id: account.id, - email: account.email, - previous_status, - outcome: AccountProbeOutcome::Dead, - status: None, - error: Some(error.to_string()), - 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(), - service.fingerprint.clone(), - ); - 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(), - installation_id: None, - session_id: None, - }) - .await -} - -async fn account_cookie_header(service: &AccountService, account_id: &str) -> Option { - let domain = request_domain(&service.config.api.base_url)?; - let request_path = primary_usage_request_path(&service.config.api.base_url); - service - .cookie_repository - .as_ref()? - .cookie_header_for_request(account_id, &domain, &request_path) - .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); - let _ = repo - .set_quota_cooldown_until(&account.id, cooldown_until) - .await; - 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); - if let Some(cookie_repo) = service.cookie_repository.as_ref() { - let _ = cookie_repo.delete_account_cookies(&account.id).await; - } - let _ = repo - .set_cloudflare_cooldown_until(&account.id, cooldown_until) - .await; - 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::CustomCa(_) => "upstream TLS configuration 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 deleted file mode 100644 index 1e9d9116a..000000000 --- a/src/codex/accounts/service/import.rs +++ /dev/null @@ -1,349 +0,0 @@ -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::gateway::oauth::RefreshFailure, -}; - -use super::{ - pool_sync::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_expires_at = entry - .access_token_expires_at - .as_deref() - .map(parse_import_datetime) - .transpose() - .map_err(StoreImportAccountError::Invalid)?; - 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, - status, - quota_limit_reached: false, - quota_verify_required: false, - quota_cooldown_until: None, - cloudflare_cooldown_until: None, - request_count: 0, - empty_response_count: 0, - image_input_tokens: 0, - image_output_tokens: 0, - image_request_count: 0, - image_request_failed_count: 0, - window_request_count: 0, - window_input_tokens: 0, - window_output_tokens: 0, - window_cached_tokens: 0, - window_image_input_tokens: 0, - window_image_output_tokens: 0, - window_image_request_count: 0, - window_image_request_failed_count: 0, - window_started_at: None, - window_reset_at: None, - limit_window_seconds: 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, - 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()) -} - -fn parse_import_datetime(value: &str) -> Result, String> { - DateTime::parse_from_rfc3339(value) - .map(|value| value.with_timezone(&Utc)) - .map_err(|_| "Invalid accessTokenExpiresAt".to_string()) -} - -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/lifecycle.rs b/src/codex/accounts/service/lifecycle.rs deleted file mode 100644 index 2f1cdd587..000000000 --- a/src/codex/accounts/service/lifecycle.rs +++ /dev/null @@ -1,173 +0,0 @@ -use crate::codex::accounts::{model::AccountStatus, repository::AccountRepository}; - -use super::{ - pool_sync::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); - self.websocket_pool.evict_account(account_id).await; - 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); - self.websocket_pool.evict_account(&account_id).await; - 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), - } - self.websocket_pool.evict_account(account_id).await; - return Ok(()); - } - self.account_pool.lock().await.remove(account_id); - self.websocket_pool.evict_account(account_id).await; - 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/mod.rs b/src/codex/accounts/service/mod.rs deleted file mode 100644 index a68a9c370..000000000 --- a/src/codex/accounts/service/mod.rs +++ /dev/null @@ -1,363 +0,0 @@ -pub mod cookies; -pub mod health; -pub mod import; -pub mod lifecycle; -pub mod pool_sync; -pub mod quota; -pub mod refresh; - -use std::sync::Arc; - -use serde_json::Value; -use tokio::sync::Mutex; - -use crate::{ - codex::accounts::cookies::repository::CookieRepository, - codex::accounts::{ - model::AccountStatus, - pool::AccountPool, - repository::{ - AccountRepository, AccountUsageRepository, StoredAccount, StoredAccountMetadata, - }, - }, - codex::gateway::{ - fingerprint::model::Fingerprint, oauth::TokenRefresher, - transport::websocket::CodexWebSocketPool, - }, - config::AppConfig, - utils::pagination::Page, -}; - -#[derive(Clone)] -pub struct AccountService { - config: Arc, - repository: Option, - usage_repository: Option, - cookie_repository: Option, - token_refresher: Option>, - account_pool: Arc>, - websocket_pool: Arc, - fingerprint: Fingerprint, -} - -pub struct AccountServiceDependencies { - pub repository: Option, - pub usage_repository: Option, - pub cookie_repository: Option, - pub token_refresher: Option>, - pub account_pool: Arc>, - pub websocket_pool: Arc, - pub fingerprint: Fingerprint, -} - -#[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, - QuotaWarnings, -} - -#[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 access_token_expires_at: 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, Clone, PartialEq)] -pub struct AccountQuotaWarnings { - pub warnings: Vec, - pub updated_at: Option>, -} - -#[derive(Debug, Clone, PartialEq)] -pub struct AccountQuotaWarning { - pub account_id: String, - pub email: Option, - pub window: QuotaWarningWindow, - pub level: QuotaWarningLevel, - pub used_percent: f64, - pub reset_at: Option, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum QuotaWarningWindow { - Primary, - Secondary, -} - -impl QuotaWarningWindow { - pub fn as_str(self) -> &'static str { - match self { - Self::Primary => "primary", - Self::Secondary => "secondary", - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum QuotaWarningLevel { - Warning, - Critical, -} - -impl QuotaWarningLevel { - pub fn as_str(self) -> &'static str { - match self { - Self::Warning => "warning", - Self::Critical => "critical", - } - } -} - -#[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, - LeaseAcquire, - TokenRefresherUnavailable, - StoreStatus, - StoreRefreshed, -} - -impl std::fmt::Display for RefreshAccountError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::RepositoryUnavailable => write!(f, "repository unavailable"), - Self::Load => write!(f, "failed to load account"), - Self::NotFound => write!(f, "account not found"), - Self::LeaseAcquire => write!(f, "failed to acquire refresh lease"), - Self::TokenRefresherUnavailable => write!(f, "token refresher unavailable"), - Self::StoreStatus => write!(f, "failed to store account status"), - Self::StoreRefreshed => write!(f, "failed to store refreshed tokens"), - } - } -} - -#[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: Arc, dependencies: AccountServiceDependencies) -> Self { - let AccountServiceDependencies { - repository, - usage_repository, - cookie_repository, - token_refresher, - account_pool, - websocket_pool, - fingerprint, - } = dependencies; - Self { - config, - repository, - usage_repository, - cookie_repository, - token_refresher, - account_pool, - websocket_pool, - fingerprint, - } - } - - pub fn has_repository(&self) -> bool { - self.repository.is_some() - } - - pub fn cookie_cleanup_repository(&self) -> Option { - self.cookie_repository.clone() - } - - 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), - } - } - - /// 列出所有账户用于刷新调度器 - pub async fn list_all_for_refresh(&self) -> Result, AccountServiceError> { - self.repository()? - .list_all() - .await - .map_err(|_| AccountServiceError::List) - } - - /// 列出所有配额锁定的账户ID（用于主动配额刷新） - pub async fn list_quota_locked_accounts(&self) -> Vec { - let pool = self.account_pool.lock().await; - pool.list_quota_locked_accounts() - } -} diff --git a/src/codex/accounts/service/pool_sync.rs b/src/codex/accounts/service/pool_sync.rs deleted file mode 100644 index 465386663..000000000 --- a/src/codex/accounts/service/pool_sync.rs +++ /dev/null @@ -1,78 +0,0 @@ -use chrono::Utc; -use secrecy::ExposeSecret; - -use crate::codex::accounts::{ - model::Account, - pool::AccountCapacitySummary, - 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 async fn runtime_capacity_summary(&self) -> AccountCapacitySummary { - self.account_pool.lock().await.capacity_summary(Utc::now()) - } -} - -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_verify_required: false, - quota_cooldown_until: None, - cloudflare_cooldown_until: None, - request_count: 0, - empty_response_count: 0, - image_input_tokens: 0, - image_output_tokens: 0, - image_request_count: 0, - image_request_failed_count: 0, - window_request_count: 0, - window_input_tokens: 0, - window_output_tokens: 0, - window_cached_tokens: 0, - window_image_input_tokens: 0, - window_image_output_tokens: 0, - window_image_request_count: 0, - window_image_request_failed_count: 0, - window_started_at: None, - window_reset_at: None, - limit_window_seconds: None, - added_at: account.added_at.to_rfc3339(), - last_used_at: None, - } -} diff --git a/src/codex/accounts/service/refresh.rs b/src/codex/accounts/service/refresh.rs deleted file mode 100644 index 7d45bb10b..000000000 --- a/src/codex/accounts/service/refresh.rs +++ /dev/null @@ -1,294 +0,0 @@ -use std::sync::Arc; - -use chrono::{Duration, Utc}; -use secrecy::{ExposeSecret, SecretString}; -use uuid::Uuid; - -use crate::{ - codex::accounts::{ - model::{Account, AccountStatus}, - repository::{AccountRepository, StoredAccount, TokenUpdate}, - }, - codex::gateway::oauth::{RefreshFailure, TokenPair, TokenRefresher}, -}; - -use super::{ - health::skipped_probe_result, pool_sync::pool_account_from_stored, AccountProbeOutcome, - AccountProbeResult, AccountService, RefreshAccountError, -}; - -const REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60; - -#[derive(Debug, Clone, Copy)] -enum RefreshFailurePolicy { - PersistStatus, - ReportOnly, - ScheduledProbe, -} - -impl AccountService { - pub async fn refresh_account( - &self, - account_id: &str, - ) -> Result { - self.refresh_account_with_failure_policy(account_id, RefreshFailurePolicy::PersistStatus) - .await - } - - pub async fn probe_account_refresh( - &self, - account_id: &str, - ) -> Result { - self.refresh_account_with_failure_policy(account_id, RefreshFailurePolicy::ReportOnly) - .await - } - - pub async fn probe_scheduled_account_refresh( - &self, - account_id: &str, - ) -> Result { - self.refresh_account_with_failure_policy(account_id, RefreshFailurePolicy::ScheduledProbe) - .await - } - - async fn refresh_account_with_failure_policy( - &self, - account_id: &str, - failure_policy: RefreshFailurePolicy, - ) -> 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")); - } - if account.refresh_token.is_none() { - return Ok(skipped_probe_result(&account, "no refresh token")); - } - let Some(refresher) = self.token_refresher.as_ref().cloned() else { - return Err(RefreshAccountError::TokenRefresherUnavailable); - }; - let lease_owner = refresh_lease_owner(); - let lease_until = Utc::now() + Duration::seconds(REFRESH_LEASE_TTL_SECONDS); - match repo - .try_acquire_refresh_lease(account_id, &lease_owner, lease_until) - .await - { - Ok(true) => {} - Ok(false) => { - return Ok(skipped_probe_result(&account, "refresh lease held")); - } - Err(_) => return Err(RefreshAccountError::LeaseAcquire), - } - - let result = self - .refresh_account_with_lease(repo, account_id, refresher, failure_policy) - .await; - if let Err(error) = repo.release_refresh_lease(account_id, &lease_owner).await { - tracing::warn!( - error = %error, - account_id = %account_id, - "释放 refresh lease 失败" - ); - } - result - } - - async fn refresh_account_with_lease( - &self, - repo: &AccountRepository, - account_id: &str, - refresher: Arc, - failure_policy: RefreshFailurePolicy, - ) -> Result { - 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 started_at = std::time::Instant::now(); - let previous_status = account.status; - if matches!(failure_policy, RefreshFailurePolicy::ScheduledProbe) { - self.mark_scheduled_refreshing(repo, account_id).await?; - } - match refresh_with_latest_disk_token_retry(repo, &account, refresh_token, &refresher).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 = match failure_policy { - RefreshFailurePolicy::PersistStatus => { - self.apply_refresh_failure_status(repo, &account, failure) - .await - } - RefreshFailurePolicy::ReportOnly | RefreshFailurePolicy::ScheduledProbe => None, - }; - 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 mark_scheduled_refreshing( - &self, - repo: &AccountRepository, - account_id: &str, - ) -> Result<(), RefreshAccountError> { - repo.set_status(account_id, AccountStatus::Refreshing) - .await - .map_err(|_| RefreshAccountError::StoreStatus)?; - self.account_pool - .lock() - .await - .set_status(account_id, AccountStatus::Refreshing); - self.websocket_pool.evict_account(account_id).await; - Ok(()) - } - - 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()); - self.websocket_pool.evict_account(account_id).await; - 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); - self.websocket_pool.evict_account(&account.id).await; - Some(status) - } -} - -async fn refresh_with_latest_disk_token_retry( - repo: &AccountRepository, - account: &StoredAccount, - refresh_token: &SecretString, - refresher: &Arc, -) -> Result { - let used_token = refresh_token.expose_secret().to_string(); - match refresher.refresh(&used_token).await { - Err(RefreshFailure::InvalidGrant) => { - if let Some(latest_token) = latest_disk_refresh_token(repo, account, &used_token).await - { - tracing::info!( - account_id = %account.id, - "检测到磁盘 refresh_token 已更新，使用最新值重试刷新" - ); - return refresher.refresh(&latest_token).await; - } - Err(RefreshFailure::InvalidGrant) - } - result => result, - } -} - -async fn latest_disk_refresh_token( - repo: &AccountRepository, - account: &StoredAccount, - used_token: &str, -) -> Option { - match repo.get(&account.id).await { - Ok(Some(latest)) => latest - .refresh_token - .as_ref() - .map(|token| token.expose_secret().to_string()) - .filter(|token| token != used_token), - Ok(None) => None, - Err(error) => { - tracing::warn!( - error = %error, - account_id = %account.id, - "刷新失败后重新读取磁盘 refresh_token 失败" - ); - None - } - } -} - -fn refresh_lease_owner() -> String { - format!("{}:{}", std::process::id(), Uuid::new_v4()) -} - -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/usage_snapshots.rs b/src/codex/accounts/usage_snapshots.rs deleted file mode 100644 index cf7732c4a..000000000 --- a/src/codex/accounts/usage_snapshots.rs +++ /dev/null @@ -1,290 +0,0 @@ -use chrono::{DateTime, Duration, Timelike, Utc}; -use serde::{Deserialize, Serialize}; -use std::collections::VecDeque; - -/// 使用历史快照 - 时间序列统计 -/// -/// 记录账户使用情况的快照，支持多种时间粒度的聚合： -/// - raw: 原始快照（每次记录） -/// - five_min: 5分钟聚合 -/// - hourly: 小时聚合 -/// - daily: 天聚合 - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct UsageSnapshot { - pub timestamp: DateTime, - pub input_tokens: i64, - pub output_tokens: i64, - pub cached_tokens: i64, - pub request_count: i64, - pub active_accounts: i32, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum Granularity { - Raw, - FiveMin, - Hourly, - Daily, -} - -impl Granularity { - pub fn bucket_duration(&self) -> Duration { - match self { - Granularity::Raw => Duration::zero(), - Granularity::FiveMin => Duration::minutes(5), - Granularity::Hourly => Duration::hours(1), - Granularity::Daily => Duration::days(1), - } - } - - /// 将时间戳对齐到桶边界 - pub fn align_timestamp(&self, ts: DateTime) -> DateTime { - match self { - Granularity::Raw => ts, - Granularity::FiveMin => { - let minutes = (ts.minute() / 5) * 5; - align_utc_time(ts, ts.hour(), minutes, 0) - } - Granularity::Hourly => align_utc_time(ts, ts.hour(), 0, 0), - Granularity::Daily => align_utc_time(ts, 0, 0, 0), - } - } -} - -fn align_utc_time(ts: DateTime, hour: u32, minute: u32, second: u32) -> DateTime { - ts.date_naive() - .and_hms_opt(hour, minute, second) - .map(|value| value.and_utc()) - .unwrap_or(ts) -} - -#[derive(Debug, Clone)] -pub struct UsageSnapshotStore { - snapshots: VecDeque, - max_snapshots: usize, - retention_days: u64, -} - -impl UsageSnapshotStore { - pub fn new(max_snapshots: usize, retention_days: u64) -> Self { - Self { - snapshots: VecDeque::with_capacity(max_snapshots), - max_snapshots, - retention_days, - } - } - - /// 添加新的快照 - pub fn push(&mut self, snapshot: UsageSnapshot) { - // 清理过期快照 - self.cleanup_expired(); - - // 添加新快照 - self.snapshots.push_back(snapshot); - - // 限制最大数量 - while self.snapshots.len() > self.max_snapshots { - self.snapshots.pop_front(); - } - } - - /// 清理过期的快照 - fn cleanup_expired(&mut self) { - let cutoff = Utc::now() - Duration::days(self.retention_days as i64); - while let Some(snapshot) = self.snapshots.front() { - if snapshot.timestamp < cutoff { - self.snapshots.pop_front(); - } else { - break; - } - } - } - - /// 获取指定时间范围和粒度的快照 - pub fn get_range( - &self, - start: DateTime, - end: DateTime, - granularity: Granularity, - ) -> Vec { - if granularity == Granularity::Raw { - // 原始快照 - 直接过滤 - return self - .snapshots - .iter() - .filter(|s| s.timestamp >= start && s.timestamp <= end) - .cloned() - .collect(); - } - - // 聚合快照 - self.aggregate(start, end, granularity) - } - - /// 聚合快照到指定粒度 - fn aggregate( - &self, - start: DateTime, - end: DateTime, - granularity: Granularity, - ) -> Vec { - use std::collections::HashMap; - - let mut buckets: HashMap, UsageSnapshot> = HashMap::new(); - - for snapshot in self.snapshots.iter() { - if snapshot.timestamp < start || snapshot.timestamp > end { - continue; - } - - let bucket_key = granularity.align_timestamp(snapshot.timestamp); - let entry = buckets.entry(bucket_key).or_insert_with(|| UsageSnapshot { - timestamp: bucket_key, - input_tokens: 0, - output_tokens: 0, - cached_tokens: 0, - request_count: 0, - active_accounts: 0, - }); - - entry.input_tokens += snapshot.input_tokens; - entry.output_tokens += snapshot.output_tokens; - entry.cached_tokens += snapshot.cached_tokens; - entry.request_count += snapshot.request_count; - entry.active_accounts = entry.active_accounts.max(snapshot.active_accounts); - } - - let mut result: Vec<_> = buckets.into_values().collect(); - result.sort_by_key(|s| s.timestamp); - result - } - - /// 获取所有快照 - pub fn all(&self) -> Vec { - self.snapshots.iter().cloned().collect() - } - - /// 获取快照数量 - pub fn len(&self) -> usize { - self.snapshots.len() - } - - /// 检查是否为空 - pub fn is_empty(&self) -> bool { - self.snapshots.is_empty() - } - - /// 清空所有快照 - pub fn clear(&mut self) { - self.snapshots.clear(); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_granularity_alignment() { - let ts = Utc::now() - .date_naive() - .and_hms_opt(14, 37, 42) - .unwrap() - .and_utc(); - - // 5分钟对齐：14:37:42 -> 14:35:00 - let aligned = Granularity::FiveMin.align_timestamp(ts); - assert_eq!(aligned.minute(), 35); - assert_eq!(aligned.second(), 0); - - // 小时对齐：14:37:42 -> 14:00:00 - let aligned = Granularity::Hourly.align_timestamp(ts); - assert_eq!(aligned.hour(), 14); - assert_eq!(aligned.minute(), 0); - - // 天对齐：14:37:42 -> 00:00:00 - let aligned = Granularity::Daily.align_timestamp(ts); - assert_eq!(aligned.hour(), 0); - assert_eq!(aligned.minute(), 0); - } - - #[test] - fn test_snapshot_store() { - let mut store = UsageSnapshotStore::new(100, 30); - - let now = Utc::now(); - store.push(UsageSnapshot { - timestamp: now, - input_tokens: 100, - output_tokens: 50, - cached_tokens: 10, - request_count: 1, - active_accounts: 5, - }); - - assert_eq!(store.len(), 1); - assert!(!store.is_empty()); - } - - #[test] - fn test_aggregation() { - let mut store = UsageSnapshotStore::new(1000, 30); - let base = Utc::now() - .date_naive() - .and_hms_opt(10, 0, 0) - .unwrap() - .and_utc(); - - // 添加多个快照 - for i in 0..10 { - store.push(UsageSnapshot { - timestamp: base + Duration::minutes(i * 3), - input_tokens: 100, - output_tokens: 50, - cached_tokens: 10, - request_count: 1, - active_accounts: 5, - }); - } - - // 获取5分钟聚合 - let aggregated = store.get_range(base, base + Duration::minutes(30), Granularity::FiveMin); - - // 应该有多个5分钟桶 - assert!(!aggregated.is_empty()); - } - - #[test] - fn test_retention() { - let mut store = UsageSnapshotStore::new(100, 1); // 1天保留期 - - let now = Utc::now(); - let old = now - Duration::days(2); // 2天前 - - // 添加过期快照 - store.push(UsageSnapshot { - timestamp: old, - input_tokens: 100, - output_tokens: 50, - cached_tokens: 10, - request_count: 1, - active_accounts: 5, - }); - - // 添加新快照（会触发清理） - store.push(UsageSnapshot { - timestamp: now, - input_tokens: 200, - output_tokens: 100, - cached_tokens: 20, - request_count: 2, - active_accounts: 10, - }); - - // 应该只剩新快照 - assert_eq!(store.len(), 1); - assert_eq!(store.all()[0].input_tokens, 200); - } -} diff --git a/src/codex/events/mod.rs b/src/codex/events/mod.rs deleted file mode 100644 index 46e84672a..000000000 --- a/src/codex/events/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub mod event; -pub mod repository; -pub mod service; diff --git a/src/codex/events/repository.rs b/src/codex/events/repository.rs deleted file mode 100644 index 3f6e97aff..000000000 --- a/src/codex/events/repository.rs +++ /dev/null @@ -1,213 +0,0 @@ -use sqlx::{QueryBuilder, Row, Sqlite, SqlitePool}; - -use crate::{ - codex::events::event::{EventLevel, EventLog}, - utils::pagination::{clamp_limit, decode_cursor, encode_cursor, Page}, -}; - -#[derive(Clone)] -pub struct EventLogRepository { - pool: SqlitePool, -} - -#[derive(Debug, Default, Clone)] -pub struct EventLogFilters { - pub kind: Option, - pub level: Option, - pub request_id: Option, - pub account_id: Option, - pub route: Option, - pub model: Option, - pub status_code: Option, - pub search: Option, -} - -impl EventLogRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - pub async fn insert(&self, event: EventLog) -> Result<(), sqlx::Error> { - let metadata_json = serde_json::to_string(&event.metadata).unwrap_or_else(|_| "{}".into()); - sqlx::query( - "insert into event_logs (id, request_id, kind, level, account_id, route, model, status_code, latency_ms, message, metadata_json, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind(event.id) - .bind(event.request_id) - .bind(event.kind) - .bind(event.level.as_str()) - .bind(event.account_id) - .bind(event.route) - .bind(event.model) - .bind(event.status_code) - .bind(event.latency_ms) - .bind(event.message) - .bind(metadata_json) - .bind(event.created_at) - .execute(&self.pool) - .await?; - Ok(()) - } - - pub async fn trim_to_capacity(&self, capacity: u32) -> Result { - sqlx::query( - "delete from event_logs where id in ( - select id from event_logs order by created_at desc, id desc limit -1 offset ? - )", - ) - .bind(i64::from(capacity)) - .execute(&self.pool) - .await - .map(|result| result.rows_affected()) - } - - pub async fn list( - &self, - cursor: Option, - limit: u32, - ) -> Result, sqlx::Error> { - self.list_filtered(EventLogFilters::default(), cursor, limit) - .await - } - - pub async fn list_filtered( - &self, - filters: EventLogFilters, - cursor: Option, - limit: u32, - ) -> Result, sqlx::Error> { - let limit = clamp_limit(limit); - let mut builder = QueryBuilder::::new( - "select id, request_id, kind, level, account_id, route, model, status_code, latency_ms, message, metadata_json, created_at from event_logs", - ); - let mut has_condition = false; - append_optional_text_filter(&mut builder, &mut has_condition, "kind", filters.kind); - append_optional_text_filter(&mut builder, &mut has_condition, "level", filters.level); - append_optional_text_filter( - &mut builder, - &mut has_condition, - "request_id", - filters.request_id, - ); - append_optional_text_filter( - &mut builder, - &mut has_condition, - "account_id", - filters.account_id, - ); - append_optional_text_filter(&mut builder, &mut has_condition, "route", filters.route); - append_optional_text_filter(&mut builder, &mut has_condition, "model", filters.model); - if let Some(status_code) = filters.status_code { - append_condition_prefix(&mut builder, &mut has_condition); - builder.push("status_code = "); - builder.push_bind(status_code); - } - if let Some(search) = filters.search { - append_condition_prefix(&mut builder, &mut has_condition); - let pattern = format!("%{search}%"); - builder.push("(message like "); - builder.push_bind(pattern.clone()); - builder.push(" or metadata_json like "); - builder.push_bind(pattern); - builder.push(")"); - } - if let Some(cursor) = cursor.and_then(|c| decode_cursor(&c)) { - append_condition_prefix(&mut builder, &mut has_condition); - builder.push("(created_at < "); - builder.push_bind(cursor.0.clone()); - builder.push(" or (created_at = "); - builder.push_bind(cursor.0); - builder.push(" and id < "); - builder.push_bind(cursor.1); - builder.push("))"); - } - builder.push(" order by created_at desc, id desc limit "); - builder.push_bind(i64::from(limit + 1)); - let mut rows = builder.build().fetch_all(&self.pool).await?; - - let has_next = rows.len() > limit as usize; - if has_next { - rows.truncate(limit as usize); - } - - let items = rows.into_iter().map(row_to_event_log).collect::>(); - let next_cursor = if has_next { - items.last().map(|e| encode_cursor(&e.created_at, &e.id)) - } else { - None - }; - Ok(Page { items, next_cursor }) - } - - pub async fn get(&self, id: &str) -> Result, sqlx::Error> { - sqlx::query( - "select id, request_id, kind, level, account_id, route, model, status_code, latency_ms, message, metadata_json, created_at from event_logs where id = ?", - ) - .bind(id) - .fetch_optional(&self.pool) - .await - .map(|row| row.map(row_to_event_log)) - } - - pub async fn count(&self) -> Result { - let row = sqlx::query("select count(*) as count from event_logs") - .fetch_one(&self.pool) - .await?; - let count = row.get::("count"); - Ok(count.max(0) as u64) - } - - pub async fn clear(&self) -> Result { - sqlx::query("delete from event_logs") - .execute(&self.pool) - .await - .map(|result| result.rows_affected()) - } -} - -fn append_optional_text_filter( - builder: &mut QueryBuilder, - has_condition: &mut bool, - column: &'static str, - value: Option, -) { - let Some(value) = value else { - return; - }; - append_condition_prefix(builder, has_condition); - builder.push(column); - builder.push(" = "); - builder.push_bind(value); -} - -fn append_condition_prefix(builder: &mut QueryBuilder, has_condition: &mut bool) { - if *has_condition { - builder.push(" and "); - } else { - builder.push(" where "); - *has_condition = true; - } -} - -fn row_to_event_log(row: sqlx::sqlite::SqliteRow) -> EventLog { - EventLog { - id: row.get("id"), - request_id: row.get("request_id"), - kind: row.get("kind"), - level: match row.get::("level").as_str() { - "debug" => EventLevel::Debug, - "warn" => EventLevel::Warn, - "error" => EventLevel::Error, - _ => EventLevel::Info, - }, - account_id: row.get("account_id"), - route: row.get("route"), - model: row.get("model"), - status_code: row.get("status_code"), - latency_ms: row.get("latency_ms"), - message: row.get("message"), - metadata: serde_json::from_str(&row.get::("metadata_json")) - .unwrap_or_else(|_| serde_json::json!({})), - created_at: row.get("created_at"), - } -} diff --git a/src/codex/events/service.rs b/src/codex/events/service.rs deleted file mode 100644 index 64feaf692..000000000 --- a/src/codex/events/service.rs +++ /dev/null @@ -1,203 +0,0 @@ -use std::sync::Arc; - -use tokio::sync::RwLock; - -use crate::{ - codex::events::{ - event::EventLog, - repository::{EventLogFilters, EventLogRepository}, - }, - config::LoggingConfig, - utils::pagination::Page, -}; - -#[derive(Clone)] -pub struct LogService { - config: Arc>, - repository: Option, -} - -#[derive(Debug, thiserror::Error)] -pub enum LogServiceError { - #[error("事件日志仓储未初始化")] - RepositoryUnavailable, - #[error("查询事件日志失败")] - List, - #[error("读取事件日志失败")] - Get, - #[error("统计事件日志失败")] - Count, - #[error("清空事件日志失败")] - Clear, - #[error("写入事件日志失败")] - Write, - #[error("日志容量必须大于 0")] - InvalidCapacity, -} - -#[derive(Debug)] -pub struct LogState { - pub enabled: bool, - pub capacity: u32, - pub capture_body: bool, - pub stored_count: u64, -} - -#[derive(Debug)] -pub struct ClearLogs { - pub cleared: u64, -} - -#[derive(Debug, Default)] -pub struct LogStateUpdate { - pub enabled: Option, - pub capacity: Option, - pub capture_body: Option, -} - -#[derive(Debug, Default)] -pub struct LogListFilter { - pub kind: Option, - pub level: Option, - pub request_id: Option, - pub account_id: Option, - pub route: Option, - pub model: Option, - pub status_code: Option, - pub search: Option, -} - -impl LogService { - pub fn new(config: LoggingConfig, repository: Option) -> Self { - Self { - config: Arc::new(RwLock::new(config)), - repository, - } - } - - pub async fn list( - &self, - cursor: Option, - limit: u32, - filter: LogListFilter, - ) -> Result, LogServiceError> { - self.repository()? - .list_filtered(EventLogFilters::from(filter), cursor, limit) - .await - .map_err(|_| LogServiceError::List) - } - - pub async fn state(&self) -> Result { - let stored_count = self - .repository()? - .count() - .await - .map_err(|_| LogServiceError::Count)?; - let config = self.config.read().await; - Ok(LogState { - enabled: config.enabled, - capacity: config.capacity, - capture_body: config.capture_body, - stored_count, - }) - } - - pub async fn update_state(&self, update: LogStateUpdate) -> Result { - if matches!(update.capacity, Some(0)) { - return Err(LogServiceError::InvalidCapacity); - } - - { - let mut config = self.config.write().await; - if let Some(enabled) = update.enabled { - config.enabled = enabled; - } - if let Some(capacity) = update.capacity { - config.capacity = capacity; - } - if let Some(capture_body) = update.capture_body { - config.capture_body = capture_body; - } - } - if let Some(capacity) = update.capacity { - self.repository()? - .trim_to_capacity(capacity) - .await - .map_err(|_| LogServiceError::Write)?; - } - - self.state().await - } - - pub async fn record(&self, mut event: EventLog) -> Result { - let config = self.config.read().await.clone(); - if !config.enabled { - return Ok(false); - } - apply_capture_body_policy(&mut event, config.capture_body); - let repository = self.repository()?; - repository - .insert(event) - .await - .map_err(|_| LogServiceError::Write)?; - repository - .trim_to_capacity(config.capacity) - .await - .map_err(|_| LogServiceError::Write)?; - Ok(true) - } - - pub async fn get(&self, id: &str) -> Result, LogServiceError> { - self.repository()? - .get(id) - .await - .map_err(|_| LogServiceError::Get) - } - - pub async fn clear(&self) -> Result { - self.repository()? - .clear() - .await - .map(|cleared| ClearLogs { cleared }) - .map_err(|_| LogServiceError::Clear) - } - - fn repository(&self) -> Result<&EventLogRepository, LogServiceError> { - self.repository - .as_ref() - .ok_or(LogServiceError::RepositoryUnavailable) - } -} - -fn apply_capture_body_policy(event: &mut EventLog, capture_body: bool) { - if capture_body { - return; - } - let Some(metadata) = event.metadata.as_object_mut() else { - return; - }; - for key in [ - "body", - "rawBody", - "requestBody", - "responseBody", - "upstreamBody", - ] { - metadata.remove(key); - } -} - -impl From for EventLogFilters { - fn from(filter: LogListFilter) -> Self { - Self { - kind: filter.kind, - level: filter.level, - request_id: filter.request_id, - account_id: filter.account_id, - route: filter.route, - model: filter.model, - status_code: filter.status_code, - search: filter.search, - } - } -} diff --git a/src/codex/gateway/fingerprint/mod.rs b/src/codex/gateway/fingerprint/mod.rs deleted file mode 100644 index 809e5a444..000000000 --- a/src/codex/gateway/fingerprint/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -pub mod model; -pub mod repository; -pub mod update_checker; -pub mod updater; diff --git a/src/codex/gateway/fingerprint/repository.rs b/src/codex/gateway/fingerprint/repository.rs deleted file mode 100644 index b5b26a124..000000000 --- a/src/codex/gateway/fingerprint/repository.rs +++ /dev/null @@ -1,122 +0,0 @@ -use chrono::Utc; -use sqlx::{Row, SqlitePool}; -use uuid::Uuid; - -use crate::codex::gateway::fingerprint::{ - model::Fingerprint, - updater::{FingerprintUpdate, CODEX_DESKTOP_UPDATE_SOURCE}, -}; - -const AUTO_UPDATED_FINGERPRINT_ID: &str = "auto_updated"; -const AUTO_UPDATE_SOURCE: &str = "auto_update"; -const AUTO_UPDATE_PLATFORM: &str = "darwin"; -const AUTO_UPDATE_ARCH: &str = "arm64"; -const AUTO_UPDATE_USER_AGENT_TEMPLATE: &str = "Codex Desktop/{version} ({platform}; {arch})"; -const DEFAULT_AUTO_UPDATE_CHROMIUM_VERSION: &str = "146"; - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct StoredFingerprint { - pub app_version: String, - pub build_number: String, - pub source: String, -} - -#[derive(Clone)] -pub struct FingerprintRepository { - pool: SqlitePool, -} - -impl FingerprintRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - pub async fn insert_update(&self, update: &FingerprintUpdate) -> Result<(), sqlx::Error> { - let mut fp = Fingerprint::default_codex_desktop(); - fp.app_version.clone_from(&update.app_version); - fp.build_number.clone_from(&update.build_number); - sqlx::query( - "insert into fingerprints (id, app_version, build_number, platform, arch, chromium_version, user_agent_template, source, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind(Uuid::new_v4().to_string()) - .bind(fp.app_version) - .bind(fp.build_number) - .bind(fp.platform) - .bind(fp.arch) - .bind(fp.chromium_version) - .bind(fp.user_agent_template) - .bind(CODEX_DESKTOP_UPDATE_SOURCE) - .bind(Utc::now().to_rfc3339()) - .execute(&self.pool) - .await?; - Ok(()) - } - - pub async fn upsert_auto_update( - &self, - app_version: &str, - build_number: &str, - chromium_version: Option<&str>, - ) -> Result<(), sqlx::Error> { - sqlx::query( - "insert into fingerprints (id, app_version, build_number, platform, arch, chromium_version, user_agent_template, source, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?) \ - on conflict(id) do update set app_version = excluded.app_version, build_number = excluded.build_number, platform = excluded.platform, arch = excluded.arch, chromium_version = excluded.chromium_version, user_agent_template = excluded.user_agent_template, source = excluded.source, created_at = excluded.created_at", - ) - .bind(AUTO_UPDATED_FINGERPRINT_ID) - .bind(app_version) - .bind(build_number) - .bind(AUTO_UPDATE_PLATFORM) - .bind(AUTO_UPDATE_ARCH) - .bind(chromium_version.unwrap_or(DEFAULT_AUTO_UPDATE_CHROMIUM_VERSION)) - .bind(AUTO_UPDATE_USER_AGENT_TEMPLATE) - .bind(AUTO_UPDATE_SOURCE) - .bind(Utc::now().to_rfc3339()) - .execute(&self.pool) - .await?; - Ok(()) - } - - pub async fn latest(&self) -> Result, sqlx::Error> { - let row = sqlx::query( - "select app_version, build_number, source from fingerprints order by created_at desc, id desc limit 1", - ) - .fetch_optional(&self.pool) - .await?; - Ok(row.map(|row| StoredFingerprint { - app_version: row.get("app_version"), - build_number: row.get("build_number"), - source: row.get("source"), - })) - } - - pub async fn load_latest_auto_updated(&self) -> Result, sqlx::Error> { - let row = sqlx::query( - r#" - select app_version, build_number, platform, arch, chromium_version, user_agent_template - from fingerprints - where source = ? - order by created_at desc - limit 1 - "#, - ) - .bind(AUTO_UPDATE_SOURCE) - .fetch_optional(&self.pool) - .await?; - - let Some(row) = row else { - return Ok(None); - }; - - Ok(Some(Fingerprint { - originator: "Codex Desktop".to_string(), - app_version: row.get("app_version"), - build_number: row.get("build_number"), - platform: row.get("platform"), - arch: row.get("arch"), - chromium_version: row.get("chromium_version"), - user_agent_template: row.get("user_agent_template"), - default_headers: Fingerprint::default_headers(), - header_order: Fingerprint::default_header_order(), - })) - } -} diff --git a/src/codex/gateway/fingerprint/update_checker.rs b/src/codex/gateway/fingerprint/update_checker.rs deleted file mode 100644 index eb73838dd..000000000 --- a/src/codex/gateway/fingerprint/update_checker.rs +++ /dev/null @@ -1,246 +0,0 @@ -use std::sync::Arc; -use std::time::Duration; - -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; -use sqlx::SqlitePool; -use tokio::sync::Mutex; -use tokio::time::interval; -use tracing::{info, warn}; - -use crate::codex::gateway::fingerprint::repository::FingerprintRepository; -use crate::platform::storage::paths; - -const APPCAST_URL: &str = "https://persistent.oaistatic.com/codex-app-prod/appcast.xml"; -const POLL_INTERVAL: Duration = Duration::from_secs(3 * 24 * 60 * 60); // 3 天 -type AppcastFields = (Option, Option, Option); - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct UpdateState { - pub last_check: DateTime, - pub latest_version: Option, - pub latest_build: Option, - pub download_url: Option, - pub update_available: bool, - pub current_version: String, - pub current_build: String, -} - -#[derive(Clone)] -pub struct UpdateChecker { - repository: Option, - state: Arc>, -} - -struct InternalState { - current_version: String, - current_build: String, -} - -impl UpdateChecker { - pub fn new(db: Option, current_version: String, current_build: String) -> Self { - Self { - repository: db.map(FingerprintRepository::new), - state: Arc::new(Mutex::new(InternalState { - current_version, - current_build, - })), - } - } - - pub async fn check_for_update(&self) -> Result { - let state = self.state.lock().await; - let current_version = state.current_version.clone(); - let current_build = state.current_build.clone(); - drop(state); - - let xml = fetch_appcast().await?; - let (version, build, download_url) = parse_appcast(&xml)?; - - let update_available = version.as_ref().is_some_and(|v| v != ¤t_version) - || build.as_ref().is_some_and(|b| b != ¤t_build); - - let update_state = UpdateState { - last_check: Utc::now(), - latest_version: version, - latest_build: build, - download_url, - update_available, - current_version, - current_build, - }; - - persist_update_state(&update_state)?; - - Ok(update_state) - } - - pub fn start_background_checker(self) -> tokio::task::JoinHandle<()> { - tokio::spawn(async move { - let mut ticker = interval(POLL_INTERVAL); - - info!( - interval_secs = POLL_INTERVAL.as_secs(), - "fingerprint 后台版本检查器已启动" - ); - - // 立即执行首次检查 - if let Err(e) = self.check_and_apply_update().await { - warn!(error = %e, "fingerprint 首次版本检查失败"); - } - - loop { - ticker.tick().await; - if let Err(e) = self.check_and_apply_update().await { - warn!(error = %e, "fingerprint 定期版本检查失败"); - } - } - }) - } - - async fn check_and_apply_update(&self) -> Result<(), UpdateError> { - let update_state = self.check_for_update().await?; - - if update_state.update_available { - let (Some(version), Some(build)) = ( - update_state.latest_version.as_ref(), - update_state.latest_build.as_ref(), - ) else { - return Ok(()); - }; - - let state = self.state.lock().await; - let current_version = state.current_version.clone(); - let current_build = state.current_build.clone(); - drop(state); - - info!( - version = %version, - build = %build, - current_version = %current_version, - current_build = %current_build, - "发现新的 fingerprint 版本" - ); - - self.apply_version_update(version, build).await?; - - let mut state = self.state.lock().await; - state.current_version = version.clone(); - state.current_build = build.clone(); - - info!( - version = %version, - build = %build, - "已自动应用 fingerprint 版本" - ); - } - - Ok(()) - } - - async fn apply_version_update(&self, version: &str, build: &str) -> Result<(), UpdateError> { - let chromium_version = load_matching_chromium_version(version, build); - - if let Some(repository) = &self.repository { - repository - .upsert_auto_update(version, build, chromium_version.as_deref()) - .await?; - } - - Ok(()) - } -} - -async fn fetch_appcast() -> Result { - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(30)) - .build()?; - - let response = client.get(APPCAST_URL).send().await?; - - if !response.status().is_success() { - return Err(UpdateError::AppcastFetch(response.status().as_u16())); - } - - let xml = response.text().await?; - Ok(xml) -} - -fn parse_appcast(xml: &str) -> Result { - let item_start = xml.find("").ok_or(UpdateError::AppcastParse)?; - let item_end = xml.find("").ok_or(UpdateError::AppcastParse)?; - let item = &xml[item_start..item_end]; - - let version = extract_sparkle_field(item, "shortVersionString"); - let build = extract_sparkle_field(item, "version"); - let download_url = item.find("url=\"").and_then(|pos| { - let start = pos + 5; - item[start..] - .find('"') - .map(|end| item[start..start + end].to_string()) - }); - - Ok((version, build, download_url)) -} - -fn extract_sparkle_field(item: &str, field: &str) -> Option { - // 支持属性语法: sparkle:version="X" - let attr_pattern = format!("sparkle:{}=\"", field); - if let Some(pos) = item.find(&attr_pattern) { - let start = pos + attr_pattern.len(); - if let Some(end) = item[start..].find('"') { - return Some(item[start..start + end].to_string()); - } - } - - // 支持元素语法: X - let elem_start = format!("", field); - let elem_end = format!("", field); - if let Some(start_pos) = item.find(&elem_start) { - let content_start = start_pos + elem_start.len(); - if let Some(end_pos) = item[content_start..].find(&elem_end) { - return Some(item[content_start..content_start + end_pos].to_string()); - } - } - - None -} - -fn persist_update_state(state: &UpdateState) -> Result<(), UpdateError> { - let data_dir = paths::ensure_data_dir()?; - let state_path = data_dir.join("update-state.json"); - let json = serde_json::to_string_pretty(state)?; - std::fs::write(state_path, json)?; - Ok(()) -} - -fn load_matching_chromium_version(version: &str, build: &str) -> Option { - let extracted_path = paths::data_dir().join("extracted-fingerprint.json"); - let content = std::fs::read_to_string(extracted_path).ok()?; - let parsed: serde_json::Value = serde_json::from_str(&content).ok()?; - - let extracted_version = parsed.get("app_version")?.as_str()?; - let extracted_build = parsed.get("build_number")?.as_str()?; - - if extracted_version == version && extracted_build == build { - parsed.get("chromium_version")?.as_str().map(String::from) - } else { - None - } -} - -#[derive(Debug, thiserror::Error)] -pub enum UpdateError { - #[error("HTTP 请求失败: {0}")] - Http(#[from] reqwest::Error), - #[error("获取 appcast 失败，状态码: {0}")] - AppcastFetch(u16), - #[error("解析 appcast 失败")] - AppcastParse, - #[error("JSON 序列化失败: {0}")] - Json(#[from] serde_json::Error), - #[error("文件操作失败: {0}")] - Io(#[from] std::io::Error), - #[error("数据库操作失败: {0}")] - Database(#[from] sqlx::Error), -} diff --git a/src/codex/gateway/fingerprint/updater.rs b/src/codex/gateway/fingerprint/updater.rs deleted file mode 100644 index ba9a4881b..000000000 --- a/src/codex/gateway/fingerprint/updater.rs +++ /dev/null @@ -1,72 +0,0 @@ -use serde::Deserialize; -use thiserror::Error; - -use crate::codex::gateway::fingerprint::repository::FingerprintRepository; - -pub const CODEX_DESKTOP_UPDATE_SOURCE: &str = "codex_desktop_update_source"; - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct FingerprintUpdate { - pub app_version: String, - pub build_number: String, -} - -#[derive(Debug, Error)] -pub enum FingerprintError { - #[error("invalid update manifest: {0}")] - InvalidManifest(#[from] serde_json::Error), - #[error("failed to fetch update manifest: {0}")] - Http(#[from] reqwest::Error), - #[error("failed to persist fingerprint update: {0}")] - Database(#[from] sqlx::Error), -} - -#[derive(Deserialize)] -struct Manifest { - version: String, - build_number: String, -} - -pub fn parse_update_manifest(input: &str) -> Result { - // 自动更新只同步桌面端指纹字段，不把远端配置当作运行时业务配置执行。 - let manifest: Manifest = serde_json::from_str(input)?; - Ok(FingerprintUpdate { - app_version: manifest.version, - build_number: manifest.build_number, - }) -} - -#[derive(Clone)] -pub struct FingerprintUpdater { - client: reqwest::Client, - repository: FingerprintRepository, - update_url: String, -} - -impl FingerprintUpdater { - pub fn new( - client: reqwest::Client, - repository: FingerprintRepository, - update_url: impl Into, - ) -> Self { - Self { - client, - repository, - update_url: update_url.into(), - } - } - - pub async fn poll_once(&self) -> Result { - let manifest = self - .client - .get(&self.update_url) - .send() - .await? - .error_for_status()? - .text() - .await?; - let update = parse_update_manifest(&manifest)?; - self.repository.insert_update(&update).await?; - Ok(update) - } -} diff --git a/src/codex/gateway/installation_id.rs b/src/codex/gateway/installation_id.rs deleted file mode 100644 index 44e528189..000000000 --- a/src/codex/gateway/installation_id.rs +++ /dev/null @@ -1,112 +0,0 @@ -//! Installation ID 管理 -//! -//! 稳定的设备级 UUID，作为 `x-codex-installation-id` header 和 `client_metadata` 发送到上游。 -//! 用于路由/亲和性提示，使上游可以将客户端固定到同一后端实例以保持提示缓存温暖。 -//! -//! 查找顺序： -//! 1. `~/.codex/installation_id`（真实 Codex Desktop 位置） -//! 2. `/installation_id`（之前持久化的） -//! 3. 生成新 UUID 并持久化到 `/installation_id` - -use std::{ - fs, - path::{Path, PathBuf}, - sync::OnceLock, -}; - -use uuid::Uuid; - -static INSTALLATION_ID: OnceLock = OnceLock::new(); - -/// 获取 installation ID（首次调用后缓存） -pub fn get_installation_id(database_path: Option<&str>) -> String { - INSTALLATION_ID - .get_or_init(|| resolve_installation_id(database_path)) - .clone() -} - -fn resolve_installation_id(database_path: Option<&str>) -> String { - // 优先尝试 ~/.codex/installation_id（真实 Codex Desktop 位置） - if let Some(home) = dirs::home_dir() { - let codex_home = home.join(".codex").join("installation_id"); - if let Some(id) = read_uuid_file(&codex_home) { - return id; - } - } - - // 从数据库路径推导数据目录（例如 "data/proxy.db" -> "data/"） - let data_dir = database_path - .and_then(|path| Path::new(path).parent()) - .map(|p| p.to_path_buf()); - - if let Some(dir) = data_dir { - let our_file = dir.join("installation_id"); - if let Some(id) = read_uuid_file(&our_file) { - return id; - } - - // 生成并持久化新 UUID - let generated = Uuid::new_v4().to_string(); - if let Err(e) = persist_uuid(&our_file, &generated) { - tracing::warn!( - error = %e, - path = %our_file.display(), - "持久化 installation_id 失败" - ); - } - return generated; - } - - // 降级：生成不持久化 - Uuid::new_v4().to_string() -} - -fn read_uuid_file(path: &PathBuf) -> Option { - let content = fs::read_to_string(path).ok()?; - let trimmed = content.trim(); - - // 验证 UUID 格式 - Uuid::parse_str(trimmed).ok()?; - - Some(trimmed.to_string()) -} - -fn persist_uuid(path: &PathBuf, uuid: &str) -> std::io::Result<()> { - if let Some(parent) = path.parent() { - fs::create_dir_all(parent)?; - } - fs::write(path, uuid) -} - -#[cfg(test)] -mod tests { - use super::*; - use std::fs; - use tempfile::TempDir; - - #[test] - fn test_generate_and_persist() { - let temp = TempDir::new().unwrap(); - let db_path = temp.path().join("proxy.db"); - let db_path_str = db_path.to_str().unwrap(); - - let id1 = resolve_installation_id(Some(db_path_str)); - assert!(Uuid::parse_str(&id1).is_ok()); - - let id2 = resolve_installation_id(Some(db_path_str)); - assert_eq!(id1, id2); - } - - #[test] - fn test_read_from_codex_home() { - let temp = TempDir::new().unwrap(); - let codex_dir = temp.path().join(".codex"); - fs::create_dir_all(&codex_dir).unwrap(); - - let expected_id = Uuid::new_v4().to_string(); - fs::write(codex_dir.join("installation_id"), &expected_id).unwrap(); - - let read_id = read_uuid_file(&codex_dir.join("installation_id")); - assert_eq!(read_id, Some(expected_id)); - } -} diff --git a/src/codex/gateway/mod.rs b/src/codex/gateway/mod.rs deleted file mode 100644 index e65e13dcc..000000000 --- a/src/codex/gateway/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -pub mod conversation_identity; -pub mod fingerprint; -pub mod installation_id; -pub mod oauth; -pub mod protocol; -pub mod transport; diff --git a/src/codex/gateway/oauth/client.rs b/src/codex/gateway/oauth/client.rs deleted file mode 100644 index 575ff14de..000000000 --- a/src/codex/gateway/oauth/client.rs +++ /dev/null @@ -1,419 +0,0 @@ -use std::collections::BTreeMap; - -use async_trait::async_trait; -use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; -use chrono::{DateTime, Duration, Utc}; -use rand::Rng; -use reqwest::{Client, StatusCode}; -use serde::Deserialize; -use sha2::{Digest, Sha256}; -use thiserror::Error; - -use crate::config::AuthConfig; - -use super::{ - refresh::{RefreshFailure, TokenRefresher}, - token::TokenPair, -}; - -const OPENAI_OAUTH_DEVICE_CODE_ENDPOINT: &str = "https://auth.openai.com/oauth/device/code"; -const OAUTH_CALLBACK_PORT: u16 = 1455; -const OAUTH_CALLBACK_PATH: &str = "/auth/callback"; - -#[derive(Debug, Clone)] -pub struct OAuthConfig { - pub client_id: String, - pub auth_endpoint: String, - pub device_code_endpoint: String, - pub token_endpoint: String, -} - -impl OAuthConfig { - pub fn from_auth_config(config: &AuthConfig) -> Self { - Self { - client_id: config.oauth_client_id.clone(), - auth_endpoint: config.oauth_auth_endpoint.clone(), - device_code_endpoint: OPENAI_OAUTH_DEVICE_CODE_ENDPOINT.to_string(), - token_endpoint: config.oauth_token_endpoint.clone(), - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] -pub struct DeviceCode { - pub device_code: String, - pub user_code: String, - pub verification_uri: String, - pub verification_uri_complete: String, - pub expires_in: u64, - pub interval: u64, -} - -#[derive(Debug, Clone, PartialEq, Eq, Error)] -pub enum OAuthError { - #[error("authorization pending")] - AuthorizationPending, - #[error("slow down")] - SlowDown, - #[error("OAuth request rejected: {0}")] - Rejected(String), - #[error("OAuth transport failed")] - Transport, -} - -impl OAuthError { - pub fn pending_code(&self) -> Option<&'static str> { - match self { - Self::AuthorizationPending => Some("authorization_pending"), - Self::SlowDown => Some("slow_down"), - Self::Rejected(_) | Self::Transport => None, - } - } -} - -#[async_trait] -pub trait OAuthClient: Send + Sync + 'static { - async fn exchange_code( - &self, - code: &str, - code_verifier: &str, - redirect_uri: &str, - ) -> Result; - - async fn request_device_code(&self) -> Result; - - async fn poll_device_token(&self, device_code: &str) -> Result; -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct PkceLogin { - pub auth_url: String, - pub state: String, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct PkceSession { - pub code_verifier: String, - pub redirect_uri: String, - pub return_host: String, -} - -#[derive(Debug, Clone)] -struct PendingPkceSession { - session: PkceSession, - created_at: DateTime, - exchanging: bool, -} - -#[derive(Debug)] -pub struct PkceSessionStore { - pending: BTreeMap, - completed: BTreeMap>, - ttl: Duration, -} - -impl Default for PkceSessionStore { - fn default() -> Self { - Self { - pending: BTreeMap::new(), - completed: BTreeMap::new(), - ttl: Duration::minutes(5), - } - } -} - -impl PkceSessionStore { - pub fn start_login(&mut self, return_host: &str, config: &OAuthConfig) -> PkceLogin { - self.cleanup(Utc::now()); - let (code_verifier, code_challenge) = generate_pkce_pair(); - let state = random_state(); - let redirect_uri = format!("http://localhost:{OAUTH_CALLBACK_PORT}{OAUTH_CALLBACK_PATH}"); - let auth_url = build_auth_url(config, &redirect_uri, &state, &code_challenge); - self.pending.insert( - state.clone(), - PendingPkceSession { - session: PkceSession { - code_verifier, - redirect_uri, - return_host: return_host.to_string(), - }, - created_at: Utc::now(), - exchanging: false, - }, - ); - PkceLogin { auth_url, state } - } - - pub fn try_acquire(&mut self, state: &str) -> Option { - self.cleanup(Utc::now()); - let pending = self.pending.get_mut(state)?; - if pending.exchanging { - return None; - } - pending.exchanging = true; - Some(pending.session.clone()) - } - - pub fn release(&mut self, state: &str) { - if let Some(pending) = self.pending.get_mut(state) { - pending.exchanging = false; - } - } - - pub fn complete(&mut self, state: &str) { - self.pending.remove(state); - self.completed.insert(state.to_string(), Utc::now()); - } - - pub fn is_completed_or_exchanging(&mut self, state: &str) -> bool { - self.cleanup(Utc::now()); - self.completed.contains_key(state) - || self - .pending - .get(state) - .is_some_and(|pending| pending.exchanging) - } - - fn cleanup(&mut self, now: DateTime) { - self.pending - .retain(|_, pending| now - pending.created_at <= self.ttl); - self.completed - .retain(|_, completed_at| now - *completed_at <= self.ttl); - } -} - -#[derive(Clone)] -pub struct OpenAiOAuthRefresher { - client: Client, - config: OAuthConfig, -} - -impl OpenAiOAuthRefresher { - pub fn new(client: Client, config: OAuthConfig) -> Self { - Self { client, config } - } -} - -#[derive(Deserialize)] -struct RefreshTokenResponse { - access_token: String, - refresh_token: Option, -} - -#[derive(Deserialize)] -struct OAuthErrorResponse { - error: Option, - error_description: Option, -} - -#[async_trait] -impl OAuthClient for OpenAiOAuthRefresher { - async fn exchange_code( - &self, - code: &str, - code_verifier: &str, - redirect_uri: &str, - ) -> Result { - let response = self - .client - .post(&self.config.token_endpoint) - .form(&[ - ("grant_type", "authorization_code"), - ("client_id", self.config.client_id.as_str()), - ("code", code), - ("redirect_uri", redirect_uri), - ("code_verifier", code_verifier), - ]) - .send() - .await - .map_err(|_| OAuthError::Transport)?; - let status = response.status(); - let body = response.text().await.map_err(|_| OAuthError::Transport)?; - if !status.is_success() { - return Err(OAuthError::Rejected(format!( - "Token exchange failed ({status}): {body}" - ))); - } - parse_token_pair(&body).map_err(|_| OAuthError::Transport) - } - - async fn request_device_code(&self) -> Result { - let response = self - .client - .post(&self.config.device_code_endpoint) - .form(&[ - ("client_id", self.config.client_id.as_str()), - ("scope", "openid profile email offline_access"), - ]) - .send() - .await - .map_err(|_| OAuthError::Transport)?; - let status = response.status(); - let body = response.text().await.map_err(|_| OAuthError::Transport)?; - if !status.is_success() { - return Err(OAuthError::Rejected(format!( - "Device code request failed ({status}): {body}" - ))); - } - serde_json::from_str::(&body).map_err(|_| OAuthError::Transport) - } - - async fn poll_device_token(&self, device_code: &str) -> Result { - let response = self - .client - .post(&self.config.token_endpoint) - .form(&[ - ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"), - ("device_code", device_code), - ("client_id", self.config.client_id.as_str()), - ]) - .send() - .await - .map_err(|_| OAuthError::Transport)?; - let status = response.status(); - let body = response.text().await.map_err(|_| OAuthError::Transport)?; - if !status.is_success() { - return Err(classify_device_poll_failure(status, &body)); - } - parse_token_pair(&body).map_err(|_| OAuthError::Transport) - } -} - -#[async_trait] -impl TokenRefresher for OpenAiOAuthRefresher { - async fn refresh(&self, refresh_token: &str) -> Result { - // 🔒 一次性刷新令牌保护： - // OAuth2 refresh tokens 可能是一次性的（用完即失效）。 - // 如果请求中途失败（网络错误、超时等），我们不能自动重试， - // 因为令牌可能已被服务器消耗。 - // - // 调用者应该： - // 1. 在调用前保存原始 refresh_token - // 2. 只在 is_safe_to_retry_refresh(&error) 返回 true 时重试 - // 3. 任何其他错误都应标记为 permanent failure - let response = self - .client - .post(&self.config.token_endpoint) - .form(&[ - ("grant_type", "refresh_token"), - ("client_id", self.config.client_id.as_str()), - ("refresh_token", refresh_token), - ]) - .send() - .await - .map_err(|_| RefreshFailure::Transport)?; - let status = response.status(); - let body = response - .text() - .await - .map_err(|_| RefreshFailure::Transport)?; - if !status.is_success() { - return Err(classify_refresh_failure(status, &body)); - } - parse_token_pair(&body).map_err(|_| RefreshFailure::Transport) - } -} - -fn parse_token_pair(body: &str) -> Result { - let tokens = serde_json::from_str::(body).map_err(|_| ())?; - if tokens.access_token.trim().is_empty() { - return Err(()); - } - Ok(TokenPair { - access_token: tokens.access_token, - refresh_token: tokens.refresh_token, - }) -} - -fn classify_device_poll_failure(status: StatusCode, body: &str) -> OAuthError { - let parsed = serde_json::from_str::(body).ok(); - match parsed.as_ref().and_then(|data| data.error.as_deref()) { - Some("authorization_pending") => OAuthError::AuthorizationPending, - Some("slow_down") => OAuthError::SlowDown, - Some(code) => OAuthError::Rejected( - parsed - .as_ref() - .and_then(|data| data.error_description.clone()) - .unwrap_or_else(|| code.to_string()), - ), - None => OAuthError::Rejected(format!("Device token poll failed ({status}): {body}")), - } -} - -fn generate_pkce_pair() -> (String, String) { - let mut bytes = [0u8; 32]; - rand::rng().fill_bytes(&mut bytes); - let code_verifier = URL_SAFE_NO_PAD.encode(bytes); - let digest = Sha256::digest(code_verifier.as_bytes()); - let code_challenge = URL_SAFE_NO_PAD.encode(digest); - (code_verifier, code_challenge) -} - -fn random_state() -> String { - let mut bytes = [0u8; 16]; - rand::rng().fill_bytes(&mut bytes); - bytes.iter().map(|byte| format!("{byte:02x}")).collect() -} - -fn build_auth_url( - config: &OAuthConfig, - redirect_uri: &str, - state: &str, - code_challenge: &str, -) -> String { - let params = [ - ("response_type", "code"), - ("client_id", config.client_id.as_str()), - ("redirect_uri", redirect_uri), - ("scope", "openid profile email offline_access"), - ("code_challenge", code_challenge), - ("code_challenge_method", "S256"), - ("id_token_add_organizations", "true"), - ("codex_cli_simplified_flow", "true"), - ("state", state), - ("originator", "codex_cli_rs"), - ]; - let query = params - .iter() - .map(|(key, value)| format!("{}={}", percent_encode(key), percent_encode(value))) - .collect::>() - .join("&"); - format!("{}?{query}", config.auth_endpoint) -} - -fn percent_encode(value: &str) -> String { - let mut encoded = String::new(); - for byte in value.bytes() { - match byte { - b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' => { - encoded.push(char::from(byte)); - } - _ => encoded.push_str(&format!("%{byte:02X}")), - } - } - encoded -} - -fn classify_refresh_failure(status: StatusCode, body: &str) -> RefreshFailure { - let lower = body.to_ascii_lowercase(); - if lower.contains("quota") { - return RefreshFailure::QuotaExhausted; - } - if lower.contains("account has been deactivated") - || lower.contains("refresh_token_reused") - || lower.contains("banned") - { - return RefreshFailure::Banned; - } - if status == StatusCode::BAD_REQUEST - || status == StatusCode::UNAUTHORIZED - || lower.contains("invalid_grant") - || lower.contains("invalid_token") - || lower.contains("access_denied") - || lower.contains("refresh_token_expired") - || lower.contains("token_revoked") - { - return RefreshFailure::InvalidGrant; - } - RefreshFailure::Transport -} diff --git a/src/codex/gateway/oauth/codex_cli.rs b/src/codex/gateway/oauth/codex_cli.rs deleted file mode 100644 index e9fb8e7c4..000000000 --- a/src/codex/gateway/oauth/codex_cli.rs +++ /dev/null @@ -1,93 +0,0 @@ -use std::{ - env, fs, io, - path::{Path, PathBuf}, -}; - -use secrecy::{ExposeSecret, SecretString}; -use serde::Deserialize; -use thiserror::Error; - -#[derive(Debug, Error)] -pub enum CliAuthImportError { - #[error("CODEX_HOME is not set and HOME is unavailable")] - MissingHome, - #[error("failed to read CLI auth file {path}: {source}")] - Read { path: PathBuf, source: io::Error }, - #[error("failed to parse CLI auth.json: {0}")] - Parse(#[from] serde_json::Error), - #[error("CLI auth.json does not contain access_token")] - MissingAccessToken, -} - -#[derive(Debug, Clone)] -pub struct CliAuth { - access_token: SecretString, - refresh_token: Option, - id_token: Option, - expires_at: Option, -} - -impl CliAuth { - pub fn access_token(&self) -> &str { - self.access_token.expose_secret() - } - - pub fn refresh_token(&self) -> Option<&str> { - self.refresh_token - .as_ref() - .map(|token| token.expose_secret()) - } - - pub fn id_token(&self) -> Option<&str> { - self.id_token.as_ref().map(|token| token.expose_secret()) - } - - pub fn expires_at(&self) -> Option { - self.expires_at - } -} - -#[derive(Debug, Deserialize)] -struct CliAuthJson { - access_token: Option, - refresh_token: Option, - id_token: Option, - expires_at: Option, -} - -pub fn default_codex_home() -> Result { - if let Some(codex_home) = env::var_os("CODEX_HOME") { - return Ok(PathBuf::from(codex_home)); - } - let home = env::var_os("HOME").ok_or(CliAuthImportError::MissingHome)?; - Ok(PathBuf::from(home).join(".codex")) -} - -pub fn read_cli_auth_from_home(codex_home: &Path) -> Result { - let auth_path = codex_home.join("auth.json"); - let raw = fs::read_to_string(&auth_path).map_err(|source| CliAuthImportError::Read { - path: auth_path, - source, - })?; - parse_cli_auth_json(&raw) -} - -pub fn parse_cli_auth_json(raw: &str) -> Result { - let parsed = serde_json::from_str::(raw)?; - let access_token = - non_empty_secret(parsed.access_token).ok_or(CliAuthImportError::MissingAccessToken)?; - - Ok(CliAuth { - access_token, - refresh_token: non_empty_secret(parsed.refresh_token), - id_token: non_empty_secret(parsed.id_token), - expires_at: parsed.expires_at, - }) -} - -fn non_empty_secret(value: Option) -> Option { - value - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) - .map(|value| SecretString::new(value.into())) -} diff --git a/src/codex/gateway/oauth/mod.rs b/src/codex/gateway/oauth/mod.rs deleted file mode 100644 index 7d3d0ddab..000000000 --- a/src/codex/gateway/oauth/mod.rs +++ /dev/null @@ -1,9 +0,0 @@ -pub mod client; -pub mod codex_cli; -pub mod refresh; -pub mod token; - -pub use client::*; -pub use codex_cli::*; -pub use refresh::*; -pub use token::*; diff --git a/src/codex/gateway/oauth/refresh.rs b/src/codex/gateway/oauth/refresh.rs deleted file mode 100644 index 3bd8afa14..000000000 --- a/src/codex/gateway/oauth/refresh.rs +++ /dev/null @@ -1,142 +0,0 @@ -use std::sync::Arc; - -use async_trait::async_trait; -use chrono::{DateTime, Duration, Utc}; -use thiserror::Error; -use tokio::sync::Semaphore; - -use crate::codex::accounts::model::{Account, AccountStatus}; - -use super::token::TokenPair; - -#[derive(Debug, Clone, Copy)] -pub struct RefreshPolicy { - pub refresh_margin_seconds: u64, - pub refresh_concurrency: u32, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum RefreshTrigger { - BeforeExpiry, - Unauthorized, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] -pub enum RefreshFailure { - #[error("refresh token is invalid or expired")] - InvalidGrant, - #[error("account quota is exhausted")] - QuotaExhausted, - #[error("account is banned")] - Banned, - #[error("account is disabled")] - Disabled, - #[error("refresh transport failed")] - Transport, -} - -#[derive(Debug, Error)] -pub enum RefreshError { - #[error("refresh task semaphore closed")] - ConcurrencyClosed, - #[error("refresh transport failed")] - Transport, -} - -#[async_trait] -pub trait TokenRefresher: Send + Sync + 'static { - async fn refresh(&self, refresh_token: &str) -> Result; -} - -#[derive(Clone)] -pub struct RefreshScheduler { - policy: RefreshPolicy, - client: Arc, - semaphore: Arc, -} - -impl RefreshScheduler -where - C: TokenRefresher, -{ - pub fn new(policy: RefreshPolicy, client: C) -> Self { - let concurrency = policy.refresh_concurrency.max(1) as usize; - Self { - policy, - client: Arc::new(client), - semaphore: Arc::new(Semaphore::new(concurrency)), - } - } - - pub async fn refresh_account_at( - &self, - account: &Account, - trigger: RefreshTrigger, - now: DateTime, - ) -> Result { - if !self.should_refresh(account, trigger, now) { - return Ok(account.clone()); - } - let Some(refresh_token) = account.refresh_token.as_deref() else { - let mut expired = account.clone(); - expired.status = AccountStatus::Expired; - return Ok(expired); - }; - - let _permit = self - .semaphore - .acquire() - .await - .map_err(|_| RefreshError::ConcurrencyClosed)?; - // 刷新响应未返回 refresh_token 时必须保留旧值，避免把账号永久刷新能力清空。 - match self.client.refresh(refresh_token).await { - Ok(token_pair) => Ok(apply_token_pair(account, token_pair)), - Err(RefreshFailure::Transport) => Err(RefreshError::Transport), - Err(error) => Ok(apply_refresh_failure(account, error)), - } - } - - fn should_refresh( - &self, - account: &Account, - trigger: RefreshTrigger, - now: DateTime, - ) -> bool { - if account.status != AccountStatus::Active { - return false; - } - match trigger { - RefreshTrigger::Unauthorized => true, - RefreshTrigger::BeforeExpiry => account - .access_token_expires_at - .is_some_and(|expires_at| expires_at <= now + self.refresh_margin()), - } - } - - fn refresh_margin(&self) -> Duration { - let seconds = self.policy.refresh_margin_seconds.min(86_400 * 7) as i64; - Duration::seconds(seconds) - } -} - -fn apply_token_pair(account: &Account, token_pair: TokenPair) -> Account { - let mut refreshed = account.clone(); - refreshed.access_token = token_pair.access_token; - if let Some(refresh_token) = token_pair.refresh_token { - refreshed.refresh_token = Some(refresh_token); - } - refreshed.status = AccountStatus::Active; - refreshed -} - -fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Account { - let mut updated = account.clone(); - updated.status = match failure { - RefreshFailure::InvalidGrant => AccountStatus::Expired, - RefreshFailure::QuotaExhausted => AccountStatus::QuotaExhausted, - RefreshFailure::Banned => AccountStatus::Banned, - RefreshFailure::Disabled => AccountStatus::Disabled, - RefreshFailure::Transport => AccountStatus::Active, - }; - updated -} diff --git a/src/codex/gateway/oauth/token.rs b/src/codex/gateway/oauth/token.rs deleted file mode 100644 index a650e65e5..000000000 --- a/src/codex/gateway/oauth/token.rs +++ /dev/null @@ -1,5 +0,0 @@ -#[derive(Debug, Clone)] -pub struct TokenPair { - pub access_token: String, - pub refresh_token: Option, -} diff --git a/src/codex/gateway/protocol/codex_to_openai.rs b/src/codex/gateway/protocol/codex_to_openai.rs deleted file mode 100644 index ec0b70cf2..000000000 --- a/src/codex/gateway/protocol/codex_to_openai.rs +++ /dev/null @@ -1,567 +0,0 @@ -use std::collections::{BTreeMap, BTreeSet}; - -use chrono::Utc; -use serde_json::{json, Value}; -use uuid::Uuid; - -use crate::codex::gateway::{ - protocol::tuple_schema::reconvert_tuple_values, - transport::sse::{parse_sse_events, SseError}, -}; - -pub fn openai_error(message: &str, code: &str) -> Value { - json!({ - "error": { - "message": message, - "type": "server_error", - "param": null, - "code": code - } - }) -} - -pub fn chat_completion_from_codex_sse( - body: &str, - model: &str, - include_reasoning: bool, - tuple_schema: Option<&Value>, -) -> Result, SseError> { - let events = parse_sse_events(body)?; - let mut content = String::new(); - let mut reasoning = String::new(); - let mut function_call_items = BTreeMap::new(); - let mut finished_call_ids = BTreeSet::new(); - let mut tool_calls = Vec::new(); - let mut completed_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) { - 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.output_item.added") => { - if let Some((item_id, call_id, name)) = function_call_item(value.get("item")) { - function_call_items.insert(item_id, FunctionCallInfo { call_id, name }); - } - } - Some("response.function_call_arguments.done") => { - if let Some(tool_call) = - tool_call_from_done_event(&value, &function_call_items, &mut finished_call_ids) - { - tool_calls.push(tool_call); - } - } - Some("response.output_item.done") => { - if let Some(tool_call) = tool_call_from_output_item( - value.get("item"), - &function_call_items, - &mut finished_call_ids, - ) { - tool_calls.push(tool_call); - } - } - Some("response.completed") => { - completed_response = value.get("response").cloned(); - } - _ => {} - } - } - let Some(response) = completed_response else { - return Ok(None); - }; - if content.is_empty() { - content = output_text_from_response(&response); - } - if let Some(tuple_schema) = tuple_schema { - content = reconvert_tuple_text(&content, tuple_schema).unwrap_or(content); - } - let mut message = json!({ - "role": "assistant", - "content": if content.is_empty() && !tool_calls.is_empty() { - Value::Null - } else { - Value::String(content) - }, - }); - if include_reasoning && !reasoning.is_empty() { - message["reasoning_content"] = Value::String(reasoning); - } - let has_tool_calls = !tool_calls.is_empty(); - if has_tool_calls { - message["tool_calls"] = Value::Array(tool_calls); - } - - Ok(Some(json!({ - "id": format!("chatcmpl-{}", Uuid::new_v4().simple()), - "object": "chat.completion", - "created": Utc::now().timestamp(), - "model": model, - "choices": [{ - "index": 0, - "message": message, - "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" }, - }], - "usage": openai_usage(response.get("usage")), - }))) -} - -pub fn chat_completion_stream_from_codex_sse( - body: &str, - model: &str, - include_reasoning: bool, - tuple_schema: Option<&Value>, -) -> Result, SseError> { - let events = parse_sse_events(body)?; - let chunk_id = format!("chatcmpl-{}", Uuid::new_v4().simple()); - let created = Utc::now().timestamp(); - let mut output = String::new(); - let mut completed_response = None; - let mut has_tool_calls = false; - let mut function_call_items = BTreeMap::new(); - let mut tool_call_indices = BTreeMap::new(); - let mut call_ids_with_deltas = BTreeSet::new(); - let mut next_tool_call_index = 0usize; - let mut tuple_text_buffer = tuple_schema.map(|_| String::new()); - - push_sse_data( - &mut output, - json!({ - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [{ - "index": 0, - "delta": {"role": "assistant"}, - "finish_reason": null, - }], - }), - ); - - 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) { - if let Some(buffer) = tuple_text_buffer.as_mut() { - buffer.push_str(delta); - } else { - push_chat_content_delta(&mut output, &chunk_id, created, model, delta); - } - } - } - Some("response.reasoning_summary_text.delta") if include_reasoning => { - if let Some(delta) = value.get("delta").and_then(Value::as_str) { - push_sse_data( - &mut output, - json!({ - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [{ - "index": 0, - "delta": {"reasoning_content": delta}, - "finish_reason": null, - }], - }), - ); - } - } - Some("response.output_item.added") => { - if let Some((item_id, call_id, name)) = function_call_item(value.get("item")) { - has_tool_calls = true; - let index = next_tool_call_index; - next_tool_call_index = next_tool_call_index.saturating_add(1); - function_call_items.insert( - item_id, - FunctionCallInfo { - call_id: call_id.clone(), - name: name.clone(), - }, - ); - tool_call_indices.insert(call_id.clone(), index); - push_sse_data( - &mut output, - json!({ - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [{ - "index": 0, - "delta": { - "tool_calls": [{ - "index": index, - "id": call_id, - "type": "function", - "function": { - "name": name, - "arguments": "", - }, - }], - }, - "finish_reason": null, - }], - }), - ); - } - } - Some("response.function_call_arguments.delta") => { - if let Some((call_id, index, delta)) = - stream_tool_call_delta(&value, &function_call_items, &tool_call_indices) - { - has_tool_calls = true; - call_ids_with_deltas.insert(call_id); - push_sse_data( - &mut output, - json!({ - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [{ - "index": 0, - "delta": { - "tool_calls": [{ - "index": index, - "function": { - "arguments": delta, - }, - }], - }, - "finish_reason": null, - }], - }), - ); - } - } - Some("response.function_call_arguments.done") => { - if let Some((call_id, index, arguments)) = stream_tool_call_done( - &value, - &function_call_items, - &tool_call_indices, - &call_ids_with_deltas, - ) { - has_tool_calls = true; - push_sse_data( - &mut output, - json!({ - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [{ - "index": 0, - "delta": { - "tool_calls": [{ - "index": index, - "function": { - "arguments": arguments, - }, - }], - }, - "finish_reason": null, - }], - }), - ); - 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.completed") => { - completed_response = value.get("response").cloned(); - } - _ => {} - } - } - - let Some(response) = completed_response else { - return Ok(None); - }; - - if let (Some(tuple_schema), Some(buffer)) = (tuple_schema, tuple_text_buffer.as_ref()) { - if !buffer.is_empty() { - let content = - reconvert_tuple_text(buffer, tuple_schema).unwrap_or_else(|| buffer.clone()); - push_chat_content_delta(&mut output, &chunk_id, created, model, &content); - } - } - - push_sse_data( - &mut output, - json!({ - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [{ - "index": 0, - "delta": {}, - "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" }, - }], - "usage": openai_usage(response.get("usage")), - }), - ); - output.push_str("data: [DONE]\n\n"); - - Ok(Some(output)) -} - -fn push_chat_content_delta( - output: &mut String, - chunk_id: &str, - created: i64, - model: &str, - content: &str, -) { - push_sse_data( - output, - json!({ - "id": chunk_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [{ - "index": 0, - "delta": {"content": content}, - "finish_reason": null, - }], - }), - ); -} - -fn reconvert_tuple_text(text: &str, tuple_schema: &Value) -> Option { - let parsed = serde_json::from_str::(text).ok()?; - Some(reconvert_tuple_values(parsed, tuple_schema).to_string()) -} - -fn output_text_from_response(response: &Value) -> String { - response - .get("output_text") - .and_then(Value::as_str) - .map(ToString::to_string) - .unwrap_or_else(|| { - response - .get("output") - .and_then(Value::as_array) - .map(|items| { - items - .iter() - .filter_map(output_text_from_item) - .collect::>() - .join("\n\n") - }) - .unwrap_or_default() - }) -} - -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) -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct FunctionCallInfo { - call_id: String, - name: String, -} - -fn function_call_item(item: Option<&Value>) -> Option<(String, String, String)> { - let item = item?; - if item.get("type").and_then(Value::as_str) != Some("function_call") { - return None; - } - let item_id = item.get("id").and_then(Value::as_str)?; - let call_id = item.get("call_id").and_then(Value::as_str)?; - let name = item.get("name").and_then(Value::as_str)?; - Some((item_id.to_string(), call_id.to_string(), name.to_string())) -} - -fn tool_call_from_done_event( - event: &Value, - function_call_items: &BTreeMap, - finished_call_ids: &mut BTreeSet, -) -> Option { - let event_id = event - .get("call_id") - .or_else(|| event.get("item_id")) - .and_then(Value::as_str)?; - let info = function_call_items.get(event_id); - let call_id = info - .map(|info| info.call_id.as_str()) - .unwrap_or(event_id) - .to_string(); - if !finished_call_ids.insert(call_id.clone()) { - return None; - } - let name = event - .get("name") - .and_then(Value::as_str) - .filter(|name| !name.is_empty()) - .or_else(|| info.map(|info| info.name.as_str())) - .unwrap_or("unknown"); - let arguments = event - .get("arguments") - .and_then(Value::as_str) - .unwrap_or_default(); - Some(openai_tool_call(call_id, name, arguments)) -} - -fn tool_call_from_output_item( - item: Option<&Value>, - function_call_items: &BTreeMap, - finished_call_ids: &mut BTreeSet, -) -> Option { - let item = item?; - if item.get("type").and_then(Value::as_str) != Some("function_call") { - return None; - } - let item_id = item.get("id").and_then(Value::as_str); - let info = item_id.and_then(|item_id| function_call_items.get(item_id)); - let call_id = item - .get("call_id") - .and_then(Value::as_str) - .or_else(|| info.map(|info| info.call_id.as_str())) - .or(item_id)?; - if !finished_call_ids.insert(call_id.to_string()) { - return None; - } - let name = item - .get("name") - .and_then(Value::as_str) - .filter(|name| !name.is_empty()) - .or_else(|| info.map(|info| info.name.as_str())) - .unwrap_or("unknown"); - let arguments = item - .get("arguments") - .and_then(Value::as_str) - .unwrap_or_default(); - Some(openai_tool_call(call_id.to_string(), name, arguments)) -} - -fn openai_tool_call(call_id: String, name: &str, arguments: &str) -> Value { - json!({ - "id": call_id, - "type": "function", - "function": { - "name": name, - "arguments": arguments, - }, - }) -} - -fn stream_tool_call_delta( - event: &Value, - function_call_items: &BTreeMap, - tool_call_indices: &BTreeMap, -) -> Option<(String, usize, String)> { - let event_id = event - .get("call_id") - .or_else(|| event.get("item_id")) - .and_then(Value::as_str)?; - let info = function_call_items.get(event_id); - let call_id = info - .map(|info| info.call_id.as_str()) - .unwrap_or(event_id) - .to_string(); - let index = *tool_call_indices.get(&call_id).unwrap_or(&0); - let delta = event.get("delta").and_then(Value::as_str)?.to_string(); - Some((call_id, index, delta)) -} - -fn stream_tool_call_done( - event: &Value, - function_call_items: &BTreeMap, - tool_call_indices: &BTreeMap, - call_ids_with_deltas: &BTreeSet, -) -> Option<(String, usize, String)> { - let event_id = event - .get("call_id") - .or_else(|| event.get("item_id")) - .and_then(Value::as_str)?; - let info = function_call_items.get(event_id); - let call_id = info - .map(|info| info.call_id.as_str()) - .unwrap_or(event_id) - .to_string(); - if call_ids_with_deltas.contains(&call_id) { - return None; - } - let index = *tool_call_indices.get(&call_id).unwrap_or(&0); - let arguments = event - .get("arguments") - .and_then(Value::as_str) - .unwrap_or_default() - .to_string(); - Some((call_id, index, arguments)) -} - -fn push_sse_data(output: &mut String, value: Value) { - output.push_str("data: "); - output.push_str(&value.to_string()); - output.push_str("\n\n"); -} - -fn is_function_call_item(item: Option<&Value>) -> bool { - item.and_then(|item| item.get("type")) - .and_then(Value::as_str) - == Some("function_call") -} - -fn openai_usage(usage: Option<&Value>) -> Value { - let prompt_tokens = usage - .and_then(|usage| usage.get("input_tokens")) - .and_then(Value::as_u64) - .unwrap_or_default(); - let completion_tokens = usage - .and_then(|usage| usage.get("output_tokens")) - .and_then(Value::as_u64) - .unwrap_or_default(); - let cached_tokens = usage - .and_then(|usage| usage.get("input_tokens_details")) - .and_then(|details| details.get("cached_tokens")) - .and_then(Value::as_u64) - .unwrap_or_default(); - let reasoning_tokens = usage - .and_then(|usage| usage.get("output_tokens_details")) - .and_then(|details| details.get("reasoning_tokens")) - .and_then(Value::as_u64); - let mut openai_usage = json!({ - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "total_tokens": prompt_tokens.saturating_add(completion_tokens), - "prompt_tokens_details": { - "cached_tokens": cached_tokens, - }, - }); - if let Some(reasoning_tokens) = reasoning_tokens { - openai_usage["completion_tokens_details"] = json!({ - "reasoning_tokens": reasoning_tokens, - }); - } - openai_usage -} diff --git a/src/codex/gateway/protocol/error.rs b/src/codex/gateway/protocol/error.rs deleted file mode 100644 index 8034c4a17..000000000 --- a/src/codex/gateway/protocol/error.rs +++ /dev/null @@ -1,70 +0,0 @@ -use axum::{ - http::StatusCode, - response::{IntoResponse, Response}, - Json, -}; -use serde::Serialize; -use thiserror::Error; - -pub type AppResult = Result; - -#[derive(Debug, Error)] -pub enum AppError { - #[error("configuration error: {0}")] - Config(String), - #[error("database error: {0}")] - Database(#[from] sqlx::Error), - #[error("upstream error: {0}")] - Upstream(String), - #[error("unauthorized")] - Unauthorized, - #[error("forbidden")] - Forbidden, - #[error("bad request: {0}")] - BadRequest(String), -} - -#[derive(Serialize)] -struct ErrorBody<'a> { - error: ErrorMessage<'a>, -} - -#[derive(Serialize)] -struct ErrorMessage<'a> { - message: &'a str, - #[serde(rename = "type")] - kind: &'a str, - code: &'a str, -} - -impl IntoResponse for AppError { - fn into_response(self) -> Response { - let (status, code, kind) = match &self { - Self::BadRequest(_) | Self::Config(_) => ( - StatusCode::BAD_REQUEST, - "bad_request", - "invalid_request_error", - ), - Self::Unauthorized => ( - StatusCode::UNAUTHORIZED, - "unauthorized", - "invalid_request_error", - ), - Self::Forbidden => (StatusCode::FORBIDDEN, "forbidden", "invalid_request_error"), - Self::Database(_) | Self::Upstream(_) => { - (StatusCode::BAD_GATEWAY, "upstream_error", "server_error") - } - }; - - let message = self.to_string(); - let body = ErrorBody { - error: ErrorMessage { - message: &message, - kind, - code, - }, - }; - - (status, Json(body)).into_response() - } -} diff --git a/src/codex/gateway/protocol/mod.rs b/src/codex/gateway/protocol/mod.rs deleted file mode 100644 index f1c8f873c..000000000 --- a/src/codex/gateway/protocol/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub mod codex_to_openai; -pub mod error; -pub mod openai_to_codex; -pub mod schema; -pub(crate) mod tuple_schema; diff --git a/src/codex/gateway/protocol/openai_to_codex.rs b/src/codex/gateway/protocol/openai_to_codex.rs deleted file mode 100644 index 7be38786b..000000000 --- a/src/codex/gateway/protocol/openai_to_codex.rs +++ /dev/null @@ -1,311 +0,0 @@ -use serde::{Deserialize, Serialize}; -use serde_json::{json, Value}; - -use crate::{ - codex::gateway::protocol::error::{AppError, AppResult}, - codex::gateway::protocol::tuple_schema::prepare_schema, - codex::gateway::transport::types::CodexResponsesRequest, -}; - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ChatCompletionRequest { - pub model: String, - #[serde(default)] - pub stream: bool, - pub messages: Vec, - #[serde(default)] - pub reasoning_effort: Option, - #[serde(default)] - pub service_tier: Option, - #[serde(default)] - pub tools: Option>, - #[serde(default)] - pub tool_choice: Option, - #[serde(default)] - pub parallel_tool_calls: Option, - #[serde(default)] - pub functions: Option>, - #[serde(default)] - pub response_format: Option, - #[serde(default)] - pub user: Option, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ChatMessage { - pub role: String, - #[serde(default)] - pub content: Option, - #[serde(default)] - pub name: Option, - #[serde(default)] - pub tool_calls: Option>, - #[serde(default)] - pub tool_call_id: Option, - #[serde(default)] - pub function_call: Option, -} - -pub fn translate_chat_to_codex(req: ChatCompletionRequest) -> AppResult { - if req.messages.is_empty() { - return Err(AppError::BadRequest( - "messages must not be empty".to_string(), - )); - } - - let client_conversation_id = req - .user - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToString::to_string); - let instructions = chat_instructions(&req.messages); - let mut input = chat_input(req.messages); - if input.is_empty() { - input.push(json!({"role": "user", "content": ""})); - } - - let mut request = CodexResponsesRequest::new_http_sse(req.model, instructions, input); - request.force_http_sse = true; - request.tools = codex_tools(req.tools, req.functions); - request.tool_choice = req.tool_choice; - request.parallel_tool_calls = req.parallel_tool_calls; - let response_format = response_format_text(req.response_format); - request.text = response_format.text; - request.tuple_schema = response_format.tuple_schema; - request.service_tier = req.service_tier; - request.prompt_cache_key = client_conversation_id.clone(); - request.client_conversation_id = client_conversation_id; - if let Some(effort) = req.reasoning_effort { - request.reasoning = Some(json!({"effort": effort, "summary": "auto"})); - } - Ok(request) -} - -fn chat_instructions(messages: &[ChatMessage]) -> String { - let instructions = messages - .iter() - .filter(|message| message.role == "system" || message.role == "developer") - .map(|message| extract_text(&message.content)) - .filter(|text| !text.is_empty()) - .collect::>() - .join("\n\n"); - if instructions.is_empty() { - "You are a helpful assistant.".to_string() - } else { - instructions - } -} - -fn chat_input(messages: Vec) -> Vec { - let mut input = Vec::new(); - for message in messages { - match message.role.as_str() { - "system" | "developer" => {} - "assistant" => push_assistant_message(&mut input, message), - "tool" => input.push(json!({ - "type": "function_call_output", - "call_id": message.tool_call_id.unwrap_or_else(|| "unknown".to_string()), - "output": extract_text(&message.content), - })), - "function" => input.push(json!({ - "type": "function_call_output", - "call_id": format!("fc_{}", message.name.unwrap_or_else(|| "unknown".to_string())), - "output": extract_text(&message.content), - })), - _ => input.push(json!({ - "role": "user", - "content": extract_content(&message.content), - })), - } - } - input -} - -fn push_assistant_message(input: &mut Vec, message: ChatMessage) { - let text = extract_text(&message.content); - let has_tool_calls = message - .tool_calls - .as_ref() - .is_some_and(|calls| !calls.is_empty()); - if !text.is_empty() || (!has_tool_calls && message.function_call.is_none()) { - input.push(json!({"role": "assistant", "content": text})); - } - - if let Some(tool_calls) = message.tool_calls { - for tool_call in tool_calls { - let call_id = tool_call - .get("id") - .and_then(Value::as_str) - .unwrap_or("unknown"); - let function = tool_call.get("function").unwrap_or(&Value::Null); - let name = function - .get("name") - .and_then(Value::as_str) - .unwrap_or("unknown"); - let arguments = function - .get("arguments") - .and_then(Value::as_str) - .unwrap_or(""); - input.push(json!({ - "type": "function_call", - "call_id": call_id, - "name": name, - "arguments": arguments, - })); - } - } - - if let Some(function_call) = message.function_call { - let name = function_call - .get("name") - .and_then(Value::as_str) - .unwrap_or("unknown"); - let arguments = function_call - .get("arguments") - .and_then(Value::as_str) - .unwrap_or(""); - input.push(json!({ - "type": "function_call", - "call_id": format!("fc_{name}"), - "name": name, - "arguments": arguments, - })); - } -} - -fn extract_text(content: &Option) -> String { - let Some(content) = content else { - return String::new(); - }; - if let Some(text) = content.as_str() { - return text.to_string(); - } - let Some(parts) = content.as_array() else { - return String::new(); - }; - parts - .iter() - .filter_map(|part| { - (part.get("type").and_then(Value::as_str) == Some("text")) - .then(|| part.get("text").and_then(Value::as_str)) - .flatten() - }) - .collect::>() - .join("\n") -} - -fn extract_content(content: &Option) -> Value { - let Some(content) = content else { - return Value::String(String::new()); - }; - if content.as_str().is_some() { - return content.clone(); - } - let Some(parts) = content.as_array() else { - return Value::String(String::new()); - }; - let has_image = parts - .iter() - .any(|part| part.get("type").and_then(Value::as_str) == Some("image_url")); - if !has_image { - return Value::String(extract_text(&Some(content.clone()))); - } - let codex_parts = parts - .iter() - .filter_map(codex_content_part) - .collect::>(); - if codex_parts.is_empty() { - Value::String(String::new()) - } else { - Value::Array(codex_parts) - } -} - -fn codex_content_part(part: &Value) -> Option { - match part.get("type").and_then(Value::as_str)? { - "text" => part - .get("text") - .and_then(Value::as_str) - .map(|text| json!({"type": "input_text", "text": text})), - "image_url" => image_url(part).map(|url| json!({"type": "input_image", "image_url": url})), - _ => None, - } -} - -fn image_url(part: &Value) -> Option<&str> { - let image = part.get("image_url")?; - image - .as_str() - .or_else(|| image.get("url").and_then(Value::as_str)) -} - -fn codex_tools(tools: Option>, functions: Option>) -> Option> { - if let Some(tools) = tools.filter(|tools| !tools.is_empty()) { - return Some(tools); - } - functions - .filter(|functions| !functions.is_empty()) - .map(|functions| { - functions - .into_iter() - .map(|function| json!({"type": "function", "function": function})) - .collect() - }) -} - -struct PreparedResponseFormat { - text: Option, - tuple_schema: Option, -} - -fn response_format_text(response_format: Option) -> PreparedResponseFormat { - let Some(format) = response_format else { - return PreparedResponseFormat { - text: None, - tuple_schema: None, - }; - }; - let Some(kind) = format.get("type").and_then(Value::as_str) else { - return PreparedResponseFormat { - text: None, - tuple_schema: None, - }; - }; - match kind { - "json_object" => PreparedResponseFormat { - text: Some(json!({"format": {"type": "json_object"}})), - tuple_schema: None, - }, - "json_schema" => { - let Some(schema) = format.get("json_schema") else { - return PreparedResponseFormat { - text: None, - tuple_schema: None, - }; - }; - let name = schema - .get("name") - .and_then(Value::as_str) - .unwrap_or("response"); - let prepared_schema = - prepare_schema(schema.get("schema").cloned().unwrap_or_else(|| json!({}))); - let mut codex_format = json!({ - "type": "json_schema", - "name": name, - "schema": prepared_schema.schema, - }); - if let Some(strict) = schema.get("strict").and_then(Value::as_bool) { - codex_format["strict"] = Value::Bool(strict); - } - PreparedResponseFormat { - text: Some(json!({"format": codex_format})), - tuple_schema: prepared_schema.original_schema, - } - } - _ => PreparedResponseFormat { - text: None, - tuple_schema: None, - }, - } -} diff --git a/src/codex/gateway/protocol/schema.rs b/src/codex/gateway/protocol/schema.rs deleted file mode 100644 index 7767e0ada..000000000 --- a/src/codex/gateway/protocol/schema.rs +++ /dev/null @@ -1,6 +0,0 @@ -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ResponseFormat { - Text, - JsonObject, - JsonSchema, -} diff --git a/src/codex/gateway/transport/custom_ca.rs b/src/codex/gateway/transport/custom_ca.rs deleted file mode 100644 index f7fc5fc78..000000000 --- a/src/codex/gateway/transport/custom_ca.rs +++ /dev/null @@ -1,361 +0,0 @@ -use std::{env, fs, io, path::PathBuf, sync::Arc}; - -use rustls::{ClientConfig, RootCertStore}; -use rustls_pki_types::{ - pem::{self, PemObject, SectionKind}, - CertificateDer, -}; -use thiserror::Error; - -pub const CODEX_CA_CERT_ENV: &str = "CODEX_CA_CERTIFICATE"; -pub const SSL_CERT_FILE_ENV: &str = "SSL_CERT_FILE"; - -const CA_CERT_HINT: &str = "If you set CODEX_CA_CERTIFICATE or SSL_CERT_FILE, ensure it points to a PEM file containing one or more CERTIFICATE blocks, or unset it to use system roots."; - -type PemSection = (SectionKind, Vec); - -#[derive(Debug, Error)] -pub enum CustomCaError { - #[error( - "Failed to read CA certificate file {} selected by {}: {source}. {hint}", - path.display(), - source_env, - hint = CA_CERT_HINT - )] - ReadCaFile { - source_env: &'static str, - path: PathBuf, - source: io::Error, - }, - #[error( - "Failed to load CA certificates from {} selected by {}: {detail}. {hint}", - path.display(), - source_env, - hint = CA_CERT_HINT - )] - InvalidCaFile { - source_env: &'static str, - path: PathBuf, - detail: String, - }, - #[error( - "Failed to parse certificate #{certificate_index} from {} selected by {}: {source}. {hint}", - path.display(), - source_env, - hint = CA_CERT_HINT - )] - RegisterCertificate { - source_env: &'static str, - path: PathBuf, - certificate_index: usize, - source: reqwest::Error, - }, - #[error( - "Failed to register certificate #{certificate_index} from {} selected by {} in rustls root store: {source}. {hint}", - path.display(), - source_env, - hint = CA_CERT_HINT - )] - RegisterRustlsCertificate { - source_env: &'static str, - path: PathBuf, - certificate_index: usize, - source: rustls::Error, - }, - #[error("Failed to build HTTP client while using CA bundle from {} ({}): {source}", source_env, path.display())] - BuildClientWithCustomCa { - source_env: &'static str, - path: PathBuf, - source: reqwest::Error, - }, - #[error("Failed to build HTTP client while using system root certificates: {0}")] - BuildClientWithSystemRoots(reqwest::Error), - #[error("Failed to load native root certificates for custom CA transport: {0}")] - LoadNativeRoots(io::Error), -} - -impl From for io::Error { - fn from(error: CustomCaError) -> Self { - match error { - CustomCaError::ReadCaFile { ref source, .. } => io::Error::new(source.kind(), error), - CustomCaError::InvalidCaFile { .. } - | CustomCaError::RegisterCertificate { .. } - | CustomCaError::RegisterRustlsCertificate { .. } => { - io::Error::new(io::ErrorKind::InvalidData, error) - } - CustomCaError::BuildClientWithCustomCa { .. } - | CustomCaError::BuildClientWithSystemRoots(_) - | CustomCaError::LoadNativeRoots(_) => io::Error::other(error), - } - } -} - -pub fn build_reqwest_client_with_custom_ca( - builder: reqwest::ClientBuilder, -) -> Result { - build_reqwest_client_with_env(&ProcessEnv, builder) -} - -pub fn custom_ca_env_cache_key() -> Option { - ProcessEnv - .configured_ca_bundle() - .map(|bundle| format!("{}={}", bundle.source_env, bundle.path.display())) -} - -pub fn maybe_build_rustls_client_config_with_custom_ca( -) -> Result>, CustomCaError> { - maybe_build_rustls_client_config_with_env(&ProcessEnv) -} - -fn build_reqwest_client_with_env( - env_source: &dyn EnvSource, - mut builder: reqwest::ClientBuilder, -) -> Result { - let Some(bundle) = env_source.configured_ca_bundle() else { - return builder - .build() - .map_err(CustomCaError::BuildClientWithSystemRoots); - }; - - builder = builder.use_rustls_tls(); - for (idx, cert) in bundle.load_certificates()?.iter().enumerate() { - let certificate = reqwest::Certificate::from_der(cert.as_ref()).map_err(|source| { - CustomCaError::RegisterCertificate { - source_env: bundle.source_env, - path: bundle.path.clone(), - certificate_index: idx + 1, - source, - } - })?; - builder = builder.add_root_certificate(certificate); - } - builder - .build() - .map_err(|source| CustomCaError::BuildClientWithCustomCa { - source_env: bundle.source_env, - path: bundle.path, - source, - }) -} - -fn maybe_build_rustls_client_config_with_env( - env_source: &dyn EnvSource, -) -> Result>, CustomCaError> { - let Some(bundle) = env_source.configured_ca_bundle() else { - return Ok(None); - }; - - let mut root_store = native_root_store().map_err(CustomCaError::LoadNativeRoots)?; - for (idx, cert) in bundle.load_certificates()?.into_iter().enumerate() { - root_store - .add(cert) - .map_err(|source| CustomCaError::RegisterRustlsCertificate { - source_env: bundle.source_env, - path: bundle.path.clone(), - certificate_index: idx + 1, - source, - })?; - } - - Ok(Some(Arc::new( - ClientConfig::builder() - .with_root_certificates(root_store) - .with_no_client_auth(), - ))) -} - -pub fn native_root_store() -> Result { - let mut root_store = RootCertStore::empty(); - let rustls_native_certs::CertificateResult { certs, errors, .. } = - rustls_native_certs::load_native_certs(); - if !errors.is_empty() { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("failed to load native root certificates: {errors:?}"), - )); - } - - let (added, _) = root_store.add_parsable_certificates(certs); - if added == 0 { - return Err(io::Error::new( - io::ErrorKind::NotFound, - "no native root certificates found", - )); - } - - Ok(root_store) -} - -trait EnvSource { - fn var(&self, key: &str) -> Option; - - fn non_empty_path(&self, key: &str) -> Option { - self.var(key) - .filter(|value| !value.is_empty()) - .map(PathBuf::from) - } - - fn configured_ca_bundle(&self) -> Option { - self.non_empty_path(CODEX_CA_CERT_ENV) - .map(|path| ConfiguredCaBundle { - source_env: CODEX_CA_CERT_ENV, - path, - }) - .or_else(|| { - self.non_empty_path(SSL_CERT_FILE_ENV) - .map(|path| ConfiguredCaBundle { - source_env: SSL_CERT_FILE_ENV, - path, - }) - }) - } -} - -struct ProcessEnv; - -impl EnvSource for ProcessEnv { - fn var(&self, key: &str) -> Option { - env::var(key).ok() - } -} - -struct ConfiguredCaBundle { - source_env: &'static str, - path: PathBuf, -} - -impl ConfiguredCaBundle { - fn load_certificates(&self) -> Result>, CustomCaError> { - let pem_data = fs::read(&self.path).map_err(|source| CustomCaError::ReadCaFile { - source_env: self.source_env, - path: self.path.clone(), - source, - })?; - let normalized = normalize_trusted_certificate_labels(&pem_data); - let mut certificates = Vec::new(); - for section in PemSection::pem_slice_iter(normalized.as_bytes()) { - let (kind, der) = section.map_err(|error| self.pem_parse_error(&error))?; - if kind == SectionKind::Certificate { - certificates.push(CertificateDer::from(der)); - } - } - if certificates.is_empty() { - return Err(self.pem_parse_error(&pem::Error::NoItemsFound)); - } - Ok(certificates) - } - - fn pem_parse_error(&self, error: &pem::Error) -> CustomCaError { - let detail = match error { - pem::Error::NoItemsFound => "no certificates found in PEM file".to_string(), - _ => format!("failed to parse PEM file: {error}"), - }; - CustomCaError::InvalidCaFile { - source_env: self.source_env, - path: self.path.clone(), - detail, - } - } -} - -fn normalize_trusted_certificate_labels(pem_data: &[u8]) -> String { - String::from_utf8_lossy(pem_data) - .replace("BEGIN TRUSTED CERTIFICATE", "BEGIN CERTIFICATE") - .replace("END TRUSTED CERTIFICATE", "END CERTIFICATE") -} - -#[cfg(test)] -mod tests { - use std::{collections::HashMap, fs, path::PathBuf}; - - use tempfile::TempDir; - - use super::*; - - const TEST_CERT: &str = include_str!("../../../../tests/fixtures/test-ca.pem"); - - struct MapEnv { - values: HashMap, - } - - impl EnvSource for MapEnv { - fn var(&self, key: &str) -> Option { - self.values.get(key).cloned() - } - } - - fn map_env(pairs: &[(&str, &str)]) -> MapEnv { - MapEnv { - values: pairs - .iter() - .map(|(key, value)| ((*key).to_string(), (*value).to_string())) - .collect(), - } - } - - fn write_cert_file(temp_dir: &TempDir, name: &str, contents: &str) -> PathBuf { - let path = temp_dir.path().join(name); - fs::write(&path, contents).unwrap(); - path - } - - #[test] - fn ca_path_prefers_codex_env() { - let env = map_env(&[ - (CODEX_CA_CERT_ENV, "/tmp/codex.pem"), - (SSL_CERT_FILE_ENV, "/tmp/fallback.pem"), - ]); - - assert_eq!( - env.configured_ca_bundle().map(|bundle| bundle.path), - Some(PathBuf::from("/tmp/codex.pem")) - ); - } - - #[test] - fn ca_path_falls_back_to_ssl_cert_file() { - let env = map_env(&[(SSL_CERT_FILE_ENV, "/tmp/fallback.pem")]); - - assert_eq!( - env.configured_ca_bundle().map(|bundle| bundle.path), - Some(PathBuf::from("/tmp/fallback.pem")) - ); - } - - #[test] - fn ca_path_ignores_empty_values() { - let env = map_env(&[ - (CODEX_CA_CERT_ENV, ""), - (SSL_CERT_FILE_ENV, "/tmp/fallback.pem"), - ]); - - assert_eq!( - env.configured_ca_bundle().map(|bundle| bundle.path), - Some(PathBuf::from("/tmp/fallback.pem")) - ); - } - - #[test] - fn rustls_config_uses_custom_ca_bundle_when_configured() { - let temp_dir = TempDir::new().unwrap(); - let cert_path = write_cert_file(&temp_dir, "ca.pem", TEST_CERT); - let env = map_env(&[(CODEX_CA_CERT_ENV, cert_path.to_string_lossy().as_ref())]); - - let config = maybe_build_rustls_client_config_with_env(&env) - .expect("rustls config") - .expect("custom CA config should be present"); - - assert!(config.enable_sni); - } - - #[test] - fn rustls_config_reports_invalid_ca_file() { - let temp_dir = TempDir::new().unwrap(); - let cert_path = write_cert_file(&temp_dir, "empty.pem", ""); - let env = map_env(&[(CODEX_CA_CERT_ENV, cert_path.to_string_lossy().as_ref())]); - - let error = maybe_build_rustls_client_config_with_env(&env).unwrap_err(); - - assert!(matches!(error, CustomCaError::InvalidCaFile { .. })); - } -} diff --git a/src/codex/gateway/transport/endpoints.rs b/src/codex/gateway/transport/endpoints.rs deleted file mode 100644 index f77c71154..000000000 --- a/src/codex/gateway/transport/endpoints.rs +++ /dev/null @@ -1,61 +0,0 @@ -use reqwest::Url; - -pub(crate) const CODEX_RESPONSES_PATH: &str = "/codex/responses"; -pub(crate) const CODEX_RESPONSES_COMPACT_PATH: &str = "/codex/responses/compact"; -pub(crate) const CODEX_USAGE_PATH: &str = "/codex/usage"; -pub(crate) const CODEX_USAGE_API_PATH: &str = "/api/codex/usage"; -pub(crate) const WHAM_USAGE_PATH: &str = "/wham/usage"; - -pub(crate) fn endpoint_url(base_url: &str, endpoint_path: &str) -> String { - format!( - "{}/{}", - base_url.trim_end_matches('/'), - endpoint_path.trim_start_matches('/') - ) -} - -pub(crate) fn endpoint_request_path(base_url: &str, endpoint_path: &str) -> String { - let endpoint_path = endpoint_path.trim_start_matches('/'); - let base_path = Url::parse(base_url) - .ok() - .map(|url| url.path().trim_end_matches('/').to_string()) - .filter(|path| !path.is_empty()) - .unwrap_or_default(); - - if base_path.is_empty() { - format!("/{endpoint_path}") - } else { - format!("{base_path}/{endpoint_path}") - } -} - -pub(crate) fn usage_endpoint_urls(base_url: &str) -> Vec { - usage_endpoint_paths(base_url) - .into_iter() - .map(|path| endpoint_url(base_url, path)) - .collect() -} - -pub(crate) fn primary_usage_request_path(base_url: &str) -> String { - let endpoint_path = usage_endpoint_paths(base_url) - .into_iter() - .next() - .unwrap_or(CODEX_USAGE_API_PATH); - endpoint_request_path(base_url, endpoint_path) -} - -fn usage_endpoint_paths(base_url: &str) -> Vec<&'static str> { - if has_backend_api_base_path(base_url) { - vec![WHAM_USAGE_PATH, CODEX_USAGE_PATH] - } else { - vec![CODEX_USAGE_API_PATH, CODEX_USAGE_PATH] - } -} - -fn has_backend_api_base_path(base_url: &str) -> bool { - Url::parse(base_url).ok().is_some_and(|url| { - url.path() - .split('/') - .any(|segment| segment == "backend-api") - }) -} diff --git a/src/codex/gateway/transport/headers.rs b/src/codex/gateway/transport/headers.rs deleted file mode 100644 index b27916ffc..000000000 --- a/src/codex/gateway/transport/headers.rs +++ /dev/null @@ -1,92 +0,0 @@ -use indexmap::IndexMap; - -use crate::codex::gateway::fingerprint::model::Fingerprint; - -pub fn build_codex_base_headers( - fp: &Fingerprint, - access_token: &str, - account_id: Option<&str>, -) -> IndexMap { - let mut headers = IndexMap::new(); - - headers.insert( - "authorization".to_string(), - format!("Bearer {access_token}"), - ); - if let Some(id) = account_id { - headers.insert("chatgpt-account-id".to_string(), id.to_string()); - } - headers.insert("originator".to_string(), fp.originator.clone()); - headers.insert("user-agent".to_string(), fp.user_agent()); - headers.insert("sec-ch-ua".to_string(), fp.sec_ch_ua()); - - for (key, value) in &fp.default_headers { - let key_lower = key.to_lowercase(); - if !headers.contains_key(&key_lower) { - headers.insert(key_lower, value.clone()); - } - } - - headers -} - -pub fn build_codex_headers( - fp: &Fingerprint, - access_token: &str, - account_id: Option<&str>, - turn_state: Option<&str>, - request_id: &str, -) -> IndexMap { - let mut headers = build_codex_base_headers(fp, access_token, account_id); - headers.insert( - "x-openai-internal-codex-residency".to_string(), - "us".to_string(), - ); - headers.insert("x-client-request-id".to_string(), request_id.to_string()); - - if let Some(state) = turn_state { - headers.insert("x-codex-turn-state".to_string(), state.to_string()); - } - - headers.insert("accept".to_string(), "text/event-stream".to_string()); - - headers -} - -fn order_headers(headers: IndexMap, order: &[String]) -> IndexMap { - let mut ordered = IndexMap::new(); - - for key in order { - if let Some(value) = headers.get(key) { - ordered.insert(key.clone(), value.clone()); - } - } - - for (key, value) in headers { - if !ordered.contains_key(&key) { - ordered.insert(key, value); - } - } - - ordered -} - -pub fn build_ordered_codex_headers( - fp: &Fingerprint, - access_token: &str, - account_id: Option<&str>, - turn_state: Option<&str>, - request_id: &str, -) -> IndexMap { - let headers = build_codex_headers(fp, access_token, account_id, turn_state, request_id); - order_headers(headers, &fp.header_order) -} - -pub fn build_ordered_codex_base_headers( - fp: &Fingerprint, - access_token: &str, - account_id: Option<&str>, -) -> IndexMap { - let headers = build_codex_base_headers(fp, access_token, account_id); - order_headers(headers, &fp.header_order) -} diff --git a/src/codex/gateway/transport/http_client.rs b/src/codex/gateway/transport/http_client.rs deleted file mode 100644 index f2c6a25e9..000000000 --- a/src/codex/gateway/transport/http_client.rs +++ /dev/null @@ -1,960 +0,0 @@ -use std::{ - collections::HashMap, - sync::{Arc, Mutex, OnceLock}, - time::Duration, -}; - -use futures::StreamExt; -use indexmap::IndexMap; -use reqwest::{ - header::{ - HeaderMap, HeaderName, HeaderValue, ACCEPT, ACCEPT_ENCODING, CONTENT_TYPE, COOKIE, - RETRY_AFTER, SET_COOKIE, - }, - Client, Response as ReqwestResponse, StatusCode, -}; -use serde_json::{Map, Value}; -use thiserror::Error; -use uuid::Uuid; - -use crate::{ - codex::gateway::fingerprint::model::Fingerprint, - codex::gateway::transport::{ - custom_ca::{build_reqwest_client_with_custom_ca, custom_ca_env_cache_key, CustomCaError}, - endpoints::{ - endpoint_url, usage_endpoint_urls, CODEX_RESPONSES_COMPACT_PATH, CODEX_RESPONSES_PATH, - }, - headers::build_ordered_codex_base_headers, - retry_after::retry_after_seconds_from_body, - sse::SseError, - types::{CodexCompactRequest, CodexResponsesRequest}, - usage_events::{extract_sse_usage, TokenUsage}, - websocket::{ - append_rate_limit_updates, create_response_via_websocket, - create_response_via_websocket_stream, create_response_via_websocket_stream_with_pool, - latest_turn_state, transport_for_request, CodexTransport, CodexWebSocketError, - CodexWebSocketPool, CodexWebSocketPoolKey, CodexWebSocketSseStream, - CodexWebSocketStreamResponse, SharedRateLimitUpdates, SharedTurnState, - WebSocketSupportError, - }, - }, - codex::models::catalog::BackendModelEntry, -}; - -#[derive(Debug, Error)] -pub enum CodexClientError { - #[error("http transport error: {0}")] - Http(#[from] reqwest::Error), - #[error("custom CA transport error: {0}")] - CustomCa(#[from] CustomCaError), - #[error("invalid request header name: {0}")] - InvalidHeaderName(#[from] reqwest::header::InvalidHeaderName), - #[error("invalid request header value: {0}")] - InvalidHeaderValue(#[from] reqwest::header::InvalidHeaderValue), - #[error("unsupported transport: {0}")] - UnsupportedTransport(#[from] WebSocketSupportError), - #[error("websocket transport error: {0}")] - WebSocket(#[from] CodexWebSocketError), - #[error("invalid upstream SSE response: {0}")] - InvalidSse(#[from] SseError), - #[error("upstream returned status {status}: {body}")] - Upstream { - status: StatusCode, - body: String, - retry_after_seconds: Option, - }, - #[error("backend model catalog is unavailable")] - ModelsUnavailable, -} - -pub type CodexClientResult = Result; - -#[derive(Debug, Clone, Copy)] -pub struct CodexRequestContext<'a> { - pub access_token: &'a str, - pub account_id: Option<&'a str>, - pub request_id: &'a str, - pub turn_state: Option<&'a str>, - pub turn_metadata: Option<&'a str>, - pub beta_features: Option<&'a str>, - pub include_timing_metrics: Option<&'a str>, - pub version: Option<&'a str>, - pub codex_window_id: Option<&'a str>, - pub parent_thread_id: Option<&'a str>, - pub cookie_header: Option<&'a str>, - pub installation_id: Option<&'a str>, - pub session_id: Option<&'a str>, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexBackendResponse { - pub body: String, - pub usage: Option, - pub turn_state: Option, - pub set_cookie_headers: Vec, - pub rate_limit_headers: Vec<(String, String)>, -} - -#[derive(Debug, Clone, PartialEq)] -pub struct CodexCompactResponse { - pub body: Value, - pub set_cookie_headers: Vec, - pub rate_limit_headers: Vec<(String, String)>, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexConnectivityProbe { - pub endpoint: String, - pub status: StatusCode, -} - -pub struct CodexBackendStream { - pub response: ReqwestResponse, - pub turn_state: Option, - pub set_cookie_headers: Vec, - pub rate_limit_headers: Vec<(String, String)>, -} - -pub struct CodexBackendWebSocketStream { - pub body_stream: CodexWebSocketSseStream, - pub turn_state: Option, - pub set_cookie_headers: Vec, - pub rate_limit_headers: Vec<(String, String)>, - pub rate_limit_updates: SharedRateLimitUpdates, - pub turn_state_updates: SharedTurnState, -} - -const MAX_UPSTREAM_ERROR_BODY_BYTES: usize = 1024 * 1024; -type ReqwestClientCacheKey = (bool, Option); -type ReqwestClientCache = Mutex>; - -#[derive(Clone)] -pub struct CodexBackendClient { - client: Client, - base_url: String, - fingerprint: Fingerprint, - websocket_pool: Option>, - websocket_pool_account_id: Option, -} - -impl CodexBackendClient { - pub fn new(client: Client, base_url: impl Into, fingerprint: Fingerprint) -> Self { - Self { - client, - base_url: base_url.into().trim_end_matches('/').to_string(), - fingerprint, - websocket_pool: None, - websocket_pool_account_id: None, - } - } - - pub fn with_websocket_pool( - mut self, - pool: Arc, - account_id: impl Into, - ) -> Self { - self.websocket_pool = Some(pool); - self.websocket_pool_account_id = Some(account_id.into()); - self - } - - #[tracing::instrument( - skip(self, request, context), - fields( - request_id = %context.request_id, - account_id = %context.account_id.unwrap_or("unknown"), - model = %request.model, - transport = ?transport_for_request(request), - ) - )] - pub async fn create_response( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - let transport = transport_for_request(request); - let upstream_request = response_upstream_request(request, context); - if matches!( - transport, - CodexTransport::WebSocketPreferred | CodexTransport::WebSocketRequired - ) { - let headers = - self.request_headers_for_websocket_response(&upstream_request, context)?; - match self - .create_response_via_configured_websocket(&upstream_request, headers.clone()) - .await - { - Ok(response) => { - let usage = extract_sse_usage(&response.body)?; - return Ok(CodexBackendResponse { - body: response.body, - usage, - turn_state: response.turn_state, - set_cookie_headers: response.set_cookie_headers, - rate_limit_headers: response.rate_limit_headers, - }); - } - Err(error) - if transport == CodexTransport::WebSocketPreferred - && websocket_error_allows_http_sse_fallback(&error) => - { - tracing::info!( - error = %error, - "Codex WebSocket 不可用，降级到 HTTP SSE" - ); - } - Err(error) => return Err(codex_websocket_error(error)), - } - } - - let headers = self.request_headers_for_http_response(&upstream_request, context)?; - let response = self - .send_response_request(&upstream_request, headers) - .await?; - let status = response.status(); - let turn_state = turn_state(&response); - let set_cookie_headers = set_cookie_headers(&response); - let rate_limit_headers = rate_limit_headers(response.headers()); - let retry_after_seconds = retry_after_seconds(response.headers(), None); - if !status.is_success() { - let body = read_capped_error_body(response).await?; - return Err(CodexClientError::Upstream { - status, - retry_after_seconds: retry_after_seconds - .or_else(|| retry_after_seconds_from_body(&body)), - body, - }); - } - let body = response.text().await?; - let usage = extract_sse_usage(&body)?; - - Ok(CodexBackendResponse { - body, - usage, - turn_state, - set_cookie_headers, - rate_limit_headers, - }) - } - - #[tracing::instrument( - skip(self, request, context), - fields( - request_id = %context.request_id, - account_id = %context.account_id.unwrap_or("unknown"), - model = %request.model, - transport = "http_json", - ) - )] - pub async fn create_compact_response( - &self, - request: &CodexCompactRequest, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - let headers = self.compact_request_headers(context)?; - let response = self.send_compact_request(request, headers).await?; - let status = response.status(); - let set_cookie_headers = set_cookie_headers(&response); - let rate_limit_headers = rate_limit_headers(response.headers()); - let retry_after_seconds = retry_after_seconds(response.headers(), None); - let body = response.text().await?; - - if !status.is_success() { - return Err(CodexClientError::Upstream { - status, - retry_after_seconds: retry_after_seconds - .or_else(|| retry_after_seconds_from_body(&body)), - body, - }); - } - - let parsed = - serde_json::from_str::(&body).map_err(|_| CodexClientError::Upstream { - status: StatusCode::BAD_GATEWAY, - retry_after_seconds: None, - body: format!( - "Compact response is not valid JSON: {}", - truncate_for_error(&body) - ), - })?; - - Ok(CodexCompactResponse { - body: parsed, - set_cookie_headers, - rate_limit_headers, - }) - } - - #[tracing::instrument( - skip(self, request, context), - fields( - request_id = %context.request_id, - account_id = %context.account_id.unwrap_or("unknown"), - model = %request.model, - transport = "http_sse", - ) - )] - pub async fn stream_response( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - let upstream_request = response_upstream_request(request, context); - let headers = self.request_headers_for_http_response(&upstream_request, context)?; - let response = self - .send_response_request(&upstream_request, headers) - .await?; - let status = response.status(); - let turn_state = turn_state(&response); - let set_cookie_headers = set_cookie_headers(&response); - let rate_limit_headers = rate_limit_headers(response.headers()); - if !status.is_success() { - let retry_after_seconds = retry_after_seconds(response.headers(), None); - let body = read_capped_error_body(response).await?; - return Err(CodexClientError::Upstream { - status, - retry_after_seconds: retry_after_seconds - .or_else(|| retry_after_seconds_from_body(&body)), - body, - }); - } - - Ok(CodexBackendStream { - response, - turn_state, - set_cookie_headers, - rate_limit_headers, - }) - } - - #[tracing::instrument( - skip(self, request, context), - fields( - request_id = %context.request_id, - account_id = %context.account_id.unwrap_or("unknown"), - model = %request.model, - transport = "websocket", - ) - )] - pub async fn websocket_stream_response( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - let upstream_request = response_upstream_request(request, context); - let headers = self.request_headers_for_websocket_response(&upstream_request, context)?; - let response = self - .create_stream_via_configured_websocket(&upstream_request, headers) - .await - .map_err(codex_websocket_error)?; - - Ok(CodexBackendWebSocketStream { - body_stream: response.body_stream, - turn_state: response.turn_state, - set_cookie_headers: response.set_cookie_headers, - rate_limit_headers: response.rate_limit_headers, - rate_limit_updates: response.rate_limit_updates, - turn_state_updates: response.turn_state_updates, - }) - } - - async fn create_response_via_configured_websocket( - &self, - request: &CodexResponsesRequest, - headers: HeaderMap, - ) -> Result< - crate::codex::gateway::transport::websocket::CodexWebSocketResponse, - CodexWebSocketError, - > { - let Some((pool, pool_key)) = self.websocket_pool_key(request) else { - return create_response_via_websocket(&self.base_url, request, headers).await; - }; - let response = create_response_via_websocket_stream_with_pool( - &self.base_url, - request, - headers, - pool, - pool_key, - ) - .await?; - collect_websocket_stream_response(response).await - } - - async fn create_stream_via_configured_websocket( - &self, - request: &CodexResponsesRequest, - headers: HeaderMap, - ) -> Result { - let Some((pool, pool_key)) = self.websocket_pool_key(request) else { - return create_response_via_websocket_stream(&self.base_url, request, headers).await; - }; - create_response_via_websocket_stream_with_pool( - &self.base_url, - request, - headers, - pool, - pool_key, - ) - .await - } - - fn websocket_pool_key( - &self, - request: &CodexResponsesRequest, - ) -> Option<(Arc, CodexWebSocketPoolKey)> { - let pool = self.websocket_pool.clone()?; - let account_id = self.websocket_pool_account_id.as_deref()?; - let conversation_id = request - .prompt_cache_key - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty())?; - Some(( - pool, - CodexWebSocketPoolKey::new(&self.base_url, account_id, conversation_id), - )) - } - - pub async fn fetch_models( - &self, - context: CodexRequestContext<'_>, - ) -> CodexClientResult> { - let endpoints = [ - format!( - "{}/codex/models?client_version={}", - self.base_url, self.fingerprint.app_version - ), - format!("{}/models", self.base_url), - format!("{}/sentinel/chat-requirements", self.base_url), - ]; - - for endpoint in endpoints { - let headers = self.auxiliary_request_headers(context)?; - let response = self.client.get(endpoint).headers(headers).send().await?; - if !response.status().is_success() { - continue; - } - let parsed = response.json::().await?; - let models = extract_backend_model_entries(&parsed); - if !models.is_empty() { - return Ok(models); - } - } - - Err(CodexClientError::ModelsUnavailable) - } - - pub async fn probe_models_endpoint( - &self, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - let endpoint = format!( - "{}/codex/models?client_version={}", - self.base_url, self.fingerprint.app_version - ); - let headers = self.auxiliary_request_headers(context)?; - let response = self.client.get(&endpoint).headers(headers).send().await?; - - Ok(CodexConnectivityProbe { - endpoint, - status: response.status(), - }) - } - - pub async fn fetch_usage(&self, context: CodexRequestContext<'_>) -> CodexClientResult { - let mut last_invalid_body = None; - for endpoint in self.usage_endpoints() { - let headers = self.auxiliary_request_headers(context)?; - let response = self.client.get(endpoint).headers(headers).send().await?; - let status = response.status(); - let retry_after_seconds = retry_after_seconds(response.headers(), None); - if status == StatusCode::NOT_FOUND { - let body = read_capped_error_body(response).await?; - last_invalid_body = Some(body); - continue; - } - if !status.is_success() { - let body = read_capped_error_body(response).await?; - return Err(CodexClientError::Upstream { - status, - retry_after_seconds: retry_after_seconds - .or_else(|| retry_after_seconds_from_body(&body)), - body, - }); - } - let body = response.text().await?; - match serde_json::from_str::(&body) { - Ok(parsed) if parsed.get("rate_limit").is_some() => return Ok(parsed), - _ => last_invalid_body = Some(body), - } - } - - Err(CodexClientError::Upstream { - status: StatusCode::BAD_GATEWAY, - retry_after_seconds: None, - body: last_invalid_body - .map(|body| format!("invalid usage response: {}", truncate_for_error(&body))) - .unwrap_or_else(|| "usage endpoint is unavailable".to_string()), - }) - } - - async fn send_response_request( - &self, - request: &CodexResponsesRequest, - headers: HeaderMap, - ) -> CodexClientResult { - let url = endpoint_url(&self.base_url, CODEX_RESPONSES_PATH); - Ok(self - .client - .post(url) - .headers(headers) - .json(request) - .send() - .await?) - } - - async fn send_compact_request( - &self, - request: &CodexCompactRequest, - headers: HeaderMap, - ) -> CodexClientResult { - let url = endpoint_url(&self.base_url, CODEX_RESPONSES_COMPACT_PATH); - Ok(self - .client - .post(url) - .headers(headers) - .json(request) - .send() - .await?) - } - - fn usage_endpoints(&self) -> Vec { - usage_endpoint_urls(&self.base_url) - } - - fn request_headers(&self, context: CodexRequestContext<'_>) -> CodexClientResult { - let request_id = context.session_id.unwrap_or(context.request_id); - let ordered_headers = build_ordered_codex_base_headers( - &self.fingerprint, - context.access_token, - context.account_id, - ); - - let mut headers = HeaderMap::new(); - insert_ordered_headers(&mut headers, &ordered_headers)?; - headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); - insert_optional_header(&mut headers, "cookie", context.cookie_header)?; - headers.insert(ACCEPT, HeaderValue::from_static("text/event-stream")); - headers.insert( - HeaderName::from_static("openai-beta"), - HeaderValue::from_static("responses_websockets=2026-02-06"), - ); - headers.insert( - HeaderName::from_static("x-openai-internal-codex-residency"), - HeaderValue::from_static("us"), - ); - headers.insert( - HeaderName::from_static("x-client-request-id"), - HeaderValue::from_str(request_id)?, - ); - insert_optional_header( - &mut headers, - "x-codex-installation-id", - context.installation_id, - )?; - insert_optional_header(&mut headers, "session_id", context.session_id)?; - insert_optional_header(&mut headers, "x-codex-window-id", context.codex_window_id)?; - insert_optional_header(&mut headers, "x-codex-turn-state", context.turn_state)?; - insert_optional_header(&mut headers, "x-codex-turn-metadata", context.turn_metadata)?; - insert_optional_header(&mut headers, "x-codex-beta-features", context.beta_features)?; - insert_optional_header( - &mut headers, - "x-responsesapi-include-timing-metrics", - context.include_timing_metrics, - )?; - insert_optional_header(&mut headers, "version", context.version)?; - insert_optional_header( - &mut headers, - "x-codex-parent-thread-id", - context.parent_thread_id, - )?; - - Ok(headers) - } - - fn auxiliary_request_headers( - &self, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - let ordered_headers = build_ordered_codex_base_headers( - &self.fingerprint, - context.access_token, - context.account_id, - ); - - let mut headers = HeaderMap::new(); - - insert_ordered_headers(&mut headers, &ordered_headers)?; - - if let Some(cookie_header) = context.cookie_header { - headers.insert(COOKIE, HeaderValue::from_str(cookie_header)?); - } - headers.insert(ACCEPT, HeaderValue::from_static("application/json")); - headers.insert(ACCEPT_ENCODING, HeaderValue::from_static("gzip, deflate")); - - Ok(headers) - } - - fn compact_request_headers( - &self, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - let ordered_headers = build_ordered_codex_base_headers( - &self.fingerprint, - context.access_token, - context.account_id, - ); - - let mut headers = HeaderMap::new(); - insert_ordered_headers(&mut headers, &ordered_headers)?; - headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); - insert_optional_header(&mut headers, "cookie", context.cookie_header)?; - headers.insert( - HeaderName::from_static("openai-beta"), - HeaderValue::from_static("responses_websockets=2026-02-06"), - ); - headers.insert( - HeaderName::from_static("x-openai-internal-codex-residency"), - HeaderValue::from_static("us"), - ); - headers.insert( - HeaderName::from_static("x-client-request-id"), - HeaderValue::from_str(&Uuid::new_v4().to_string())?, - ); - insert_optional_header( - &mut headers, - "x-codex-installation-id", - context.installation_id, - )?; - - Ok(headers) - } - - fn request_headers_for_http_response( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - let mut headers = self.request_headers(context)?; - if let Some(subagent) = openai_subagent_from_metadata(request.client_metadata.as_ref()) { - headers.insert( - HeaderName::from_static("x-openai-subagent"), - HeaderValue::from_str(&subagent)?, - ); - } - Ok(headers) - } - - fn request_headers_for_websocket_response( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - let mut headers = self.request_headers_for_http_response(request, context)?; - headers.remove(CONTENT_TYPE); - headers.remove(ACCEPT); - Ok(headers) - } -} - -fn response_upstream_request( - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, -) -> CodexResponsesRequest { - let mut upstream = request.clone(); - if let Some(session_id) = context.session_id { - upstream.prompt_cache_key = Some(session_id.to_string()); - } - upstream.client_metadata = response_client_metadata(request.client_metadata.as_ref(), context); - upstream -} - -async fn read_capped_error_body(response: ReqwestResponse) -> Result { - let mut stream = response.bytes_stream(); - let mut body = Vec::new(); - while let Some(chunk) = stream.next().await { - let chunk = chunk?; - let remaining = MAX_UPSTREAM_ERROR_BODY_BYTES.saturating_sub(body.len()); - if remaining == 0 { - break; - } - if chunk.len() <= remaining { - body.extend_from_slice(&chunk); - } else { - body.extend_from_slice(&chunk[..remaining]); - break; - } - } - Ok(String::from_utf8_lossy(&body).into_owned()) -} - -async fn collect_websocket_stream_response( - response: CodexWebSocketStreamResponse, -) -> Result -{ - let CodexWebSocketStreamResponse { - mut body_stream, - turn_state, - set_cookie_headers, - mut rate_limit_headers, - rate_limit_updates, - turn_state_updates, - } = response; - let mut body = String::new(); - while let Some(chunk) = body_stream.next().await { - body.push_str(&chunk?); - } - if body.is_empty() { - return Err(CodexWebSocketError::EmptyResponse); - } - append_rate_limit_updates(&mut rate_limit_headers, &rate_limit_updates).await; - let turn_state = latest_turn_state(&turn_state_updates).await.or(turn_state); - Ok( - crate::codex::gateway::transport::websocket::CodexWebSocketResponse { - body, - turn_state, - set_cookie_headers, - rate_limit_headers, - }, - ) -} - -fn response_client_metadata( - client_metadata: Option<&Value>, - context: CodexRequestContext<'_>, -) -> Option { - let mut metadata = Map::new(); - if let Some(Value::Object(input)) = client_metadata { - for (key, value) in input { - if let Some(value) = value.as_str() { - metadata.insert(key.clone(), Value::String(value.to_string())); - } - } - } - - // 上游安全链路要求 body metadata 与 header 中的派生身份保持一致。 - insert_metadata_string( - &mut metadata, - "x-codex-installation-id", - context.installation_id, - ); - insert_metadata_string(&mut metadata, "x-codex-window-id", context.codex_window_id); - insert_metadata_string( - &mut metadata, - "x-codex-turn-metadata", - context.turn_metadata, - ); - insert_metadata_string( - &mut metadata, - "x-codex-parent-thread-id", - context.parent_thread_id, - ); - - if metadata.is_empty() { - None - } else { - Some(Value::Object(metadata)) - } -} - -fn insert_metadata_string(metadata: &mut Map, key: &str, value: Option<&str>) { - if let Some(value) = value.filter(|value| !value.trim().is_empty()) { - metadata.insert(key.to_string(), Value::String(value.to_string())); - } -} - -fn openai_subagent_from_metadata(client_metadata: Option<&Value>) -> Option { - let value = client_metadata? - .as_object()? - .get("x-openai-subagent")? - .as_str()? - .trim(); - if matches!( - value, - "review" | "compact" | "memory_consolidation" | "collab_spawn" - ) { - Some(value.to_string()) - } else { - None - } -} - -fn extract_backend_model_entries(value: &Value) -> Vec { - let Some(models) = value - .pointer("/chat_models/models") - .or_else(|| value.get("models")) - .or_else(|| value.get("data")) - .or_else(|| value.get("categories")) - .and_then(Value::as_array) - else { - return Vec::new(); - }; - - let mut entries = Vec::new(); - for model in models { - if let Some(nested) = model.get("models").and_then(Value::as_array) { - entries.extend(nested.iter().filter_map(parse_backend_model_entry)); - } else if let Some(entry) = parse_backend_model_entry(model) { - entries.push(entry); - } - } - entries -} - -fn parse_backend_model_entry(value: &Value) -> Option { - let entry = serde_json::from_value::(value.clone()).ok()?; - (entry.slug.is_some() || entry.id.is_some() || entry.name.is_some()).then_some(entry) -} - -fn insert_optional_header( - headers: &mut HeaderMap, - name: &'static str, - value: Option<&str>, -) -> CodexClientResult<()> { - let Some(value) = value else { - return Ok(()); - }; - headers.insert(HeaderName::from_static(name), HeaderValue::from_str(value)?); - Ok(()) -} - -fn insert_ordered_headers( - headers: &mut HeaderMap, - ordered_headers: &IndexMap, -) -> CodexClientResult<()> { - for (name, value) in ordered_headers { - headers.insert( - HeaderName::from_bytes(name.as_bytes())?, - HeaderValue::from_str(value)?, - ); - } - Ok(()) -} - -pub fn build_reqwest_client(force_http11: bool) -> Result { - let cache_key = (force_http11, custom_ca_env_cache_key()); - static CLIENTS: OnceLock = OnceLock::new(); - let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new())); - let mut clients = cache - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - if let Some(client) = clients.get(&cache_key) { - return Ok(client.clone()); - } - - let builder = Client::builder() - .use_rustls_tls() - .no_proxy() - .pool_max_idle_per_host(4) - .tcp_keepalive(Duration::from_secs(30)) - .gzip(true) - .brotli(true) - .zstd(true) - .deflate(true); - let builder = if force_http11 { - builder.http1_only() - } else { - builder - }; - let client = build_reqwest_client_with_custom_ca(builder)?; - clients.insert(cache_key, client.clone()); - Ok(client) -} - -fn turn_state(response: &ReqwestResponse) -> Option { - response - .headers() - .get("x-codex-turn-state") - .and_then(|value| value.to_str().ok()) - .map(ToString::to_string) -} - -fn set_cookie_headers(response: &ReqwestResponse) -> Vec { - response - .headers() - .get_all(SET_COOKIE) - .iter() - .filter_map(|value| value.to_str().ok().map(ToString::to_string)) - .collect() -} - -fn rate_limit_headers(headers: &HeaderMap) -> Vec<(String, String)> { - headers - .iter() - .filter(|(name, _)| is_rate_limit_header(name.as_str())) - .filter_map(|(name, value)| { - value - .to_str() - .ok() - .map(|value| (name.as_str().to_string(), value.to_string())) - }) - .collect() -} - -fn is_rate_limit_header(name: &str) -> bool { - let name = name.to_ascii_lowercase(); - name == "retry-after" - || name.contains("ratelimit") - || name.contains("rate-limit") - || name.starts_with("x-codex-primary-") - || name.starts_with("x-codex-secondary-") - || name.starts_with("x-codex-code-review-") - || name.starts_with("x-codex-review-") - || name.starts_with("x-code-review-") -} - -fn retry_after_seconds(headers: &HeaderMap, body: Option<&str>) -> Option { - headers - .get(RETRY_AFTER) - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.trim().parse::().ok()) - .filter(|seconds| *seconds > 0) - .or_else(|| body.and_then(retry_after_seconds_from_body)) -} - -fn truncate_for_error(body: &str) -> String { - body.chars().take(200).collect() -} - -fn websocket_error_allows_http_sse_fallback(error: &CodexWebSocketError) -> bool { - match error { - CodexWebSocketError::Upstream { status, .. } => *status == StatusCode::UPGRADE_REQUIRED, - CodexWebSocketError::InvalidRequest(_) - | CodexWebSocketError::Encode(_) - | CodexWebSocketError::Transport(_) - | CodexWebSocketError::OpenTimeout { .. } - | CodexWebSocketError::SendIdleTimeout { .. } - | CodexWebSocketError::ReceiveIdleTimeout { .. } - | CodexWebSocketError::EmptyResponse - | CodexWebSocketError::UnexpectedBinaryEvent - | CodexWebSocketError::ClosedByServerBeforeCompleted - | CodexWebSocketError::StreamClosedBeforeCompleted - | CodexWebSocketError::IncompleteResponse { .. } - | CodexWebSocketError::InvalidCompletedResponse { .. } => false, - } -} - -fn codex_websocket_error(error: CodexWebSocketError) -> CodexClientError { - match error { - CodexWebSocketError::Upstream { - status, - body, - retry_after_seconds, - } => CodexClientError::Upstream { - status, - body, - retry_after_seconds, - }, - error => CodexClientError::WebSocket(error), - } -} diff --git a/src/codex/gateway/transport/mod.rs b/src/codex/gateway/transport/mod.rs deleted file mode 100644 index 47174a405..000000000 --- a/src/codex/gateway/transport/mod.rs +++ /dev/null @@ -1,10 +0,0 @@ -pub mod custom_ca; -pub mod endpoints; -pub mod headers; -pub mod http_client; -pub mod rate_limits; -mod retry_after; -pub mod sse; -pub mod types; -pub mod usage_events; -pub mod websocket; diff --git a/src/codex/gateway/transport/retry_after.rs b/src/codex/gateway/transport/retry_after.rs deleted file mode 100644 index 0d437e7df..000000000 --- a/src/codex/gateway/transport/retry_after.rs +++ /dev/null @@ -1,143 +0,0 @@ -use serde_json::Value; - -pub(crate) fn retry_after_seconds_from_body(body: &str) -> Option { - let value = serde_json::from_str::(body).ok()?; - let error = value - .pointer("/response/error") - .or_else(|| value.get("error")) - .unwrap_or(&value); - if let Some(seconds) = error - .get("resets_in_seconds") - .and_then(Value::as_u64) - .filter(|seconds| *seconds > 0) - { - return Some(seconds); - } - retry_after_seconds_from_resets_at(error) - .or_else(|| retry_after_seconds_from_rate_limit_message(error)) -} - -fn retry_after_seconds_from_resets_at(error: &Value) -> Option { - let resets_at = error.get("resets_at").and_then(Value::as_u64)?; - let now = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .ok()? - .as_secs(); - (resets_at > now).then_some(resets_at - now) -} - -fn retry_after_seconds_from_rate_limit_message(error: &Value) -> Option { - let code = error - .get("code") - .or_else(|| error.get("type")) - .and_then(Value::as_str)?; - if code != "rate_limit_exceeded" { - return None; - } - let message = error.get("message").and_then(Value::as_str)?; - parse_try_again_delay_seconds(message) -} - -fn parse_try_again_delay_seconds(message: &str) -> Option { - let lower = message.to_ascii_lowercase(); - let marker = "try again in"; - let offset = lower.find(marker)? + marker.len(); - let remainder = message.get(offset..)?.trim_start(); - let number_end = number_prefix_len(remainder)?; - let value = remainder.get(..number_end)?.parse::().ok()?; - if !value.is_finite() || value <= 0.0 { - return None; - } - let unit_text = remainder - .get(number_end..)? - .trim_start() - .to_ascii_lowercase(); - let unit = unit_token(&unit_text)?; - match unit { - "ms" => positive_seconds_ceil(value / 1000.0), - "s" | "second" | "seconds" => positive_seconds_ceil(value), - _ => None, - } -} - -fn number_prefix_len(input: &str) -> Option { - let mut seen_digit = false; - let mut seen_dot = false; - let mut end = 0; - for (index, ch) in input.char_indices() { - if ch.is_ascii_digit() { - seen_digit = true; - end = index + ch.len_utf8(); - } else if ch == '.' && !seen_dot { - seen_dot = true; - end = index + ch.len_utf8(); - } else { - break; - } - } - seen_digit.then_some(end) -} - -fn unit_token(input: &str) -> Option<&str> { - let end = input - .char_indices() - .find_map(|(index, ch)| (!ch.is_ascii_alphabetic()).then_some(index)) - .unwrap_or(input.len()); - (end > 0).then_some(&input[..end]) -} - -fn positive_seconds_ceil(seconds: f64) -> Option { - if !seconds.is_finite() || seconds <= 0.0 || seconds > u64::MAX as f64 { - return None; - } - Some(seconds.ceil() as u64) -} - -#[cfg(test)] -mod tests { - use super::retry_after_seconds_from_body; - use serde_json::json; - - #[test] - fn retry_after_seconds_from_body_should_parse_rate_limit_message_seconds() { - let body = json!({ - "response": { - "error": { - "code": "rate_limit_exceeded", - "message": "Rate limit reached. Please try again in 11.054s." - } - } - }) - .to_string(); - - assert_eq!(retry_after_seconds_from_body(&body), Some(12)); - } - - #[test] - fn retry_after_seconds_from_body_should_parse_rate_limit_message_milliseconds() { - let body = json!({ - "error": { - "code": "rate_limit_exceeded", - "message": "Rate limit reached. Please try again in 28ms." - } - }) - .to_string(); - - assert_eq!(retry_after_seconds_from_body(&body), Some(1)); - } - - #[test] - fn retry_after_seconds_from_body_should_ignore_retry_message_for_other_codes() { - let body = json!({ - "response": { - "error": { - "code": "upstream_transient_error", - "message": "Try again in 35 seconds." - } - } - }) - .to_string(); - - assert_eq!(retry_after_seconds_from_body(&body), None); - } -} diff --git a/src/codex/gateway/transport/usage_events.rs b/src/codex/gateway/transport/usage_events.rs deleted file mode 100644 index 7989d1137..000000000 --- a/src/codex/gateway/transport/usage_events.rs +++ /dev/null @@ -1,122 +0,0 @@ -use serde::{Deserialize, Serialize}; -use serde_json::Value; - -use crate::codex::gateway::transport::sse::{parse_sse_events, SseError}; - -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct TokenUsage { - pub input_tokens: u64, - pub output_tokens: u64, - pub cached_tokens: u64, - pub image_input_tokens: u64, - pub image_output_tokens: u64, - pub total_tokens: u64, -} - -impl TokenUsage { - pub fn new(input_tokens: u64, output_tokens: u64, cached_tokens: u64) -> Self { - Self { - input_tokens, - output_tokens, - cached_tokens, - image_input_tokens: 0, - image_output_tokens: 0, - total_tokens: input_tokens + output_tokens, - } - } - - pub fn merged(self, other: Self) -> Self { - Self { - input_tokens: self.input_tokens + other.input_tokens, - output_tokens: self.output_tokens + other.output_tokens, - cached_tokens: self.cached_tokens + other.cached_tokens, - image_input_tokens: self.image_input_tokens + other.image_input_tokens, - image_output_tokens: self.image_output_tokens + other.image_output_tokens, - total_tokens: self.total_tokens + other.total_tokens, - } - } -} - -pub fn extract_usage(body: &Value) -> Option { - let usage = body.get("usage").unwrap_or(body); - if !usage.is_object() { - return None; - } - - let input_tokens = number_field(usage, "input_tokens") - .or_else(|| number_field(usage, "prompt_tokens")) - .unwrap_or_default(); - let output_tokens = number_field(usage, "output_tokens") - .or_else(|| number_field(usage, "completion_tokens")) - .unwrap_or_default(); - let cached_tokens = nested_number_field(usage, &["input_tokens_details", "cached_tokens"]) - .or_else(|| nested_number_field(usage, &["prompt_tokens_details", "cached_tokens"])) - .or_else(|| number_field(usage, "cached_tokens")) - .unwrap_or_default(); - let image_input_tokens = - nested_number_field(body, &["tool_usage", "image_gen", "input_tokens"]).unwrap_or_default(); - let image_output_tokens = - nested_number_field(body, &["tool_usage", "image_gen", "output_tokens"]) - .unwrap_or_default(); - let total_tokens = number_field(usage, "total_tokens").unwrap_or(input_tokens + output_tokens); - - let has_usage = [ - "input_tokens", - "output_tokens", - "prompt_tokens", - "completion_tokens", - "cached_tokens", - "total_tokens", - ] - .iter() - .any(|field| usage.get(*field).is_some()) - || usage.get("input_tokens_details").is_some() - || usage.get("prompt_tokens_details").is_some() - || image_input_tokens > 0 - || image_output_tokens > 0; - has_usage.then_some(TokenUsage { - input_tokens, - output_tokens, - cached_tokens, - image_input_tokens, - image_output_tokens, - total_tokens, - }) -} - -pub fn extract_sse_usage(body: &str) -> Result, SseError> { - let events = parse_sse_events(body)?; - let mut fallback_usage: Option = None; - for event in events { - if event.data == "[DONE]" { - continue; - } - let Ok(value) = serde_json::from_str::(&event.data) else { - continue; - }; - let event_usage = value - .get("response") - .and_then(extract_usage) - .or_else(|| extract_usage(&value)); - if let Some(event_usage) = event_usage { - if event.event.as_deref() == Some("response.completed") { - return Ok(Some(event_usage)); - } - fallback_usage = Some(event_usage); - } - } - Ok(fallback_usage) -} - -fn number_field(value: &Value, field: &str) -> Option { - value.get(field)?.as_u64() -} - -fn nested_number_field(value: &Value, path: &[&str]) -> Option { - let mut current = value; - for segment in path { - current = current.get(*segment)?; - } - current.as_u64() -} diff --git a/src/codex/gateway/transport/websocket/audit.rs b/src/codex/gateway/transport/websocket/audit.rs deleted file mode 100644 index c47d5d5f5..000000000 --- a/src/codex/gateway/transport/websocket/audit.rs +++ /dev/null @@ -1,229 +0,0 @@ -use std::{ - io, - path::{Path, PathBuf}, -}; - -use chrono::Utc; -use serde_json::{json, Value}; -use uuid::Uuid; - -use crate::codex::gateway::transport::types::CodexResponsesRequest; - -use super::{ - codec::websocket_payload_audit_snapshot_from_request_body, http_sse_fallback_allowed, - transport_for_request, CodexTransport, OpeningAuditSnapshot, PayloadAuditSnapshot, -}; - -pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR"; - -#[derive(Debug, Clone, PartialEq, serde::Serialize)] -pub struct WebSocketAuditArtifact { - pub transport_mode: String, - pub fallback_allowed: bool, - pub opening: Option, - pub payload: Option, - pub error: Option, -} - -#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)] -pub struct WebSocketAuditErrorSnapshot { - pub classification: String, - pub message: String, -} - -#[derive(Debug, Clone, PartialEq, serde::Serialize)] -pub struct WebSocketParityDiff { - pub differences: Vec, -} - -#[derive(Debug, Clone, PartialEq, serde::Serialize)] -pub struct WebSocketParityDifference { - pub path: String, - pub current: Value, - pub reference: Value, -} - -impl WebSocketAuditArtifact { - pub fn from_attempt( - request: &CodexResponsesRequest, - opening: OpeningAuditSnapshot, - payload: PayloadAuditSnapshot, - ) -> Self { - Self { - transport_mode: transport_mode_name(request), - fallback_allowed: http_sse_fallback_allowed(request), - opening: Some(opening), - payload: Some(payload), - error: None, - } - } -} - -pub fn websocket_parity_diff( - current: &WebSocketAuditArtifact, - reference: &WebSocketAuditArtifact, -) -> WebSocketParityDiff { - let mut differences = Vec::new(); - push_difference( - &mut differences, - "transport_mode", - json!(current.transport_mode), - json!(reference.transport_mode), - ); - push_difference( - &mut differences, - "fallback_allowed", - json!(current.fallback_allowed), - json!(reference.fallback_allowed), - ); - push_difference( - &mut differences, - "opening.request_line", - opening_request_line(¤t.opening), - opening_request_line(&reference.opening), - ); - push_difference( - &mut differences, - "opening.header_order", - opening_header_order(¤t.opening), - opening_header_order(&reference.opening), - ); - push_difference( - &mut differences, - "opening.sec_websocket_extensions", - opening_header_value(¤t.opening, "Sec-WebSocket-Extensions"), - opening_header_value(&reference.opening, "Sec-WebSocket-Extensions"), - ); - push_difference( - &mut differences, - "payload.top_level_keys", - payload_top_level_keys(¤t.payload), - payload_top_level_keys(&reference.payload), - ); - push_difference( - &mut differences, - "error.classification", - error_classification(¤t.error), - error_classification(&reference.error), - ); - WebSocketParityDiff { differences } -} - -fn push_difference( - differences: &mut Vec, - path: &str, - current: Value, - reference: Value, -) { - if current == reference { - return; - } - differences.push(WebSocketParityDifference { - path: path.to_string(), - current, - reference, - }); -} - -fn opening_request_line(opening: &Option) -> Value { - opening - .as_ref() - .map(|opening| json!(opening.request_line)) - .unwrap_or(Value::Null) -} - -fn opening_header_order(opening: &Option) -> Value { - opening - .as_ref() - .map(|opening| { - json!(opening - .headers - .iter() - .map(|header| header.name.as_str()) - .collect::>()) - }) - .unwrap_or(Value::Null) -} - -fn opening_header_value(opening: &Option, name: &str) -> Value { - opening - .as_ref() - .and_then(|opening| { - opening - .headers - .iter() - .find(|header| header.name.eq_ignore_ascii_case(name)) - }) - .map(|header| json!(header.value)) - .unwrap_or(Value::Null) -} - -fn payload_top_level_keys(payload: &Option) -> Value { - payload - .as_ref() - .map(|payload| json!(payload.top_level_keys)) - .unwrap_or(Value::Null) -} - -fn error_classification(error: &Option) -> Value { - error - .as_ref() - .map(|error| json!(error.classification)) - .unwrap_or(Value::Null) -} - -pub async fn record_websocket_audit_attempt_from_env( - request: &CodexResponsesRequest, - opening: &OpeningAuditSnapshot, - payload: &Value, -) { - let Some(dir) = websocket_audit_dir_from_env() else { - return; - }; - let artifact = WebSocketAuditArtifact::from_attempt( - request, - opening.clone(), - websocket_payload_audit_snapshot_from_request_body(request, payload), - ); - if let Err(error) = write_websocket_audit_artifact_for_dir(Some(&dir), &artifact).await { - tracing::warn!( - error = %error, - "写入 Codex WebSocket audit artifact 失败" - ); - } -} - -pub async fn write_websocket_audit_artifact_for_dir( - dir: Option<&Path>, - artifact: &WebSocketAuditArtifact, -) -> io::Result> { - let Some(dir) = dir.filter(|dir| !dir.as_os_str().is_empty()) else { - return Ok(None); - }; - - tokio::fs::create_dir_all(dir).await?; - let path = dir.join(audit_file_name()); - let body = serde_json::to_vec_pretty(artifact).map_err(io::Error::other)?; - tokio::fs::write(&path, body).await?; - Ok(Some(path)) -} - -fn websocket_audit_dir_from_env() -> Option { - std::env::var_os(WS_AUDIT_DIR_ENV) - .filter(|value| !value.as_os_str().is_empty()) - .map(PathBuf::from) -} - -fn audit_file_name() -> String { - let timestamp = Utc::now().format("%Y%m%dT%H%M%S%.3fZ"); - format!("codex-ws-audit-{timestamp}-{}.json", Uuid::new_v4()) -} - -fn transport_mode_name(request: &CodexResponsesRequest) -> String { - match transport_for_request(request) { - CodexTransport::HttpSse => "http_sse", - CodexTransport::WebSocketPreferred => "websocket_preferred", - CodexTransport::WebSocketRequired => "websocket_required", - } - .to_string() -} diff --git a/src/codex/gateway/transport/websocket/mod.rs b/src/codex/gateway/transport/websocket/mod.rs deleted file mode 100644 index 7778680e0..000000000 --- a/src/codex/gateway/transport/websocket/mod.rs +++ /dev/null @@ -1,896 +0,0 @@ -use std::{ - future::Future, - pin::Pin, - sync::Arc, - time::{Duration, Instant}, -}; - -use thiserror::Error; -use tokio::{sync::Mutex, time::timeout}; -use tokio_tungstenite::tungstenite::{ - client::IntoClientRequest, http::Request as WsRequest, Error as WsError, Message, -}; - -use futures::{channel::mpsc, SinkExt, Stream, StreamExt}; -use reqwest::{header::HeaderMap, StatusCode}; - -use crate::codex::gateway::transport::{ - endpoints::{endpoint_url, CODEX_RESPONSES_PATH}, - rate_limits::{parse_rate_limits_event_raw, rate_limits_to_header_pairs, ParsedRateLimits}, - retry_after::retry_after_seconds_from_body, - types::CodexResponsesRequest, -}; - -mod audit; -mod codec; -mod deflate; -mod opening; -mod pool; - -pub use audit::{ - websocket_parity_diff, write_websocket_audit_artifact_for_dir, WebSocketAuditArtifact, - WebSocketAuditErrorSnapshot, WebSocketParityDiff, WebSocketParityDifference, WS_AUDIT_DIR_ENV, -}; -use codec::{ - agent_message_output_item_event_invalid_required_fields, classify_ws_error_frame, - codex_websocket_transport_error, compaction_output_item_event_invalid_required_fields, - custom_tool_call_output_item_event_invalid_required_fields, - custom_tool_call_output_payload_item_event_invalid_required_fields, - delta_event_missing_official_required_fields, - function_call_output_item_event_invalid_required_fields, - function_call_output_payload_item_event_invalid_required_fields, - image_generation_call_output_item_event_invalid_required_fields, incomplete_response_reason, - is_internal_websocket_event, is_terminal_websocket_event, - local_shell_call_output_item_event_invalid_required_fields, - message_output_item_event_invalid_required_fields, metadata_turn_state, - output_item_event_invalid_item_type_tag, output_item_event_invalid_metadata, - output_item_event_missing_item, output_item_event_non_object_item, - reasoning_output_item_event_invalid_required_fields, - reasoning_summary_part_added_missing_summary_index, response_completed_missing_response, - response_completed_parse_error, response_created_missing_response, - response_output_text_delta_missing_delta, responses_stream_event_shape_parse_error, - retry_after_seconds_from_wrapped_error_headers, - tool_search_call_output_item_event_invalid_required_fields, - tool_search_output_item_event_invalid_required_fields, - web_search_call_output_item_event_invalid_required_fields, websocket_event_type, - websocket_message_text, websocket_request_body, websocket_request_text, websocket_sse_chunk, - WebSocketErrorClassificationProfile, -}; -pub use codec::{websocket_payload_audit_snapshot, PayloadAuditSnapshot}; -pub use opening::{websocket_opening_audit_snapshot, OpeningAuditHeader, OpeningAuditSnapshot}; -use pool::{ - CodexWebSocketConnectionMetadata, CodexWsStream, PooledWebSocketConnection, - WebSocketPoolAcquire, -}; -pub use pool::{CodexWebSocketPool, CodexWebSocketPoolConfig, CodexWebSocketPoolKey}; - -const WEBSOCKET_OPEN_TIMEOUT: Duration = Duration::from_secs(20); -const WEBSOCKET_EVENT_TIMEOUT: Duration = Duration::from_secs(20); - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum CodexTransport { - HttpSse, - WebSocketPreferred, - WebSocketRequired, -} - -#[derive(Debug, Error, PartialEq, Eq)] -pub enum WebSocketSupportError { - #[error("previous_response_id requires Codex WebSocket transport")] - PreviousResponseRequiresWebSocket, - #[error("request explicitly requires Codex WebSocket transport")] - ExplicitWebSocketRequired, -} - -#[derive(Debug, Error)] -pub enum CodexWebSocketError { - #[error("invalid WebSocket request: {0}")] - InvalidRequest(#[from] tokio_tungstenite::tungstenite::http::Error), - #[error("failed to encode websocket request: {0}")] - Encode(#[source] serde_json::Error), - #[error("websocket transport error: {0}")] - Transport(#[source] WsError), - #[error("websocket handshake timed out after {timeout:?}")] - OpenTimeout { timeout: Duration }, - #[error("idle timeout sending websocket request")] - SendIdleTimeout { timeout: Duration }, - #[error("idle timeout waiting for websocket")] - ReceiveIdleTimeout { timeout: Duration }, - #[error("websocket handshake returned status {status}: {body}")] - Upstream { - status: StatusCode, - body: String, - retry_after_seconds: Option, - }, - #[error("websocket response ended before any events")] - EmptyResponse, - #[error("unexpected binary websocket event")] - UnexpectedBinaryEvent, - #[error("websocket closed by server before response.completed")] - ClosedByServerBeforeCompleted, - #[error("stream closed before response.completed")] - StreamClosedBeforeCompleted, - #[error("Incomplete response returned, reason: {reason}")] - IncompleteResponse { reason: String }, - #[error("{message}")] - InvalidCompletedResponse { message: String }, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexWebSocketResponse { - pub body: String, - pub turn_state: Option, - pub set_cookie_headers: Vec, - pub rate_limit_headers: Vec<(String, String)>, -} - -pub type CodexWebSocketSseStream = - Pin> + Send>>; -pub type SharedRateLimitUpdates = Arc>>; -pub type SharedTurnState = Arc>>; - -pub struct CodexWebSocketStreamResponse { - pub body_stream: CodexWebSocketSseStream, - pub turn_state: Option, - pub set_cookie_headers: Vec, - pub rate_limit_headers: Vec<(String, String)>, - pub rate_limit_updates: SharedRateLimitUpdates, - pub turn_state_updates: SharedTurnState, -} - -struct ActiveWebSocket { - websocket: CodexWsStream, - metadata: CodexWebSocketConnectionMetadata, - pool_return: Option, - reused: bool, - last_activity_at: Instant, -} - -struct WebSocketPoolReturn { - pool: Arc, - key: CodexWebSocketPoolKey, - created_at: Instant, -} - -impl ActiveWebSocket { - fn error_classification_profile(&self) -> WebSocketErrorClassificationProfile { - if self.pool_return.is_some() { - WebSocketErrorClassificationProfile::Pooled - } else { - WebSocketErrorClassificationProfile::OneShot - } - } - - async fn finish(self) { - let Self { - websocket, - metadata, - pool_return, - reused: _, - last_activity_at, - } = self; - let Some(pool_return) = pool_return else { - let mut websocket = websocket; - let _ = websocket.close(None).await; - return; - }; - pool_return - .pool - .put( - pool_return.key, - PooledWebSocketConnection { - websocket, - metadata, - created_at: pool_return.created_at, - last_activity_at, - last_ping_at: None, - }, - ) - .await; - } - - async fn discard(self) { - let Self { - mut websocket, - pool_return, - .. - } = self; - if let Some(pool_return) = pool_return { - pool_return.pool.discard(&pool_return.key).await; - } - let _ = websocket.close(None).await; - } -} - -pub fn transport_for_request(request: &CodexResponsesRequest) -> CodexTransport { - if request.previous_response_id.is_some() { - CodexTransport::WebSocketRequired - } else if request.force_http_sse { - CodexTransport::HttpSse - } else { - CodexTransport::WebSocketPreferred - } -} - -pub fn http_sse_fallback_allowed(request: &CodexResponsesRequest) -> bool { - matches!( - transport_for_request(request), - CodexTransport::WebSocketPreferred | CodexTransport::HttpSse - ) -} - -#[tracing::instrument( - skip(base_url, request, headers), - fields( - model = %request.model, - has_previous_response_id = request.previous_response_id.is_some(), - ) -)] -pub async fn create_response_via_websocket( - base_url: &str, - request: &CodexResponsesRequest, - headers: HeaderMap, -) -> Result { - let CodexWebSocketStreamResponse { - mut body_stream, - turn_state, - set_cookie_headers, - mut rate_limit_headers, - rate_limit_updates, - turn_state_updates, - } = create_response_via_websocket_stream(base_url, request, headers).await?; - - let mut body = String::new(); - while let Some(chunk) = body_stream.next().await { - body.push_str(&chunk?); - } - if body.is_empty() { - return Err(CodexWebSocketError::EmptyResponse); - } - append_rate_limit_updates(&mut rate_limit_headers, &rate_limit_updates).await; - let turn_state = latest_turn_state(&turn_state_updates).await.or(turn_state); - - Ok(CodexWebSocketResponse { - body, - turn_state, - set_cookie_headers, - rate_limit_headers, - }) -} - -#[tracing::instrument( - skip(base_url, request, headers), - fields( - model = %request.model, - has_previous_response_id = request.previous_response_id.is_some(), - ) -)] -pub async fn create_response_via_websocket_stream( - base_url: &str, - request: &CodexResponsesRequest, - headers: HeaderMap, -) -> Result { - create_response_via_websocket_stream_inner(base_url, request, headers, None).await -} - -#[tracing::instrument( - skip(base_url, request, headers, pool, pool_key), - fields( - model = %request.model, - has_previous_response_id = request.previous_response_id.is_some(), - ) -)] -pub async fn create_response_via_websocket_stream_with_pool( - base_url: &str, - request: &CodexResponsesRequest, - headers: HeaderMap, - pool: Arc, - pool_key: CodexWebSocketPoolKey, -) -> Result { - create_response_via_websocket_stream_inner(base_url, request, headers, Some((pool, pool_key))) - .await -} - -async fn create_response_via_websocket_stream_inner( - base_url: &str, - request: &CodexResponsesRequest, - headers: HeaderMap, - pool: Option<(Arc, CodexWebSocketPoolKey)>, -) -> Result { - let mut pool_context = pool; - let mut retry_stale_reuse = pool_context.is_some(); - let rate_limit_updates = Arc::new(Mutex::new(Vec::new())); - let turn_state_updates = Arc::new(Mutex::new(None)); - loop { - let mut active = acquire_websocket(base_url, headers.clone(), pool_context.clone()).await?; - let metadata = active.metadata.clone(); - set_latest_turn_state(&turn_state_updates, metadata.turn_state.clone()).await; - let payload = websocket_request_body(request); - let payload_text = websocket_request_text(request).map_err(CodexWebSocketError::Encode)?; - audit::record_websocket_audit_attempt_from_env( - request, - &active.metadata.opening_audit, - &payload, - ) - .await; - if let Err(error) = timeout_websocket_send( - active.websocket.send(Message::Text(payload_text.into())), - WEBSOCKET_EVENT_TIMEOUT, - ) - .await - { - let reused = active.reused; - active.discard().await; - if reused && retry_stale_reuse { - retry_stale_reuse = false; - pool_context = None; - continue; - } - return Err(error); - } - loop { - let message = match next_websocket_message(&mut active.websocket).await { - Ok(Some(message)) => message, - Ok(None) => { - let reused = active.reused; - active.discard().await; - if reused && retry_stale_reuse { - retry_stale_reuse = false; - pool_context = None; - break; - } - return Err(CodexWebSocketError::StreamClosedBeforeCompleted); - } - Err(CodexWebSocketError::Transport(error)) => { - let reused = active.reused; - active.discard().await; - if reused && retry_stale_reuse { - retry_stale_reuse = false; - pool_context = None; - break; - } - return Err(CodexWebSocketError::Transport(error)); - } - Err(error) => { - active.discard().await; - return Err(error); - } - }; - active.last_activity_at = Instant::now(); - let raw = match websocket_message_text(message) { - Ok(Some(raw)) => raw, - Ok(None) => continue, - Err(error) => { - let reused = active.reused; - active.discard().await; - if reused && retry_stale_reuse { - retry_stale_reuse = false; - pool_context = None; - break; - } - return Err(error); - } - }; - if is_internal_websocket_event(&raw) { - capture_internal_rate_limit_event(&raw, &rate_limit_updates).await; - continue; - } - let classification_profile = active.error_classification_profile(); - if let Some(classified) = classify_ws_error_frame(&raw, classification_profile) { - active.discard().await; - return Err(CodexWebSocketError::Upstream { - status: classified.status, - retry_after_seconds: retry_after_seconds_from_wrapped_error_headers(&raw) - .or_else(|| retry_after_seconds_from_body(&raw)), - body: raw, - }); - } - - let Some(first_event) = websocket_event_type(&raw) else { - continue; - }; - if responses_stream_event_shape_parse_error(&raw) { - continue; - } - if let Some(reason) = incomplete_response_reason(&raw) { - active.discard().await; - return Err(CodexWebSocketError::IncompleteResponse { reason }); - } - if let Some(message) = response_completed_parse_error(&raw) { - active.discard().await; - return Err(CodexWebSocketError::InvalidCompletedResponse { message }); - } - if response_completed_missing_response(&raw) { - continue; - } - if response_created_missing_response(&raw) { - continue; - } - if response_output_text_delta_missing_delta(&raw) { - continue; - } - if delta_event_missing_official_required_fields(&raw) { - continue; - } - if output_item_event_missing_item(&raw) { - continue; - } - if output_item_event_non_object_item(&raw) { - continue; - } - if output_item_event_invalid_item_type_tag(&raw) { - continue; - } - if output_item_event_invalid_metadata(&raw) { - continue; - } - if message_output_item_event_invalid_required_fields(&raw) { - continue; - } - if agent_message_output_item_event_invalid_required_fields(&raw) { - continue; - } - if reasoning_output_item_event_invalid_required_fields(&raw) { - continue; - } - if function_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if function_call_output_payload_item_event_invalid_required_fields(&raw) { - continue; - } - if custom_tool_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if custom_tool_call_output_payload_item_event_invalid_required_fields(&raw) { - continue; - } - if tool_search_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if tool_search_output_item_event_invalid_required_fields(&raw) { - continue; - } - if local_shell_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if web_search_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if image_generation_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if compaction_output_item_event_invalid_required_fields(&raw) { - continue; - } - if reasoning_summary_part_added_missing_summary_index(&raw) { - continue; - } - if first_event == "response.metadata" { - capture_metadata_turn_state_event(&raw, &turn_state_updates).await; - continue; - } - if first_event == "error" { - continue; - } - let first_chunk = websocket_sse_chunk(&raw, &first_event); - let terminal = is_terminal_websocket_event(&first_event); - let (tx, rx) = mpsc::unbounded(); - if tx.unbounded_send(Ok(first_chunk)).is_err() { - active.discard().await; - return Err(CodexWebSocketError::EmptyResponse); - } - if terminal { - active.finish().await; - } else { - let rate_limit_updates = rate_limit_updates.clone(); - let turn_state_updates = turn_state_updates.clone(); - tokio::spawn(async move { - forward_websocket_as_sse(active, tx, rate_limit_updates, turn_state_updates) - .await; - }); - } - return Ok(CodexWebSocketStreamResponse { - body_stream: Box::pin(rx), - turn_state: metadata.turn_state, - set_cookie_headers: metadata.set_cookie_headers, - rate_limit_headers: metadata.rate_limit_headers, - rate_limit_updates, - turn_state_updates, - }); - } - } -} - -async fn acquire_websocket( - base_url: &str, - headers: HeaderMap, - pool: Option<(Arc, CodexWebSocketPoolKey)>, -) -> Result { - let Some((pool, key)) = pool else { - let created_at = Instant::now(); - let (websocket, metadata) = connect_websocket(base_url, headers).await?; - return Ok(ActiveWebSocket { - websocket, - metadata, - pool_return: None, - reused: false, - last_activity_at: created_at, - }); - }; - - match pool.acquire(&key).await { - WebSocketPoolAcquire::Reused(connection) => { - let PooledWebSocketConnection { - websocket, - metadata, - created_at, - last_activity_at, - last_ping_at: _, - } = *connection; - return Ok(ActiveWebSocket { - websocket, - metadata, - pool_return: Some(WebSocketPoolReturn { - pool, - key, - created_at, - }), - reused: true, - last_activity_at, - }); - } - WebSocketPoolAcquire::FreshReserved => {} - WebSocketPoolAcquire::Bypass => { - let (websocket, metadata) = connect_websocket(base_url, headers).await?; - return Ok(ActiveWebSocket { - websocket, - metadata, - pool_return: None, - reused: false, - last_activity_at: Instant::now(), - }); - } - } - - let created_at = Instant::now(); - let (websocket, metadata) = match connect_websocket(base_url, headers).await { - Ok(connection) => connection, - Err(error) => { - pool.discard(&key).await; - return Err(error); - } - }; - Ok(ActiveWebSocket { - websocket, - metadata, - pool_return: Some(WebSocketPoolReturn { - pool, - key, - created_at, - }), - reused: false, - last_activity_at: created_at, - }) -} - -async fn connect_websocket( - base_url: &str, - headers: HeaderMap, -) -> Result<(CodexWsStream, CodexWebSocketConnectionMetadata), CodexWebSocketError> { - connect_websocket_with_timeout(base_url, headers, WEBSOCKET_OPEN_TIMEOUT).await -} - -async fn connect_websocket_with_timeout( - base_url: &str, - headers: HeaderMap, - open_timeout: Duration, -) -> Result<(CodexWsStream, CodexWebSocketConnectionMetadata), CodexWebSocketError> { - let ws_request = build_ws_request(base_url, headers)?; - let (websocket, handshake_response, opening_audit) = timeout( - open_timeout, - opening::connect_with_original_opening_handshake(ws_request), - ) - .await - .map_err(|_| CodexWebSocketError::OpenTimeout { - timeout: open_timeout, - })? - .map_err(codex_websocket_transport_error)?; - Ok(( - websocket, - CodexWebSocketConnectionMetadata::from_headers(handshake_response.headers(), opening_audit), - )) -} - -fn build_ws_request( - base_url: &str, - headers: HeaderMap, -) -> Result, CodexWebSocketError> { - let mut request = websocket_url(base_url) - .into_client_request() - .map_err(codex_websocket_transport_error)?; - for (name, value) in &headers { - let Ok(name) = - tokio_tungstenite::tungstenite::http::HeaderName::from_bytes(name.as_str().as_bytes()) - else { - continue; - }; - let Ok(value) = - tokio_tungstenite::tungstenite::http::HeaderValue::from_bytes(value.as_bytes()) - else { - continue; - }; - request.headers_mut().insert(name, value); - } - Ok(request) -} - -fn websocket_url(base_url: &str) -> String { - let url = endpoint_url(base_url, CODEX_RESPONSES_PATH); - if let Some(rest) = url.strip_prefix("https://") { - format!("wss://{rest}") - } else if let Some(rest) = url.strip_prefix("http://") { - format!("ws://{rest}") - } else { - url - } -} - -async fn forward_websocket_as_sse( - mut active: ActiveWebSocket, - tx: mpsc::UnboundedSender>, - rate_limit_updates: SharedRateLimitUpdates, - turn_state_updates: SharedTurnState, -) { - loop { - match next_websocket_message(&mut active.websocket).await { - Ok(Some(message)) => { - active.last_activity_at = Instant::now(); - let raw = match websocket_message_text(message) { - Ok(Some(raw)) => raw, - Ok(None) => continue, - Err(error) => { - active.discard().await; - let _ = tx.unbounded_send(Err(error)); - return; - } - }; - if is_internal_websocket_event(&raw) { - capture_internal_rate_limit_event(&raw, &rate_limit_updates).await; - continue; - } - let classification_profile = active.error_classification_profile(); - if let Some(classified) = classify_ws_error_frame(&raw, classification_profile) { - active.discard().await; - let _ = tx.unbounded_send(Err(CodexWebSocketError::Upstream { - status: classified.status, - retry_after_seconds: retry_after_seconds_from_wrapped_error_headers(&raw) - .or_else(|| retry_after_seconds_from_body(&raw)), - body: raw, - })); - return; - } - let Some(event) = websocket_event_type(&raw) else { - continue; - }; - if responses_stream_event_shape_parse_error(&raw) { - continue; - } - if let Some(reason) = incomplete_response_reason(&raw) { - active.discard().await; - let _ = - tx.unbounded_send(Err(CodexWebSocketError::IncompleteResponse { reason })); - return; - } - if let Some(message) = response_completed_parse_error(&raw) { - active.discard().await; - let _ = tx.unbounded_send(Err(CodexWebSocketError::InvalidCompletedResponse { - message, - })); - return; - } - if response_completed_missing_response(&raw) { - continue; - } - if response_created_missing_response(&raw) { - continue; - } - if response_output_text_delta_missing_delta(&raw) { - continue; - } - if delta_event_missing_official_required_fields(&raw) { - continue; - } - if output_item_event_missing_item(&raw) { - continue; - } - if output_item_event_non_object_item(&raw) { - continue; - } - if output_item_event_invalid_item_type_tag(&raw) { - continue; - } - if output_item_event_invalid_metadata(&raw) { - continue; - } - if message_output_item_event_invalid_required_fields(&raw) { - continue; - } - if agent_message_output_item_event_invalid_required_fields(&raw) { - continue; - } - if reasoning_output_item_event_invalid_required_fields(&raw) { - continue; - } - if function_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if function_call_output_payload_item_event_invalid_required_fields(&raw) { - continue; - } - if custom_tool_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if custom_tool_call_output_payload_item_event_invalid_required_fields(&raw) { - continue; - } - if tool_search_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if tool_search_output_item_event_invalid_required_fields(&raw) { - continue; - } - if local_shell_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if web_search_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if image_generation_call_output_item_event_invalid_required_fields(&raw) { - continue; - } - if compaction_output_item_event_invalid_required_fields(&raw) { - continue; - } - if reasoning_summary_part_added_missing_summary_index(&raw) { - continue; - } - if event == "response.metadata" { - capture_metadata_turn_state_event(&raw, &turn_state_updates).await; - continue; - } - if event == "error" { - continue; - } - let terminal = is_terminal_websocket_event(&event); - let chunk = websocket_sse_chunk(&raw, &event); - if tx.unbounded_send(Ok(chunk)).is_err() { - active.discard().await; - return; - } - if terminal { - active.finish().await; - return; - } - } - Ok(None) => { - let _ = tx.unbounded_send(Err(CodexWebSocketError::StreamClosedBeforeCompleted)); - active.discard().await; - return; - } - Err(error) => { - let _ = tx.unbounded_send(Err(error)); - active.discard().await; - return; - } - } - } -} - -async fn next_websocket_message( - websocket: &mut CodexWsStream, -) -> Result, CodexWebSocketError> { - match timeout(WEBSOCKET_EVENT_TIMEOUT, websocket.next()).await { - Ok(Some(Ok(message))) => Ok(Some(message)), - Ok(Some(Err(error))) => Err(CodexWebSocketError::Transport(error)), - Ok(None) => Ok(None), - Err(_) => Err(CodexWebSocketError::ReceiveIdleTimeout { - timeout: WEBSOCKET_EVENT_TIMEOUT, - }), - } -} - -async fn timeout_websocket_send( - send: F, - timeout_duration: Duration, -) -> Result<(), CodexWebSocketError> -where - F: Future>, -{ - match timeout(timeout_duration, send).await { - Ok(Ok(())) => Ok(()), - Ok(Err(error)) => Err(CodexWebSocketError::Transport(error)), - Err(_) => Err(CodexWebSocketError::SendIdleTimeout { - timeout: timeout_duration, - }), - } -} - -async fn capture_internal_rate_limit_event(raw: &str, updates: &SharedRateLimitUpdates) { - if let Some(parsed) = parse_rate_limits_event_raw(raw) { - updates.lock().await.push(parsed); - } -} - -async fn capture_metadata_turn_state_event(raw: &str, updates: &SharedTurnState) { - if let Some(turn_state) = metadata_turn_state(raw) { - set_latest_turn_state(updates, Some(turn_state)).await; - } -} - -async fn set_latest_turn_state(updates: &SharedTurnState, turn_state: Option) { - if turn_state.is_some() { - *updates.lock().await = turn_state; - } -} - -pub async fn latest_turn_state(updates: &SharedTurnState) -> Option { - updates.lock().await.clone() -} - -pub async fn append_rate_limit_updates( - headers: &mut Vec<(String, String)>, - updates: &SharedRateLimitUpdates, -) { - let updates = updates.lock().await; - for update in updates.iter() { - headers.extend(rate_limits_to_header_pairs(update)); - } -} - -#[cfg(test)] -mod tests { - use futures::future; - use reqwest::header::HeaderMap; - use tokio::{net::TcpListener, time::Duration}; - - use super::{connect_websocket_with_timeout, timeout_websocket_send, CodexWebSocketError}; - - #[tokio::test] - async fn connect_websocket_with_timeout_should_fail_when_handshake_stalls() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (_stream, _) = listener.accept().await.unwrap(); - tokio::time::sleep(Duration::from_millis(100)).await; - }); - - let result = connect_websocket_with_timeout( - &format!("http://{addr}"), - HeaderMap::new(), - Duration::from_millis(10), - ) - .await; - - let Err(CodexWebSocketError::OpenTimeout { timeout }) = result else { - panic!("expected websocket open timeout"); - }; - assert_eq!(timeout, Duration::from_millis(10)); - server.abort(); - } - - #[tokio::test(start_paused = true)] - async fn timeout_websocket_send_should_fail_when_send_future_stalls() { - let result = tokio::spawn(timeout_websocket_send( - future::pending(), - Duration::from_secs(20), - )); - tokio::task::yield_now().await; - tokio::time::advance(Duration::from_secs(20)).await; - - let Err(CodexWebSocketError::SendIdleTimeout { timeout }) = result.await.unwrap() else { - panic!("expected websocket send timeout"); - }; - assert_eq!(timeout, Duration::from_secs(20)); - } -} diff --git a/src/codex/gateway/transport/websocket/opening.rs b/src/codex/gateway/transport/websocket/opening.rs deleted file mode 100644 index 04dadd166..000000000 --- a/src/codex/gateway/transport/websocket/opening.rs +++ /dev/null @@ -1,495 +0,0 @@ -use std::sync::Arc; - -use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; -use rustls_pki_types::ServerName; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::TcpStream, -}; -use tokio_rustls::TlsConnector; -use tokio_tungstenite::{ - tungstenite::{ - error::{CapacityError, Error as WsError, ProtocolError, TlsError, UrlError}, - handshake::{ - client::{generate_key, Response as WsResponse}, - derive_accept_key, - }, - http::{Request as WsRequest, StatusCode, Uri}, - protocol::Role, - }, - MaybeTlsStream, WebSocketStream, -}; - -use super::deflate::PerMessageDeflateStream; -use super::pool::CodexWsStream; -use crate::codex::gateway::transport::custom_ca::{ - maybe_build_rustls_client_config_with_custom_ca, native_root_store, -}; - -const MAX_OPENING_RESPONSE_BYTES: usize = 64 * 1024; -const ORIGINAL_WS_PERMESSAGE_DEFLATE_EXTENSION: &str = "permessage-deflate; client_max_window_bits"; -const REDACTED_VALUE: &str = ""; - -const ORIGINAL_WS_HEADER_ORDER: &[&str] = &[ - "chatgpt-account-id", - "authorization", - "user-agent", - "originator", - "openai-beta", - "x-codex-beta-features", - "x-client-request-id", - "session_id", - "thread-id", - "x-codex-window-id", - "x-codex-turn-metadata", - "x-codex-turn-state", - "x-responsesapi-include-timing-metrics", - "version", - "x-codex-parent-thread-id", - "x-openai-subagent", -]; - -#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)] -pub struct OpeningAuditSnapshot { - pub request_line: String, - pub headers: Vec, -} - -#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)] -pub struct OpeningAuditHeader { - pub name: String, - pub value: String, -} - -struct OpeningRequestHead { - request_line: String, - headers: Vec, -} - -struct OpeningHeader { - name: String, - value: Vec, -} - -pub(super) async fn connect_with_original_opening_handshake( - mut request: WsRequest<()>, -) -> Result<(CodexWsStream, WsResponse, OpeningAuditSnapshot), WsError> { - ensure_websocket_key(&mut request)?; - let audit_snapshot = websocket_opening_audit_snapshot(&request)?; - let uri = request.uri().clone(); - let mut stream = connect_stream(&uri).await?; - let request_bytes = opening_request_bytes(&request)?; - - stream - .write_all(&request_bytes) - .await - .map_err(WsError::Io)?; - stream.flush().await.map_err(WsError::Io)?; - - let (response, buffered) = read_opening_response(&mut stream).await?; - if response.status() != StatusCode::SWITCHING_PROTOCOLS { - return Err(WsError::Http(Box::new(response))); - } - verify_accept_key(&request, &response)?; - let permessage_deflate = response_accepts_permessage_deflate(&response); - - let stream = PerMessageDeflateStream::new(stream, permessage_deflate, buffered); - let websocket = - WebSocketStream::from_partially_read(stream, Vec::new(), Role::Client, None).await; - Ok((websocket, response, audit_snapshot)) -} - -pub fn websocket_opening_audit_snapshot( - request: &WsRequest<()>, -) -> Result { - let head = opening_request_head(request)?; - Ok(OpeningAuditSnapshot { - request_line: head.request_line, - headers: head - .headers - .into_iter() - .map(|header| OpeningAuditHeader { - value: audit_header_value(&header.name, &header.value), - name: header.name, - }) - .collect(), - }) -} - -fn ensure_websocket_key(request: &mut WsRequest<()>) -> Result<(), WsError> { - if request.headers().contains_key("Sec-WebSocket-Key") { - return Ok(()); - } - - request - .headers_mut() - .insert("Sec-WebSocket-Key", HeaderValue::from_str(&generate_key())?); - Ok(()) -} - -async fn connect_stream(uri: &Uri) -> Result, WsError> { - let host = uri.host().ok_or(UrlError::NoHostName)?; - if host.is_empty() { - return Err(UrlError::EmptyHostName.into()); - } - let port = uri - .port_u16() - .or_else(|| match uri.scheme_str() { - Some("wss") => Some(443), - Some("ws") => Some(80), - _ => None, - }) - .ok_or(UrlError::UnsupportedUrlScheme)?; - - let socket = TcpStream::connect((host, port)) - .await - .map_err(WsError::Io)?; - - match uri.scheme_str() { - Some("ws") => Ok(MaybeTlsStream::Plain(socket)), - Some("wss") => { - let server_name = ServerName::try_from(host.to_string()) - .map_err(|_| WsError::Tls(TlsError::InvalidDnsName))?; - let config = match maybe_build_rustls_client_config_with_custom_ca() - .map_err(|err| WsError::Io(err.into()))? - { - Some(config) => config, - None => Arc::new( - rustls::ClientConfig::builder() - .with_root_certificates(native_root_store().map_err(WsError::Io)?) - .with_no_client_auth(), - ), - }; - let connector = TlsConnector::from(config); - let stream = connector - .connect(server_name, socket) - .await - .map_err(WsError::Io)?; - Ok(MaybeTlsStream::Rustls(stream)) - } - _ => Err(UrlError::UnsupportedUrlScheme.into()), - } -} - -fn opening_request_bytes(request: &WsRequest<()>) -> Result, WsError> { - let head = opening_request_head(request)?; - let mut bytes = Vec::new(); - bytes.extend_from_slice(head.request_line.as_bytes()); - bytes.extend_from_slice(b"\r\n"); - for header in head.headers { - write_header(&mut bytes, &header.name, &header.value); - } - bytes.extend_from_slice(b"\r\n"); - Ok(bytes) -} - -fn opening_request_head(request: &WsRequest<()>) -> Result { - let path = request - .uri() - .path_and_query() - .ok_or(UrlError::NoPathOrQuery)? - .as_str(); - let host = host_header(request.uri())?; - let key = request - .headers() - .get("Sec-WebSocket-Key") - .ok_or_else(|| ProtocolError::MissingSecWebSocketKey)? - .as_bytes(); - - let mut headers = Vec::new(); - push_header(&mut headers, "Host", host.as_bytes()); - push_header(&mut headers, "Connection", b"Upgrade"); - push_header(&mut headers, "Upgrade", b"websocket"); - push_header(&mut headers, "Sec-WebSocket-Version", b"13"); - push_header(&mut headers, "Sec-WebSocket-Key", key); - push_original_business_headers(&mut headers, request.headers()); - push_header( - &mut headers, - "sec-websocket-extensions", - ORIGINAL_WS_PERMESSAGE_DEFLATE_EXTENSION.as_bytes(), - ); - Ok(OpeningRequestHead { - request_line: format!("GET {path} HTTP/1.1"), - headers, - }) -} - -fn host_header(uri: &Uri) -> Result { - let host = uri.host().ok_or(UrlError::NoHostName)?; - if host.is_empty() { - return Err(UrlError::EmptyHostName.into()); - } - let Some(port) = uri.port_u16() else { - return Ok(host.to_string()); - }; - let is_default = matches!( - (uri.scheme_str(), port), - (Some("ws"), 80) | (Some("wss"), 443) - ); - if is_default { - Ok(host.to_string()) - } else { - Ok(format!("{host}:{port}")) - } -} - -fn push_original_business_headers(output: &mut Vec, headers: &HeaderMap) { - let mut emitted = Vec::new(); - for name in ORIGINAL_WS_HEADER_ORDER { - if *name == "session_id" { - if push_session_headers(output, headers) { - emitted.push("session_id"); - emitted.push("session-id"); - if headers.get("thread-id").is_none() { - emitted.push("thread-id"); - } - } - } else if push_named_header(output, headers, name) { - emitted.push(*name); - } - } - - for (name, value) in headers { - let lower = name.as_str(); - if emitted.iter().any(|seen| seen.eq_ignore_ascii_case(lower)) - || should_skip_websocket_business_header(lower) - { - continue; - } - push_header(output, original_header_name(lower), value.as_bytes()); - } -} - -fn push_named_header(output: &mut Vec, headers: &HeaderMap, name: &str) -> bool { - let Some(value) = headers.get(name) else { - return false; - }; - push_header(output, original_header_name(name), value.as_bytes()); - true -} - -fn push_session_headers(output: &mut Vec, headers: &HeaderMap) -> bool { - let Some(value) = headers - .get("session_id") - .or_else(|| headers.get("session-id")) - else { - return false; - }; - - push_header(output, "session-id", value.as_bytes()); - if headers.get("thread-id").is_none() { - push_header(output, "thread-id", value.as_bytes()); - } - true -} - -fn push_header(output: &mut Vec, name: &str, value: &[u8]) { - output.push(OpeningHeader { - name: name.to_string(), - value: value.to_vec(), - }); -} - -fn write_header(bytes: &mut Vec, name: &str, value: &[u8]) { - bytes.extend_from_slice(name.as_bytes()); - bytes.extend_from_slice(b": "); - bytes.extend_from_slice(value); - bytes.extend_from_slice(b"\r\n"); -} - -fn is_opening_header(name: &str) -> bool { - matches!( - name, - "host" - | "connection" - | "upgrade" - | "sec-websocket-version" - | "sec-websocket-extensions" - | "sec-websocket-key" - ) -} - -fn should_skip_websocket_business_header(name: &str) -> bool { - is_opening_header(name) - || matches!( - name, - "content-type" - | "accept" - | "cookie" - | "sec-ch-ua" - | "sec-ch-ua-mobile" - | "sec-ch-ua-platform" - | "accept-encoding" - | "accept-language" - | "sec-fetch-site" - | "sec-fetch-mode" - | "sec-fetch-dest" - | "x-openai-internal-codex-residency" - | "x-codex-installation-id" - ) -} - -fn original_header_name(name: &str) -> &str { - match name { - "authorization" => "authorization", - "chatgpt-account-id" => "chatgpt-account-id", - "user-agent" => "user-agent", - "accept-encoding" => "Accept-Encoding", - "accept-language" => "Accept-Language", - "openai-beta" => "openai-beta", - "cookie" => "Cookie", - "content-type" => "Content-Type", - "accept" => "Accept", - _ => name, - } -} - -fn audit_header_value(name: &str, value: &[u8]) -> String { - if is_sensitive_audit_header(name) { - return REDACTED_VALUE.to_string(); - } - String::from_utf8_lossy(value).into_owned() -} - -fn is_sensitive_audit_header(name: &str) -> bool { - matches!( - name.to_ascii_lowercase().as_str(), - "authorization" - | "chatgpt-account-id" - | "cookie" - | "x-client-request-id" - | "x-codex-installation-id" - | "session_id" - | "session-id" - | "thread-id" - | "x-codex-window-id" - | "x-codex-turn-state" - | "x-codex-turn-metadata" - | "x-codex-parent-thread-id" - ) -} - -async fn read_opening_response( - stream: &mut MaybeTlsStream, -) -> Result<(WsResponse, Vec), WsError> { - let mut bytes = Vec::new(); - let mut chunk = [0_u8; 1024]; - let header_end = loop { - let read = stream.read(&mut chunk).await.map_err(WsError::Io)?; - if read == 0 { - return Err(ProtocolError::HandshakeIncomplete.into()); - } - bytes.extend_from_slice(&chunk[..read]); - if bytes.len() > MAX_OPENING_RESPONSE_BYTES { - return Err(CapacityError::MessageTooLong { - size: bytes.len(), - max_size: MAX_OPENING_RESPONSE_BYTES, - } - .into()); - } - if let Some(index) = find_header_end(&bytes) { - break index; - } - }; - - let body_start = header_end + 4; - let mut body = bytes[body_start..].to_vec(); - let mut response = parse_response_head(&bytes[..header_end])?; - if response.status() != StatusCode::SWITCHING_PROTOCOLS { - read_response_body(stream, response.headers(), &mut body).await?; - *response.body_mut() = Some(body); - return Ok((response, Vec::new())); - } - - Ok((response, body)) -} - -fn find_header_end(bytes: &[u8]) -> Option { - bytes.windows(4).position(|window| window == b"\r\n\r\n") -} - -fn parse_response_head(head: &[u8]) -> Result { - let head = std::str::from_utf8(head)?; - let mut lines = head.split("\r\n"); - let status_line = lines - .next() - .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "empty response"))?; - let status = status_line - .split_whitespace() - .nth(1) - .ok_or_else(|| { - std::io::Error::new(std::io::ErrorKind::InvalidData, "missing response status") - })? - .parse::() - .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))?; - - let mut response = WsResponse::new(None); - *response.status_mut() = StatusCode::from_u16(status)?; - for line in lines { - let Some((name, value)) = line.split_once(':') else { - continue; - }; - response.headers_mut().append( - HeaderName::from_bytes(name.trim().as_bytes())?, - HeaderValue::from_bytes(value.trim_start().as_bytes())?, - ); - } - Ok(response) -} - -async fn read_response_body( - stream: &mut MaybeTlsStream, - headers: &HeaderMap, - body: &mut Vec, -) -> Result<(), WsError> { - let Some(content_length) = headers - .get("content-length") - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.parse::().ok()) - else { - return Ok(()); - }; - - while body.len() < content_length { - let remaining = content_length - body.len(); - let mut chunk = vec![0_u8; remaining.min(1024)]; - let read = stream.read(&mut chunk).await.map_err(WsError::Io)?; - if read == 0 { - break; - } - body.extend_from_slice(&chunk[..read]); - } - body.truncate(content_length); - Ok(()) -} - -fn verify_accept_key(request: &WsRequest<()>, response: &WsResponse) -> Result<(), WsError> { - let key = request - .headers() - .get("Sec-WebSocket-Key") - .ok_or(ProtocolError::MissingSecWebSocketKey)?; - let expected = derive_accept_key(key.as_bytes()); - let actual = response - .headers() - .get("Sec-WebSocket-Accept") - .ok_or(ProtocolError::SecWebSocketAcceptKeyMismatch)?; - - if actual == expected.as_str() { - Ok(()) - } else { - Err(ProtocolError::SecWebSocketAcceptKeyMismatch.into()) - } -} - -fn response_accepts_permessage_deflate(response: &WsResponse) -> bool { - response - .headers() - .get_all("sec-websocket-extensions") - .iter() - .filter_map(|value| value.to_str().ok()) - .any(|value| { - value - .split(',') - .any(|extension| extension.trim_start().starts_with("permessage-deflate")) - }) -} diff --git a/src/codex/mod.rs b/src/codex/mod.rs deleted file mode 100644 index b222d7c10..000000000 --- a/src/codex/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -pub mod accounts; -pub mod events; -pub mod gateway; -pub mod models; -pub mod serving; -pub mod tasks; -pub mod usage; diff --git a/src/codex/models/mod.rs b/src/codex/models/mod.rs deleted file mode 100644 index c01714035..000000000 --- a/src/codex/models/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub mod catalog; -pub mod repository; -pub mod service; diff --git a/src/codex/models/repository.rs b/src/codex/models/repository.rs deleted file mode 100644 index 5ec49256b..000000000 --- a/src/codex/models/repository.rs +++ /dev/null @@ -1,62 +0,0 @@ -use chrono::Utc; -use sqlx::{Row, SqlitePool}; -use thiserror::Error; - -use crate::codex::models::catalog::{CodexModelInfo, ModelPlanSnapshot}; - -#[derive(Debug, Error)] -pub enum ModelSnapshotRepositoryError { - #[error("model snapshot database error: {0}")] - Database(#[from] sqlx::Error), - #[error("model snapshot json error: {0}")] - Json(#[from] serde_json::Error), -} - -pub type ModelSnapshotRepositoryResult = Result; - -#[derive(Debug, Clone)] -pub struct ModelSnapshotRepository { - pool: SqlitePool, -} - -impl ModelSnapshotRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - pub async fn replace_plan_snapshot( - &self, - snapshot: &ModelPlanSnapshot, - ) -> ModelSnapshotRepositoryResult<()> { - let models_json = serde_json::to_string(&snapshot.models)?; - let fetched_at = Utc::now().to_rfc3339(); - sqlx::query( - "insert into model_plan_snapshots (plan_type, models_json, fetched_at) values (?, ?, ?) \ - on conflict(plan_type) do update set models_json = excluded.models_json, fetched_at = excluded.fetched_at", - ) - .bind(&snapshot.plan_type) - .bind(models_json) - .bind(fetched_at) - .execute(&self.pool) - .await?; - Ok(()) - } - - pub async fn list_plan_snapshots( - &self, - ) -> ModelSnapshotRepositoryResult> { - let rows = sqlx::query( - "select plan_type, models_json from model_plan_snapshots order by plan_type asc", - ) - .fetch_all(&self.pool) - .await?; - rows.into_iter() - .map(|row| { - let plan_type: String = row.get("plan_type"); - let models_json: String = row.get("models_json"); - let models = serde_json::from_str::>(&models_json)?; - Ok(ModelPlanSnapshot { plan_type, models }) - }) - .collect() - } -} diff --git a/src/codex/models/service.rs b/src/codex/models/service.rs deleted file mode 100644 index 1fe69f128..000000000 --- a/src/codex/models/service.rs +++ /dev/null @@ -1,191 +0,0 @@ -use std::sync::Arc; - -use tokio::sync::Mutex; - -use crate::{ - codex::accounts::{ - model::{Account, AccountStatus}, - pool::AccountPool, - repository::AccountRepository, - }, - codex::gateway::fingerprint::model::Fingerprint, - codex::gateway::transport::http_client::{ - build_reqwest_client, CodexBackendClient, CodexRequestContext, - }, - codex::models::{ - catalog::{ModelCatalog, ModelPlanSnapshot}, - repository::ModelSnapshotRepository, - }, - config::AppConfig, -}; - -#[derive(Clone)] -pub struct ModelService { - config: Arc, - snapshot_repository: Option, - account_repository: Option, - account_pool: Arc>, - fingerprint: Fingerprint, -} - -#[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: Arc, - snapshot_repository: Option, - account_repository: Option, - account_pool: Arc>, - fingerprint: Fingerprint, - ) -> Self { - Self { - config, - snapshot_repository, - account_repository, - account_pool, - fingerprint, - } - } - - 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 = %error, "加载缓存模型快照失败"); - 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(), - self.fingerprint.clone(), - ) - }) - .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, - installation_id: None, - session_id: 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 = %error, plan_type, "刷新后端模型失败"); - 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/codex/serving/chat.rs b/src/codex/serving/chat.rs deleted file mode 100644 index 1b0400fbc..000000000 --- a/src/codex/serving/chat.rs +++ /dev/null @@ -1,313 +0,0 @@ -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::events::event::EventLevel, - codex::gateway::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::models::catalog::ModelCatalog, - codex::models::service::ModelService, - codex::serving::dispatch::{ - classify_upstream_account_retry, no_available_accounts_response, - normalize_service_tier_for_upstream, CodexRequestLogContext, CodexUpstreamService, - UpstreamAccountRecoveryTransition, - }, - codex::serving::http::errors::{ - codex_client_error_message, codex_client_error_response, - codex_client_error_response_with_status_and_message, model_not_found_response, - }, - config::ModelConfig, -}; - -#[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 acquired) = self.upstream.acquire_account(&codex_request.model).await else { - return no_available_accounts_response().into_response(); - }; - let mut log_context = CodexRequestLogContext::new( - request_id, - &acquired.account.id, - &codex_request.model, - client_stream, - started_at, - ); - - let mut excluded_account_ids = Vec::new(); - let mut model_unsupported_retry_used = false; - let response = loop { - self.upstream - .stagger_request(acquired.previous_slot_at) - .await; - let response = self - .upstream - .send_codex_request_with_upstream_retries( - &codex_request, - &acquired.account, - request_id, - ) - .await; - self.upstream.release_account(&acquired.account.id).await; - - match response { - Ok(response) => break response, - Err(error) => { - if let Some(retry) = - classify_upstream_account_retry(&error, model_unsupported_retry_used) - { - if retry.is_model_unsupported() { - model_unsupported_retry_used = true; - } - let transition = self - .upstream - .apply_account_recovery_transition( - &acquired.account, - retry, - &codex_request.model, - &mut excluded_account_ids, - false, - codex_client_error_message(&error), - ) - .await; - self.upstream - .log_response( - &log_context, - retry.status(), - EventLevel::Warn, - "v1 chat completions 上游请求将使用备用账户重试", - retry.metadata(client_stream), - ) - .await; - match transition { - UpstreamAccountRecoveryTransition::Retry(fallback) => { - log_context = log_context.with_account(&fallback.account.id); - acquired = *fallback; - continue; - } - UpstreamAccountRecoveryTransition::Respond { status, message } => { - let error_response = - codex_client_error_response_with_status_and_message( - error, status, &message, - ); - self.upstream - .log_response( - &log_context, - error_response.0, - EventLevel::Error, - "v1 chat completions fallback 已耗尽", - json!({"stream": client_stream}), - ) - .await; - return error_response.into_response(); - } - } - } - if self - .upstream - .record_request_attempt(&acquired.account.id, false) - .await - .is_err() - { - self.upstream - .log_response( - &log_context, - StatusCode::OK, - EventLevel::Warn, - "v1 chat completions 记录失败请求计数失败", - json!({"stream": client_stream, "requestAttemptStoreError": true}), - ) - .await; - } - let error_response = codex_client_error_response(error); - self.upstream - .log_response( - &log_context, - error_response.0, - EventLevel::Error, - "v1 chat completions 上游请求失败", - json!({"stream": client_stream}), - ) - .await; - return error_response.into_response(); - } - } - }; - if self - .upstream - .persist_cookies(&acquired.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(&acquired.account.id, usage, false) - .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, - codex_request.tuple_schema.as_ref(), - ) { - 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, - codex_request.tuple_schema.as_ref(), - ) { - 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/codex/serving/diagnostics.rs b/src/codex/serving/diagnostics.rs deleted file mode 100644 index 835078b6a..000000000 --- a/src/codex/serving/diagnostics.rs +++ /dev/null @@ -1,143 +0,0 @@ -use std::{sync::Arc, time::Instant}; - -use reqwest::StatusCode; -use serde::Serialize; - -use crate::{ - codex::gateway::{ - fingerprint::model::Fingerprint, - transport::http_client::{ - build_reqwest_client, CodexBackendClient, CodexClientError, CodexRequestContext, - }, - }, - config::AppConfig, -}; - -#[derive(Clone)] -pub struct DiagnosticsService { - config: Arc, - fingerprint: Fingerprint, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct UpstreamProbeResult { - pub target: &'static str, - pub status: &'static str, - pub backend_base_url: String, - pub endpoint: Option, - pub reachable: bool, - pub status_code: Option, - pub authorization: &'static str, - pub duration_ms: u64, - pub error: Option, -} - -impl DiagnosticsService { - pub fn new(config: Arc, fingerprint: Fingerprint) -> Self { - Self { - config, - fingerprint, - } - } - - pub(crate) fn fingerprint(&self) -> &Fingerprint { - &self.fingerprint - } - - pub async fn probe_upstream(&self, request_id: &str) -> UpstreamProbeResult { - let started_at = Instant::now(); - let client = match build_reqwest_client(self.config.tls.force_http11) { - Ok(client) => CodexBackendClient::new( - client, - self.config.api.base_url.clone(), - self.fingerprint.clone(), - ), - Err(error) => { - return self.unreachable_probe( - None, - started_at, - format!("failed to build HTTP client: {error}"), - ); - } - }; - - match client - .probe_models_endpoint(CodexRequestContext { - access_token: "diagnostic-probe", - account_id: None, - 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, - installation_id: None, - session_id: None, - }) - .await - { - Ok(probe) => self.reachable_probe(probe.endpoint, probe.status, started_at), - Err(error) => self.unreachable_probe(None, started_at, public_probe_error(&error)), - } - } - - fn reachable_probe( - &self, - endpoint: String, - status: StatusCode, - started_at: Instant, - ) -> UpstreamProbeResult { - UpstreamProbeResult { - target: "codexModels", - status: "reachable", - backend_base_url: self.config.api.base_url.clone(), - endpoint: Some(endpoint), - reachable: true, - status_code: Some(status.as_u16()), - authorization: authorization_status(status), - duration_ms: elapsed_millis(started_at), - error: None, - } - } - - fn unreachable_probe( - &self, - endpoint: Option, - started_at: Instant, - error: String, - ) -> UpstreamProbeResult { - UpstreamProbeResult { - target: "codexModels", - status: "unreachable", - backend_base_url: self.config.api.base_url.clone(), - endpoint, - reachable: false, - status_code: None, - authorization: "unknown", - duration_ms: elapsed_millis(started_at), - error: Some(error), - } - } -} - -fn authorization_status(status: StatusCode) -> &'static str { - if status.is_success() { - "accepted" - } else if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { - "rejected" - } else { - "unknown" - } -} - -fn elapsed_millis(started_at: Instant) -> u64 { - started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64 -} - -fn public_probe_error(error: &CodexClientError) -> String { - error.to_string().chars().take(200).collect() -} diff --git a/src/codex/serving/dispatch/affinity.rs b/src/codex/serving/dispatch/affinity.rs deleted file mode 100644 index 408a58631..000000000 --- a/src/codex/serving/dispatch/affinity.rs +++ /dev/null @@ -1,698 +0,0 @@ -use std::collections::HashMap; -use std::sync::Arc; -use std::time::Duration; - -use chrono::{DateTime, Utc}; -use sha2::{Digest, Sha256}; -use sqlx::{Row, SqlitePool}; -use thiserror::Error; -use tokio::sync::RwLock; -use tokio::time::interval; - -use crate::codex::gateway::conversation_identity::derive_stable_conversation_key; -use crate::codex::gateway::transport::types::CodexResponsesRequest; - -/// 会话亲和性条目 - 记录响应ID到账户的映射 -#[derive(Debug, Clone)] -pub struct AffinityEntry { - /// 账户ID - pub account_id: String, - /// 对话ID - pub conversation_id: String, - /// 上游 turn state token - pub turn_state: Option, - /// 指令哈希（SHA-256）- 用于缓存命中检测 - pub instructions_hash: Option, - /// 输入token数 - pub input_tokens: Option, - /// 函数调用ID列表 - pub function_call_ids: Vec, - /// 变体哈希 - 用于区分并发对话分支 - pub variant_hash: Option, - /// 创建时间 - pub created_at: DateTime, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct StoredSessionAffinity { - pub response_id: String, - pub account_id: String, - pub conversation_id: String, - pub turn_state: Option, - pub instructions_hash: Option, - pub input_tokens: Option, - pub function_call_ids: Vec, - pub variant_hash: Option, - pub expires_at: DateTime, - pub created_at: DateTime, -} - -#[derive(Debug, Error)] -pub enum SessionAffinityRepositoryError { - #[error("database error: {0}")] - Database(#[from] sqlx::Error), - #[error("invalid session affinity timestamp: {0}")] - InvalidTimestamp(#[from] chrono::ParseError), - #[error("invalid session affinity function call ids: {0}")] - InvalidFunctionCallIds(#[from] serde_json::Error), -} - -pub type SessionAffinityRepositoryResult = Result; - -#[derive(Clone)] -pub struct SessionAffinityRepository { - pool: SqlitePool, -} - -/// 会话亲和性映射 - 将 previous_response_id 映射到创建它的账户 -/// -/// 功能： -/// - 支持服务器端对话历史重用 -/// - 提示缓存命中（缓存按账户存储在后端） -/// - 4小时TTL -/// - 每10分钟自动清理过期条目 -pub struct SessionAffinityMap { - map: Arc>>, - ttl: Duration, - cleanup_handle: Option>, -} - -const DEFAULT_TTL: Duration = Duration::from_secs(4 * 60 * 60); // 4小时 -const CLEANUP_INTERVAL: Duration = Duration::from_secs(10 * 60); // 10分钟 - -const UPSERT_SESSION_AFFINITY_SQL: &str = r" -insert into session_affinities ( - response_id, - account_id, - conversation_id, - turn_state, - instructions_hash, - input_tokens, - function_call_ids_json, - variant_hash, - expires_at, - created_at -) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -on conflict(response_id) do update set - account_id = excluded.account_id, - conversation_id = excluded.conversation_id, - turn_state = excluded.turn_state, - instructions_hash = excluded.instructions_hash, - input_tokens = excluded.input_tokens, - function_call_ids_json = excluded.function_call_ids_json, - variant_hash = excluded.variant_hash, - expires_at = excluded.expires_at, - created_at = excluded.created_at"; - -const LIST_ACTIVE_SESSION_AFFINITIES_SQL: &str = r" -select - response_id, - account_id, - conversation_id, - turn_state, - instructions_hash, - input_tokens, - function_call_ids_json, - variant_hash, - expires_at, - created_at -from session_affinities -where expires_at > ? -order by created_at asc, response_id asc"; - -pub(super) fn compute_variant_hash(request: &CodexResponsesRequest) -> String { - compute_variant_hash_with_identity(request, request.variant_identity.as_deref()) -} - -fn compute_variant_hash_with_identity( - request: &CodexResponsesRequest, - identity: Option<&str>, -) -> String { - let tools_json = request - .tools - .as_ref() - .map_or_else(|| "[]".to_string(), |tools| tools_json(tools)); - let mut hasher = Sha256::new(); - hasher.update(request.instructions.as_bytes()); - hasher.update(b"\0"); - hasher.update(tools_json.as_bytes()); - if let Some(identity) = identity - .map(str::trim) - .filter(|identity| !identity.is_empty()) - { - hasher.update(b"\0"); - hasher.update(identity.as_bytes()); - } - hex::encode(hasher.finalize()).chars().take(12).collect() -} - -fn tools_json(tools: &[serde_json::Value]) -> String { - serde_json::to_string(tools).unwrap_or_else(|_| "[]".to_string()) -} - -pub(super) fn prepare_variant_identity(request: &mut CodexResponsesRequest) { - request.variant_identity = build_variant_identity(request); -} - -fn build_variant_identity(request: &CodexResponsesRequest) -> Option { - let mut parts = Vec::with_capacity(2); - if let Some(window_id) = non_empty_string(request.codex_window_id.as_deref()) { - parts.push(format!("window:{window_id}")); - } - if request.explicit_prompt_cache_key - || non_empty_string(request.client_conversation_id.as_deref()).is_some() - { - if let Some(anchor) = derive_stable_conversation_key(request) { - parts.push(format!("anchor:{anchor}")); - } - } - - (!parts.is_empty()).then(|| parts.join("\0")) -} - -fn non_empty_string(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -pub(super) fn hash_instructions(instructions: Option<&str>) -> String { - let mut hasher = Sha256::new(); - hasher.update(instructions.unwrap_or_default().as_bytes()); - hex::encode(hasher.finalize()) -} - -impl SessionAffinityRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - pub async fn upsert( - &self, - response_id: &str, - entry: &AffinityEntry, - ttl: Duration, - ) -> SessionAffinityRepositoryResult<()> { - let function_call_ids_json = serde_json::to_string(&entry.function_call_ids)?; - let expires_at = entry.created_at + chrono_duration_from_std(ttl); - sqlx::query(UPSERT_SESSION_AFFINITY_SQL) - .bind(response_id) - .bind(&entry.account_id) - .bind(&entry.conversation_id) - .bind(&entry.turn_state) - .bind(&entry.instructions_hash) - .bind(entry.input_tokens.map(u64_to_i64_saturating)) - .bind(function_call_ids_json) - .bind(&entry.variant_hash) - .bind(expires_at.to_rfc3339()) - .bind(entry.created_at.to_rfc3339()) - .execute(&self.pool) - .await?; - Ok(()) - } - - pub async fn list_active( - &self, - now: DateTime, - ) -> SessionAffinityRepositoryResult> { - let rows = sqlx::query(LIST_ACTIVE_SESSION_AFFINITIES_SQL) - .bind(now.to_rfc3339()) - .fetch_all(&self.pool) - .await?; - rows.into_iter() - .map(|row| stored_session_affinity_from_row(&row)) - .collect() - } - - pub async fn delete_expired(&self, now: DateTime) -> SessionAffinityRepositoryResult { - let result = sqlx::query("delete from session_affinities where expires_at <= ?") - .bind(now.to_rfc3339()) - .execute(&self.pool) - .await?; - Ok(result.rows_affected()) - } -} - -impl SessionAffinityMap { - pub fn new(ttl: Duration) -> Self { - Self { - map: Arc::new(RwLock::new(HashMap::new())), - ttl, - cleanup_handle: None, - } - } - - pub fn with_default_ttl() -> Self { - Self::new(DEFAULT_TTL) - } - - /// 启动自动清理任务 - pub fn start_cleanup(&mut self) { - if self.cleanup_handle.is_some() { - return; // 已经启动 - } - - let map = self.map.clone(); - let ttl = self.ttl; - - let handle = tokio::spawn(async move { - let mut ticker = interval(CLEANUP_INTERVAL); - loop { - ticker.tick().await; - let now = Utc::now(); - let mut map_lock = map.write().await; - map_lock.retain(|_, entry| !entry_expired(entry.created_at, ttl, now)); - } - }); - - self.cleanup_handle = Some(handle); - } - - /// 停止自动清理任务 - pub fn stop_cleanup(&mut self) { - if let Some(handle) = self.cleanup_handle.take() { - handle.abort(); - } - } - - /// 记录响应ID到账户的映射 - #[allow(clippy::too_many_arguments)] - pub async fn record( - &self, - response_id: String, - account_id: String, - conversation_id: String, - turn_state: Option, - instructions: Option<&str>, - input_tokens: Option, - function_call_ids: Option>, - variant_hash: Option, - ) -> AffinityEntry { - let instructions_hash = instructions.map(|value| hash_instructions(Some(value))); - - let entry = AffinityEntry { - account_id, - conversation_id, - turn_state, - instructions_hash, - input_tokens, - function_call_ids: function_call_ids.unwrap_or_default(), - variant_hash, - created_at: Utc::now(), - }; - - self.map.write().await.insert(response_id, entry.clone()); - entry - } - - pub async fn restore(&self, records: Vec) -> usize { - let now = Utc::now(); - let mut restored = 0usize; - let mut map = self.map.write().await; - for record in records { - if record.expires_at <= now || entry_expired(record.created_at, self.ttl, now) { - continue; - } - map.insert( - record.response_id, - AffinityEntry { - account_id: record.account_id, - conversation_id: record.conversation_id, - turn_state: record.turn_state, - instructions_hash: record.instructions_hash, - input_tokens: record.input_tokens, - function_call_ids: record.function_call_ids, - variant_hash: record.variant_hash, - created_at: record.created_at, - }, - ); - restored = restored.saturating_add(1); - } - restored - } - - pub fn ttl(&self) -> Duration { - self.ttl - } - - /// 查找响应ID对应的账户ID - pub async fn lookup_account(&self, response_id: &str) -> Option { - self.get_entry(response_id).await.map(|e| e.account_id) - } - - /// 查找响应ID对应的对话ID - pub async fn lookup_conversation_id(&self, response_id: &str) -> Option { - self.get_entry(response_id).await.map(|e| e.conversation_id) - } - - /// 查找对话ID的最新响应ID - /// - /// 参数： - /// - conversation_id: 对话ID - /// - max_age: 最大年龄（用于避免使用已过期的提示缓存） - /// - variant_hash: 变体哈希（用于区分并发分支） - pub async fn lookup_latest_response_by_conversation( - &self, - conversation_id: &str, - max_age: Option, - variant_hash: Option<&str>, - ) -> Option { - let map = self.map.read().await; - let now = Utc::now(); - - let mut latest_response_id: Option = None; - let mut latest_created_at: Option> = None; - - for (response_id, entry) in map.iter() { - // 过滤对话ID - if entry.conversation_id != conversation_id { - continue; - } - - // 过滤变体哈希 - if let Some(vh) = variant_hash { - if entry.variant_hash.as_deref() != Some(vh) { - continue; - } - } - - // 检查TTL - if entry_expired(entry.created_at, self.ttl, now) { - continue; - } - - // 检查最大年龄 - if let Some(max) = max_age { - if entry_age_exceeds(entry.created_at, now, max) { - continue; - } - } - - // 保留最新的 - if latest_created_at.is_none_or(|latest| entry.created_at > latest) { - latest_created_at = Some(entry.created_at); - latest_response_id = Some(response_id.clone()); - } - } - - latest_response_id - } - - /// 查找响应ID对应的 turn state - pub async fn lookup_turn_state(&self, response_id: &str) -> Option { - self.get_entry(response_id).await.and_then(|e| e.turn_state) - } - - /// 查找响应ID对应的指令哈希 - pub async fn lookup_instructions_hash(&self, response_id: &str) -> Option { - self.get_entry(response_id) - .await - .and_then(|e| e.instructions_hash) - } - - /// 查找响应ID对应的输入token数 - pub async fn lookup_input_tokens(&self, response_id: &str) -> Option { - self.get_entry(response_id) - .await - .and_then(|e| e.input_tokens) - } - - /// 查找响应ID对应的函数调用ID列表 - pub async fn lookup_function_call_ids(&self, response_id: &str) -> Vec { - self.get_entry(response_id) - .await - .map(|e| e.function_call_ids) - .unwrap_or_default() - } - - /// 删除响应ID（上游返回 not-found 时调用） - pub async fn forget(&self, response_id: &str) { - self.map.write().await.remove(response_id); - } - - /// 获取映射大小 - pub async fn size(&self) -> usize { - self.map.read().await.len() - } - - /// 清空所有映射 - pub async fn clear(&self) { - self.map.write().await.clear(); - } - - /// 获取条目（检查TTL） - async fn get_entry(&self, response_id: &str) -> Option { - let map = self.map.read().await; - let entry = map.get(response_id)?; - - let now = Utc::now(); - if entry_expired(entry.created_at, self.ttl, now) { - return None; - } - - Some(entry.clone()) - } -} - -fn entry_expired(created_at: DateTime, ttl: Duration, now: DateTime) -> bool { - now.signed_duration_since(created_at) - .to_std() - .is_ok_and(|age| age >= ttl) -} - -fn entry_age_exceeds(created_at: DateTime, now: DateTime, max_age: Duration) -> bool { - now.signed_duration_since(created_at) - .to_std() - .is_ok_and(|age| age > max_age) -} - -fn chrono_duration_from_std(duration: Duration) -> chrono::Duration { - chrono::Duration::from_std(duration).unwrap_or_else(|_| chrono::Duration::seconds(i64::MAX)) -} - -fn u64_to_i64_saturating(value: u64) -> i64 { - value.min(i64::MAX as u64) as i64 -} - -fn optional_nonnegative_i64_to_u64(value: Option) -> Option { - value.and_then(|value| u64::try_from(value).ok()) -} - -fn parse_rfc3339(value: &str) -> SessionAffinityRepositoryResult> { - Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc)) -} - -fn stored_session_affinity_from_row( - row: &sqlx::sqlite::SqliteRow, -) -> SessionAffinityRepositoryResult { - let function_call_ids_json = row.get::("function_call_ids_json"); - Ok(StoredSessionAffinity { - response_id: row.get("response_id"), - account_id: row.get("account_id"), - conversation_id: row.get("conversation_id"), - turn_state: row.get("turn_state"), - instructions_hash: row.get("instructions_hash"), - input_tokens: optional_nonnegative_i64_to_u64(row.get("input_tokens")), - function_call_ids: serde_json::from_str(&function_call_ids_json)?, - variant_hash: row.get("variant_hash"), - expires_at: parse_rfc3339(&row.get::("expires_at"))?, - created_at: parse_rfc3339(&row.get::("created_at"))?, - }) -} - -impl Drop for SessionAffinityMap { - fn drop(&mut self) { - self.stop_cleanup(); - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::codex::gateway::transport::types::CodexResponsesRequest; - use serde_json::json; - - #[tokio::test] - async fn session_affinity_records_and_retrieves_account() { - let map = SessionAffinityMap::with_default_ttl(); - - map.record( - "resp_123".to_string(), - "acc_456".to_string(), - "conv_789".to_string(), - None, - None, - None, - None, - None, - ) - .await; - - let account = map.lookup_account("resp_123").await; - assert_eq!(account, Some("acc_456".to_string())); - } - - #[tokio::test] - async fn session_affinity_returns_none_for_unknown_response() { - let map = SessionAffinityMap::with_default_ttl(); - let account = map.lookup_account("unknown").await; - assert_eq!(account, None); - } - - #[tokio::test] - async fn session_affinity_expires_old_entries() { - let map = SessionAffinityMap::new(Duration::from_millis(100)); - - map.record( - "resp_123".to_string(), - "acc_456".to_string(), - "conv_789".to_string(), - None, - None, - None, - None, - None, - ) - .await; - - // 立即查找应该成功 - assert!(map.lookup_account("resp_123").await.is_some()); - - // 等待过期 - tokio::time::sleep(Duration::from_millis(150)).await; - - // 过期后应该返回 None - assert!(map.lookup_account("resp_123").await.is_none()); - } - - #[tokio::test] - async fn session_affinity_finds_latest_response_by_conversation() { - let map = SessionAffinityMap::with_default_ttl(); - - map.record( - "resp_1".to_string(), - "acc_1".to_string(), - "conv_789".to_string(), - None, - None, - None, - None, - None, - ) - .await; - - tokio::time::sleep(Duration::from_millis(10)).await; - - map.record( - "resp_2".to_string(), - "acc_1".to_string(), - "conv_789".to_string(), - None, - None, - None, - None, - None, - ) - .await; - - let latest = map - .lookup_latest_response_by_conversation("conv_789", None, None) - .await; - assert_eq!(latest, Some("resp_2".to_string())); - } - - #[tokio::test] - async fn session_affinity_forget_removes_entry() { - let map = SessionAffinityMap::with_default_ttl(); - - map.record( - "resp_123".to_string(), - "acc_456".to_string(), - "conv_789".to_string(), - None, - None, - None, - None, - None, - ) - .await; - - assert!(map.lookup_account("resp_123").await.is_some()); - - map.forget("resp_123").await; - - assert!(map.lookup_account("resp_123").await.is_none()); - } - - #[test] - fn compute_variant_hash_should_follow_original_shape_contract() { - let mut request = CodexResponsesRequest::new_http_sse( - "gpt-5.5", - "You are Codex.", - vec![json!({"role": "user", "content": "hello"})], - ); - request.tools = Some(vec![json!({"type": "function", "name": "shell"})]); - let first = compute_variant_hash(&request); - request.input = vec![ - json!({"role": "user", "content": "hello"}), - json!({"role": "assistant", "content": "cached"}), - json!({"role": "user", "content": "continue"}), - ]; - request.previous_response_id = Some("resp_dynamic".to_string()); - request.prompt_cache_key = Some("conversation_dynamic".to_string()); - request.include = Some(vec!["reasoning.encrypted_content".to_string()]); - let same_shape = compute_variant_hash(&request); - - request.instructions = "You are Codex with different instructions.".to_string(); - let different_instructions = compute_variant_hash(&request); - - assert_eq!(first, same_shape); - assert_ne!(first, different_instructions); - } - - #[test] - fn compute_variant_hash_should_include_optional_identity() { - let request = CodexResponsesRequest::new_http_sse( - "gpt-5.5", - "You are Codex.", - vec![json!({"role": "user", "content": "hello"})], - ); - let first = compute_variant_hash_with_identity(&request, Some("window:main")); - let second = compute_variant_hash_with_identity(&request, Some("window:review")); - - assert_ne!(first, second); - } - - #[test] - fn compute_variant_hash_should_include_prepared_codex_window_identity() { - let mut first = CodexResponsesRequest::new_http_sse( - "gpt-5.5", - "You are Codex.", - vec![json!({"role": "user", "content": "hello"})], - ); - first.codex_window_id = Some("main-window".to_string()); - prepare_variant_identity(&mut first); - - let mut second = first.clone(); - second.codex_window_id = Some("review-window".to_string()); - prepare_variant_identity(&mut second); - - assert_ne!(compute_variant_hash(&first), compute_variant_hash(&second)); - } - - #[test] - fn compute_variant_hash_should_include_explicit_prompt_cache_anchor() { - let mut first = CodexResponsesRequest::new_http_sse( - "gpt-5.5", - "You are Codex.", - vec![json!({"role": "user", "content": "alpha"})], - ); - first.prompt_cache_key = Some("shared-session".to_string()); - first.explicit_prompt_cache_key = true; - prepare_variant_identity(&mut first); - - let mut second = first.clone(); - second.input = vec![json!({"role": "user", "content": "beta"})]; - prepare_variant_identity(&mut second); - - assert_ne!(compute_variant_hash(&first), compute_variant_hash(&second)); - } -} diff --git a/src/codex/serving/dispatch/fallback.rs b/src/codex/serving/dispatch/fallback.rs deleted file mode 100644 index d86fa220d..000000000 --- a/src/codex/serving/dispatch/fallback.rs +++ /dev/null @@ -1,632 +0,0 @@ -use axum::http::StatusCode; -use chrono::{Duration, Utc}; -use serde_json::{json, Value}; - -use crate::codex::{ - accounts::{ - model::{Account, AccountStatus}, - pool::AccountPoolStatusSummary, - }, - gateway::transport::{http_client::CodexClientError, rate_limits::cooldown_with_jitter}, -}; - -use super::{ - stream::contains_invalid_encrypted_content_signal, usage::record_request_attempt, - CodexUpstreamDependencies, -}; - -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, PartialEq, Eq)] -pub(crate) enum UpstreamAccountRetry { - ModelUnsupported { status: StatusCode }, - RateLimited { retry_after_seconds: u64 }, - QuotaExhausted, - CloudflareChallenge { cooldown_seconds: u64 }, - CloudflarePathBlock, - TokenInvalid { account_status: AccountStatus }, - Banned, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum UpstreamRequestRecovery { - StripPreviousResponse { reason: HistoryRecoveryReason }, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum UpstreamRecoveryAction { - Request(UpstreamRequestRecovery), - Account(UpstreamAccountRetry), - RespondWithError, -} - -#[derive(Debug, Clone, Copy)] -pub(crate) struct UpstreamRecoveryState { - pub(crate) request_recovery_used: bool, - pub(crate) model_unsupported_retry_used: bool, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum HistoryRecoveryReason { - PreviousResponseNotFound, - UnansweredFunctionCall, - InvalidReasoningReplay, -} - -impl UpstreamAccountRetry { - pub(crate) fn status(self) -> StatusCode { - match self { - Self::ModelUnsupported { status } => status, - Self::RateLimited { .. } => StatusCode::TOO_MANY_REQUESTS, - Self::QuotaExhausted => StatusCode::PAYMENT_REQUIRED, - Self::CloudflareChallenge { .. } => StatusCode::BAD_GATEWAY, - Self::CloudflarePathBlock => StatusCode::BAD_GATEWAY, - Self::TokenInvalid { .. } => StatusCode::UNAUTHORIZED, - Self::Banned => StatusCode::FORBIDDEN, - } - } - - pub(crate) fn metadata(self, stream: bool) -> Value { - match self { - Self::ModelUnsupported { .. } => json!({ - "stream": stream, - "retry": true, - "reason": "modelUnsupported", - }), - 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::CloudflarePathBlock => json!({ - "stream": stream, - "retry": true, - "reason": "cloudflarePathBlock", - }), - Self::TokenInvalid { account_status } => json!({ - "stream": stream, - "retry": true, - "reason": "tokenInvalid", - "accountStatus": account_status_metadata(account_status), - }), - Self::Banned => json!({ - "stream": stream, - "retry": true, - "reason": "banned", - }), - } - } - - pub(crate) fn is_model_unsupported(self) -> bool { - matches!(self, Self::ModelUnsupported { .. }) - } - - pub(crate) fn fallback_response_message(self, upstream_message: String) -> String { - match self { - Self::CloudflareChallenge { .. } => { - "Upstream blocked the request (Cloudflare challenge)".to_string() - } - Self::CloudflarePathBlock => { - "Upstream blocked the request (Cloudflare path-block)".to_string() - } - _ => upstream_message, - } - } -} - -impl UpstreamRequestRecovery { - pub(crate) fn metadata(self, stream: bool, stale_response_id: Option<&str>) -> Value { - match self { - Self::StripPreviousResponse { reason } => json!({ - "stream": stream, - "retry": true, - "reason": reason.as_metadata_value(), - "previousResponseId": stale_response_id, - }), - } - } -} - -impl HistoryRecoveryReason { - fn as_metadata_value(self) -> &'static str { - match self { - Self::PreviousResponseNotFound => "previousResponseNotFound", - Self::UnansweredFunctionCall => "unansweredFunctionCall", - Self::InvalidReasoningReplay => "invalidReasoningReplay", - } - } -} - -pub(crate) fn classify_upstream_request_recovery( - error: &CodexClientError, - recovery_already_used: bool, -) -> Option { - if recovery_already_used { - return None; - } - if is_previous_response_not_found_error(error) { - return Some(UpstreamRequestRecovery::StripPreviousResponse { - reason: HistoryRecoveryReason::PreviousResponseNotFound, - }); - } - if is_unanswered_function_call_error(error) { - return Some(UpstreamRequestRecovery::StripPreviousResponse { - reason: HistoryRecoveryReason::UnansweredFunctionCall, - }); - } - if is_invalid_reasoning_replay_error(error) { - return Some(UpstreamRequestRecovery::StripPreviousResponse { - reason: HistoryRecoveryReason::InvalidReasoningReplay, - }); - } - None -} - -pub(crate) fn classify_upstream_recovery_action( - error: &CodexClientError, - state: UpstreamRecoveryState, -) -> UpstreamRecoveryAction { - if let Some(recovery) = classify_upstream_request_recovery(error, state.request_recovery_used) { - return UpstreamRecoveryAction::Request(recovery); - } - if let Some(retry) = classify_upstream_account_retry(error, state.model_unsupported_retry_used) - { - return UpstreamRecoveryAction::Account(retry); - } - UpstreamRecoveryAction::RespondWithError -} - -pub(crate) fn classify_upstream_account_retry( - error: &CodexClientError, - model_unsupported_retry_used: bool, -) -> Option { - match error { - CodexClientError::Upstream { status, .. } - if is_model_not_supported_error(error) && !model_unsupported_retry_used => - { - Some(UpstreamAccountRetry::ModelUnsupported { status: *status }) - } - 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) - } - } - CodexClientError::Upstream { status, body, .. } - if *status == StatusCode::NOT_FOUND && body.trim().is_empty() => - { - Some(UpstreamAccountRetry::CloudflarePathBlock) - } - CodexClientError::Upstream { status, body, .. } if *status == StatusCode::UNAUTHORIZED => { - Some(UpstreamAccountRetry::TokenInvalid { - account_status: account_status_for_unauthorized(body), - }) - } - _ => None, - } -} - -fn is_model_not_supported_error(error: &CodexClientError) -> bool { - let CodexClientError::Upstream { status, body, .. } = error else { - return false; - }; - if !status.is_client_error() || *status == StatusCode::TOO_MANY_REQUESTS { - return false; - } - let haystack = format!( - "{} {} {}", - error_code(body).unwrap_or_default(), - error_message(body).unwrap_or_default(), - body - ) - .to_ascii_lowercase(); - haystack.contains("model") - && (haystack.contains("not supported") - || haystack.contains("not_supported") - || haystack.contains("not available") - || haystack.contains("not_available")) -} - -fn account_status_for_unauthorized(body: &str) -> AccountStatus { - if body.to_ascii_lowercase().contains("deactivated") { - AccountStatus::Banned - } else { - AccountStatus::Expired - } -} - -fn is_previous_response_not_found_error(error: &CodexClientError) -> bool { - let CodexClientError::Upstream { body, .. } = error else { - return false; - }; - if error_code(body).as_deref() == Some("previous_response_not_found") { - return true; - } - let lower = body.to_ascii_lowercase(); - lower.contains("previous_response_not_found") - || (lower.contains("previous response with id") && lower.contains("not found")) -} - -fn is_unanswered_function_call_error(error: &CodexClientError) -> bool { - let CodexClientError::Upstream { status, body, .. } = error else { - return false; - }; - if *status != StatusCode::BAD_REQUEST { - return false; - } - body.to_ascii_lowercase() - .contains("no tool output found for function call") -} - -fn is_invalid_reasoning_replay_error(error: &CodexClientError) -> bool { - let CodexClientError::Upstream { body, .. } = error else { - return false; - }; - let Ok(value) = serde_json::from_str::(body) else { - return false; - }; - contains_invalid_encrypted_content_signal(&value) -} - -fn error_code(body: &str) -> Option { - let value = serde_json::from_str::(body).ok()?; - value - .pointer("/response/error/code") - .or_else(|| value.pointer("/response/error/type")) - .or_else(|| value.pointer("/error/code")) - .or_else(|| value.pointer("/error/type")) - .and_then(Value::as_str) - .map(|code| code.to_ascii_lowercase()) -} - -fn error_message(body: &str) -> Option { - let value = serde_json::from_str::(body).ok()?; - value - .pointer("/response/error/message") - .or_else(|| value.pointer("/error/message")) - .and_then(Value::as_str) - .map(ToOwned::to_owned) -} - -pub(super) fn build_account_exhaustion_detail( - summary: AccountPoolStatusSummary, - message: &str, -) -> String { - let mut parts = Vec::new(); - if summary.rate_limited > 0 { - parts.push(format!("{} rate-limited", summary.rate_limited)); - } - if summary.expired > 0 { - parts.push(format!("{} expired", summary.expired)); - } - if summary.banned > 0 { - parts.push(format!("{} banned", summary.banned)); - } - if summary.disabled > 0 { - parts.push(format!("{} disabled", summary.disabled)); - } - if summary.quota_exhausted > 0 { - parts.push(format!("{} quota-exhausted", summary.quota_exhausted)); - } - if summary.refreshing > 0 { - parts.push(format!("{} refreshing", summary.refreshing)); - } - if parts.is_empty() { - format!("No accounts available. {message}") - } else { - format!("All accounts exhausted ({}). {message}", parts.join(", ")) - } -} - -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") -} - -pub(super) async fn apply_upstream_account_retry_with_deps( - deps: &CodexUpstreamDependencies, - account: &Account, - retry: UpstreamAccountRetry, - image_generation_requested: bool, -) { - if let Err(error) = record_request_attempt(deps, &account.id, image_generation_requested).await - { - tracing::warn!( - error = ?error, - account_id = %account.id, - "记录上游失败账户请求尝试失败" - ); - } - let evict_websocket_pool = match retry { - UpstreamAccountRetry::ModelUnsupported { .. } => { - tracing::warn!( - account_id = %account.id, - "上游账号不支持当前 model，将尝试备用账号" - ); - false - } - UpstreamAccountRetry::RateLimited { - retry_after_seconds, - } => { - let cooldown_until = Utc::now() + cooldown_with_jitter(retry_after_seconds, 2_000); - if let Some(repo) = deps.account_repository.as_ref() { - if let Err(error) = repo - .set_quota_cooldown_until(&account.id, cooldown_until) - .await - { - tracing::warn!( - error = %error, - account_id = %account.id, - cooldown_until = %cooldown_until, - "持久化 quota cooldown 失败" - ); - } - } - deps.account_pool - .lock() - .await - .mark_quota_limited_until(&account.id, cooldown_until); - true - } - UpstreamAccountRetry::QuotaExhausted => { - set_account_status(deps, account, AccountStatus::QuotaExhausted).await; - true - } - UpstreamAccountRetry::CloudflareChallenge { cooldown_seconds } => { - let cooldown_until = Utc::now() + Duration::seconds(cooldown_seconds as i64); - if let Some(cookie_repo) = deps.cookie_repository.as_ref() { - if let Err(error) = cookie_repo.delete_account_cookies(&account.id).await { - tracing::warn!( - error = %error, - account_id = %account.id, - "清理 Cloudflare 阻断账户 cookies 失败" - ); - } - } - if let Some(repo) = deps.account_repository.as_ref() { - if let Err(error) = repo - .set_cloudflare_cooldown_until(&account.id, cooldown_until) - .await - { - tracing::warn!( - error = %error, - account_id = %account.id, - cooldown_until = %cooldown_until, - "持久化 Cloudflare cooldown 失败" - ); - } - } - deps.account_pool - .lock() - .await - .set_cloudflare_cooldown_until(&account.id, cooldown_until); - true - } - UpstreamAccountRetry::CloudflarePathBlock => { - if let Some(cookie_repo) = deps.cookie_repository.as_ref() { - if let Err(error) = cookie_repo.delete_account_cookies(&account.id).await { - tracing::warn!( - error = %error, - account_id = %account.id, - "清理 Cloudflare path-block 账户 cookies 失败" - ); - } - } - let block_count = deps - .cf_path_block_tracker - .record_path_block(&account.id) - .await; - if deps.cf_path_block_tracker.should_disable(&account.id).await { - set_account_status(deps, account, AccountStatus::Disabled).await; - } - tracing::warn!( - account_id = %account.id, - block_count, - "上游返回 Cloudflare path-block 404，已清理 cookies" - ); - true - } - UpstreamAccountRetry::TokenInvalid { account_status } => { - set_account_status(deps, account, account_status).await; - true - } - UpstreamAccountRetry::Banned => { - set_account_status(deps, account, AccountStatus::Banned).await; - true - } - }; - if evict_websocket_pool { - deps.websocket_pool.evict_account(&account.id).await; - } -} - -fn account_status_metadata(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", - } -} - -async fn set_account_status( - deps: &CodexUpstreamDependencies, - account: &Account, - status: AccountStatus, -) { - if let Some(repo) = deps.account_repository.as_ref() { - if let Err(error) = repo.set_status(&account.id, status).await { - tracing::warn!( - error = %error, - account_id = %account.id, - status = ?status, - "持久化上游账户状态失败" - ); - } - } - deps.account_pool - .lock() - .await - .set_status(&account.id, status); -} - -#[cfg(test)] -mod tests { - use super::*; - - fn upstream_error(status: StatusCode, code: &str, message: &str) -> CodexClientError { - CodexClientError::Upstream { - status, - retry_after_seconds: None, - body: json!({ - "error": { - "code": code, - "message": message, - } - }) - .to_string(), - } - } - - #[test] - fn classify_upstream_recovery_action_should_prioritize_request_recovery() { - let error = upstream_error( - StatusCode::TOO_MANY_REQUESTS, - "previous_response_not_found", - "Previous response with id resp_missing was not found", - ); - - let action = classify_upstream_recovery_action( - &error, - UpstreamRecoveryState { - request_recovery_used: false, - model_unsupported_retry_used: false, - }, - ); - - assert_eq!( - action, - UpstreamRecoveryAction::Request(UpstreamRequestRecovery::StripPreviousResponse { - reason: HistoryRecoveryReason::PreviousResponseNotFound, - }) - ); - } - - #[test] - fn classify_upstream_recovery_action_should_delegate_after_request_recovery_is_used() { - let error = upstream_error( - StatusCode::TOO_MANY_REQUESTS, - "previous_response_not_found", - "Previous response with id resp_missing was not found", - ); - - let action = classify_upstream_recovery_action( - &error, - UpstreamRecoveryState { - request_recovery_used: true, - model_unsupported_retry_used: false, - }, - ); - - assert_eq!( - action, - UpstreamRecoveryAction::Account(UpstreamAccountRetry::RateLimited { - retry_after_seconds: DEFAULT_RATE_LIMIT_BACKOFF_SECONDS, - }) - ); - } - - #[test] - fn classify_upstream_recovery_action_should_strip_history_after_invalid_reasoning_replay() { - let error = upstream_error( - StatusCode::SERVICE_UNAVAILABLE, - "invalid_encrypted_content", - "Invalid encrypted content in reasoning replay", - ); - - let action = classify_upstream_recovery_action( - &error, - UpstreamRecoveryState { - request_recovery_used: false, - model_unsupported_retry_used: false, - }, - ); - - assert_eq!( - action, - UpstreamRecoveryAction::Request(UpstreamRequestRecovery::StripPreviousResponse { - reason: HistoryRecoveryReason::InvalidReasoningReplay, - }) - ); - } - - #[test] - fn classify_upstream_recovery_action_should_limit_model_retry_to_once() { - let error = upstream_error( - StatusCode::BAD_REQUEST, - "model_not_supported", - "Model gpt-5.5 is not supported on this account plan", - ); - - let first_action = classify_upstream_recovery_action( - &error, - UpstreamRecoveryState { - request_recovery_used: false, - model_unsupported_retry_used: false, - }, - ); - let second_action = classify_upstream_recovery_action( - &error, - UpstreamRecoveryState { - request_recovery_used: false, - model_unsupported_retry_used: true, - }, - ); - - assert_eq!( - first_action, - UpstreamRecoveryAction::Account(UpstreamAccountRetry::ModelUnsupported { - status: StatusCode::BAD_REQUEST, - }) - ); - assert_eq!(second_action, UpstreamRecoveryAction::RespondWithError); - } -} diff --git a/src/codex/serving/dispatch/limits.rs b/src/codex/serving/dispatch/limits.rs deleted file mode 100644 index 54f92cb85..000000000 --- a/src/codex/serving/dispatch/limits.rs +++ /dev/null @@ -1,91 +0,0 @@ -use chrono::Utc; -use serde_json::Value; - -use crate::codex::gateway::transport::rate_limits::{parse_rate_limit_headers, rate_limit_quota}; - -use super::CodexUpstreamDependencies; - -pub(super) async fn apply_rate_limit_headers_with_deps( - deps: &CodexUpstreamDependencies, - account_id: &str, - plan_type: Option<&str>, - rate_limit_headers: &[(String, String)], -) { - let Some(rate_limits) = parse_rate_limit_headers(rate_limit_headers) else { - return; - }; - - let existing_quota = existing_quota_json(deps, account_id).await; - let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref()); - deps.account_pool - .lock() - .await - .set_quota_verify_required(account_id, false); - if let Some(repo) = deps.account_repository.as_ref() { - if let Err(error) = repo.update_quota_json(account_id, "a.to_string()).await { - tracing::warn!( - error = %error, - account_id = %account_id, - "被动同步 quota 缓存失败" - ); - } - if let Err(error) = repo.set_quota_verify_required(account_id, false).await { - tracing::warn!( - error = %error, - account_id = %account_id, - "清理被动 quota dirty 标记失败" - ); - } - } - - let Some(reset_at) = rate_limits.primary_reset_at() else { - return; - }; - deps.account_pool.lock().await.sync_rate_limit_window( - account_id, - reset_at, - rate_limits.primary_limit_window_seconds(), - ); - if let Some(repo) = deps.account_repository.as_ref() { - if let Err(error) = repo - .sync_rate_limit_window( - account_id, - reset_at, - rate_limits.primary_limit_window_seconds(), - ) - .await - { - tracing::warn!( - error = %error, - account_id = %account_id, - window_reset_at = %reset_at, - "持久化 rate-limit window 失败" - ); - } - } - if !rate_limits.primary_limit_reached() || reset_at <= Utc::now() { - return; - } - - if let Some(repo) = deps.account_repository.as_ref() { - if let Err(error) = repo.set_quota_cooldown_until(account_id, reset_at).await { - tracing::warn!( - error = %error, - account_id = %account_id, - cooldown_until = %reset_at, - "持久化被动 quota cooldown 失败" - ); - } - } - deps.account_pool - .lock() - .await - .mark_quota_limited_until(account_id, reset_at); - deps.websocket_pool.evict_account(account_id).await; -} - -async fn existing_quota_json(deps: &CodexUpstreamDependencies, account_id: &str) -> Option { - let repo = deps.account_repository.as_ref()?; - let raw = repo.get_quota_json(account_id).await.ok().flatten()?; - serde_json::from_str(&raw).ok() -} diff --git a/src/codex/serving/dispatch/mod.rs b/src/codex/serving/dispatch/mod.rs deleted file mode 100644 index 0ead13d32..000000000 --- a/src/codex/serving/dispatch/mod.rs +++ /dev/null @@ -1,2063 +0,0 @@ -pub mod affinity; -pub mod fallback; -mod implicit_resume; -mod limits; -// 上游调度辅助命名为 routing，避免出现 dispatch::dispatch 的模块套娃。 -mod reasoning_replay; -pub mod routing; -pub mod stream; -mod stream_audit; -mod transition; -pub mod usage; - -use std::{ - future::Future, - sync::Arc, - time::{Duration as StdDuration, Instant}, -}; - -use axum::{ - body::Body, - http::{ - header::{CACHE_CONTROL, CONTENT_TYPE}, - StatusCode, - }, - response::{IntoResponse, Response}, - Json, -}; -use chrono::{DateTime, Utc}; -use futures::{stream as futures_stream, StreamExt}; -use rand::RngExt; -use serde_json::{json, Value}; -use tokio::sync::Mutex; - -use crate::{ - codex::accounts::cookies::repository::CookieRepository, - codex::accounts::{ - cloudflare_challenge::CfPathBlockTracker, - model::Account, - pool::{AccountAcquireRequest, AccountPool, AcquiredAccount}, - repository::AccountRepository, - service::quota::quota_from_usage, - }, - codex::events::{ - event::{EventLevel, EventLog}, - service::LogService, - }, - codex::gateway::conversation_identity::{build_conversation_identity, ensure_prompt_cache_key}, - codex::gateway::fingerprint::model::Fingerprint, - codex::gateway::installation_id::get_installation_id, - codex::gateway::protocol::codex_to_openai::openai_error, - codex::gateway::transport::{ - endpoints::{ - endpoint_request_path, primary_usage_request_path, CODEX_RESPONSES_COMPACT_PATH, - CODEX_RESPONSES_PATH, - }, - http_client::{ - build_reqwest_client, CodexBackendClient, CodexBackendStream, - CodexBackendWebSocketStream, CodexClientError, CodexCompactResponse, - CodexRequestContext, - }, - types::{CodexCompactRequest, CodexResponsesRequest}, - usage_events::TokenUsage, - websocket::{transport_for_request, CodexTransport, CodexWebSocketPool}, - }, - config::AppConfig, -}; - -use crate::codex::serving::http::errors::{ - codex_client_error_message, codex_client_error_response, -}; - -pub(crate) use self::{ - fallback::{ - classify_upstream_account_retry, classify_upstream_recovery_action, UpstreamAccountRetry, - UpstreamRecoveryAction, UpstreamRecoveryState, UpstreamRequestRecovery, - }, - routing::{no_available_accounts_response, normalize_service_tier_for_upstream}, - stream::{completed_response_json, CollectedResponse}, - transition::UpstreamAccountRecoveryTransition, -}; - -use self::affinity::{ - compute_variant_hash, hash_instructions, prepare_variant_identity, SessionAffinityMap, - SessionAffinityRepository, -}; -pub(crate) use self::implicit_resume::ImplicitResumeSnapshot; -use self::implicit_resume::{continuation_input_start, implicit_resume_allowed}; -use self::reasoning_replay::ReasoningReplayCache; -use self::stream_audit::{StreamAudit, WebSocketStreamAudit}; -use self::transition::{ - execute_upstream_account_recovery_transition_with_deps, - execute_upstream_request_recovery_transition_with_deps, -}; -use self::{ - limits::apply_rate_limit_headers_with_deps, - stream::{ - completed_response_metadata, ensure_stream_metadata, has_terminal_sse_event, - premature_close_failed_event, responses_stream_error_event, TupleStreamReconverter, - }, - usage::{record_empty_response_with_deps, record_request_attempt, record_usage_with_deps}, -}; - -#[derive(Clone)] -struct CodexUpstreamDependencies { - config: Arc, - account_pool: Arc>, - account_repository: Option, - cookie_repository: Option, - logs: LogService, - fingerprint: Fingerprint, // 用于实际请求的指纹 - session_affinity: Arc, - session_affinity_repository: Option, - reasoning_replay: Arc, - websocket_pool: Arc, - cf_path_block_tracker: CfPathBlockTracker, -} - -const MAX_UPSTREAM_5XX_RETRIES: u8 = 2; -const UPSTREAM_5XX_RETRY_BASE_DELAY_MS: u64 = 1_000; -const MAX_DIRTY_QUOTA_VERIFY_ATTEMPTS: usize = 5; - -#[derive(Clone)] -pub(crate) struct CodexUpstreamService { - deps: CodexUpstreamDependencies, -} - -pub(crate) struct CodexUpstreamRepositories { - pub(crate) account: Option, - pub(crate) cookie: Option, - pub(crate) session_affinity: Option, -} - -impl CodexUpstreamService { - pub(crate) fn new( - config: Arc, - account_pool: Arc>, - repositories: CodexUpstreamRepositories, - logs: LogService, - fingerprint: Fingerprint, - websocket_pool: Arc, - ) -> Self { - Self { - deps: CodexUpstreamDependencies { - config, - account_pool, - account_repository: repositories.account, - cookie_repository: repositories.cookie, - logs, - fingerprint, - session_affinity: Arc::new(SessionAffinityMap::with_default_ttl()), - session_affinity_repository: repositories.session_affinity, - reasoning_replay: Arc::new(ReasoningReplayCache::default()), - websocket_pool, - cf_path_block_tracker: CfPathBlockTracker::new(), - }, - } - } - - pub(crate) async fn acquire_account(&self, model: &str) -> Option { - self.deps - .account_pool - .lock() - .await - .acquire_with(AccountAcquireRequest::new(model, Utc::now())) - } - - pub(crate) async fn prepare_response_session( - &self, - request: &mut CodexResponsesRequest, - ) -> Option { - prepare_variant_identity(request); - if let Some(previous_response_id) = request.previous_response_id.as_deref() { - if let Some(conversation_id) = self - .deps - .session_affinity - .lookup_conversation_id(previous_response_id) - .await - { - request.prompt_cache_key = Some(conversation_id); - } - if request - .turn_state - .as_deref() - .is_none_or(|value| value.trim().is_empty()) - { - request.turn_state = self - .deps - .session_affinity - .lookup_turn_state(previous_response_id) - .await; - } - ensure_prompt_cache_key(request); - return None; - } - ensure_prompt_cache_key(request); - apply_implicit_resume_with_deps(&self.deps, request).await - } - - pub(crate) async fn acquire_account_for_request( - &self, - request: &CodexResponsesRequest, - request_id: &str, - ) -> Option { - let preferred_account_id = match request.previous_response_id.as_deref() { - Some(previous_response_id) => { - self.deps - .session_affinity - .lookup_account(previous_response_id) - .await - } - None => None, - }; - let mut excluded_account_ids = Vec::new(); - let mut verify_attempts = 0; - - loop { - let mut acquire_request = AccountAcquireRequest::new(&request.model, Utc::now()) - .with_exclude_account_ids(excluded_account_ids.iter().cloned()); - if let Some(preferred_account_id) = preferred_account_id.as_ref() { - acquire_request = - acquire_request.with_preferred_account_id(preferred_account_id.clone()); - } - let mut acquired = self - .deps - .account_pool - .lock() - .await - .acquire_with(acquire_request)?; - if !acquired.account.quota_verify_required { - return Some(acquired); - } - - match verify_dirty_quota_with_deps(&self.deps, &acquired.account, request_id).await { - Ok(verification) if verification.limit_reached => { - let account_id = acquired.account.id.clone(); - self.deps.account_pool.lock().await.release(&account_id); - excluded_account_ids.push(account_id.clone()); - verify_attempts += 1; - tracing::warn!( - account_id = %account_id, - verify_attempts, - max_verify_attempts = MAX_DIRTY_QUOTA_VERIFY_ATTEMPTS, - "dirty quota 校验确认账号仍被限流，改用备用账号" - ); - if verify_attempts >= MAX_DIRTY_QUOTA_VERIFY_ATTEMPTS { - tracing::warn!( - max_verify_attempts = MAX_DIRTY_QUOTA_VERIFY_ATTEMPTS, - "dirty quota 校验达到最大次数，停止放大 upstream usage 请求" - ); - return None; - } - } - Ok(_) => { - acquired.account.quota_verify_required = false; - return Some(acquired); - } - Err(error) => { - tracing::warn!( - error = %error, - account_id = %acquired.account.id, - "dirty quota 校验失败，保留 dirty 标记并继续当前账号" - ); - return Some(acquired); - } - } - } - } - - pub(crate) async fn release_account(&self, account_id: &str) { - self.deps.account_pool.lock().await.release(account_id); - } - - pub(crate) async fn stagger_request(&self, previous_slot_at: Option>) { - stagger_request_with_deps(&self.deps, previous_slot_at).await; - } - - /// 获取当前使用的指纹（用于诊断） - pub(crate) fn fingerprint(&self) -> &Fingerprint { - &self.deps.fingerprint - } - - pub(crate) async fn send_codex_request_with_upstream_retries( - &self, - request: &CodexResponsesRequest, - account: &Account, - request_id: &str, - ) -> Result - { - send_codex_request_with_upstream_retries_deps(&self.deps, request, account, request_id) - .await - } - - pub(crate) async fn send_compact_request_with_upstream_retries( - &self, - request: &CodexCompactRequest, - account: &Account, - request_id: &str, - ) -> Result { - send_compact_request_with_upstream_retries_deps(&self.deps, request, account, request_id) - .await - } - - pub(crate) async fn apply_account_recovery_transition( - &self, - account: &Account, - retry: UpstreamAccountRetry, - model: &str, - excluded_account_ids: &mut Vec, - image_generation_requested: bool, - upstream_message: String, - ) -> UpstreamAccountRecoveryTransition { - execute_upstream_account_recovery_transition_with_deps( - &self.deps, - account, - retry, - model, - excluded_account_ids, - image_generation_requested, - upstream_message, - ) - .await - } - - pub(crate) async fn apply_request_recovery_transition( - &self, - request: &mut CodexResponsesRequest, - recovery: UpstreamRequestRecovery, - stream: bool, - log_context: &CodexRequestLogContext, - history_recovery_used: &mut bool, - implicit_resume: &mut Option, - ) { - execute_upstream_request_recovery_transition_with_deps( - &self.deps, - request, - recovery, - stream, - log_context, - history_recovery_used, - implicit_resume, - ) - .await; - } - - pub(crate) async fn responses_stream( - &self, - request: CodexResponsesRequest, - acquired: AcquiredAccount, - log_context: CodexRequestLogContext, - implicit_resume: Option, - ) -> Response { - let deps = self.deps.clone(); - if matches!( - transport_for_request(&request), - CodexTransport::WebSocketPreferred | CodexTransport::WebSocketRequired - ) { - return responses_websocket_stream( - deps, - request, - acquired, - log_context, - implicit_resume, - ) - .await; - } - - responses_http_sse_stream(deps, request, acquired, log_context, implicit_resume).await - } - - 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, - image_generation_requested: bool, - ) -> Result<(), ()> { - record_usage_with_deps(&self.deps, account_id, usage, image_generation_requested).await - } - - pub(crate) async fn record_empty_response( - &self, - account_id: &str, - image_generation_requested: bool, - ) -> Result<(), ()> { - record_empty_response_with_deps(&self.deps, account_id, image_generation_requested).await - } - - pub(crate) async fn record_request_attempt( - &self, - account_id: &str, - image_generation_requested: bool, - ) -> Result<(), ()> { - record_request_attempt(&self.deps, account_id, image_generation_requested).await - } - - pub(crate) async fn record_response_affinity( - &self, - request: &CodexResponsesRequest, - account_id: &str, - body: &str, - turn_state: Option<&str>, - usage: Option, - ) { - record_response_affinity_with_deps( - &self.deps, request, account_id, body, turn_state, usage, - ) - .await; - } - - pub(crate) async fn evict_reasoning_replay( - &self, - request: &CodexResponsesRequest, - account_id: &str, - ) { - evict_reasoning_replay_with_deps(&self.deps, request, account_id).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; - } - - pub(crate) async fn reload_session_affinity_from_repository( - &self, - ) -> self::affinity::SessionAffinityRepositoryResult { - let Some(repository) = self.deps.session_affinity_repository.as_ref() else { - return Ok(0); - }; - let now = Utc::now(); - repository.delete_expired(now).await?; - let records = repository.list_active(now).await?; - Ok(self.deps.session_affinity.restore(records).await) - } -} - -const IMPLICIT_RESUME_MAX_AGE: StdDuration = StdDuration::from_secs(55 * 60); - -async fn record_failed_request_attempt_with_deps( - deps: &CodexUpstreamDependencies, - account_id: &str, - image_generation_requested: bool, -) { - if let Err(error) = record_request_attempt(deps, account_id, image_generation_requested).await { - tracing::warn!( - error = ?error, - account_id = %account_id, - "记录上游最终失败请求尝试失败" - ); - } -} - -async fn upstream_cookie_header_for_endpoint( - deps: &CodexUpstreamDependencies, - account_id: &str, - endpoint_path: &str, -) -> Option { - let request_path = endpoint_request_path(&deps.config.api.base_url, endpoint_path); - upstream_cookie_header_for_request_path(deps, account_id, &request_path).await -} - -async fn upstream_cookie_header_for_request_path( - deps: &CodexUpstreamDependencies, - account_id: &str, - request_path: &str, -) -> Option { - let repo = deps.cookie_repository.as_ref()?; - let domain = routing::request_domain(&deps.config.api.base_url)?; - repo.cookie_header_for_request(account_id, &domain, request_path) - .await - .ok() - .flatten() -} - -struct DirtyQuotaVerification { - limit_reached: bool, -} - -async fn verify_dirty_quota_with_deps( - deps: &CodexUpstreamDependencies, - account: &Account, - request_id: &str, -) -> Result { - let usage = fetch_dirty_account_usage_with_deps(deps, account, request_id).await?; - let quota = quota_from_usage(&usage); - let limit_reached = quota_primary_limit_reached("a); - persist_verified_quota_with_deps(deps, account, "a).await; - if limit_reached { - apply_verified_quota_limit_with_deps(deps, account, "a).await; - } - Ok(DirtyQuotaVerification { limit_reached }) -} - -async fn fetch_dirty_account_usage_with_deps( - deps: &CodexUpstreamDependencies, - account: &Account, - request_id: &str, -) -> Result { - let usage_path = primary_usage_request_path(&deps.config.api.base_url); - let cookie_header = - upstream_cookie_header_for_request_path(deps, &account.id, &usage_path).await; - let client = CodexBackendClient::new( - build_reqwest_client(deps.config.tls.force_http11)?, - deps.config.api.base_url.clone(), - deps.fingerprint.clone(), - ); - client - .fetch_usage(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: cookie_header.as_deref(), - installation_id: None, - session_id: None, - }) - .await -} - -async fn persist_verified_quota_with_deps( - deps: &CodexUpstreamDependencies, - account: &Account, - quota: &Value, -) { - deps.account_pool - .lock() - .await - .set_quota_verify_required(&account.id, false); - if let Some(repo) = deps.account_repository.as_ref() { - if let Err(error) = repo - .update_quota_json(&account.id, "a.to_string()) - .await - { - tracing::warn!( - error = %error, - account_id = %account.id, - "持久化 dirty quota 校验结果失败" - ); - } - if let Err(error) = repo.set_quota_verify_required(&account.id, false).await { - tracing::warn!( - error = %error, - account_id = %account.id, - "清理 dirty quota 持久化标记失败" - ); - } - } -} - -async fn apply_verified_quota_limit_with_deps( - deps: &CodexUpstreamDependencies, - account: &Account, - quota: &Value, -) { - let Some(reset_at) = quota_primary_reset_at(quota) else { - return; - }; - let limit_window_seconds = quota_primary_limit_window_seconds(quota); - { - let mut pool = deps.account_pool.lock().await; - pool.sync_rate_limit_window(&account.id, reset_at, limit_window_seconds); - if reset_at > Utc::now() { - pool.mark_quota_limited_until(&account.id, reset_at); - } - } - if let Some(repo) = deps.account_repository.as_ref() { - if let Err(error) = repo - .sync_rate_limit_window(&account.id, reset_at, limit_window_seconds) - .await - { - tracing::warn!( - error = %error, - account_id = %account.id, - window_reset_at = %reset_at, - "持久化 dirty quota window 失败" - ); - } - if reset_at > Utc::now() { - if let Err(error) = repo.set_quota_cooldown_until(&account.id, reset_at).await { - tracing::warn!( - error = %error, - account_id = %account.id, - cooldown_until = %reset_at, - "持久化 dirty quota cooldown 失败" - ); - } - } - } - if reset_at > Utc::now() { - deps.websocket_pool.evict_account(&account.id).await; - } -} - -fn quota_primary_limit_reached(quota: &Value) -> bool { - quota - .pointer("/rate_limit/limit_reached") - .and_then(Value::as_bool) - .unwrap_or(false) -} - -fn quota_primary_reset_at(quota: &Value) -> Option> { - let reset_at = quota - .pointer("/rate_limit/reset_at") - .and_then(Value::as_i64)?; - DateTime::::from_timestamp(reset_at, 0) -} - -fn quota_primary_limit_window_seconds(quota: &Value) -> Option { - quota - .pointer("/rate_limit/limit_window_seconds") - .and_then(Value::as_u64) - .filter(|seconds| *seconds > 0) -} - -async fn apply_implicit_resume_with_deps( - deps: &CodexUpstreamDependencies, - request: &mut CodexResponsesRequest, -) -> Option { - if request.previous_response_id.is_some() { - return None; - } - let continuation_start = continuation_input_start(&request.input); - if continuation_start == 0 || continuation_start >= request.input.len() { - return None; - } - let conversation_id = request - .prompt_cache_key - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty())?; - - let snapshot = ImplicitResumeSnapshot::capture(request); - let variant_hash = compute_variant_hash(request); - let previous_response_id = deps - .session_affinity - .lookup_latest_response_by_conversation( - conversation_id, - Some(IMPLICIT_RESUME_MAX_AGE), - Some(&variant_hash), - ) - .await?; - let current_instructions_hash = hash_instructions(Some(&request.instructions)); - if deps - .session_affinity - .lookup_instructions_hash(&previous_response_id) - .await - .as_deref() - != Some(current_instructions_hash.as_str()) - { - return None; - } - - let stored_function_call_ids = deps - .session_affinity - .lookup_function_call_ids(&previous_response_id) - .await; - if !implicit_resume_allowed( - &request.input[continuation_start..], - &request.input, - &stored_function_call_ids, - ) { - return None; - } - - let account_id = deps - .session_affinity - .lookup_account(&previous_response_id) - .await?; - let replay_items = deps - .reasoning_replay - .lookup( - &previous_response_id, - &account_id, - conversation_id, - &variant_hash, - ) - .await; - let continuation = request.input[continuation_start..].to_vec(); - let mut input = replay_items; - input.extend(continuation); - - request.previous_response_id = Some(previous_response_id.clone()); - request.use_websocket = true; - request.force_http_sse = false; - request.input = input; - if let Some(turn_state) = deps - .session_affinity - .lookup_turn_state(&previous_response_id) - .await - { - request.turn_state = Some(turn_state); - } - Some(snapshot) -} - -async fn stagger_request_with_deps( - deps: &CodexUpstreamDependencies, - previous_slot_at: Option>, -) { - let Some(previous_slot_at) = previous_slot_at else { - return; - }; - let interval_ms = deps.config.auth.request_interval_ms; - if interval_ms == 0 { - return; - } - let target_interval_ms = jitter_request_interval_ms(interval_ms); - let elapsed_ms = Utc::now() - .signed_duration_since(previous_slot_at) - .num_milliseconds() - .max(0) as u64; - let Some(wait_ms) = target_interval_ms.checked_sub(elapsed_ms) else { - return; - }; - if wait_ms == 0 { - return; - } - tracing::debug!( - wait_ms, - request_interval_ms = interval_ms, - target_interval_ms, - "按账户请求间隔等待后发送上游请求" - ); - tokio::time::sleep(StdDuration::from_millis(wait_ms)).await; -} - -fn jitter_request_interval_ms(interval_ms: u64) -> u64 { - let mut rng = rand::rng(); - jitter_request_interval_ms_with_factor(interval_ms, rng.random_range(0.7..=1.3)) -} - -fn jitter_request_interval_ms_with_factor(interval_ms: u64, factor: f64) -> u64 { - let factor = factor.clamp(0.7, 1.3); - ((interval_ms as f64) * factor).round().min(u64::MAX as f64) as u64 -} - -async fn responses_http_sse_stream( - deps: CodexUpstreamDependencies, - mut request: CodexResponsesRequest, - mut acquired: AcquiredAccount, - mut log_context: CodexRequestLogContext, - mut implicit_resume: Option, -) -> Response { - let mut excluded_account_ids = Vec::new(); - let mut history_recovery_used = false; - let mut model_unsupported_retry_used = false; - let stream_response = loop { - stagger_request_with_deps(&deps, acquired.previous_slot_at).await; - let stream_response = send_codex_stream_request_with_upstream_retries( - &deps, - &request, - &acquired.account, - log_context.request_id.as_str(), - ) - .await; - match stream_response { - Ok(response) => break response, - Err(error) => { - deps.account_pool.lock().await.release(&acquired.account.id); - match classify_upstream_recovery_action( - &error, - UpstreamRecoveryState { - request_recovery_used: history_recovery_used, - model_unsupported_retry_used, - }, - ) { - UpstreamRecoveryAction::Request(recovery) => { - execute_upstream_request_recovery_transition_with_deps( - &deps, - &mut request, - recovery, - true, - &log_context, - &mut history_recovery_used, - &mut implicit_resume, - ) - .await; - continue; - } - UpstreamRecoveryAction::Account(retry) => { - if retry.is_model_unsupported() { - model_unsupported_retry_used = true; - } - let transition = execute_upstream_account_recovery_transition_with_deps( - &deps, - &acquired.account, - retry, - &request.model, - &mut excluded_account_ids, - request.expects_image_generation(), - codex_client_error_message(&error), - ) - .await; - log_codex_upstream_response_with_deps( - &deps, - &log_context, - retry.status(), - EventLevel::Warn, - "v1 responses stream 上游请求将使用备用账户重试", - retry.metadata(true), - ) - .await; - match transition { - UpstreamAccountRecoveryTransition::Retry(fallback) => { - log_context = log_context.with_account(&fallback.account.id); - acquired = *fallback; - continue; - } - UpstreamAccountRecoveryTransition::Respond { status, message } => { - log_codex_upstream_response_with_deps( - &deps, - &log_context, - status, - EventLevel::Error, - "v1 responses stream fallback 已耗尽", - json!({"stream": true}), - ) - .await; - return responses_stream_error_response(status, &message); - } - } - } - UpstreamRecoveryAction::RespondWithError => {} - } - record_failed_request_attempt_with_deps( - &deps, - &acquired.account.id, - request.expects_image_generation(), - ) - .await; - 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 上游请求失败", - json!({"stream": true}), - ) - .await; - return error_response.into_response(); - } - } - }; - - if persist_upstream_cookies_with_deps( - &deps, - &acquired.account.id, - &stream_response.set_cookie_headers, - ) - .await - .is_err() - { - deps.account_pool.lock().await.release(&acquired.account.id); - log_codex_upstream_response_with_deps( - &deps, - &log_context, - StatusCode::INTERNAL_SERVER_ERROR, - EventLevel::Error, - "v1 responses stream 持久化 cookie 失败", - 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 rate_limit_headers = stream_response.rate_limit_headers.clone(); - let upstream = Box::pin(stream_response.response.bytes_stream()); - let tuple_schema = request.tuple_schema.clone(); - let audit = StreamAudit::new( - deps, - log_context, - acquired.account.id, - acquired.account.plan_type, - request, - rate_limit_headers, - ); - - use tokio::time::{interval_at, Duration, Instant as TokioInstant}; - const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15); - const HEARTBEAT_CHUNK: &[u8] = b": ping\n\n"; - let heartbeat_timer = interval_at(TokioInstant::now() + HEARTBEAT_INTERVAL, HEARTBEAT_INTERVAL); - let tuple_reconverter = TupleStreamReconverter::new(tuple_schema); - - let body_stream = futures_stream::unfold( - Some(( - upstream, - Vec::new(), - audit, - heartbeat_timer, - tuple_reconverter, - )), - |state| async move { - let ( - mut upstream, - mut collected, - mut audit, - mut heartbeat_timer, - mut tuple_reconverter, - ) = state?; - - tokio::select! { - chunk_result = upstream.next() => { - match chunk_result { - Some(Ok(chunk)) => { - collected.extend_from_slice(&chunk); - let chunk = match std::str::from_utf8(&chunk) { - Ok(text) => { - axum::body::Bytes::from(tuple_reconverter.transform_chunk(text)) - } - Err(_) => chunk, - }; - Some(( - Ok::(chunk), - Some(( - upstream, - collected, - audit, - heartbeat_timer, - tuple_reconverter, - )), - )) - } - Some(Err(error)) => { - if collected_has_terminal_sse_event(&collected) { - let tail = tuple_reconverter.finish(); - audit.complete(&collected).await; - if tail.is_empty() { - None - } else { - Some(( - Ok::(tail.into()), - None, - )) - } - } else { - let detail = error.to_string(); - let mut output = tuple_reconverter.finish(); - let failure = append_premature_close_failed_event( - &mut collected, - Some(detail.as_str()), - ); - output.push_str(&failure); - audit.complete(&collected).await; - Some(( - Ok::(output.into()), - None, - )) - } - } - None => { - if collected_has_terminal_sse_event(&collected) { - let tail = tuple_reconverter.finish(); - audit.complete(&collected).await; - if tail.is_empty() { - None - } else { - Some(( - Ok::(tail.into()), - None, - )) - } - } else { - let mut output = tuple_reconverter.finish(); - let failure = - append_premature_close_failed_event(&mut collected, None); - output.push_str(&failure); - audit.complete(&collected).await; - Some(( - Ok::(output.into()), - None, - )) - } - } - } - } - _ = heartbeat_timer.tick() => { - Some(( - Ok::(HEARTBEAT_CHUNK.into()), - Some(( - upstream, - collected, - audit, - heartbeat_timer, - tuple_reconverter, - )), - )) - } - } - }, - ); - - 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() - }) -} - -fn websocket_stream_error_allows_http_sse_fallback(error: &CodexClientError) -> bool { - match error { - CodexClientError::Upstream { status, .. } => *status == StatusCode::UPGRADE_REQUIRED, - _ => false, - } -} - -fn collected_has_terminal_sse_event(collected: &[u8]) -> bool { - let body = String::from_utf8_lossy(collected); - has_terminal_sse_event(&body).unwrap_or(false) -} - -fn append_premature_close_failed_event(collected: &mut Vec, detail: Option<&str>) -> String { - let failure = premature_close_failed_event(detail); - let prefix = sse_event_boundary_prefix(collected); - collected.extend_from_slice(prefix.as_bytes()); - collected.extend_from_slice(failure.as_bytes()); - if prefix.is_empty() { - failure - } else { - let mut output = String::with_capacity(prefix.len() + failure.len()); - output.push_str(prefix); - output.push_str(&failure); - output - } -} - -fn responses_stream_error_response(status: StatusCode, message: &str) -> Response { - Response::builder() - .status(StatusCode::OK) - .header(CONTENT_TYPE, "text/event-stream") - .header(CACHE_CONTROL, "no-cache") - .body(Body::from(responses_stream_error_event(status, message))) - .unwrap_or_else(|_| { - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(openai_error( - "Failed to build stream error response", - "stream_response_error", - )), - ) - .into_response() - }) -} - -fn sse_event_boundary_prefix(collected: &[u8]) -> &'static str { - if collected.is_empty() || collected_ends_with_sse_event_boundary(collected) { - "" - } else if collected.ends_with(b"\n") { - "\n" - } else { - "\n\n" - } -} - -fn collected_ends_with_sse_event_boundary(collected: &[u8]) -> bool { - collected.ends_with(b"\n\n") - || collected.ends_with(b"\r\n\r\n") - || collected.ends_with(b"\n\r\n") -} - -async fn send_codex_request_with_upstream_retries_deps( - deps: &CodexUpstreamDependencies, - request: &CodexResponsesRequest, - account: &Account, - request_id: &str, -) -> Result { - let result = retry_upstream_5xx( - || send_codex_request(deps, request, account, request_id), - request_id, - &request.model, - ) - .await; - if let Ok(response) = &result { - apply_rate_limit_headers_with_deps( - deps, - &account.id, - account.plan_type.as_deref(), - &response.rate_limit_headers, - ) - .await; - } - result -} - -async fn send_compact_request_with_upstream_retries_deps( - deps: &CodexUpstreamDependencies, - request: &CodexCompactRequest, - account: &Account, - request_id: &str, -) -> Result { - let result = retry_upstream_5xx( - || send_compact_request(deps, request, account, request_id), - request_id, - &request.model, - ) - .await; - if let Ok(response) = &result { - apply_rate_limit_headers_with_deps( - deps, - &account.id, - account.plan_type.as_deref(), - &response.rate_limit_headers, - ) - .await; - } - result -} - -async fn send_codex_request( - deps: &CodexUpstreamDependencies, - request: &CodexResponsesRequest, - account: &Account, - request_id: &str, -) -> Result { - let cookie_header = - upstream_cookie_header_for_endpoint(deps, &account.id, CODEX_RESPONSES_PATH).await; - - let account_scope = &account.id; - let identity = build_conversation_identity( - request.prompt_cache_key.as_deref(), - request.codex_window_id.as_deref(), - account_scope, - ); - - let installation_id = get_installation_id(Some(&deps.config.database.url)); - - let client = CodexBackendClient::new( - build_reqwest_client(deps.config.tls.force_http11)?, - deps.config.api.base_url.clone(), - deps.fingerprint.clone(), - ) - .with_websocket_pool(deps.websocket_pool.clone(), account.id.clone()); - 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: identity.window_id.as_deref(), - parent_thread_id: request.parent_thread_id.as_deref(), - cookie_header: cookie_header.as_deref(), - installation_id: Some(&installation_id), - session_id: identity.conversation_id.as_deref(), - }, - ) - .await -} - -async fn send_compact_request( - deps: &CodexUpstreamDependencies, - request: &CodexCompactRequest, - account: &Account, - request_id: &str, -) -> Result { - let cookie_header = - upstream_cookie_header_for_endpoint(deps, &account.id, CODEX_RESPONSES_COMPACT_PATH).await; - - let installation_id = get_installation_id(Some(&deps.config.database.url)); - let client = CodexBackendClient::new( - build_reqwest_client(deps.config.tls.force_http11)?, - deps.config.api.base_url.clone(), - deps.fingerprint.clone(), - ); - - client - .create_compact_response( - request, - 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: cookie_header.as_deref(), - installation_id: Some(&installation_id), - session_id: None, - }, - ) - .await -} - -async fn send_codex_stream_request_with_upstream_retries( - deps: &CodexUpstreamDependencies, - request: &CodexResponsesRequest, - account: &Account, - request_id: &str, -) -> Result { - retry_upstream_5xx( - || send_codex_stream_request(deps, request, account, request_id), - request_id, - &request.model, - ) - .await -} - -async fn send_codex_stream_request( - deps: &CodexUpstreamDependencies, - request: &CodexResponsesRequest, - account: &Account, - request_id: &str, -) -> Result { - let cookie_header = - upstream_cookie_header_for_endpoint(deps, &account.id, CODEX_RESPONSES_PATH).await; - - // Build conversation identity for session affinity - // Use account.id as the scope since entry_id doesn't exist - let account_scope = &account.id; - let identity = build_conversation_identity( - request.prompt_cache_key.as_deref(), - request.codex_window_id.as_deref(), - account_scope, - ); - - // Get installation ID (cached after first call) - let installation_id = get_installation_id(Some(&deps.config.database.url)); - - let client = CodexBackendClient::new( - build_reqwest_client(deps.config.tls.force_http11)?, - deps.config.api.base_url.clone(), - deps.fingerprint.clone(), - ) - .with_websocket_pool(deps.websocket_pool.clone(), account.id.clone()); - 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: identity.window_id.as_deref(), - parent_thread_id: request.parent_thread_id.as_deref(), - cookie_header: cookie_header.as_deref(), - installation_id: Some(&installation_id), - session_id: identity.conversation_id.as_deref(), - }, - ) - .await -} - -async fn send_codex_websocket_stream_request_with_upstream_retries( - deps: &CodexUpstreamDependencies, - request: &CodexResponsesRequest, - account: &Account, - request_id: &str, -) -> Result { - retry_upstream_5xx( - || send_codex_websocket_stream_request(deps, request, account, request_id), - request_id, - &request.model, - ) - .await -} - -async fn retry_upstream_5xx( - mut operation: F, - request_id: &str, - model: &str, -) -> Result -where - F: FnMut() -> Fut, - Fut: Future>, -{ - let mut attempt = 0; - loop { - match operation().await { - Err(error) - if is_retryable_upstream_5xx(&error) && attempt < MAX_UPSTREAM_5XX_RETRIES => - { - let delay = upstream_5xx_retry_delay(attempt); - tracing::warn!( - error = %error, - request_id = %request_id, - model = %model, - retry_attempt = attempt + 1, - max_retries = MAX_UPSTREAM_5XX_RETRIES, - delay_ms = delay.as_millis(), - "Codex 上游 5xx，按原版策略同账户重试" - ); - tokio::time::sleep(delay).await; - attempt += 1; - } - result => return result, - } - } -} - -fn is_retryable_upstream_5xx(error: &CodexClientError) -> bool { - let CodexClientError::Upstream { status, body, .. } = error else { - return false; - }; - if !status.is_server_error() { - return false; - } - let Ok(value) = serde_json::from_str::(body) else { - return true; - }; - !stream::contains_invalid_encrypted_content_signal(&value) -} - -fn upstream_5xx_retry_delay(attempt: u8) -> StdDuration { - StdDuration::from_millis(UPSTREAM_5XX_RETRY_BASE_DELAY_MS * (1_u64 << u32::from(attempt))) -} - -async fn send_codex_websocket_stream_request( - deps: &CodexUpstreamDependencies, - request: &CodexResponsesRequest, - account: &Account, - request_id: &str, -) -> Result { - let cookie_header = - upstream_cookie_header_for_endpoint(deps, &account.id, CODEX_RESPONSES_PATH).await; - - let account_scope = &account.id; - let identity = build_conversation_identity( - request.prompt_cache_key.as_deref(), - request.codex_window_id.as_deref(), - account_scope, - ); - let installation_id = get_installation_id(Some(&deps.config.database.url)); - - let client = CodexBackendClient::new( - build_reqwest_client(deps.config.tls.force_http11)?, - deps.config.api.base_url.clone(), - deps.fingerprint.clone(), - ) - .with_websocket_pool(deps.websocket_pool.clone(), account.id.clone()); - client - .websocket_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: identity.window_id.as_deref(), - parent_thread_id: request.parent_thread_id.as_deref(), - cookie_header: cookie_header.as_deref(), - installation_id: Some(&installation_id), - session_id: identity.conversation_id.as_deref(), - }, - ) - .await -} - -async fn record_response_affinity_with_deps( - deps: &CodexUpstreamDependencies, - request: &CodexResponsesRequest, - account_id: &str, - body: &str, - turn_state: Option<&str>, - usage: Option, -) { - let Some(conversation_id) = request - .prompt_cache_key - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - else { - return; - }; - let metadata = match completed_response_metadata(body) { - Ok(Some(metadata)) => metadata, - Ok(None) => return, - Err(error) => { - tracing::warn!( - error = %error, - "解析已完成响应 metadata 用于 session affinity 失败" - ); - return; - } - }; - - let variant_hash = compute_variant_hash(request); - let response_id = metadata.response_id; - let entry = deps - .session_affinity - .record( - response_id.clone(), - account_id.to_string(), - conversation_id.to_string(), - turn_state - .filter(|value| !value.trim().is_empty()) - .map(ToString::to_string) - .or_else(|| request.turn_state.clone()), - Some(&request.instructions), - usage.map(|usage| usage.input_tokens), - Some(metadata.function_call_ids), - Some(variant_hash.clone()), - ) - .await; - deps.reasoning_replay - .record( - &response_id, - account_id, - conversation_id, - &variant_hash, - &metadata.replay_items, - ) - .await; - if let Some(repository) = deps.session_affinity_repository.as_ref() { - if let Err(error) = repository - .upsert(&response_id, &entry, deps.session_affinity.ttl()) - .await - { - tracing::warn!( - error = %error, - response_id = %response_id, - account_id = %account_id, - "持久化 session affinity 失败" - ); - } - } -} - -async fn evict_reasoning_replay_with_deps( - deps: &CodexUpstreamDependencies, - request: &CodexResponsesRequest, - account_id: &str, -) { - let Some(conversation_id) = request - .prompt_cache_key - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - else { - return; - }; - let variant_hash = compute_variant_hash(request); - let evicted = deps - .reasoning_replay - .evict_by_identity(account_id, conversation_id, &variant_hash) - .await; - if evicted > 0 { - tracing::info!( - account_id = %account_id, - conversation_id = %conversation_id, - variant_hash = %variant_hash, - evicted, - "已驱逐无效 encrypted content 对应的 reasoning replay" - ); - } -} - -async fn responses_websocket_stream( - deps: CodexUpstreamDependencies, - mut request: CodexResponsesRequest, - mut acquired: AcquiredAccount, - mut log_context: CodexRequestLogContext, - mut implicit_resume: Option, -) -> Response { - let mut excluded_account_ids = Vec::new(); - let mut history_recovery_used = false; - let mut model_unsupported_retry_used = false; - let stream_response = loop { - stagger_request_with_deps(&deps, acquired.previous_slot_at).await; - let response = send_codex_websocket_stream_request_with_upstream_retries( - &deps, - &request, - &acquired.account, - log_context.request_id.as_str(), - ) - .await; - - match response { - Ok(response) => break response, - Err(error) => { - if transport_for_request(&request) == CodexTransport::WebSocketPreferred - && websocket_stream_error_allows_http_sse_fallback(&error) - { - acquired.previous_slot_at = None; - return responses_http_sse_stream( - deps, - request, - acquired, - log_context, - implicit_resume, - ) - .await; - } - deps.account_pool.lock().await.release(&acquired.account.id); - match classify_upstream_recovery_action( - &error, - UpstreamRecoveryState { - request_recovery_used: history_recovery_used, - model_unsupported_retry_used, - }, - ) { - UpstreamRecoveryAction::Request(recovery) => { - execute_upstream_request_recovery_transition_with_deps( - &deps, - &mut request, - recovery, - true, - &log_context, - &mut history_recovery_used, - &mut implicit_resume, - ) - .await; - continue; - } - UpstreamRecoveryAction::Account(retry) => { - if retry.is_model_unsupported() { - model_unsupported_retry_used = true; - } - let transition = execute_upstream_account_recovery_transition_with_deps( - &deps, - &acquired.account, - retry, - &request.model, - &mut excluded_account_ids, - request.expects_image_generation(), - codex_client_error_message(&error), - ) - .await; - log_codex_upstream_response_with_deps( - &deps, - &log_context, - retry.status(), - EventLevel::Warn, - "v1 responses WebSocket 上游请求将使用备用账户重试", - retry.metadata(true), - ) - .await; - match transition { - UpstreamAccountRecoveryTransition::Retry(fallback) => { - log_context = log_context.with_account(&fallback.account.id); - acquired = *fallback; - continue; - } - UpstreamAccountRecoveryTransition::Respond { status, message } => { - log_codex_upstream_response_with_deps( - &deps, - &log_context, - status, - EventLevel::Error, - "v1 responses WebSocket stream fallback 已耗尽", - json!({"stream": true, "transport": "websocket"}), - ) - .await; - return responses_stream_error_response(status, &message); - } - } - } - UpstreamRecoveryAction::RespondWithError => {} - } - record_failed_request_attempt_with_deps( - &deps, - &acquired.account.id, - request.expects_image_generation(), - ) - .await; - 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 上游请求失败", - json!({"stream": true, "transport": "websocket"}), - ) - .await; - return error_response.into_response(); - } - } - }; - - if persist_upstream_cookies_with_deps( - &deps, - &acquired.account.id, - &stream_response.set_cookie_headers, - ) - .await - .is_err() - { - deps.account_pool.lock().await.release(&acquired.account.id); - log_codex_upstream_response_with_deps( - &deps, - &log_context, - StatusCode::INTERNAL_SERVER_ERROR, - EventLevel::Error, - "v1 responses WebSocket stream 持久化 cookie 失败", - 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 tuple_schema = request.tuple_schema.clone(); - let audit = WebSocketStreamAudit::new( - deps, - log_context, - acquired.account.id, - acquired.account.plan_type, - request, - stream_response.turn_state, - stream_response.rate_limit_headers, - stream_response.rate_limit_updates, - stream_response.turn_state_updates, - ); - let upstream = stream_response.body_stream; - - use tokio::time::{interval_at, Duration, Instant as TokioInstant}; - const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15); - const HEARTBEAT_CHUNK: &str = ": ping\n\n"; - let heartbeat_timer = interval_at(TokioInstant::now() + HEARTBEAT_INTERVAL, HEARTBEAT_INTERVAL); - let tuple_reconverter = TupleStreamReconverter::new(tuple_schema); - - let body_stream = futures_stream::unfold( - Some(( - upstream, - Vec::new(), - audit, - heartbeat_timer, - tuple_reconverter, - )), - |state| async move { - let ( - mut upstream, - mut collected, - mut audit, - mut heartbeat_timer, - mut tuple_reconverter, - ) = state?; - - tokio::select! { - chunk_result = upstream.next() => { - match chunk_result { - Some(Ok(chunk)) => { - collected.extend_from_slice(chunk.as_bytes()); - let chunk = tuple_reconverter.transform_chunk(&chunk); - Some(( - Ok::(chunk), - Some(( - upstream, - collected, - audit, - heartbeat_timer, - tuple_reconverter, - )), - )) - } - Some(Err(error)) => { - if collected_has_terminal_sse_event(&collected) { - let tail = tuple_reconverter.finish(); - audit.complete(&collected).await; - if tail.is_empty() { - None - } else { - Some((Ok::(tail), None)) - } - } else { - let detail = error.to_string(); - let mut output = tuple_reconverter.finish(); - let failure = append_premature_close_failed_event( - &mut collected, - Some(detail.as_str()), - ); - output.push_str(&failure); - audit.complete(&collected).await; - Some((Ok::(output), None)) - } - } - None => { - if collected_has_terminal_sse_event(&collected) { - let tail = tuple_reconverter.finish(); - audit.complete(&collected).await; - if tail.is_empty() { - None - } else { - Some((Ok::(tail), None)) - } - } else { - let mut output = tuple_reconverter.finish(); - let failure = - append_premature_close_failed_event(&mut collected, None); - output.push_str(&failure); - audit.complete(&collected).await; - Some((Ok::(output), None)) - } - } - } - } - _ = heartbeat_timer.tick() => { - Some(( - Ok::(HEARTBEAT_CHUNK.to_string()), - Some(( - upstream, - collected, - audit, - heartbeat_timer, - tuple_reconverter, - )), - )) - } - } - }, - ); - - 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 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, - } - } -} - -async fn log_codex_upstream_response_with_deps( - deps: &CodexUpstreamDependencies, - context: &CodexRequestLogContext, - status: StatusCode, - level: EventLevel, - message: &str, - mut metadata: Value, -) { - 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) = deps.logs.record(event).await { - tracing::warn!(error = %error, "写入 v1 response 事件日志失败"); - } -} - -#[cfg(test)] -mod tests { - use std::{collections::BTreeMap, sync::Arc}; - - use axum::http::StatusCode; - use serde_json::json; - use tokio::sync::Mutex; - use wiremock::{ - matchers::{header, method, path}, - Mock, MockServer, ResponseTemplate, - }; - - use super::{ - append_premature_close_failed_event, jitter_request_interval_ms_with_factor, - CodexUpstreamRepositories, CodexUpstreamService, - }; - use crate::{ - codex::{ - accounts::{ - model::{Account, AccountStatus}, - pool::AccountPool, - }, - events::service::LogService, - gateway::{ - fingerprint::model::Fingerprint, - transport::{ - http_client::CodexClientError, types::CodexResponsesRequest, - websocket::CodexWebSocketPool, - }, - }, - }, - config::{ - AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, - ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, - TlsConfig, UsageStatsConfig, - }, - }; - - #[test] - fn jitter_request_interval_ms_with_factor_should_match_original_bounds() { - assert_eq!(jitter_request_interval_ms_with_factor(300, 0.7), 210); - assert_eq!(jitter_request_interval_ms_with_factor(300, 1.3), 390); - assert_eq!(jitter_request_interval_ms_with_factor(300, 0.1), 210); - assert_eq!(jitter_request_interval_ms_with_factor(300, 2.0), 390); - } - - #[test] - fn append_premature_close_failed_event_should_close_partial_sse_event_before_failure() { - let mut collected = - b"event: response.output_text.delta\ndata: {\"delta\":\"partial\"}\n".to_vec(); - - append_premature_close_failed_event(&mut collected, None); - - let body = String::from_utf8(collected).unwrap(); - let failure = super::stream::responses_sse_failure(&body).unwrap(); - let metadata = failure.unwrap().metadata(true); - assert_eq!(metadata["upstreamCode"], "stream_disconnected"); - } - - #[test] - fn is_retryable_upstream_5xx_should_exclude_invalid_reasoning_replay() { - let error = CodexClientError::Upstream { - status: StatusCode::SERVICE_UNAVAILABLE, - retry_after_seconds: None, - body: json!({ - "type": "response.failed", - "response": { - "error": { - "code": "invalid_encrypted_content", - "message": "Invalid encrypted content in reasoning replay" - } - } - }) - .to_string(), - }; - - assert!(!super::is_retryable_upstream_5xx(&error)); - } - - #[tokio::test] - async fn acquire_account_for_request_should_verify_dirty_quota_and_try_next_account() { - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/api/codex/usage")) - .and(header("authorization", "Bearer dirty-access")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "plan_type": "plus", - "rate_limit": { - "allowed": true, - "limit_reached": true, - "primary_window": { - "used_percent": 100, - "reset_at": chrono::Utc::now().timestamp() + 300, - "limit_window_seconds": 300 - } - } - }))) - .mount(&server) - .await; - - let mut dirty = Account::test("acct_a", AccountStatus::Active); - dirty.access_token = "dirty-access".to_string(); - dirty.account_id = Some("chatgpt-dirty".to_string()); - dirty.quota_verify_required = true; - let mut clean = Account::test("acct_b", AccountStatus::Active); - clean.access_token = "clean-access".to_string(); - clean.account_id = Some("chatgpt-clean".to_string()); - let mut pool = AccountPool::default(); - pool.insert(dirty); - pool.insert(clean); - let service = test_upstream_service(server.uri(), pool); - let request = CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new()); - - let acquired = service - .acquire_account_for_request(&request, "req_dirty_quota") - .await - .unwrap(); - - assert_eq!(acquired.account.id, "acct_b"); - } - - fn test_upstream_service(base_url: String, account_pool: AccountPool) -> CodexUpstreamService { - CodexUpstreamService::new( - Arc::new(test_app_config(base_url)), - Arc::new(Mutex::new(account_pool)), - CodexUpstreamRepositories { - account: None, - cookie: None, - session_affinity: None, - }, - LogService::new(test_logging_config(), None), - Fingerprint::default_for_tests(), - Arc::new(CodexWebSocketPool::default()), - ) - } - - fn test_app_config(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(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - 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: "sqlite::memory:".to_string(), - }, - 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, - }, - ws_pool: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - session_cleanup_interval_secs: 3600, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - }, - logging: test_logging_config(), - } - } - - fn test_logging_config() -> LoggingConfig { - LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - capacity: 2_000, - capture_body: false, - } - } -} diff --git a/src/codex/serving/dispatch/routing.rs b/src/codex/serving/dispatch/routing.rs deleted file mode 100644 index da28a0bf5..000000000 --- a/src/codex/serving/dispatch/routing.rs +++ /dev/null @@ -1,29 +0,0 @@ -use axum::{http::StatusCode, Json}; -use reqwest::Url; -use serde_json::Value; - -use crate::codex::gateway::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/serving/dispatch/stream.rs b/src/codex/serving/dispatch/stream.rs deleted file mode 100644 index a216d0a32..000000000 --- a/src/codex/serving/dispatch/stream.rs +++ /dev/null @@ -1,638 +0,0 @@ -use std::collections::BTreeSet; - -use axum::http::StatusCode; -use serde::Serialize; -use serde_json::{json, Value}; -use uuid::Uuid; - -use crate::codex::gateway::{ - protocol::tuple_schema::reconvert_tuple_values, - transport::{ - http_client::CodexClientError, - sse::{encode_sse_event, parse_sse_events}, - }, -}; - -pub(crate) enum CollectedResponse { - Completed(Value), - Failed(ResponsesSseFailure), - MissingCompleted, - Empty, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct CompletedResponseMetadata { - pub response_id: String, - pub function_call_ids: Vec, - pub replay_items: Vec, -} - -#[derive(Debug, Clone)] -pub(crate) struct ResponsesSseFailure { - event: String, - message: String, - upstream_code: Option, - invalid_reasoning_replay: bool, -} - -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), - invalid_reasoning_replay: contains_invalid_encrypted_content_signal(value), - } - } - - pub(crate) fn invalid_reasoning_replay(&self) -> bool { - self.invalid_reasoning_replay - } - - pub(crate) fn upstream_error(&self) -> CodexClientError { - let code = self.upstream_code.as_deref().unwrap_or("error"); - CodexClientError::Upstream { - status: status_for_failure_code(code), - retry_after_seconds: None, - body: json!({ - "error": { - "type": code, - "code": code, - "message": self.message, - } - }) - .to_string(), - } - } - - 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)); - } - } -} - -fn status_for_failure_code(code: &str) -> StatusCode { - let lower = code.to_ascii_lowercase(); - if lower.contains("model") - && (lower.contains("not_supported") || lower.contains("not_available")) - { - return StatusCode::BAD_REQUEST; - } - if lower.contains("invalid_request") || lower.contains("not_found") { - return StatusCode::BAD_REQUEST; - } - if lower.contains("rate_limit") || lower.contains("usage_limit") { - return StatusCode::TOO_MANY_REQUESTS; - } - if lower.contains("unauthorized") - || lower.contains("invalid_api_key") - || lower == "token_invalid" - || lower == "token_expired" - || lower == "account_deactivated" - { - return StatusCode::UNAUTHORIZED; - } - if lower.contains("forbidden") || lower.contains("banned") { - return StatusCode::FORBIDDEN; - } - if lower.contains("payment") || lower.contains("quota") { - return StatusCode::PAYMENT_REQUIRED; - } - StatusCode::BAD_GATEWAY -} - -pub(crate) fn completed_response_json( - body: &str, - tuple_schema: Option<&Value>, -) -> 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); - }; - - // 检测空响应：没有输出文本且没有工具调用 - if is_empty_response(&response, &output_text, &output_items) { - return Ok(CollectedResponse::Empty); - } - - ensure_completed_response_output(&mut response, &output_items, &output_text); - reconvert_completed_response_tuple_values(&mut response, tuple_schema); - sync_output_text_from_output(&mut response); - Ok(CollectedResponse::Completed(response)) -} - -/// 检测是否为空响应 -fn is_empty_response(response: &Value, output_text: &str, output_items: &[Value]) -> bool { - // 如果有输出文本，不是空响应 - if !output_text.trim().is_empty() { - return false; - } - - // 如果有工具调用，不是空响应 - if !output_items.is_empty() { - return false; - } - - // 检查 output_tokens 是否为 0 - let output_tokens = response - .pointer("/usage/output_tokens") - .and_then(Value::as_u64) - .unwrap_or(0); - - output_tokens == 0 -} - -pub(super) struct TupleStreamReconverter { - tuple_schema: Option, - pending: String, - tuple_text_buffer: String, -} - -impl TupleStreamReconverter { - pub(super) fn new(tuple_schema: Option) -> Self { - Self { - tuple_schema, - pending: String::new(), - tuple_text_buffer: String::new(), - } - } - - pub(super) fn transform_chunk(&mut self, chunk: &str) -> String { - if self.tuple_schema.is_none() { - return chunk.to_string(); - } - self.pending.push_str(chunk); - let mut output = String::new(); - while let Some(delimiter_index) = self.pending.find("\n\n") { - let block = self.pending[..delimiter_index].to_string(); - self.pending.drain(..delimiter_index + 2); - output.push_str(&self.transform_block(&block)); - } - output - } - - pub(super) fn finish(&mut self) -> String { - if self.tuple_schema.is_none() || self.pending.trim().is_empty() { - self.pending.clear(); - return String::new(); - } - let block = std::mem::take(&mut self.pending); - self.transform_block(&block) - } - - fn transform_block(&mut self, block: &str) -> String { - if block.trim().is_empty() { - return "\n\n".to_string(); - } - let frame = format!("{block}\n\n"); - let Ok(mut events) = parse_sse_events(&frame) else { - return frame; - }; - if events.len() != 1 { - return frame; - } - let event = events.remove(0); - let Some(event_name) = event.event.as_deref() else { - return frame; - }; - match event_name { - "response.output_text.delta" => self.transform_output_text_delta(&event.data, &frame), - "response.completed" => { - self.transform_response_completed(&event.data, event_name, &frame) - } - _ => frame, - } - } - - fn transform_output_text_delta(&mut self, data: &str, original_frame: &str) -> String { - let Ok(value) = serde_json::from_str::(data) else { - return original_frame.to_string(); - }; - let Some(delta) = value.get("delta").and_then(Value::as_str) else { - return original_frame.to_string(); - }; - self.tuple_text_buffer.push_str(delta); - String::new() - } - - fn transform_response_completed( - &mut self, - data: &str, - event_name: &str, - original_frame: &str, - ) -> String { - let Some(tuple_schema) = self.tuple_schema.as_ref() else { - return original_frame.to_string(); - }; - let Ok(mut value) = serde_json::from_str::(data) else { - return original_frame.to_string(); - }; - - let mut output = String::new(); - if !self.tuple_text_buffer.is_empty() { - let reconverted_text = reconvert_tuple_text(&self.tuple_text_buffer, tuple_schema) - .unwrap_or_else(|| self.tuple_text_buffer.clone()); - self.tuple_text_buffer.clear(); - output.push_str(&encode_sse_event( - "response.output_text.delta", - &json!({ - "type": "response.output_text.delta", - "delta": reconverted_text, - }) - .to_string(), - )); - } - - if let Some(response) = value.get_mut("response") { - reconvert_completed_response_tuple_values(response, Some(tuple_schema)); - sync_output_text_from_output(response); - } - output.push_str(&encode_sse_event(event_name, &value.to_string())); - output - } -} - -pub(super) fn responses_sse_failure( - body: &str, -) -> Result, crate::codex::gateway::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 completed_response_metadata( - body: &str, -) -> Result, crate::codex::gateway::transport::sse::SseError> { - let events = parse_sse_events(body)?; - let mut response_id = None; - let mut function_call_ids = BTreeSet::new(); - let mut replay_items = Vec::new(); - - for event in events { - let Ok(value) = serde_json::from_str::(&event.data) else { - continue; - }; - match event.event.as_deref() { - Some("response.output_item.done") => { - if let Some(item) = value.get("item") { - collect_response_replay_items(item, &mut replay_items); - } - if let Some(call_id) = value.pointer("/item/call_id").and_then(Value::as_str) { - function_call_ids.insert(call_id.to_string()); - } - } - Some("response.completed") => { - response_id = value - .pointer("/response/id") - .and_then(Value::as_str) - .map(ToString::to_string); - collect_response_function_call_ids(&value, &mut function_call_ids); - if let Some(output) = value.pointer("/response/output").and_then(Value::as_array) { - for item in output { - collect_response_replay_items(item, &mut replay_items); - } - } - } - _ => {} - } - } - - Ok(response_id.map(|response_id| CompletedResponseMetadata { - response_id, - function_call_ids: function_call_ids.into_iter().collect(), - replay_items, - })) -} - -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)); -} - -pub(super) fn has_terminal_sse_event( - body: &str, -) -> Result { - parse_sse_events(body).map(|events| { - events.iter().any(|event| { - matches!( - event.event.as_deref(), - Some("response.completed" | "response.failed" | "error") - ) - }) - }) -} - -pub(super) fn premature_close_failed_event(detail: Option<&str>) -> String { - let message = match detail.filter(|value| !value.trim().is_empty()) { - Some(detail) => format!("Upstream stream closed before response.completed: {detail}"), - None => "Upstream stream closed before response.completed".to_string(), - }; - response_failed_event(ResponsesStreamError { - kind: "server_error", - code: "stream_disconnected", - message: &message, - }) -} - -pub(super) fn responses_stream_error_event(status: StatusCode, message: &str) -> String { - let clean_message = message - .strip_prefix("Codex API error (") - .and_then(|tail| tail.split_once("): ")) - .map_or(message, |(_, clean)| clean); - let lower = clean_message.to_ascii_lowercase(); - let error = if status == StatusCode::TOO_MANY_REQUESTS { - ResponsesStreamError { - kind: "rate_limit_error", - code: "rate_limit_exceeded", - message: clean_message, - } - } else if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { - ResponsesStreamError { - kind: "invalid_request_error", - code: "authentication_error", - message: clean_message, - } - } else if lower.contains("error sending request") { - ResponsesStreamError { - kind: "server_error", - code: "upstream_transport_error", - message: clean_message, - } - } else { - ResponsesStreamError { - kind: if status.is_client_error() { - "invalid_request_error" - } else { - "server_error" - }, - code: "codex_api_error", - message: clean_message, - } - }; - response_failed_event(error) -} - -#[derive(Clone, Copy, Serialize)] -struct ResponsesStreamError<'a> { - #[serde(rename = "type")] - kind: &'static str, - code: &'static str, - message: &'a str, -} - -#[derive(Serialize)] -struct ResponseFailedData<'a> { - #[serde(rename = "type")] - event_type: &'static str, - response: ResponseFailedResponse<'a>, - error: ResponsesStreamError<'a>, -} - -#[derive(Serialize)] -struct ResponseFailedResponse<'a> { - id: String, - status: &'static str, - error: ResponsesStreamError<'a>, -} - -fn response_failed_event(error: ResponsesStreamError<'_>) -> String { - let response_id_suffix: String = Uuid::new_v4() - .simple() - .to_string() - .chars() - .take(24) - .collect(); - let data = ResponseFailedData { - event_type: "response.failed", - response: ResponseFailedResponse { - id: format!("resp_proxy_{response_id_suffix}"), - status: "failed", - error, - }, - error, - }; - encode_sse_event( - "response.failed", - &serde_json::to_string(&data).expect("response.failed event should serialize"), - ) -} - -fn collect_response_function_call_ids(value: &Value, function_call_ids: &mut BTreeSet) { - let Some(output) = value.pointer("/response/output").and_then(Value::as_array) else { - return; - }; - for item in output { - if let Some(call_id) = item.get("call_id").and_then(Value::as_str) { - function_call_ids.insert(call_id.to_string()); - } - } -} - -fn collect_response_replay_items(item: &Value, replay_items: &mut Vec) { - if matches!( - item.get("type").and_then(Value::as_str), - Some("reasoning" | "function_call") - ) { - replay_items.push(item.clone()); - } -} - -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) -} - -pub(crate) fn contains_invalid_encrypted_content_signal(value: &Value) -> bool { - visit_invalid_encrypted_content(value, 0) -} - -fn visit_invalid_encrypted_content(value: &Value, depth: usize) -> bool { - if depth > 5 { - return false; - } - match value { - Value::String(value) => string_contains_invalid_encrypted_content(value), - Value::Array(values) => values - .iter() - .any(|value| visit_invalid_encrypted_content(value, depth + 1)), - Value::Object(object) => object.iter().any(|(key, value)| { - string_contains_invalid_encrypted_content(key) - || visit_invalid_encrypted_content(value, depth + 1) - }), - _ => false, - } -} - -fn string_contains_invalid_encrypted_content(value: &str) -> bool { - let lower = value.to_ascii_lowercase(); - lower.contains("invalid_encrypted_content") - || (lower.contains("invalid") && lower.contains("encrypted") && lower.contains("content")) -} - -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 reconvert_completed_response_tuple_values(response: &mut Value, tuple_schema: Option<&Value>) { - let Some(tuple_schema) = tuple_schema else { - return; - }; - let Some(items) = response.get_mut("output").and_then(Value::as_array_mut) else { - return; - }; - for item in items { - let Some(content) = item.get_mut("content").and_then(Value::as_array_mut) else { - continue; - }; - for part in content { - let part_type = part.get("type").and_then(Value::as_str); - if !matches!(part_type, Some("output_text" | "text")) { - continue; - } - let Some(text) = part.get("text").and_then(Value::as_str) else { - continue; - }; - let Ok(parsed) = serde_json::from_str::(text) else { - continue; - }; - let reconverted = reconvert_tuple_values(parsed, tuple_schema); - part["text"] = Value::String(reconverted.to_string()); - } - } -} - -fn reconvert_tuple_text(text: &str, tuple_schema: &Value) -> Option { - let parsed = serde_json::from_str::(text).ok()?; - Some(reconvert_tuple_values(parsed, tuple_schema).to_string()) -} - -fn sync_output_text_from_output(response: &mut Value) { - 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/serving/dispatch/stream_audit.rs b/src/codex/serving/dispatch/stream_audit.rs deleted file mode 100644 index ed97430ec..000000000 --- a/src/codex/serving/dispatch/stream_audit.rs +++ /dev/null @@ -1,337 +0,0 @@ -use std::sync::Arc; - -use axum::http::StatusCode; -use serde_json::json; -use tokio::sync::Mutex; - -use crate::{ - codex::gateway::transport::types::CodexResponsesRequest, - codex::{ - accounts::pool::AccountPool, - events::event::EventLevel, - gateway::transport::{ - usage_events::extract_sse_usage, - websocket::{ - append_rate_limit_updates, latest_turn_state, SharedRateLimitUpdates, - SharedTurnState, - }, - }, - }, -}; - -use super::{ - evict_reasoning_replay_with_deps, - limits::apply_rate_limit_headers_with_deps, - log_codex_upstream_response_with_deps, record_response_affinity_with_deps, - record_usage_with_deps, - stream::{ensure_stream_metadata, responses_sse_failure}, - CodexRequestLogContext, CodexUpstreamDependencies, -}; - -pub(super) struct StreamAudit { - deps: CodexUpstreamDependencies, - context: CodexRequestLogContext, - account_slot: AccountSlotGuard, - account_plan_type: Option, - request: CodexResponsesRequest, - rate_limit_headers: Vec<(String, String)>, -} - -impl StreamAudit { - pub(super) fn new( - deps: CodexUpstreamDependencies, - context: CodexRequestLogContext, - account_id: String, - account_plan_type: Option, - request: CodexResponsesRequest, - rate_limit_headers: Vec<(String, String)>, - ) -> Self { - let account_slot = AccountSlotGuard::new(deps.account_pool.clone(), account_id); - Self { - deps, - context, - account_slot, - account_plan_type, - request, - rate_limit_headers, - } - } - - pub(super) async fn complete(&mut self, body: &[u8]) { - apply_rate_limit_headers_with_deps( - &self.deps, - &self.context.account_id, - self.account_plan_type.as_deref(), - &self.rate_limit_headers, - ) - .await; - let body = String::from_utf8_lossy(body); - let mut status = StatusCode::OK; - let mut level = EventLevel::Info; - let mut message = "v1 responses stream 已完成"; - 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, - self.request.expects_image_generation(), - ) - .await - .is_err() - { - level = EventLevel::Warn; - message = "v1 responses stream 已完成但 usage 存储失败"; - 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 已完成但 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 上游 SSE 失败"; - if failure.invalid_reasoning_replay() { - evict_reasoning_replay_with_deps( - &self.deps, - &self.request, - &self.context.account_id, - ) - .await; - } - failure.extend_metadata(&mut metadata); - } - Ok(None) => {} - Err(error) => { - level = EventLevel::Warn; - message = "v1 responses stream 已完成但 SSE 失败 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; - } -} - -pub(super) struct WebSocketStreamAudit { - deps: CodexUpstreamDependencies, - context: CodexRequestLogContext, - account_slot: AccountSlotGuard, - account_plan_type: Option, - request: CodexResponsesRequest, - turn_state: Option, - rate_limit_headers: Vec<(String, String)>, - rate_limit_updates: SharedRateLimitUpdates, - turn_state_updates: SharedTurnState, -} - -impl WebSocketStreamAudit { - #[expect( - clippy::too_many_arguments, - reason = "stream audit captures one immutable snapshot from the upstream response boundary" - )] - pub(super) fn new( - deps: CodexUpstreamDependencies, - context: CodexRequestLogContext, - account_id: String, - account_plan_type: Option, - request: CodexResponsesRequest, - turn_state: Option, - rate_limit_headers: Vec<(String, String)>, - rate_limit_updates: SharedRateLimitUpdates, - turn_state_updates: SharedTurnState, - ) -> Self { - let account_slot = AccountSlotGuard::new(deps.account_pool.clone(), account_id); - Self { - deps, - context, - account_slot, - account_plan_type, - request, - turn_state, - rate_limit_headers, - rate_limit_updates, - turn_state_updates, - } - } - - pub(super) async fn complete(&mut self, body: &[u8]) { - let mut rate_limit_headers = self.rate_limit_headers.clone(); - append_rate_limit_updates(&mut rate_limit_headers, &self.rate_limit_updates).await; - apply_rate_limit_headers_with_deps( - &self.deps, - &self.context.account_id, - self.account_plan_type.as_deref(), - &rate_limit_headers, - ) - .await; - let body = String::from_utf8_lossy(body).into_owned(); - let usage_result = extract_sse_usage(&body); - let response_usage = match &usage_result { - Ok(usage) => *usage, - Err(_) => None, - }; - let turn_state = latest_turn_state(&self.turn_state_updates) - .await - .or_else(|| self.turn_state.clone()); - record_response_affinity_with_deps( - &self.deps, - &self.request, - &self.context.account_id, - &body, - turn_state.as_deref(), - response_usage, - ) - .await; - - let mut status = StatusCode::OK; - let mut level = EventLevel::Info; - let mut message = "v1 responses WebSocket stream 已完成"; - let mut metadata = match usage_result { - Ok(Some(usage)) => { - if record_usage_with_deps( - &self.deps, - &self.context.account_id, - usage, - self.request.expects_image_generation(), - ) - .await - .is_err() - { - level = EventLevel::Warn; - message = "v1 responses WebSocket stream 已完成但 usage 存储失败"; - json!({ - "stream": true, - "transport": "websocket", - "usage": usage, - "rateLimitHeaders": self.rate_limit_headers.clone(), - "usageStoreError": true, - }) - } else { - json!({ - "stream": true, - "transport": "websocket", - "usage": usage, - "rateLimitHeaders": self.rate_limit_headers.clone(), - }) - } - } - Ok(None) => json!({ - "stream": true, - "transport": "websocket", - "usage": null, - "rateLimitHeaders": self.rate_limit_headers.clone(), - }), - Err(error) => { - level = EventLevel::Warn; - message = "v1 responses WebSocket stream 已完成但 SSE usage 无效"; - json!({ - "stream": true, - "transport": "websocket", - "rateLimitHeaders": self.rate_limit_headers.clone(), - "sseParseError": error.to_string(), - }) - } - }; - match responses_sse_failure(&body) { - Ok(Some(failure)) => { - // SSE 响应头已经发给客户端，HTTP 状态不能回滚，只能在审计中标记上游失败。 - status = StatusCode::BAD_GATEWAY; - level = EventLevel::Error; - message = "v1 responses WebSocket stream 上游 SSE 失败"; - if failure.invalid_reasoning_replay() { - evict_reasoning_replay_with_deps( - &self.deps, - &self.request, - &self.context.account_id, - ) - .await; - } - failure.extend_metadata(&mut metadata); - } - Ok(None) => {} - Err(error) => { - level = EventLevel::Warn; - message = "v1 responses WebSocket stream 已完成但 SSE 失败 metadata 无效"; - metadata = json!({ - "stream": true, - "transport": "websocket", - "rateLimitHeaders": self.rate_limit_headers.clone(), - "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; - } -} - -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); - }); - } -} diff --git a/src/codex/serving/dispatch/transition.rs b/src/codex/serving/dispatch/transition.rs deleted file mode 100644 index 916d6a6a1..000000000 --- a/src/codex/serving/dispatch/transition.rs +++ /dev/null @@ -1,92 +0,0 @@ -use axum::http::StatusCode; -use chrono::Utc; - -use crate::codex::{ - accounts::{ - model::Account, - pool::{AccountAcquireRequest, AcquiredAccount}, - }, - events::event::EventLevel, - gateway::transport::types::CodexResponsesRequest, -}; - -use super::{ - fallback::{ - apply_upstream_account_retry_with_deps, build_account_exhaustion_detail, - UpstreamAccountRetry, UpstreamRequestRecovery, - }, - log_codex_upstream_response_with_deps, CodexRequestLogContext, CodexUpstreamDependencies, - ImplicitResumeSnapshot, -}; - -pub(crate) enum UpstreamAccountRecoveryTransition { - Retry(Box), - Respond { status: StatusCode, message: String }, -} - -pub(super) async fn execute_upstream_account_recovery_transition_with_deps( - deps: &CodexUpstreamDependencies, - account: &Account, - retry: UpstreamAccountRetry, - model: &str, - excluded_account_ids: &mut Vec, - image_generation_requested: bool, - upstream_message: String, -) -> UpstreamAccountRecoveryTransition { - apply_upstream_account_retry_with_deps(deps, account, retry, image_generation_requested).await; - excluded_account_ids.push(account.id.clone()); - let fallback = deps.account_pool.lock().await.acquire_with( - AccountAcquireRequest::new(model, Utc::now()) - .with_exclude_account_ids(excluded_account_ids.iter().cloned()), - ); - - match fallback { - Some(account) => UpstreamAccountRecoveryTransition::Retry(Box::new(account)), - None => { - let retry_message = retry.fallback_response_message(upstream_message); - let message = fallback_exhausted_message_with_deps(deps, &retry_message).await; - UpstreamAccountRecoveryTransition::Respond { - status: retry.status(), - message, - } - } - } -} - -pub(super) async fn execute_upstream_request_recovery_transition_with_deps( - deps: &CodexUpstreamDependencies, - request: &mut CodexResponsesRequest, - recovery: UpstreamRequestRecovery, - stream: bool, - log_context: &CodexRequestLogContext, - history_recovery_used: &mut bool, - implicit_resume: &mut Option, -) { - *history_recovery_used = true; - let stale_response_id = request.previous_response_id.clone(); - if let Some(response_id) = stale_response_id.as_deref() { - deps.session_affinity.forget(response_id).await; - } - if let Some(snapshot) = implicit_resume.take() { - snapshot.restore(request); - } - request.previous_response_id = None; - request.turn_state = None; - log_codex_upstream_response_with_deps( - deps, - log_context, - StatusCode::BAD_REQUEST, - EventLevel::Warn, - "v1 responses 上游历史失效，去除 previous_response_id 后重试", - recovery.metadata(stream, stale_response_id.as_deref()), - ) - .await; -} - -pub(super) async fn fallback_exhausted_message_with_deps( - deps: &CodexUpstreamDependencies, - message: &str, -) -> String { - let summary = deps.account_pool.lock().await.status_summary(Utc::now()); - build_account_exhaustion_detail(summary, message) -} diff --git a/src/codex/serving/dispatch/usage.rs b/src/codex/serving/dispatch/usage.rs deleted file mode 100644 index 145db47b7..000000000 --- a/src/codex/serving/dispatch/usage.rs +++ /dev/null @@ -1,108 +0,0 @@ -use crate::codex::{ - accounts::{pool::AccountWindowUsageDelta, repository::UsageDelta}, - gateway::transport::usage_events::TokenUsage, -}; - -use super::CodexUpstreamDependencies; - -pub(super) async fn record_usage_with_deps( - deps: &CodexUpstreamDependencies, - account_id: &str, - usage: TokenUsage, - image_generation_requested: bool, -) -> Result<(), ()> { - let image_request_succeeded = image_generation_requested && usage.image_output_tokens > 0; - let image_request_failed = image_generation_requested && !image_request_succeeded; - deps.account_pool.lock().await.record_window_token_usage( - account_id, - AccountWindowUsageDelta { - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cached_tokens: usage.cached_tokens, - image_input_tokens: usage.image_input_tokens, - image_output_tokens: usage.image_output_tokens, - image_request_succeeded, - image_request_failed, - }, - ); - 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), - image_input_tokens: u64_to_i64_saturating(usage.image_input_tokens), - image_output_tokens: u64_to_i64_saturating(usage.image_output_tokens), - image_request_count: bool_to_i64(image_request_succeeded), - image_request_failed_count: bool_to_i64(image_request_failed), - empty_response_count: 0, - }, - ) - .await - .map_err(|_| ()) -} - -pub(super) async fn record_request_attempt( - deps: &CodexUpstreamDependencies, - account_id: &str, - image_generation_requested: bool, -) -> 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, - image_input_tokens: 0, - image_output_tokens: 0, - image_request_count: 0, - image_request_failed_count: bool_to_i64(image_generation_requested), - empty_response_count: 0, - }, - ) - .await - .map_err(|_| ()) -} - -pub(super) async fn record_empty_response_with_deps( - deps: &CodexUpstreamDependencies, - account_id: &str, - image_generation_requested: bool, -) -> 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, - image_input_tokens: 0, - image_output_tokens: 0, - image_request_count: 0, - image_request_failed_count: bool_to_i64(image_generation_requested), - empty_response_count: 1, - }, - ) - .await - .map_err(|_| ()) -} - -fn u64_to_i64_saturating(value: u64) -> i64 { - value.min(i64::MAX as u64) as i64 -} - -fn bool_to_i64(value: bool) -> i64 { - if value { - 1 - } else { - 0 - } -} diff --git a/src/codex/serving/http/auth.rs b/src/codex/serving/http/auth.rs deleted file mode 100644 index 4ab948d0d..000000000 --- a/src/codex/serving/http/auth.rs +++ /dev/null @@ -1,15 +0,0 @@ -use axum::http::HeaderMap; - -use crate::{platform::http::auth::client_api_key, runtime::state::AppState}; - -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/codex/serving/http/chat.rs b/src/codex/serving/http/chat.rs deleted file mode 100644 index 0c29f4b9b..000000000 --- a/src/codex/serving/http/chat.rs +++ /dev/null @@ -1,46 +0,0 @@ -use axum::{ - body::Bytes, - extract::State, - http::{HeaderMap, StatusCode}, - response::{IntoResponse, Response}, - Extension, Json, -}; -use std::time::Instant; - -use crate::{ - codex::gateway::protocol::{ - codex_to_openai::openai_error, openai_to_codex::ChatCompletionRequest, - }, - platform::http::request_id::RequestId, - runtime::state::AppState, -}; - -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/codex/serving/http/diagnostics.rs b/src/codex/serving/http/diagnostics.rs deleted file mode 100644 index 78fae9b89..000000000 --- a/src/codex/serving/http/diagnostics.rs +++ /dev/null @@ -1,296 +0,0 @@ -use std::sync::Arc; - -use axum::{ - extract::{Extension, State}, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Json, -}; -use serde::Serialize; -use serde_json::json; - -use crate::{ - admin::session::service::AdminAuthPoolSummary, - codex::{accounts::pool::AccountCapacitySummary, gateway::fingerprint::model::Fingerprint}, - config::AppConfig, - platform::http::request_id::RequestId, - runtime::state::AppState, -}; - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct DiagnosticsData { - status: &'static str, - runtime: RuntimeDiagnostics, - paths: PathDiagnostics, - transport: TransportDiagnostics, - accounts: AccountDiagnostics, - settings: SettingsDiagnostics, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct RuntimeDiagnostics { - package_name: &'static str, - package_version: &'static str, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct PathDiagnostics { - config: &'static str, - local_config: String, - database_url: String, - logs_directory: String, - master_key_file: String, - api_key_pepper_file: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct TransportDiagnostics { - backend_base_url: String, - tls: TlsDiagnostics, - fingerprint: FingerprintDiagnostics, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct TlsDiagnostics { - force_http11: bool, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct FingerprintDiagnostics { - source: &'static str, - originator: String, - app_version: String, - build_number: String, - platform: String, - arch: String, - chromium_version: String, - user_agent: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountDiagnostics { - repository_available: bool, - authenticated_state: bool, - pool: AccountPoolDiagnostics, - capacity: CapacityDiagnostics, -} - -#[derive(Debug, Default, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountPoolDiagnostics { - total: usize, - active: usize, - expired: usize, - quota_exhausted: usize, - refreshing: usize, - disabled: usize, - banned: usize, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct CapacityDiagnostics { - max_concurrent_per_account: usize, - total_slots: usize, - used_slots: usize, - available_slots: usize, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct SettingsDiagnostics { - default_model: String, - refresh_enabled: bool, - rotation_strategy: String, - quota_skip_exhausted: bool, - logs_enabled: bool, -} - -pub async fn diagnostics(State(state): State, headers: HeaderMap) -> impl IntoResponse { - if !is_local_debug_request(&headers) { - return ( - StatusCode::FORBIDDEN, - Json(json!({ "error": "debug endpoint is local-only" })), - ) - .into_response(); - } - - (StatusCode::OK, Json(diagnostics_data(&state).await)).into_response() -} - -pub(crate) async fn diagnostics_data(state: &AppState) -> DiagnosticsData { - let config = state.services.settings.current().await; - let auth_status = state.services.admin_auth.status().await.ok(); - let capacity = state.services.accounts.runtime_capacity_summary().await; - let fingerprint = state.services.diagnostics.fingerprint().clone(); - - DiagnosticsData { - status: "ok", - runtime: RuntimeDiagnostics { - package_name: env!("CARGO_PKG_NAME"), - package_version: env!("CARGO_PKG_VERSION"), - }, - paths: paths_diagnostics(state, &config), - transport: transport_diagnostics(&config, fingerprint), - accounts: AccountDiagnostics { - repository_available: state.services.accounts.has_repository(), - authenticated_state: auth_status - .as_ref() - .is_some_and(|status| status.authenticated), - pool: auth_status - .map(|status| AccountPoolDiagnostics::from(status.pool)) - .unwrap_or_default(), - capacity: CapacityDiagnostics::from(capacity), - }, - settings: SettingsDiagnostics::from(&config), - } -} - -pub async fn debug_fingerprint( - State(state): State, - headers: HeaderMap, -) -> impl IntoResponse { - if !is_local_debug_request(&headers) { - return ( - StatusCode::FORBIDDEN, - Json(json!({ "error": "debug endpoint is local-only" })), - ) - .into_response(); - } - - // 从实际服务中获取指纹 - let fingerprint = state.services.diagnostics.fingerprint(); - - ( - StatusCode::OK, - Json(fingerprint_diagnostics(fingerprint.clone())), - ) - .into_response() -} - -pub async fn debug_upstream( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, -) -> impl IntoResponse { - if !is_local_debug_request(&headers) { - return ( - StatusCode::FORBIDDEN, - Json(json!({ "error": "debug endpoint is local-only" })), - ) - .into_response(); - } - - ( - StatusCode::OK, - Json( - state - .services - .diagnostics - .probe_upstream(request_id.as_str()) - .await, - ), - ) - .into_response() -} - -fn paths_diagnostics(state: &AppState, config: &AppConfig) -> PathDiagnostics { - PathDiagnostics { - config: "config.yaml", - local_config: state - .services - .settings - .local_config_path() - .display() - .to_string(), - database_url: config.database.url.clone(), - logs_directory: config.logging.directory.clone(), - master_key_file: config.security.master_key_file.clone(), - api_key_pepper_file: config.security.api_key_pepper_file.clone(), - } -} - -fn transport_diagnostics(config: &AppConfig, fingerprint: Fingerprint) -> TransportDiagnostics { - TransportDiagnostics { - backend_base_url: config.api.base_url.clone(), - tls: TlsDiagnostics { - force_http11: config.tls.force_http11, - }, - fingerprint: fingerprint_diagnostics(fingerprint), - } -} - -fn fingerprint_diagnostics(fingerprint: Fingerprint) -> FingerprintDiagnostics { - let user_agent = fingerprint.user_agent(); - FingerprintDiagnostics { - source: "runtime", - originator: fingerprint.originator, - app_version: fingerprint.app_version, - build_number: fingerprint.build_number, - platform: fingerprint.platform, - arch: fingerprint.arch, - chromium_version: fingerprint.chromium_version, - user_agent, - } -} - -fn is_local_debug_request(headers: &HeaderMap) -> bool { - forwarded_header_is_local(headers, "x-forwarded-for") - && forwarded_header_is_local(headers, "x-real-ip") -} - -fn forwarded_header_is_local(headers: &HeaderMap, name: &str) -> bool { - let Some(value) = headers.get(name).and_then(|value| value.to_str().ok()) else { - return true; - }; - value.split(',').next().is_some_and(is_local_host) -} - -fn is_local_host(host: &str) -> bool { - let host = host.trim().trim_start_matches('[').trim_end_matches(']'); - host == "localhost" || host == "::1" || host.starts_with("127.") -} - -impl From for AccountPoolDiagnostics { - fn from(summary: AdminAuthPoolSummary) -> Self { - Self { - total: summary.total, - active: summary.active, - expired: summary.expired, - quota_exhausted: summary.quota_exhausted, - refreshing: summary.refreshing, - disabled: summary.disabled, - banned: summary.banned, - } - } -} - -impl From for CapacityDiagnostics { - fn from(summary: AccountCapacitySummary) -> Self { - Self { - max_concurrent_per_account: summary.max_concurrent_per_account, - total_slots: summary.total_slots, - used_slots: summary.used_slots, - available_slots: summary.available_slots, - } - } -} - -impl From<&Arc> for SettingsDiagnostics { - fn from(config: &Arc) -> Self { - Self { - default_model: config.model.default_model.clone(), - refresh_enabled: config.auth.refresh_enabled, - rotation_strategy: config.auth.rotation_strategy.clone(), - quota_skip_exhausted: config.quota.skip_exhausted, - logs_enabled: config.logging.enabled, - } - } -} diff --git a/src/codex/serving/http/errors.rs b/src/codex/serving/http/errors.rs deleted file mode 100644 index c35a3317e..000000000 --- a/src/codex/serving/http/errors.rs +++ /dev/null @@ -1,119 +0,0 @@ -use axum::{http::StatusCode, Json}; -use serde_json::{json, Value}; - -use crate::{ - codex::gateway::protocol::codex_to_openai::openai_error, - codex::gateway::transport::http_client::CodexClientError, -}; - -pub(crate) fn codex_client_error_response(error: CodexClientError) -> (StatusCode, Json) { - let status = codex_client_error_status(&error); - let code = codex_client_error_code(&error); - let message = codex_client_error_message(&error); - (status, Json(openai_error(&message, code))) -} - -pub(crate) fn codex_client_error_response_with_status_and_message( - error: CodexClientError, - status: StatusCode, - message: &str, -) -> (StatusCode, Json) { - let code = codex_client_error_code(&error); - (status, Json(openai_error(message, code))) -} - -pub(crate) fn responses_codex_client_error_response( - error: CodexClientError, -) -> (StatusCode, Json) { - let status = codex_client_error_status(&error); - let message = codex_client_error_message(&error); - responses_error_response(status, &message) -} - -pub(crate) fn responses_error_response( - status: StatusCode, - message: &str, -) -> (StatusCode, Json) { - let (error_type, code) = if status == StatusCode::TOO_MANY_REQUESTS { - ("rate_limit_error", "rate_limit_exceeded") - } else { - ("server_error", "codex_api_error") - }; - ( - status, - Json(json!({ - "type": "error", - "error": { - "type": error_type, - "code": code, - "message": message, - } - })), - ) -} - -pub(crate) fn responses_no_available_accounts_response() -> (StatusCode, Json) { - ( - StatusCode::SERVICE_UNAVAILABLE, - Json(json!({ - "type": "error", - "error": { - "type": "server_error", - "code": "no_available_accounts", - "message": "No available accounts. All accounts are expired or rate-limited.", - } - })), - ) -} - -pub(crate) fn codex_client_error_message(error: &CodexClientError) -> String { - match error { - CodexClientError::UnsupportedTransport(_) => { - "previous_response_id requires Codex WebSocket transport".to_string() - } - CodexClientError::Upstream { status, body, .. } - if *status == StatusCode::NOT_FOUND && body.trim().is_empty() => - { - "Upstream blocked the request (Cloudflare path-block)".to_string() - } - CodexClientError::Upstream { body, .. } => format!( - "Codex upstream error: {}", - body.chars().take(300).collect::() - ), - _ => "Codex upstream request failed".to_string(), - } -} - -fn codex_client_error_status(error: &CodexClientError) -> StatusCode { - match error { - CodexClientError::UnsupportedTransport(_) => StatusCode::BAD_REQUEST, - CodexClientError::Upstream { status, body, .. } - if *status == StatusCode::NOT_FOUND && body.trim().is_empty() => - { - StatusCode::BAD_GATEWAY - } - CodexClientError::Upstream { status, .. } => *status, - _ => StatusCode::BAD_GATEWAY, - } -} - -fn codex_client_error_code(error: &CodexClientError) -> &'static str { - match error { - CodexClientError::UnsupportedTransport(_) => "websocket_required", - _ => "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/codex/serving/http/responses.rs b/src/codex/serving/http/responses.rs deleted file mode 100644 index 2ca4fa176..000000000 --- a/src/codex/serving/http/responses.rs +++ /dev/null @@ -1,70 +0,0 @@ -use std::time::Instant; - -use axum::{ - body::Bytes, - extract::State, - http::{HeaderMap, HeaderName, HeaderValue}, - response::{IntoResponse, Response}, - Extension, -}; - -use crate::{platform::http::request_id::RequestId, runtime::state::AppState}; - -use super::{auth::authorize_client_api_key, errors::missing_client_api_key_response}; - -const OPENAI_SUBAGENT_HEADER: HeaderName = HeaderName::from_static("x-openai-subagent"); - -pub async fn responses( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - body: Bytes, -) -> Response { - handle_responses(state, request_id, headers, body).await -} - -pub async fn review_responses( - State(state): State, - Extension(request_id): Extension, - mut headers: HeaderMap, - body: Bytes, -) -> Response { - headers.insert(OPENAI_SUBAGENT_HEADER, HeaderValue::from_static("review")); - handle_responses(state, request_id, headers, body).await -} - -pub async fn compact_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_compact(request_id.as_str(), body, started_at) - .await -} - -async fn handle_responses( - state: AppState, - request_id: RequestId, - 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/codex/serving/mod.rs b/src/codex/serving/mod.rs deleted file mode 100644 index e36684a2c..000000000 --- a/src/codex/serving/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub mod chat; -pub mod diagnostics; -pub mod dispatch; -pub mod http; -pub mod responses; diff --git a/src/codex/serving/responses.rs b/src/codex/serving/responses.rs deleted file mode 100644 index 9e25408ee..000000000 --- a/src/codex/serving/responses.rs +++ /dev/null @@ -1,1313 +0,0 @@ -use std::time::Instant; - -use axum::{ - body::Bytes, - http::{HeaderMap, StatusCode}, - response::{IntoResponse, Response}, - Json, -}; -use serde_json::{json, Map, Value}; - -use crate::{ - codex::events::event::EventLevel, - codex::gateway::fingerprint::model::Fingerprint, - codex::gateway::protocol::codex_to_openai::openai_error, - codex::gateway::protocol::tuple_schema::prepare_schema, - codex::gateway::transport::http_client::CodexClientError, - codex::gateway::transport::types::{CodexCompactRequest, CodexResponsesRequest}, - codex::models::service::ModelService, - codex::serving::dispatch::{ - affinity::SessionAffinityRepositoryResult, classify_upstream_account_retry, - classify_upstream_recovery_action, completed_response_json, - normalize_service_tier_for_upstream, stream::contains_invalid_encrypted_content_signal, - CodexRequestLogContext, CodexUpstreamService, CollectedResponse, - UpstreamAccountRecoveryTransition, UpstreamRecoveryAction, UpstreamRecoveryState, - }, - codex::serving::http::errors::{ - codex_client_error_message, responses_codex_client_error_response, - responses_error_response, responses_no_available_accounts_response, - }, - config::ModelConfig, -}; - -const OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent"; - -#[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 fn upstream_fingerprint(&self) -> &Fingerprint { - self.upstream.fingerprint() - } - - pub async fn reload_session_affinity_from_repository( - &self, - ) -> SessionAffinityRepositoryResult { - self.upstream - .reload_session_affinity_from_repository() - .await - } - - 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 = match parse_responses_body(&body, default_model.clone()) { - Some(body) => body, - None => return invalid_responses_body_response(), - }; - 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.generate = body.generate; - codex_request.tuple_schema = body.tuple_schema; - codex_request.explicit_prompt_cache_key = body - .prompt_cache_key - .as_deref() - .is_some_and(|value| !value.trim().is_empty()); - codex_request.prompt_cache_key = body.prompt_cache_key; - codex_request.include = body.include; - let openai_subagent = - normalize_openai_subagent(header_string(&headers, OPENAI_SUBAGENT_HEADER).as_deref()) - .or_else(|| openai_subagent_from_metadata(body.client_metadata.as_ref())); - codex_request.client_metadata = - client_metadata_with_openai_subagent(body.client_metadata, openai_subagent); - 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 = body.version.or_else(|| 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")); - match body.use_websocket { - Some(true) => codex_request.use_websocket = true, - Some(false) => { - // 显式传输控制不覆盖历史续链；previous_response_id 会在 transport 层强制 WebSocket。 - codex_request.force_http_sse = true; - } - None => {} - } - let mut implicit_resume = self - .upstream - .prepare_response_session(&mut codex_request) - .await; - let Some(mut acquired) = self - .upstream - .acquire_account_for_request(&codex_request, request_id) - .await - else { - return responses_no_available_accounts_response().into_response(); - }; - let mut log_context = CodexRequestLogContext::new( - request_id, - &acquired.account.id, - &codex_request.model, - client_stream, - started_at, - ); - - if client_stream { - return self - .upstream - .responses_stream(codex_request, acquired, log_context, implicit_resume) - .await; - } - - let mut excluded_account_ids = Vec::new(); - let mut empty_response_retries = 0; - let mut history_recovery_used = false; - let mut model_unsupported_retry_used = false; - const MAX_EMPTY_RETRIES: u8 = 2; - - let (response, collected_response) = loop { - self.upstream - .stagger_request(acquired.previous_slot_at) - .await; - let response = self - .upstream - .send_codex_request_with_upstream_retries( - &codex_request, - &acquired.account, - request_id, - ) - .await; - self.upstream.release_account(&acquired.account.id).await; - - match response { - Ok(response) => { - let collected_response = completed_response_json( - &response.body, - codex_request.tuple_schema.as_ref(), - ); - match &collected_response { - Ok(CollectedResponse::Empty) => { - if self - .upstream - .record_empty_response( - &acquired.account.id, - codex_request.expects_image_generation(), - ) - .await - .is_err() - { - self.upstream - .log_response( - &log_context, - StatusCode::OK, - EventLevel::Warn, - "v1 responses 记录空响应计数失败", - json!({ - "stream": false, - "emptyResponse": true, - "emptyResponseStoreError": true, - }), - ) - .await; - } - empty_response_retries += 1; - if empty_response_retries <= MAX_EMPTY_RETRIES { - self.upstream - .log_response( - &log_context, - StatusCode::OK, - EventLevel::Warn, - "v1 responses 空响应，准备重试", - json!({ - "stream": false, - "emptyResponse": true, - "retryAttempt": empty_response_retries, - }), - ) - .await; - continue; - } - self.upstream - .log_response( - &log_context, - StatusCode::BAD_GATEWAY, - EventLevel::Error, - "v1 responses 空响应重试次数已耗尽", - json!({ - "stream": false, - "emptyResponse": true, - "retriesExhausted": true, - }), - ) - .await; - return ( - StatusCode::BAD_GATEWAY, - Json(openai_error( - "Codex response was empty after retries", - "empty_upstream_response", - )), - ) - .into_response(); - } - Ok(CollectedResponse::Failed(failure)) => { - if failure.invalid_reasoning_replay() { - self.upstream - .evict_reasoning_replay(&codex_request, &acquired.account.id) - .await; - } - let error = failure.upstream_error(); - match classify_upstream_recovery_action( - &error, - UpstreamRecoveryState { - request_recovery_used: history_recovery_used, - model_unsupported_retry_used, - }, - ) { - UpstreamRecoveryAction::Request(recovery) => { - self.upstream - .apply_request_recovery_transition( - &mut codex_request, - recovery, - false, - &log_context, - &mut history_recovery_used, - &mut implicit_resume, - ) - .await; - continue; - } - UpstreamRecoveryAction::Account(retry) => { - if retry.is_model_unsupported() { - model_unsupported_retry_used = true; - } - let transition = self - .upstream - .apply_account_recovery_transition( - &acquired.account, - retry, - &codex_request.model, - &mut excluded_account_ids, - codex_request.expects_image_generation(), - codex_client_error_message(&error), - ) - .await; - self.upstream - .log_response( - &log_context, - retry.status(), - EventLevel::Warn, - "v1 responses 上游 SSE 失败将使用备用账户重试", - retry.metadata(false), - ) - .await; - match transition { - UpstreamAccountRecoveryTransition::Retry(fallback) => { - log_context = - log_context.with_account(&fallback.account.id); - acquired = *fallback; - continue; - } - UpstreamAccountRecoveryTransition::Respond { - status, - message, - } => { - let error_response = - responses_error_response(status, &message); - self.upstream - .log_response( - &log_context, - error_response.0, - EventLevel::Error, - "v1 responses 上游 SSE fallback 已耗尽", - failure.metadata(false), - ) - .await; - return error_response.into_response(); - } - } - } - UpstreamRecoveryAction::RespondWithError => {} - } - if self - .upstream - .record_request_attempt( - &acquired.account.id, - codex_request.expects_image_generation(), - ) - .await - .is_err() - { - self.upstream - .log_response( - &log_context, - StatusCode::OK, - EventLevel::Warn, - "v1 responses 记录失败请求计数失败", - json!({"stream": false, "requestAttemptStoreError": true}), - ) - .await; - } - let error_response = responses_codex_client_error_response(error); - self.upstream - .log_response( - &log_context, - error_response.0, - EventLevel::Error, - "v1 responses 上游 SSE 失败", - failure.metadata(false), - ) - .await; - return error_response.into_response(); - } - _ => {} - } - break (response, collected_response); - } - Err(error) => { - if codex_client_error_invalid_reasoning_replay(&error) { - self.upstream - .evict_reasoning_replay(&codex_request, &acquired.account.id) - .await; - } - match classify_upstream_recovery_action( - &error, - UpstreamRecoveryState { - request_recovery_used: history_recovery_used, - model_unsupported_retry_used, - }, - ) { - UpstreamRecoveryAction::Request(recovery) => { - self.upstream - .apply_request_recovery_transition( - &mut codex_request, - recovery, - false, - &log_context, - &mut history_recovery_used, - &mut implicit_resume, - ) - .await; - continue; - } - UpstreamRecoveryAction::Account(retry) => { - if retry.is_model_unsupported() { - model_unsupported_retry_used = true; - } - let transition = self - .upstream - .apply_account_recovery_transition( - &acquired.account, - retry, - &codex_request.model, - &mut excluded_account_ids, - codex_request.expects_image_generation(), - codex_client_error_message(&error), - ) - .await; - self.upstream - .log_response( - &log_context, - retry.status(), - EventLevel::Warn, - "v1 responses 上游请求将使用备用账户重试", - retry.metadata(false), - ) - .await; - match transition { - UpstreamAccountRecoveryTransition::Retry(fallback) => { - log_context = log_context.with_account(&fallback.account.id); - acquired = *fallback; - continue; - } - UpstreamAccountRecoveryTransition::Respond { status, message } => { - let error_response = responses_error_response(status, &message); - self.upstream - .log_response( - &log_context, - error_response.0, - EventLevel::Error, - "v1 responses 上游请求 fallback 已耗尽", - json!({"stream": false}), - ) - .await; - return error_response.into_response(); - } - } - } - UpstreamRecoveryAction::RespondWithError => {} - } - if self - .upstream - .record_request_attempt( - &acquired.account.id, - codex_request.expects_image_generation(), - ) - .await - .is_err() - { - self.upstream - .log_response( - &log_context, - StatusCode::OK, - EventLevel::Warn, - "v1 responses 记录失败请求计数失败", - json!({"stream": false, "requestAttemptStoreError": true}), - ) - .await; - } - let error_response = responses_codex_client_error_response(error); - self.upstream - .log_response( - &log_context, - error_response.0, - EventLevel::Error, - "v1 responses 上游请求失败", - json!({"stream": false}), - ) - .await; - return error_response.into_response(); - } - } - }; - if self - .upstream - .persist_cookies(&acquired.account.id, &response.set_cookie_headers) - .await - .is_err() - { - self.upstream - .log_response( - &log_context, - StatusCode::INTERNAL_SERVER_ERROR, - EventLevel::Error, - "v1 responses 持久化 cookie 失败", - 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( - &acquired.account.id, - usage, - codex_request.expects_image_generation(), - ) - .await - .is_err() - { - self.upstream - .log_response( - &log_context, - StatusCode::INTERNAL_SERVER_ERROR, - EventLevel::Error, - "v1 responses 持久化 usage 失败", - 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 collected_response { - Ok(CollectedResponse::Completed(body)) => { - self.upstream - .record_response_affinity( - &codex_request, - &acquired.account.id, - &response.body, - response.turn_state.as_deref(), - response.usage, - ) - .await; - self.upstream - .log_response( - &log_context, - StatusCode::OK, - EventLevel::Info, - "v1 responses 已完成", - json!({"stream": false, "usage": response.usage}), - ) - .await; - (StatusCode::OK, Json(body)).into_response() - } - Ok(CollectedResponse::Empty) => { - if self - .upstream - .record_empty_response( - &acquired.account.id, - codex_request.expects_image_generation(), - ) - .await - .is_err() - { - self.upstream - .log_response( - &log_context, - StatusCode::OK, - EventLevel::Warn, - "v1 responses 记录空响应计数失败", - json!({ - "stream": false, - "emptyResponse": true, - "emptyResponseStoreError": true, - }), - ) - .await; - } - self.upstream - .log_response( - &log_context, - StatusCode::BAD_GATEWAY, - EventLevel::Error, - "v1 responses 空响应", - json!({"stream": false, "emptyResponse": true}), - ) - .await; - ( - StatusCode::BAD_GATEWAY, - Json(openai_error( - "Codex response was empty", - "empty_upstream_response", - )), - ) - .into_response() - } - Ok(CollectedResponse::Failed(failure)) => { - if failure.invalid_reasoning_replay() { - self.upstream - .evict_reasoning_replay(&codex_request, &acquired.account.id) - .await; - } - let error_response = - responses_codex_client_error_response(failure.upstream_error()); - self.upstream - .log_response( - &log_context, - error_response.0, - EventLevel::Error, - "v1 responses 上游 SSE 失败", - failure.metadata(false), - ) - .await; - error_response.into_response() - } - Ok(CollectedResponse::MissingCompleted) => { - self.upstream - .log_response( - &log_context, - StatusCode::BAD_GATEWAY, - EventLevel::Warn, - "v1 responses 缺少 completed 事件", - 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 SSE 响应无效", - 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 handle_compact( - &self, - request_id: &str, - body: Bytes, - started_at: Instant, - ) -> Response { - let default_model = self.model_config.default_model.clone(); - let body = match parse_compact_body(&body, default_model) { - Some(body) => body, - None => return invalid_responses_body_response(), - }; - let catalog = self.models.catalog().await; - if !catalog.is_recognized_model_name(&body.model) { - return ( - StatusCode::NOT_FOUND, - Json(openai_error("Model not found", "model_not_found")), - ) - .into_response(); - } - let parsed_model = catalog.parse_model_name(&body.model); - let compact_request = CodexCompactRequest { - model: parsed_model.model_id, - input: body.input, - instructions: body.instructions, - tools: body.tools, - parallel_tool_calls: body.parallel_tool_calls, - reasoning: body.reasoning, - text: body.text, - }; - - let Some(mut acquired) = self.upstream.acquire_account(&compact_request.model).await else { - return responses_no_available_accounts_response().into_response(); - }; - let mut log_context = CodexRequestLogContext::new( - request_id, - &acquired.account.id, - &compact_request.model, - false, - started_at, - ); - let mut excluded_account_ids = Vec::new(); - let mut model_unsupported_retry_used = false; - let mut attempts = 0_u8; - const MAX_COMPACT_RETRIES: u8 = 8; - - loop { - attempts += 1; - self.upstream - .stagger_request(acquired.previous_slot_at) - .await; - let response = self - .upstream - .send_compact_request_with_upstream_retries( - &compact_request, - &acquired.account, - request_id, - ) - .await; - self.upstream.release_account(&acquired.account.id).await; - - match response { - Ok(response) => { - if self - .upstream - .persist_cookies(&acquired.account.id, &response.set_cookie_headers) - .await - .is_err() - { - self.upstream - .log_response( - &log_context, - StatusCode::INTERNAL_SERVER_ERROR, - EventLevel::Error, - "v1 responses compact 持久化 cookie 失败", - json!({"stream": false, "compact": true, "cookieStoreError": true}), - ) - .await; - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(openai_error( - "Failed to persist upstream cookies", - "cookie_store_error", - )), - ) - .into_response(); - } - self.upstream - .log_response( - &log_context, - StatusCode::OK, - EventLevel::Info, - "v1 responses compact 已完成", - json!({"stream": false, "compact": true}), - ) - .await; - return (StatusCode::OK, Json(response.body)).into_response(); - } - Err(error) => { - if attempts < MAX_COMPACT_RETRIES { - if let Some(retry) = - classify_upstream_account_retry(&error, model_unsupported_retry_used) - { - if retry.is_model_unsupported() { - model_unsupported_retry_used = true; - } - let transition = self - .upstream - .apply_account_recovery_transition( - &acquired.account, - retry, - &compact_request.model, - &mut excluded_account_ids, - false, - codex_client_error_message(&error), - ) - .await; - self.upstream - .log_response( - &log_context, - retry.status(), - EventLevel::Warn, - "v1 responses compact 上游请求将使用备用账户重试", - retry.metadata(false), - ) - .await; - match transition { - UpstreamAccountRecoveryTransition::Retry(fallback) => { - log_context = log_context.with_account(&fallback.account.id); - acquired = *fallback; - continue; - } - UpstreamAccountRecoveryTransition::Respond { status, message } => { - let error_response = responses_error_response(status, &message); - self.upstream - .log_response( - &log_context, - error_response.0, - EventLevel::Error, - "v1 responses compact fallback 已耗尽", - json!({"stream": false, "compact": true}), - ) - .await; - return error_response.into_response(); - } - } - } - } - if self - .upstream - .record_request_attempt(&acquired.account.id, false) - .await - .is_err() - { - self.upstream - .log_response( - &log_context, - StatusCode::OK, - EventLevel::Warn, - "v1 responses compact 记录失败请求计数失败", - json!({ - "stream": false, - "compact": true, - "requestAttemptStoreError": true, - }), - ) - .await; - } - let error_response = responses_codex_client_error_response(error); - self.upstream - .log_response( - &log_context, - error_response.0, - EventLevel::Error, - "v1 responses compact 上游请求失败", - json!({"stream": false, "compact": true}), - ) - .await; - return error_response.into_response(); - } - } - } - } -} - -fn codex_client_error_invalid_reasoning_replay(error: &CodexClientError) -> bool { - let CodexClientError::Upstream { body, .. } = error else { - return false; - }; - let Ok(value) = serde_json::from_str::(body) else { - return false; - }; - contains_invalid_encrypted_content_signal(&value) -} - -struct ResponsesBody { - model: Option, - input: Option>, - instructions: Option, - reasoning: Option, - tools: Option>, - service_tier: Option, - tool_choice: Option, - parallel_tool_calls: Option, - text: Option, - generate: Option, - tuple_schema: Option, - prompt_cache_key: Option, - include: Option>, - client_metadata: Option, - previous_response_id: Option, - stream: Option, - turn_state: Option, - turn_metadata: Option, - beta_features: Option, - version: Option, - include_timing_metrics: Option, - codex_window_id: Option, - parent_thread_id: Option, - use_websocket: Option, -} - -struct CompactBody { - model: String, - input: Vec, - instructions: String, - reasoning: Option, - tools: Option>, - parallel_tool_calls: Option, - text: Option, -} - -fn parse_responses_body(body: &[u8], default_model: String) -> Option { - let Value::Object(mut body) = serde_json::from_slice::(body).ok()? else { - return None; - }; - - let mut parsed = default_body(default_model); - parsed.model = take_string(&mut body, "model").or(parsed.model); - parsed.input = take_array(&mut body, "input").map(sanitize_codex_input_items); - parsed.instructions = take_string(&mut body, "instructions").or(parsed.instructions); - parsed.reasoning = take_object_value(&mut body, "reasoning"); - parsed.tools = take_non_empty_array(&mut body, "tools"); - parsed.service_tier = take_string(&mut body, "service_tier"); - parsed.tool_choice = body.remove("tool_choice"); - parsed.parallel_tool_calls = take_bool(&mut body, "parallel_tool_calls"); - let prepared_text = take_prepared_text_format(&mut body); - parsed.text = prepared_text.as_ref().map(|text| text.text.clone()); - parsed.tuple_schema = prepared_text.and_then(|text| text.tuple_schema); - parsed.generate = take_bool(&mut body, "generate"); - parsed.prompt_cache_key = take_string(&mut body, "prompt_cache_key"); - parsed.include = take_string_array(&mut body, "include"); - parsed.client_metadata = sanitize_client_metadata(body.remove("client_metadata")); - parsed.previous_response_id = take_string(&mut body, "previous_response_id"); - parsed.stream = take_bool(&mut body, "stream"); - parsed.turn_state = take_non_empty_string(&mut body, "turnState"); - parsed.turn_metadata = take_non_empty_string(&mut body, "turnMetadata"); - parsed.beta_features = take_non_empty_string(&mut body, "betaFeatures"); - parsed.version = take_non_empty_string(&mut body, "version"); - parsed.include_timing_metrics = take_non_empty_string(&mut body, "includeTimingMetrics"); - parsed.codex_window_id = take_non_empty_string(&mut body, "codexWindowId"); - parsed.parent_thread_id = take_non_empty_string(&mut body, "parentThreadId"); - parsed.use_websocket = take_bool(&mut body, "use_websocket"); - Some(parsed) -} - -fn parse_compact_body(body: &[u8], default_model: String) -> Option { - let Value::Object(mut body) = serde_json::from_slice::(body).ok()? else { - return None; - }; - - Some(CompactBody { - model: take_string(&mut body, "model").unwrap_or(default_model), - input: take_array(&mut body, "input") - .map(sanitize_codex_input_items) - .unwrap_or_default(), - instructions: take_string(&mut body, "instructions").unwrap_or_default(), - reasoning: sanitize_compact_reasoning(take_object_value(&mut body, "reasoning")), - tools: take_non_empty_array(&mut body, "tools"), - parallel_tool_calls: take_bool(&mut body, "parallel_tool_calls"), - text: take_text_format(&mut body), - }) -} - -fn invalid_responses_body_response() -> Response { - ( - StatusCode::BAD_REQUEST, - Json(openai_error( - "Request body must be a JSON object", - "invalid_request", - )), - ) - .into_response() -} - -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, - generate: None, - tuple_schema: 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, - version: None, - include_timing_metrics: None, - codex_window_id: None, - parent_thread_id: None, - use_websocket: None, - } -} - -fn take_string(body: &mut Map, key: &str) -> Option { - body.remove(key).and_then(|value| match value { - Value::String(value) => Some(value), - _ => None, - }) -} - -fn take_non_empty_string(body: &mut Map, key: &str) -> Option { - take_string(body, key).and_then(|value| { - let trimmed = value.trim(); - (!trimmed.is_empty()).then(|| trimmed.to_string()) - }) -} - -fn take_bool(body: &mut Map, key: &str) -> Option { - body.remove(key).and_then(|value| value.as_bool()) -} - -fn take_array(body: &mut Map, key: &str) -> Option> { - body.remove(key).and_then(|value| match value { - Value::Array(values) => Some(values), - _ => None, - }) -} - -fn take_non_empty_array(body: &mut Map, key: &str) -> Option> { - take_array(body, key).filter(|values| !values.is_empty()) -} - -fn take_object_value(body: &mut Map, key: &str) -> Option { - body.remove(key).and_then(|value| match value { - Value::Object(_) => Some(value), - _ => None, - }) -} - -fn take_string_array(body: &mut Map, key: &str) -> Option> { - let values = take_array(body, key)?; - values - .into_iter() - .map(|value| match value { - Value::String(value) => Some(value), - _ => None, - }) - .collect() -} - -struct PreparedTextFormat { - text: Value, - tuple_schema: Option, -} - -fn take_prepared_text_format(body: &mut Map) -> Option { - take_text_format_inner(body, true) -} - -fn take_text_format(body: &mut Map) -> Option { - take_text_format_inner(body, false).map(|prepared| prepared.text) -} - -fn take_text_format_inner( - body: &mut Map, - prepare_tuple_schema: bool, -) -> Option { - let Value::Object(text) = body.remove("text")? else { - return None; - }; - let Value::Object(format) = text.get("format")? else { - return None; - }; - let format_type = format.get("type")?.as_str()?; - let mut sanitized_format = Map::new(); - sanitized_format.insert("type".to_string(), Value::String(format_type.to_string())); - let mut tuple_schema = None; - if let Some(name) = format.get("name").and_then(Value::as_str) { - sanitized_format.insert("name".to_string(), Value::String(name.to_string())); - } - if let Some(Value::Object(schema)) = format.get("schema") { - let schema = Value::Object(schema.clone()); - let schema = if prepare_tuple_schema { - let prepared = prepare_schema(schema); - tuple_schema = prepared.original_schema; - prepared.schema - } else { - schema - }; - sanitized_format.insert("schema".to_string(), schema); - } - if let Some(strict) = format.get("strict").and_then(Value::as_bool) { - sanitized_format.insert("strict".to_string(), Value::Bool(strict)); - } - - let mut sanitized_text = Map::new(); - sanitized_text.insert("format".to_string(), Value::Object(sanitized_format)); - Some(PreparedTextFormat { - text: Value::Object(sanitized_text), - tuple_schema, - }) -} - -fn sanitize_client_metadata(client_metadata: Option) -> Option { - let Value::Object(input) = client_metadata? else { - return None; - }; - let metadata: Map = input - .into_iter() - .filter_map(|(key, value)| { - value - .as_str() - .map(|value| (key, Value::String(value.to_string()))) - }) - .collect(); - (!metadata.is_empty()).then_some(Value::Object(metadata)) -} - -fn sanitize_compact_reasoning(reasoning: Option) -> Option { - let Value::Object(input) = reasoning? else { - return None; - }; - let mut output = Map::new(); - if let Some(effort) = input.get("effort").and_then(Value::as_str) { - output.insert("effort".to_string(), Value::String(effort.to_string())); - } - if let Some(summary) = input.get("summary").and_then(Value::as_str) { - output.insert("summary".to_string(), Value::String(summary.to_string())); - } - (!output.is_empty()).then_some(Value::Object(output)) -} - -fn sanitize_codex_input_items(input: Vec) -> Vec { - input - .into_iter() - .filter_map(|item| { - let Value::Object(object) = item else { - return Some(item); - }; - match object.get("type").and_then(Value::as_str) { - Some("reasoning") => sanitize_reasoning_item(&object), - Some("compaction") => sanitize_compaction_item(&object), - _ => Some(Value::Object(object)), - } - }) - .collect() -} - -fn sanitize_reasoning_item(item: &Map) -> Option { - let id = non_empty_string(item.get("id"))?; - let summary = sanitize_summary(item.get("summary"))?; - let mut sanitized = Map::new(); - sanitized.insert("type".to_string(), Value::String("reasoning".to_string())); - sanitized.insert("id".to_string(), Value::String(id.to_string())); - sanitized.insert("summary".to_string(), Value::Array(summary)); - if let Some(status) = item - .get("status") - .and_then(Value::as_str) - .filter(|status| matches!(*status, "in_progress" | "completed" | "incomplete")) - { - sanitized.insert("status".to_string(), Value::String(status.to_string())); - } - if let Some(encrypted_content) = non_empty_string(item.get("encrypted_content")) { - sanitized.insert( - "encrypted_content".to_string(), - Value::String(encrypted_content.to_string()), - ); - } - if let Some(content) = sanitize_reasoning_content(item.get("content")) { - sanitized.insert("content".to_string(), Value::Array(content)); - } - Some(Value::Object(sanitized)) -} - -fn sanitize_summary(value: Option<&Value>) -> Option> { - let Value::Array(parts) = value? else { - return None; - }; - let summary: Vec = parts - .iter() - .filter_map(|part| { - let Value::Object(part) = part else { - return None; - }; - if part.get("type").and_then(Value::as_str) != Some("summary_text") { - return None; - } - let text = part.get("text").and_then(Value::as_str)?; - Some(json!({"type": "summary_text", "text": text})) - }) - .collect(); - Some(summary) -} - -fn sanitize_reasoning_content(value: Option<&Value>) -> Option> { - let Value::Array(parts) = value? else { - return None; - }; - let content: Vec = parts - .iter() - .filter_map(|part| { - let Value::Object(part) = part else { - return None; - }; - if part.get("type").and_then(Value::as_str) != Some("reasoning_text") { - return None; - } - let text = part.get("text").and_then(Value::as_str)?; - Some(json!({"type": "reasoning_text", "text": text})) - }) - .collect(); - (!content.is_empty()).then_some(content) -} - -fn sanitize_compaction_item(item: &Map) -> Option { - let encrypted_content = non_empty_string(item.get("encrypted_content"))?; - let mut sanitized = Map::new(); - sanitized.insert("type".to_string(), Value::String("compaction".to_string())); - sanitized.insert( - "encrypted_content".to_string(), - Value::String(encrypted_content.to_string()), - ); - if let Some(id) = non_empty_string(item.get("id")) { - sanitized.insert("id".to_string(), Value::String(id.to_string())); - } - Some(Value::Object(sanitized)) -} - -fn non_empty_string(value: Option<&Value>) -> Option<&str> { - let value = value?.as_str()?; - (!value.trim().is_empty()).then_some(value) -} - -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; - } - if request - .include - .as_ref() - .is_some_and(|include| !include.is_empty()) - { - return; - } - request.include = Some(vec!["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) -} - -fn client_metadata_with_openai_subagent( - client_metadata: Option, - openai_subagent: Option, -) -> Option { - let mut metadata = sanitize_client_metadata(client_metadata) - .and_then(|value| match value { - Value::Object(metadata) => Some(metadata), - _ => None, - }) - .unwrap_or_default(); - metadata.remove(OPENAI_SUBAGENT_HEADER); - if let Some(openai_subagent) = openai_subagent { - metadata.insert( - OPENAI_SUBAGENT_HEADER.to_string(), - Value::String(openai_subagent), - ); - } - - if metadata.is_empty() { - None - } else { - Some(Value::Object(metadata)) - } -} - -fn openai_subagent_from_metadata(client_metadata: Option<&Value>) -> Option { - normalize_openai_subagent( - client_metadata? - .as_object()? - .get(OPENAI_SUBAGENT_HEADER)? - .as_str(), - ) -} - -fn normalize_openai_subagent(value: Option<&str>) -> Option { - let value = value?.trim(); - if matches!( - value, - "review" | "compact" | "memory_consolidation" | "collab_spawn" - ) { - Some(value.to_string()) - } else { - None - } -} diff --git a/src/codex/tasks/mod.rs b/src/codex/tasks/mod.rs deleted file mode 100644 index f9a9de114..000000000 --- a/src/codex/tasks/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -pub mod cookie_cleanup; -pub mod model_refresh; -pub mod quota_refresh; -pub mod token_refresh; diff --git a/src/codex/tasks/quota_refresh.rs b/src/codex/tasks/quota_refresh.rs deleted file mode 100644 index c1c5f7d08..000000000 --- a/src/codex/tasks/quota_refresh.rs +++ /dev/null @@ -1,117 +0,0 @@ -use std::{collections::HashMap, sync::Arc, time::Duration}; -use tokio::time::{interval, sleep, Instant}; -use tracing::{debug, info, warn}; - -use crate::{codex::accounts::service::AccountService, runtime::tasks::types::SchedulerHandle}; - -/// 主动配额刷新调度器 -/// -/// 定期扫描并刷新配额锁定（limit_reached）的账户，防止账户被"永久锁死" -/// -/// 功能： -/// - 每15分钟扫描一次所有账户 -/// - 识别配额锁定的账户 -/// - 主动调用上游 getUsage() 刷新配额 -/// - 每个账户最少30分钟刷新一次（防滥用） -pub struct QuotaRefresher { - account_service: Arc, - interval_secs: u64, - min_refresh_interval_secs: u64, -} - -const DEFAULT_INTERVAL_SECS: u64 = 15 * 60; // 15分钟 -const MIN_REFRESH_INTERVAL_SECS: u64 = 30 * 60; // 30分钟 - -impl QuotaRefresher { - pub fn new(account_service: Arc) -> Self { - Self { - account_service, - interval_secs: DEFAULT_INTERVAL_SECS, - min_refresh_interval_secs: MIN_REFRESH_INTERVAL_SECS, - } - } - - /// 启动配额刷新调度器 - pub fn start(self) -> SchedulerHandle { - let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); - - tokio::spawn(async move { - let mut ticker = interval(Duration::from_secs(self.interval_secs)); - let mut last_refreshed: HashMap = HashMap::new(); - - info!(interval_secs = self.interval_secs, "quota 刷新器已启动"); - - loop { - tokio::select! { - _ = ticker.tick() => { - self.tick(&mut last_refreshed).await; - } - _ = shutdown_rx.recv() => { - info!("quota 刷新器正在关闭"); - break; - } - } - } - }); - - SchedulerHandle::new(shutdown_tx) - } - - async fn tick(&self, last_refreshed: &mut HashMap) { - let now = Instant::now(); - let min_interval = Duration::from_secs(self.min_refresh_interval_secs); - - // 获取所有配额锁定的账户 - let accounts_to_refresh = self.account_service.list_quota_locked_accounts().await; - - if accounts_to_refresh.is_empty() { - debug!("未找到 quota 锁定账户"); - return; - } - - info!( - count = accounts_to_refresh.len(), - "发现需要刷新 quota 的锁定账户" - ); - - for account_id in accounts_to_refresh { - // 检查是否最近刷新过 - if let Some(&last_time) = last_refreshed.get(&account_id) { - if now.duration_since(last_time) < min_interval { - debug!( - account_id = %account_id, - "账户最近已刷新，跳过 quota 刷新" - ); - continue; - } - } - - info!(account_id = %account_id, "正在主动刷新 quota"); - - // 记录刷新时间 - last_refreshed.insert(account_id.clone(), now); - - // 调用配额刷新（使用 health check 的 quota 方法） - let request_id = uuid::Uuid::new_v4().to_string(); - match self - .account_service - .account_quota(&account_id, &request_id) - .await - { - Ok(_) => { - info!(account_id = %account_id, "quota 刷新成功"); - } - Err(e) => { - warn!( - account_id = %account_id, - error = ?e, - "quota 刷新失败" - ); - } - } - - // 错开请求，避免突发 - sleep(Duration::from_millis(3000)).await; - } - } -} diff --git a/src/codex/tasks/token_refresh.rs b/src/codex/tasks/token_refresh.rs deleted file mode 100644 index 47058ad49..000000000 --- a/src/codex/tasks/token_refresh.rs +++ /dev/null @@ -1,882 +0,0 @@ -use rand::RngExt; -use std::{collections::HashMap, sync::Arc, time::Duration}; -use tokio::{ - sync::{RwLock, Semaphore}, - task::JoinHandle, - time::Instant, -}; -use tracing::{debug, error, info, warn}; - -use crate::{ - codex::accounts::{ - model::AccountStatus, - service::{AccountProbeOutcome, AccountService}, - }, - config::AppConfig, - runtime::tasks::types::{SchedulerError, SchedulerHandle}, -}; - -pub type SchedulerResult = Result; - -/// OAuth 刷新调度器 - 在 JWT 过期前自动刷新访问令牌 -/// -/// 功能： -/// - 在 exp - margin 时刻调度刷新 -/// - 指数退避（5次尝试：5s → 15s → 45s → 135s → 300s） -/// - 永久失败检测（invalid_grant / invalid_token） -/// - 临时失败的恢复调度（10分钟） -/// - 崩溃恢复：refreshing → 立即重试，expired + refreshToken → 延迟重试 -#[derive(Clone)] -pub struct RefreshScheduler { - account_service: Arc, - config: AppConfig, - timers: Arc>>>, - in_flight: Arc>>, - refresh_permits: Arc, - destroyed: Arc>, -} - -const MAX_ATTEMPTS: u32 = 5; -const BASE_DELAY_MS: u64 = 5_000; -const RECOVERY_DELAY_MS: u64 = 10 * 60 * 1000; // 10分钟 -const PERMANENT_THRESHOLD: u32 = 2; - -// 表示账户被封禁/停用的上游错误 -const BAN_ERRORS: &[&str] = &["account has been deactivated", "refresh_token_reused"]; - -// 表示刷新令牌无效但不一定是封禁的错误 -const EXPIRED_ERRORS: &[&str] = &[ - "invalid_grant", - "invalidgrant", - "invalid_token", - "access_denied", - "refresh_token_expired", -]; - -impl RefreshScheduler { - pub fn new(account_service: Arc, config: AppConfig) -> Self { - let refresh_permits = Arc::new(Semaphore::new(refresh_concurrency_limit( - config.auth.refresh_concurrency, - ))); - Self { - account_service, - config, - timers: Arc::new(RwLock::new(HashMap::new())), - in_flight: Arc::new(RwLock::new(HashMap::new())), - refresh_permits, - destroyed: Arc::new(RwLock::new(false)), - } - } - - /// 检查账户是否正在刷新（用于健康检查避免竞争） - pub async fn is_refreshing(&self, account_id: &str) -> bool { - self.in_flight.read().await.contains_key(account_id) - } - - /// 为所有账户调度刷新 - pub async fn schedule_all(&self) -> SchedulerResult<()> { - if !self.config.auth.refresh_enabled { - info!("自动刷新已关闭（refresh_enabled = false）"); - return Ok(()); - } - - let accounts = self - .account_service - .list_all_for_refresh() - .await - .map_err(|_| SchedulerError::AccountNotFound("failed to list accounts".to_string()))?; - - let mut expired_index = 0; - for account in accounts { - // 跳过没有刷新令牌的账户 - if account.refresh_token.is_none() { - continue; - } - - // 跳过永久禁用/封禁的账户 - if matches!( - account.status, - AccountStatus::Disabled | AccountStatus::Banned - ) { - continue; - } - - match account.status { - AccountStatus::Refreshing => { - // 崩溃恢复：进程终止时正在刷新 - info!(account_id = %account.id, "正在从 refreshing 状态恢复"); - self.do_refresh(account.id.clone()).await; - } - AccountStatus::Expired => { - // 恢复尝试 - 每个账户间隔2秒以避免突发 - let delay = Duration::from_millis(30_000 + expired_index * 2_000); - expired_index += 1; - info!( - account_id = %account.id, - delay_secs = delay.as_secs(), - "账户已过期，已调度恢复尝试" - ); - self.schedule_recovery(account.id.clone(), delay).await; - } - _ => { - // active / quota_exhausted - 在令牌过期时调度刷新 - use secrecy::ExposeSecret; - let token_str = account.access_token.expose_secret().to_string(); - self.schedule_one(account.id.clone(), token_str).await?; - } - } - } - - Ok(()) - } - - /// 触发账户的立即刷新（用于401触发的刷新） - pub async fn trigger_refresh_now(&self, account_id: String) -> SchedulerResult<()> { - // 检查账户是否正在刷新 - if self.is_refreshing(&account_id).await { - debug!(account_id = %account_id, "账户正在刷新，跳过本次触发"); - return Ok(()); - } - - // 取消现有定时器 - self.clear_one(&account_id).await; - - // 执行刷新 - self.do_refresh(account_id).await; - Ok(()) - } - - /// 为单个账户调度刷新 - async fn schedule_one(&self, account_id: String, access_token: String) -> SchedulerResult<()> { - // 清除现有定时器 - self.clear_one(&account_id).await; - - // 解析 JWT 获取过期时间 - let exp = match self.parse_jwt_exp(&access_token) { - Some(exp) => exp, - None => { - warn!(account_id = %account_id, "解析 JWT exp 失败，跳过刷新调度"); - return Ok(()); - } - }; - - let margin = self.config.auth.refresh_margin_seconds; - let now = chrono::Utc::now().timestamp() as u64; - let refresh_at = exp.saturating_sub(margin); - - if refresh_at <= now { - // 已经过了刷新时间 - 立即尝试刷新 - debug!(account_id = %account_id, "token 已超过刷新时间，立即刷新"); - self.do_refresh(account_id).await; - return Ok(()); - } - - let delay = jitter_delay(Duration::from_secs(refresh_at - now)); - info!( - account_id = %account_id, - delay_secs = delay.as_secs(), - "已调度 token 刷新" - ); - - let handle = self.spawn_delayed_refresh(account_id.clone(), delay); - - self.timers.write().await.insert(account_id, handle); - Ok(()) - } - - /// 调度下次刷新（在刷新成功后调用，避免递归） - async fn schedule_next_refresh( - &self, - account_id: String, - access_token: String, - ) -> SchedulerResult<()> { - // 清除现有定时器 - self.clear_one(&account_id).await; - - // 解析 JWT 获取过期时间 - let exp = match self.parse_jwt_exp(&access_token) { - Some(exp) => exp, - None => { - warn!(account_id = %account_id, "解析 JWT exp 失败，跳过刷新调度"); - return Ok(()); - } - }; - - let margin = self.config.auth.refresh_margin_seconds; - let now = chrono::Utc::now().timestamp() as u64; - let refresh_at = exp.saturating_sub(margin); - - if refresh_at <= now { - // 已经过了刷新时间 - debug!(account_id = %account_id, "刷新后 token 仍已超过下次刷新时间"); - return Ok(()); - } - - let delay = jitter_delay(Duration::from_secs(refresh_at - now)); - info!( - account_id = %account_id, - delay_secs = delay.as_secs(), - "已调度下一次 token 刷新" - ); - - let handle = self.spawn_delayed_refresh(account_id.clone(), delay); - - self.timers.write().await.insert(account_id, handle); - Ok(()) - } - - /// 调度恢复尝试（用于临时失败） - async fn schedule_recovery(&self, account_id: String, delay: Duration) { - let handle = self.spawn_delayed_refresh(account_id.clone(), jitter_delay(delay)); - - self.timers.write().await.insert(account_id, handle); - } - - fn spawn_delayed_refresh(&self, account_id: String, delay: Duration) -> JoinHandle<()> { - let scheduler = self.clone(); - tokio::spawn(async move { - tokio::time::sleep(delay).await; - // 定时器触发后先移除自身，避免刷新成功后调度下一次时 abort 当前任务。 - scheduler.timers.write().await.remove(&account_id); - scheduler.do_refresh(account_id).await; - }) - } - - /// 取消单个账户的定时器 - async fn clear_one(&self, account_id: &str) { - if let Some(handle) = self.timers.write().await.remove(account_id) { - handle.abort(); - } - } - - /// 执行刷新（带并发控制） - async fn do_refresh(&self, account_id: String) { - // 检查是否已销毁 - if *self.destroyed.read().await { - return; - } - - // 检查是否已在刷新中 - { - let mut in_flight = self.in_flight.write().await; - if in_flight.contains_key(&account_id) { - debug!(account_id = %account_id, "刷新任务已在执行，跳过"); - return; - } - in_flight.insert(account_id.clone(), Instant::now()); - } - - let _permit = match self.refresh_permits.clone().acquire_owned().await { - Ok(permit) => permit, - Err(error) => { - self.in_flight.write().await.remove(&account_id); - error!( - account_id = %account_id, - error = %error, - "refresh 并发控制已关闭" - ); - return; - } - }; - - // 执行刷新逻辑 - let result = self.do_refresh_inner(&account_id).await; - - // 清理 in_flight 标记 - self.in_flight.write().await.remove(&account_id); - - if let Err(e) = result { - error!(account_id = %account_id, error = %e, "token 刷新失败"); - } - } - - /// 内部刷新逻辑（带重试和错误处理） - #[tracing::instrument(skip(self), fields(account_id = %account_id))] - async fn do_refresh_inner(&self, account_id: &str) -> SchedulerResult<()> { - info!(account_id = %account_id, "开始刷新 token"); - - let mut permanent_hits = 0; - - for attempt in 1..=MAX_ATTEMPTS { - let error = match self - .account_service - .probe_scheduled_account_refresh(account_id) - .await - { - Ok(result) if matches!(result.outcome, AccountProbeOutcome::Alive) => { - info!(account_id = %account_id, "token 刷新成功"); - // 获取新令牌并重新调度下次刷新 - if let Ok(accounts) = self - .account_service - .export(vec![account_id.to_string()]) - .await - { - if let Some(account) = accounts.first() { - // 从 SecretString 中提取字符串 - use secrecy::ExposeSecret; - let token_str = account.access_token.expose_secret().to_string(); - // 不调用 schedule_one 避免递归，直接创建定时器 - let _ = self - .schedule_next_refresh(account_id.to_string(), token_str) - .await; - } - } - return Ok(()); - } - Ok(result) if matches!(result.outcome, AccountProbeOutcome::Skipped) => { - info!( - account_id = %account_id, - reason = result.error.as_deref().unwrap_or("skipped"), - "token 刷新已跳过" - ); - return Ok(()); - } - Ok(result) => result - .error - .unwrap_or_else(|| "token refresh failed".to_string()), - Err(error) => error.to_string(), - }; - - let error_msg = error.to_lowercase(); - - // 检查是否是永久错误 - let is_ban = BAN_ERRORS.iter().any(|&err| error_msg.contains(err)); - let is_expired = EXPIRED_ERRORS.iter().any(|&err| error_msg.contains(err)); - let is_permanent = is_ban || is_expired; - - if is_permanent { - permanent_hits += 1; - if permanent_hits >= PERMANENT_THRESHOLD { - error!( - account_id = %account_id, - hits = permanent_hits, - "检测到永久刷新失败" - ); - let status = if is_ban { "banned" } else { "expired" }; - if let Err(status_error) = - self.account_service.update_status(account_id, status).await - { - error!( - account_id = %account_id, - status, - error = ?status_error, - "写入永久刷新失败状态失败" - ); - } - return Err(SchedulerError::AccountNotFound(format!( - "Permanent failure: {}", - error - ))); - } - warn!( - account_id = %account_id, - hits = permanent_hits, - threshold = PERMANENT_THRESHOLD, - "检测到永久错误，继续重试确认" - ); - } - - if attempt < MAX_ATTEMPTS { - // 指数退避：5s, 15s, 45s, 135s, 300s（有上限） - let backoff = BASE_DELAY_MS.saturating_mul(3_u64.pow(attempt - 1)); - let backoff = backoff.min(300_000); - let delay = Duration::from_millis(backoff); - - warn!( - account_id = %account_id, - attempt, - max_attempts = MAX_ATTEMPTS, - delay_secs = delay.as_secs(), - error = %error, - "token 刷新尝试失败，准备重试" - ); - tokio::time::sleep(delay).await; - } else { - error!( - account_id = %account_id, - attempts = MAX_ATTEMPTS, - error = %error, - "token 刷新重试次数已耗尽" - ); - self.restore_account_for_recovery(account_id).await; - self.schedule_recovery( - account_id.to_string(), - Duration::from_millis(RECOVERY_DELAY_MS), - ) - .await; - return Err(SchedulerError::AccountNotFound(format!( - "Max attempts reached: {}", - error - ))); - } - } - - Ok(()) - } - - async fn restore_account_for_recovery(&self, account_id: &str) { - match self - .account_service - .update_status(account_id, "active") - .await - { - Ok(Some(_)) => { - info!( - account_id = %account_id, - "临时 refresh 失败后已恢复账号状态" - ); - } - Ok(None) => { - warn!( - account_id = %account_id, - "临时 refresh 失败后恢复账号状态失败：账号不存在" - ); - } - Err(error) => { - warn!( - account_id = %account_id, - error = ?error, - "临时 refresh 失败后恢复账号状态失败" - ); - } - } - } - - /// 解析 JWT 获取过期时间 - fn parse_jwt_exp(&self, token: &str) -> Option { - let parts: Vec<&str> = token.split('.').collect(); - if parts.len() != 3 { - return None; - } - - let payload = parts[1]; - use base64::Engine; - let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD - .decode(payload) - .ok()?; - let json: serde_json::Value = serde_json::from_slice(&decoded).ok()?; - json.get("exp")?.as_u64() - } - - /// 启动调度器并返回句柄 - pub async fn start(self) -> SchedulerHandle { - let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); - let scheduler = Arc::new(self); - - // 初始调度所有账户 - let _ = scheduler.schedule_all().await; - - // 启动后台任务监听关闭信号 - let scheduler_clone = scheduler.clone(); - tokio::spawn(async move { - shutdown_rx.recv().await; - *scheduler_clone.destroyed.write().await = true; - - // 取消所有定时器 - let mut timers = scheduler_clone.timers.write().await; - for handle in timers.values() { - handle.abort(); - } - timers.clear(); - - info!("token 刷新调度器已关闭"); - }); - - SchedulerHandle::new(shutdown_tx) - } - - /// 销毁调度器 - pub async fn destroy(&self) { - *self.destroyed.write().await = true; - - // 取消所有定时器 - let mut timers = self.timers.write().await; - for handle in timers.values() { - handle.abort(); - } - timers.clear(); - - // 清理 in_flight - self.in_flight.write().await.clear(); - } -} - -fn refresh_concurrency_limit(configured: u32) -> usize { - configured.max(1) as usize -} - -fn jitter_delay(delay: Duration) -> Duration { - let mut rng = rand::rng(); - jitter_delay_with_factor(delay, rng.random_range(0.8..=1.2)) -} - -fn jitter_delay_with_factor(delay: Duration, factor: f64) -> Duration { - let factor = factor.clamp(0.8, 1.2); - let millis = (delay.as_millis() as f64 * factor) - .round() - .min(u64::MAX as f64); - Duration::from_millis(millis as u64) -} - -#[cfg(test)] -mod tests { - use super::*; - use async_trait::async_trait; - use secrecy::SecretString; - use std::sync::{ - atomic::{AtomicUsize, Ordering}, - Arc, - }; - use tokio::sync::Mutex; - - use crate::{ - codex::{ - accounts::{ - pool::AccountPool, - repository::{AccountRepository, AccountUsageRepository, NewAccount}, - service::{AccountService, AccountServiceDependencies}, - }, - gateway::{ - fingerprint::model::Fingerprint, - oauth::{RefreshFailure, TokenPair, TokenRefresher}, - transport::websocket::CodexWebSocketPool, - }, - }, - config::{ - AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, - QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, - }, - platform::{crypto::SecretBox, storage::db::connect_sqlite}, - }; - - #[test] - fn refresh_concurrency_limit_should_never_be_zero() { - assert_eq!(refresh_concurrency_limit(0), 1); - } - - #[test] - fn jitter_delay_with_factor_should_keep_delay_inside_original_range() { - assert_eq!( - jitter_delay_with_factor(Duration::from_secs(100), 0.8), - Duration::from_secs(80) - ); - assert_eq!( - jitter_delay_with_factor(Duration::from_secs(100), 1.2), - Duration::from_secs(120) - ); - } - - #[tokio::test] - async fn do_refresh_inner_should_mark_expired_only_after_second_permanent_failure() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("refresh-scheduler.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([41u8; 32])); - repo.insert(NewAccount { - id: "acct_scheduler".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-old".to_string().into()), - refresh_token: Some(SecretString::new("refresh-old".to_string().into())), - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - let observed_statuses = Arc::new(Mutex::new(Vec::new())); - let refresher = AlwaysInvalidGrantRefresher { - repo: repo.clone(), - account_id: "acct_scheduler".to_string(), - calls: Arc::new(AtomicUsize::new(0)), - observed_statuses: observed_statuses.clone(), - }; - let config = test_config(url); - let account_service = Arc::new(AccountService::new( - Arc::new(config.clone()), - AccountServiceDependencies { - repository: Some(repo.clone()), - usage_repository: Some(AccountUsageRepository::new(pool)), - cookie_repository: None, - token_refresher: Some(Arc::new(refresher.clone())), - account_pool: Arc::new(Mutex::new(AccountPool::default())), - websocket_pool: Arc::new(CodexWebSocketPool::with_default_max_age()), - fingerprint: Fingerprint::default_for_tests(), - }, - )); - let scheduler = RefreshScheduler::new(account_service, config); - - let result = scheduler.do_refresh_inner("acct_scheduler").await; - - assert!(result.is_err()); - assert_eq!(refresher.calls.load(Ordering::SeqCst), 2); - assert_eq!( - observed_statuses.lock().await.as_slice(), - [AccountStatus::Refreshing, AccountStatus::Refreshing] - ); - let stored = repo.get("acct_scheduler").await.unwrap().unwrap(); - assert_eq!(stored.status, AccountStatus::Expired); - } - - #[tokio::test] - async fn do_refresh_inner_should_persist_refreshing_during_refresh_attempt() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("refresh-scheduler-refreshing.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([43u8; 32])); - repo.insert(NewAccount { - id: "acct_refreshing".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-old".to_string().into()), - refresh_token: Some(SecretString::new("refresh-old".to_string().into())), - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - let observed_statuses = Arc::new(Mutex::new(Vec::new())); - let refresher = StatusObservingSuccessRefresher { - repo: repo.clone(), - account_id: "acct_refreshing".to_string(), - observed_statuses: observed_statuses.clone(), - }; - let config = test_config(url); - let account_service = Arc::new(AccountService::new( - Arc::new(config.clone()), - AccountServiceDependencies { - repository: Some(repo.clone()), - usage_repository: Some(AccountUsageRepository::new(pool)), - cookie_repository: None, - token_refresher: Some(Arc::new(refresher)), - account_pool: Arc::new(Mutex::new(AccountPool::default())), - websocket_pool: Arc::new(CodexWebSocketPool::with_default_max_age()), - fingerprint: Fingerprint::default_for_tests(), - }, - )); - let scheduler = RefreshScheduler::new(account_service, config); - - let result = scheduler.do_refresh_inner("acct_refreshing").await; - - assert!(result.is_ok()); - assert_eq!( - observed_statuses.lock().await.as_slice(), - [AccountStatus::Refreshing] - ); - let stored = repo.get("acct_refreshing").await.unwrap().unwrap(); - assert_eq!(stored.status, AccountStatus::Active); - } - - #[tokio::test] - async fn do_refresh_inner_should_restore_refreshing_account_after_transient_failures() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("refresh-scheduler-recovery.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([42u8; 32])); - repo.insert(NewAccount { - id: "acct_recovery".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-old".to_string().into()), - refresh_token: Some(SecretString::new("refresh-old".to_string().into())), - access_token_expires_at: None, - status: AccountStatus::Refreshing, - }) - .await - .unwrap(); - let refresher = AlwaysTransportFailingRefresher::default(); - let config = test_config(url); - let account_service = Arc::new(AccountService::new( - Arc::new(config.clone()), - AccountServiceDependencies { - repository: Some(repo.clone()), - usage_repository: Some(AccountUsageRepository::new(pool)), - cookie_repository: None, - token_refresher: Some(Arc::new(refresher.clone())), - account_pool: Arc::new(Mutex::new(AccountPool::default())), - websocket_pool: Arc::new(CodexWebSocketPool::with_default_max_age()), - fingerprint: Fingerprint::default_for_tests(), - }, - )); - let scheduler = RefreshScheduler::new(account_service, config); - tokio::time::pause(); - let running_scheduler = scheduler.clone(); - let refresh = - tokio::spawn(async move { running_scheduler.do_refresh_inner("acct_recovery").await }); - - for expected_calls in 1..=MAX_ATTEMPTS as usize { - wait_for_refresh_calls(&refresher.calls, expected_calls).await; - } - - let result = refresh.await.unwrap(); - - assert!(result.is_err()); - assert_eq!( - refresher.calls.load(Ordering::SeqCst), - MAX_ATTEMPTS as usize - ); - let stored = repo.get("acct_recovery").await.unwrap().unwrap(); - assert_eq!(stored.status, AccountStatus::Active); - assert!(scheduler.timers.read().await.contains_key("acct_recovery")); - scheduler.destroy().await; - } - - async fn wait_for_refresh_calls(calls: &AtomicUsize, expected: usize) { - for _ in 0..160 { - if calls.load(Ordering::SeqCst) >= expected { - return; - } - tokio::task::yield_now().await; - if calls.load(Ordering::SeqCst) >= expected { - return; - } - std::thread::sleep(Duration::from_millis(1)); - if calls.load(Ordering::SeqCst) >= expected { - return; - } - tokio::time::advance(Duration::from_secs(5)).await; - } - let actual = calls.load(Ordering::SeqCst); - assert!( - actual >= expected, - "expected at least {expected} refresh calls, got {actual}" - ); - } - - #[derive(Clone)] - struct AlwaysInvalidGrantRefresher { - repo: AccountRepository, - account_id: String, - calls: Arc, - observed_statuses: Arc>>, - } - - #[async_trait] - impl TokenRefresher for AlwaysInvalidGrantRefresher { - async fn refresh(&self, _refresh_token: &str) -> Result { - self.calls.fetch_add(1, Ordering::SeqCst); - let status = self - .repo - .get(&self.account_id) - .await - .unwrap() - .unwrap() - .status; - self.observed_statuses.lock().await.push(status); - Err(RefreshFailure::InvalidGrant) - } - } - - #[derive(Clone)] - struct StatusObservingSuccessRefresher { - repo: AccountRepository, - account_id: String, - observed_statuses: Arc>>, - } - - #[async_trait] - impl TokenRefresher for StatusObservingSuccessRefresher { - async fn refresh(&self, _refresh_token: &str) -> Result { - let status = self - .repo - .get(&self.account_id) - .await - .unwrap() - .unwrap() - .status; - self.observed_statuses.lock().await.push(status); - Ok(TokenPair { - access_token: "access-new".to_string(), - refresh_token: None, - }) - } - } - - #[derive(Clone, Default)] - struct AlwaysTransportFailingRefresher { - calls: Arc, - } - - #[async_trait] - impl TokenRefresher for AlwaysTransportFailingRefresher { - async fn refresh(&self, _refresh_token: &str) -> Result { - self.calls.fetch_add(1, Ordering::SeqCst); - Err(RefreshFailure::Transport) - } - } - - 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: Default::default(), - }, - 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(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - 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, - }, - ws_pool: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - capacity: 2_000, - capture_body: false, - }, - } - } -} diff --git a/src/codex/usage/mod.rs b/src/codex/usage/mod.rs deleted file mode 100644 index 1f278a4d5..000000000 --- a/src/codex/usage/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod service; diff --git a/src/codex/usage/service.rs b/src/codex/usage/service.rs deleted file mode 100644 index 4edaa1ba0..000000000 --- a/src/codex/usage/service.rs +++ /dev/null @@ -1,47 +0,0 @@ -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/config/loader.rs b/src/config/loader.rs deleted file mode 100644 index 3dd9bb937..000000000 --- a/src/config/loader.rs +++ /dev/null @@ -1,19 +0,0 @@ -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 deleted file mode 100644 index 447e78c72..000000000 --- a/src/config/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -mod loader; -mod types; - -pub use types::*; - -pub type ConfigResult = Result; diff --git a/src/lib.rs b/src/lib.rs deleted file mode 100644 index c75032fff..000000000 --- a/src/lib.rs +++ /dev/null @@ -1,8 +0,0 @@ -pub mod admin; -pub mod codex; -pub mod runtime; -pub mod web; - -pub mod config; -pub mod platform; -pub mod utils; diff --git a/src/main.rs b/src/main.rs deleted file mode 100644 index c93618e0d..000000000 --- a/src/main.rs +++ /dev/null @@ -1,42 +0,0 @@ -use codex_proxy_rs::{ - config::AppConfig, - platform::logging::rotation::{init_tracing, RotationConfig}, - runtime::{bootstrap::build_state, build_router, tasks::start_background_tasks}, -}; -use tokio::signal; - -#[tokio::main] -async fn main() -> anyhow::Result<()> { - let config = AppConfig::load()?; - - let _log_guard = init_tracing(RotationConfig::new( - &config.logging.directory, - config.logging.retention_days, - ))?; - - let host = config.server.host.clone(); - let port = config.server.port; - let (state, db_pool, restored_accounts) = build_state(config.clone()).await?; - tracing::info!(account_count = restored_accounts, "账户池已从 SQLite 恢复"); - - let background_tasks = start_background_tasks(&state, db_pool.clone(), &config).await; - - let app = build_router(state); - let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?; - tracing::info!(host, port, "codex-proxy-rs 已开始监听"); - - tokio::select! { - result = axum::serve(listener, app) => { - if let Err(e) = result { - tracing::error!(error = %e, "服务器运行失败"); - } - } - _ = signal::ctrl_c() => { - tracing::info!("收到关闭信号"); - } - } - - background_tasks.shutdown().await; - - Ok(()) -} diff --git a/src/platform/crypto/mod.rs b/src/platform/crypto/mod.rs deleted file mode 100644 index a02525874..000000000 --- a/src/platform/crypto/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub mod secret_box; - -pub use secret_box::{CryptoError, CryptoResult, SecretBox}; diff --git a/src/platform/http/auth.rs b/src/platform/http/auth.rs deleted file mode 100644 index 6eb213df1..000000000 --- a/src/platform/http/auth.rs +++ /dev/null @@ -1,24 +0,0 @@ -use axum::http::HeaderMap; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ClientApiKey<'a>(&'a str); - -impl<'a> ClientApiKey<'a> { - pub fn as_str(self) -> &'a str { - self.0 - } -} - -pub fn client_api_key(headers: &HeaderMap) -> Option> { - let auth = headers.get("authorization")?.to_str().ok()?.trim(); - let token = auth.strip_prefix("Bearer ")?.trim(); - token.starts_with("cpr_").then_some(ClientApiKey(token)) -} - -pub fn admin_session_id(headers: &HeaderMap) -> Option<&str> { - let cookie = headers.get("cookie")?.to_str().ok()?; - cookie.split(';').map(str::trim).find_map(|part| { - let (name, value) = part.split_once('=')?; - (name == "cpr_admin_session" && !value.is_empty()).then_some(value) - }) -} diff --git a/src/platform/http/health.rs b/src/platform/http/health.rs deleted file mode 100644 index b6cee6ee9..000000000 --- a/src/platform/http/health.rs +++ /dev/null @@ -1,6 +0,0 @@ -use axum::Json; -use serde_json::{json, Value}; - -pub async fn health() -> Json { - Json(json!({ "status": "ok" })) -} diff --git a/src/platform/http/mod.rs b/src/platform/http/mod.rs deleted file mode 100644 index 52c7290a8..000000000 --- a/src/platform/http/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub mod auth; -pub mod health; -pub mod request_id; diff --git a/src/platform/identity/client_key_repository.rs b/src/platform/identity/client_key_repository.rs deleted file mode 100644 index 4082d396f..000000000 --- a/src/platform/identity/client_key_repository.rs +++ /dev/null @@ -1,219 +0,0 @@ -use chrono::Utc; -use sqlx::{Row, SqlitePool}; -use thiserror::Error; -use uuid::Uuid; - -use crate::{ - platform::identity::client_key::{ApiKeyHasher, GeneratedClientApiKey}, - utils::pagination::{decode_cursor, encode_cursor, Page}, -}; - -#[derive(Debug, Error)] -pub enum ClientApiKeyRepositoryError { - #[error("database error: {0}")] - Database(#[from] sqlx::Error), - #[error("api key verification error: {0}")] - Auth(#[from] crate::platform::identity::error::AuthError), - #[error("invalid pagination cursor")] - InvalidCursor, -} - -pub type ClientApiKeyRepositoryResult = Result; - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct StoredClientApiKey { - 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(Clone)] -pub struct ClientApiKeyRepository { - pool: SqlitePool, -} - -impl ClientApiKeyRepository { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - pub async fn insert_generated( - &self, - name: &str, - generated: &GeneratedClientApiKey, - ) -> ClientApiKeyRepositoryResult { - self.insert_generated_with_metadata(name, None, true, generated) - .await - } - - pub async fn insert_generated_with_metadata( - &self, - name: &str, - label: Option<&str>, - enabled: bool, - generated: &GeneratedClientApiKey, - ) -> ClientApiKeyRepositoryResult { - let id = format!("key_{}", Uuid::new_v4().simple()); - let now = Utc::now().to_rfc3339(); - sqlx::query( - "insert into client_api_keys (id, name, label, prefix, key_hash, enabled, created_at, last_used_at) values (?, ?, ?, ?, ?, ?, ?, null)", - ) - .bind(&id) - .bind(name) - .bind(label) - .bind(&generated.prefix) - .bind(&generated.key_hash) - .bind(if enabled { 1_i64 } else { 0_i64 }) - .bind(&now) - .execute(&self.pool) - .await?; - Ok(StoredClientApiKey { - id, - name: name.to_string(), - label: label.map(ToString::to_string), - prefix: generated.prefix.clone(), - enabled, - created_at: now, - last_used_at: None, - }) - } - - pub async fn list( - &self, - cursor: Option, - limit: u32, - ) -> ClientApiKeyRepositoryResult> { - let fetch_limit = i64::from(limit) + 1; - let rows = if let Some(cursor) = cursor { - let (created_at, id) = - decode_cursor(&cursor).ok_or(ClientApiKeyRepositoryError::InvalidCursor)?; - sqlx::query( - "select id, name, label, prefix, enabled, created_at, last_used_at from client_api_keys where created_at < ? or (created_at = ? and id < ?) order by created_at desc, id desc limit ?", - ) - .bind(&created_at) - .bind(created_at) - .bind(id) - .bind(fetch_limit) - .fetch_all(&self.pool) - .await? - } else { - sqlx::query( - "select id, name, label, prefix, enabled, created_at, last_used_at from client_api_keys order by created_at desc, id desc limit ?", - ) - .bind(fetch_limit) - .fetch_all(&self.pool) - .await? - }; - - let has_next = rows.len() > limit as usize; - let take_count = rows.len().min(limit as usize); - let items = rows - .into_iter() - .take(take_count) - .map(|row| key_from_row(&row)) - .collect::>(); - let next_cursor = if has_next { - items - .last() - .map(|key| encode_cursor(&key.created_at, &key.id)) - } else { - None - }; - Ok(Page { items, next_cursor }) - } - - pub async fn get(&self, id: &str) -> ClientApiKeyRepositoryResult> { - let row = sqlx::query( - "select id, name, label, prefix, enabled, created_at, last_used_at from client_api_keys where id = ?", - ) - .bind(id) - .fetch_optional(&self.pool) - .await?; - Ok(row.map(|row| key_from_row(&row))) - } - - pub async fn list_all(&self) -> ClientApiKeyRepositoryResult> { - let rows = sqlx::query( - "select id, name, label, prefix, enabled, created_at, last_used_at from client_api_keys order by created_at desc, id desc", - ) - .fetch_all(&self.pool) - .await?; - Ok(rows.into_iter().map(|row| key_from_row(&row)).collect()) - } - - pub async fn verify_and_touch( - &self, - plaintext: &str, - hasher: &ApiKeyHasher, - ) -> ClientApiKeyRepositoryResult { - let prefix = plaintext.chars().take(12).collect::(); - let rows = sqlx::query( - "select id, key_hash from client_api_keys where prefix = ? and enabled = 1", - ) - .bind(prefix) - .fetch_all(&self.pool) - .await?; - for row in rows { - let key_hash = row.get::("key_hash"); - if hasher.verify_client_api_key(plaintext, &key_hash)? { - let id = row.get::("id"); - sqlx::query("update client_api_keys set last_used_at = ? where id = ?") - .bind(Utc::now().to_rfc3339()) - .bind(id) - .execute(&self.pool) - .await?; - return Ok(true); - } - } - Ok(false) - } - - pub async fn set_enabled(&self, id: &str, enabled: bool) -> ClientApiKeyRepositoryResult { - let result = sqlx::query("update client_api_keys set enabled = ? where id = ?") - .bind(if enabled { 1_i64 } else { 0_i64 }) - .bind(id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - pub async fn set_label( - &self, - id: &str, - label: Option, - ) -> ClientApiKeyRepositoryResult> { - let result = sqlx::query("update client_api_keys set label = ? where id = ?") - .bind(label) - .bind(id) - .execute(&self.pool) - .await?; - if result.rows_affected() == 0 { - return Ok(None); - } - self.get(id).await - } - - pub async fn delete(&self, id: &str) -> ClientApiKeyRepositoryResult { - let result = sqlx::query("delete from client_api_keys where id = ?") - .bind(id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } -} - -fn key_from_row(row: &sqlx::sqlite::SqliteRow) -> StoredClientApiKey { - StoredClientApiKey { - id: row.get("id"), - name: row.get("name"), - label: row.get("label"), - prefix: row.get("prefix"), - enabled: row.get::("enabled") != 0, - created_at: row.get("created_at"), - last_used_at: row.get("last_used_at"), - } -} diff --git a/src/platform/identity/mod.rs b/src/platform/identity/mod.rs deleted file mode 100644 index d230d844f..000000000 --- a/src/platform/identity/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -pub mod admin_session; -pub mod client_key; -pub mod client_key_repository; -pub mod error; diff --git a/src/platform/logging/mod.rs b/src/platform/logging/mod.rs deleted file mode 100644 index a72110a3e..000000000 --- a/src/platform/logging/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod rotation; diff --git a/src/platform/mod.rs b/src/platform/mod.rs deleted file mode 100644 index 12b2bb3ab..000000000 --- a/src/platform/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub mod crypto; -pub mod http; -pub mod identity; -pub mod logging; -pub mod storage; diff --git a/src/platform/storage/mod.rs b/src/platform/storage/mod.rs deleted file mode 100644 index d533cee53..000000000 --- a/src/platform/storage/mod.rs +++ /dev/null @@ -1,2 +0,0 @@ -pub mod db; -pub mod paths; diff --git a/src/platform/storage/paths.rs b/src/platform/storage/paths.rs deleted file mode 100644 index d5fd4fa3b..000000000 --- a/src/platform/storage/paths.rs +++ /dev/null @@ -1,14 +0,0 @@ -use std::path::PathBuf; - -pub fn data_dir() -> PathBuf { - dirs::data_local_dir() - .unwrap_or_else(|| PathBuf::from(".")) - .join("codex-proxy-rs") - .join("data") -} - -pub fn ensure_data_dir() -> std::io::Result { - let dir = data_dir(); - std::fs::create_dir_all(&dir)?; - Ok(dir) -} diff --git a/src/runtime/bootstrap.rs b/src/runtime/bootstrap.rs deleted file mode 100644 index a0824b2df..000000000 --- a/src/runtime/bootstrap.rs +++ /dev/null @@ -1,120 +0,0 @@ -use thiserror::Error; - -use crate::{ - admin::session::repository::AdminAuthRepository, - codex::accounts::repository::AccountRepositoryError, - codex::gateway::fingerprint::{model::Fingerprint, repository::FingerprintRepository}, - codex::gateway::oauth::{OAuthConfig, OpenAiOAuthRefresher}, - codex::serving::dispatch::affinity::SessionAffinityRepositoryError, - config::AppConfig, - platform::{ - crypto::{CryptoError, SecretBox}, - identity::{ - admin_session::hash_admin_password, client_key::ApiKeyHasher, error::AuthError, - }, - storage::db::connect_sqlite, - }, -}; - -use super::state::AppState; -use sqlx::SqlitePool; - -#[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), - #[error("session affinity repository error: {0}")] - SessionAffinityRepository(#[from] SessionAffinityRepositoryError), -} - -pub type BootstrapResult = Result; - -pub async fn build_state(config: AppConfig) -> BootstrapResult<(AppState, SqlitePool, 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_config = OAuthConfig::from_auth_config(&config.auth); - let oauth_client = OpenAiOAuthRefresher::new( - reqwest::Client::builder() - .use_rustls_tls() - .no_proxy() - .build()?, - oauth_config, - ); - let pool = connect_sqlite(&config.database.url).await?; - - // 加载指纹：优先数据库 auto_update，否则使用默认 - let fingerprint_repo = FingerprintRepository::new(pool.clone()); - let fingerprint = match fingerprint_repo.load_latest_auto_updated().await { - Ok(Some(fp)) => { - tracing::info!( - version = %fp.app_version, - build = %fp.build_number, - source = "database", - "已从数据库加载请求 fingerprint" - ); - fp - } - Ok(None) => { - let fp = Fingerprint::default_codex_desktop(); - tracing::info!( - version = %fp.app_version, - build = %fp.build_number, - source = "default", - "使用默认请求 fingerprint" - ); - fp - } - Err(e) => { - tracing::warn!( - error = %e, - "从数据库加载 fingerprint 失败，使用默认值" - ); - Fingerprint::default_codex_desktop() - } - }; - - let admin_auth_repo = AdminAuthRepository::new(pool.clone()); - ensure_default_admin_exists(&admin_auth_repo, &config).await?; - - let state = AppState::with_pool_secret_api_key_hasher_oauth_client_and_fingerprint( - config, - pool.clone(), - secret_box, - api_key_hasher, - oauth_client, - fingerprint, - ); - let restored_accounts = state.reload_account_pool_from_repository().await?; - let restored_affinities = state.reload_session_affinity_from_repository().await?; - tracing::info!( - account_count = restored_accounts, - session_affinity_count = restored_affinities, - "已从 SQLite 恢复运行时状态" - ); - Ok((state, pool, restored_accounts)) -} - -/// 确保默认管理员账号存在（首次启动时创建） -async fn ensure_default_admin_exists( - repository: &AdminAuthRepository, - config: &AppConfig, -) -> Result<(), sqlx::Error> { - let password_hash = hash_admin_password(&config.admin.default_password) - .map_err(|e| sqlx::Error::Protocol(format!("failed to hash password: {e}")))?; - if repository.ensure_default_admin(&password_hash).await? { - tracing::info!( - username = %config.admin.default_username, - "已创建默认管理员用户" - ); - } - - Ok(()) -} diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs deleted file mode 100644 index d6641e904..000000000 --- a/src/runtime/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -pub mod bootstrap; -pub mod router; -pub mod state; -pub mod tasks; - -pub use router::{build_router, http_trace_layer}; diff --git a/src/runtime/router.rs b/src/runtime/router.rs deleted file mode 100644 index de8d5e002..000000000 --- a/src/runtime/router.rs +++ /dev/null @@ -1,120 +0,0 @@ -use std::time::Duration; - -use axum::{ - http::{Request, Response}, - middleware::from_fn, - routing::get, - Router, -}; -use tower_http::{ - classify::ServerErrorsFailureClass, - trace::{HttpMakeClassifier, MakeSpan, OnFailure, OnRequest, OnResponse, TraceLayer}, -}; -use tracing::Span; - -use crate::{ - admin::api as admin_api, - codex::serving::http::{ - diagnostics::{debug_fingerprint, debug_upstream, diagnostics}, - router as serving_http, - }, - platform::http::{ - health::health, - request_id::{attach_request_id, RequestId}, - }, -}; - -use super::state::AppState; - -pub type HttpTraceLayer = TraceLayer< - HttpMakeClassifier, - HttpMakeSpan, - HttpOnRequest, - HttpOnResponse, - (), - (), - HttpOnFailure, ->; - -#[derive(Debug, Clone, Copy)] -pub struct HttpMakeSpan; - -#[derive(Debug, Clone, Copy)] -pub struct HttpOnRequest; - -#[derive(Debug, Clone, Copy)] -pub struct HttpOnResponse; - -#[derive(Debug, Clone, Copy)] -pub struct HttpOnFailure; - -pub fn build_router(state: AppState) -> Router { - Router::new() - .route("/health", get(health)) - .route("/debug/diagnostics", get(diagnostics)) - .route("/debug/fingerprint", get(debug_fingerprint)) - .route("/debug/upstream", get(debug_upstream)) - .merge(serving_http::router()) - .merge(admin_api::router()) - .with_state(state) - .layer(http_trace_layer()) - .layer(from_fn(attach_request_id)) -} - -pub fn http_trace_layer() -> HttpTraceLayer { - TraceLayer::new_for_http() - .make_span_with(HttpMakeSpan) - .on_request(HttpOnRequest) - .on_response(HttpOnResponse) - .on_body_chunk(()) - .on_eos(()) - .on_failure(HttpOnFailure) -} - -impl MakeSpan for HttpMakeSpan { - fn make_span(&mut self, request: &Request) -> Span { - let request_id = request - .extensions() - .get::() - .map(RequestId::as_str) - .unwrap_or("missing"); - - tracing::info_span!( - "http_request", - request_id = %request_id, - method = %request.method(), - uri = %request.uri(), - status = tracing::field::Empty, - latency_ms = tracing::field::Empty, - ) - } -} - -impl OnRequest for HttpOnRequest { - fn on_request(&mut self, _request: &Request, _span: &Span) { - tracing::info!("收到 HTTP 请求"); - } -} - -impl OnResponse for HttpOnResponse { - fn on_response(self, response: &Response, latency: Duration, span: &Span) { - let status = response.status().as_u16(); - let latency_ms = latency.as_millis() as u64; - span.record("status", status); - span.record("latency_ms", latency_ms); - tracing::info!(status, latency_ms, "HTTP 请求完成"); - } -} - -impl OnFailure for HttpOnFailure { - fn on_failure( - &mut self, - failure_class: ServerErrorsFailureClass, - latency: Duration, - span: &Span, - ) { - let latency_ms = latency.as_millis() as u64; - span.record("latency_ms", latency_ms); - tracing::warn!(error = %failure_class, latency_ms, "HTTP 请求失败"); - } -} diff --git a/src/runtime/state.rs b/src/runtime/state.rs deleted file mode 100644 index 4af46a322..000000000 --- a/src/runtime/state.rs +++ /dev/null @@ -1,489 +0,0 @@ -use std::{path::PathBuf, sync::Arc, time::Duration}; - -use sqlx::SqlitePool; -use tokio::sync::Mutex; - -use crate::admin::{ - client_keys::service::ApiKeyService, - session::{ - repository::AdminAuthRepository, - service::{AdminAuthDependencies, AdminAuthService}, - }, - settings::SettingsService, -}; -use crate::codex::accounts::cookies::repository::CookieRepository; -use crate::codex::accounts::{ - pool::{AccountPool, AccountPoolOptions, RotationStrategy}, - repository::{AccountRepository, AccountRepositoryResult, AccountUsageRepository}, - service::{AccountService, AccountServiceDependencies}, -}; -use crate::codex::gateway::fingerprint::model::Fingerprint; -use crate::codex::gateway::oauth::{OAuthClient, OAuthConfig, PkceSessionStore, TokenRefresher}; -use crate::codex::gateway::transport::websocket::{CodexWebSocketPool, CodexWebSocketPoolConfig}; -use crate::codex::models::repository::ModelSnapshotRepository; -use crate::codex::models::service::ModelService; -use crate::codex::serving::dispatch::affinity::{ - SessionAffinityRepository, SessionAffinityRepositoryResult, -}; -use crate::codex::serving::dispatch::{CodexUpstreamRepositories, CodexUpstreamService}; -use crate::codex::serving::{ - chat::ChatService, diagnostics::DiagnosticsService, responses::ResponsesService, -}; -use crate::codex::usage::service::UsageService; -use crate::codex::{events::repository::EventLogRepository, events::service::LogService}; -use crate::config::AppConfig; -use crate::platform::crypto::SecretBox; -use crate::platform::identity::{ - client_key::ApiKeyHasher, client_key_repository::ClientApiKeyRepository, -}; - -#[derive(Clone)] -pub struct AppState { - pub services: Arc, -} - -pub struct AppServices { - pub config: Arc, - pub admin_auth: AdminAuthService, - pub accounts: AccountService, - pub api_keys: ApiKeyService, - pub logs: LogService, - pub usage: UsageService, - pub settings: SettingsService, - pub diagnostics: DiagnosticsService, - 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>, - local_config_path: Option, - fingerprint: 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_local_config_path( - config: AppConfig, - pool: SqlitePool, - local_config_path: impl Into, - ) -> Self { - Self::from_dependencies( - config, - AppStateDependencies { - pool: Some(pool), - local_config_path: Some(local_config_path.into()), - ..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), - ..AppStateDependencies::default() - }, - ) - } - - pub fn with_pool_secret_api_key_hasher_and_fingerprint( - config: AppConfig, - pool: SqlitePool, - secret_box: SecretBox, - api_key_hasher: ApiKeyHasher, - fingerprint: Fingerprint, - ) -> Self { - Self::from_dependencies( - config, - AppStateDependencies { - pool: Some(pool), - secret_box: Some(secret_box), - api_key_hasher: Some(api_key_hasher), - fingerprint: Some(fingerprint), - ..AppStateDependencies::default() - }, - ) - } - - pub fn with_pool_secret_api_key_hasher_oauth_client_and_fingerprint( - config: AppConfig, - pool: SqlitePool, - secret_box: SecretBox, - api_key_hasher: ApiKeyHasher, - oauth_client: C, - fingerprint: Fingerprint, - ) -> 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), - fingerprint: Some(fingerprint), - ..AppStateDependencies::default() - }, - ) - } - - 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 - } - - pub async fn reload_session_affinity_from_repository( - &self, - ) -> SessionAffinityRepositoryResult { - self.services - .responses - .reload_session_affinity_from_repository() - .await - } - - /// 获取 AccountService 的引用（用于调度器） - pub fn account_service(&self) -> Arc { - Arc::new(self.services.accounts.clone()) - } - - /// 获取 ModelService 的引用（用于调度器） - pub fn model_service(&self) -> Arc { - Arc::new(self.services.models.clone()) - } - - fn from_dependencies(config: AppConfig, dependencies: AppStateDependencies) -> Self { - let AppStateDependencies { - pool, - secret_box, - api_key_hasher, - token_refresher, - oauth_client, - local_config_path, - fingerprint, - } = dependencies; - let pool_ref = pool.as_ref(); - let secret_box_ref = secret_box.as_ref(); - let account_pool = account_pool_from_config(&config); - let websocket_pool = websocket_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 config = Arc::new(config); - let logs = log_service(&config, pool_ref); - let usage = usage_service(pool_ref); - let fingerprint = fingerprint.unwrap_or_else(Fingerprint::default_codex_desktop); - let settings = SettingsService::new( - config.clone(), - local_config_path.unwrap_or_else(|| PathBuf::from("local.yaml")), - ); - let diagnostics = DiagnosticsService::new(config.clone(), fingerprint.clone()); - let service_context = ServiceBuildContext { - config: config.clone(), - pool: pool_ref, - secret_box: secret_box_ref, - account_pool, - websocket_pool, - }; - let accounts = account_service( - &service_context, - token_refresher.clone(), - fingerprint.clone(), - ); - let admin_auth = admin_auth_service( - &service_context, - oauth_client, - oauth_sessions, - accounts.clone(), - ); - let V1Services { - chat, - responses, - models, - } = v1_services(service_context, logs.clone(), fingerprint); - Self { - services: Arc::new(AppServices { - config, - admin_auth, - accounts, - api_keys, - logs, - usage, - settings, - diagnostics, - 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(), - }))) -} - -fn websocket_pool_from_config(config: &AppConfig) -> Arc { - Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { - enabled: config.ws_pool.enabled, - max_age: Duration::from_millis(config.ws_pool.max_age_ms), - max_per_account: config.ws_pool.max_per_account, - ..CodexWebSocketPoolConfig::default() - })) -} - -struct V1Services { - chat: ChatService, - responses: ResponsesService, - models: ModelService, -} - -#[derive(Clone)] -struct ServiceBuildContext<'a> { - config: Arc, - pool: Option<&'a SqlitePool>, - secret_box: Option<&'a SecretBox>, - account_pool: Arc>, - websocket_pool: Arc, -} - -fn v1_services( - context: ServiceBuildContext<'_>, - logs: LogService, - fingerprint: Fingerprint, -) -> V1Services { - let upstream = codex_upstream_service(&context, logs, fingerprint.clone()); - let models = model_service( - context.config.clone(), - context.pool, - context.secret_box, - context.account_pool, - fingerprint, - ); - let model_config = context.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: Arc, - pool: Option<&SqlitePool>, - secret_box: Option<&SecretBox>, - account_pool: Arc>, - fingerprint: Fingerprint, -) -> ModelService { - ModelService::new( - config, - pool.cloned().map(ModelSnapshotRepository::new), - account_repository(pool, secret_box), - account_pool, - fingerprint, - ) -} - -fn admin_auth_service( - context: &ServiceBuildContext<'_>, - oauth_client: Option>, - oauth_sessions: Arc>, - accounts: AccountService, -) -> AdminAuthService { - let config = context.config.clone(); - AdminAuthService::new( - config, - AdminAuthDependencies { - auth_repository: context.pool.cloned().map(AdminAuthRepository::new), - account_repository: account_repository(context.pool, context.secret_box), - account_pool: context.account_pool.clone(), - oauth_config: OAuthConfig::from_auth_config(&context.config.auth), - oauth_client, - oauth_sessions, - accounts, - }, - ) -} - -fn account_service( - context: &ServiceBuildContext<'_>, - token_refresher: Option>, - fingerprint: Fingerprint, -) -> AccountService { - AccountService::new( - context.config.clone(), - AccountServiceDependencies { - repository: account_repository(context.pool, context.secret_box), - usage_repository: context.pool.cloned().map(AccountUsageRepository::new), - cookie_repository: cookie_repository(context.pool, context.secret_box), - token_refresher, - account_pool: context.account_pool.clone(), - websocket_pool: context.websocket_pool.clone(), - fingerprint, - }, - ) -} - -fn codex_upstream_service( - context: &ServiceBuildContext<'_>, - logs: LogService, - fingerprint: Fingerprint, -) -> CodexUpstreamService { - CodexUpstreamService::new( - context.config.clone(), - context.account_pool.clone(), - CodexUpstreamRepositories { - account: account_repository(context.pool, context.secret_box), - cookie: cookie_repository(context.pool, context.secret_box), - session_affinity: context.pool.cloned().map(SessionAffinityRepository::new), - }, - logs, - fingerprint, - context.websocket_pool.clone(), - ) -} - -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(config: &AppConfig, pool: Option<&SqlitePool>) -> LogService { - LogService::new( - config.logging.clone(), - pool.cloned().map(EventLogRepository::new), - ) -} - -fn usage_service(pool: Option<&SqlitePool>) -> UsageService { - UsageService::new(pool.cloned().map(AccountUsageRepository::new)) -} diff --git a/src/runtime/tasks/coordinator.rs b/src/runtime/tasks/coordinator.rs deleted file mode 100644 index 752210c56..000000000 --- a/src/runtime/tasks/coordinator.rs +++ /dev/null @@ -1,116 +0,0 @@ -use sqlx::SqlitePool; - -use crate::{ - admin::{ - session::repository::AdminAuthRepository, tasks::session_cleanup::SessionCleanupScheduler, - }, - codex::{ - gateway::fingerprint::{ - model::Fingerprint, repository::FingerprintRepository, update_checker::UpdateChecker, - }, - tasks::{ - cookie_cleanup::CookieCleanupScheduler, model_refresh::ModelRefresher, - quota_refresh::QuotaRefresher, token_refresh::RefreshScheduler, - }, - }, - config::AppConfig, - runtime::state::AppState, -}; - -use super::types::SchedulerHandle; - -#[derive(Default)] -pub struct BackgroundTaskCoordinator { - handles: Vec<(&'static str, SchedulerHandle)>, -} - -impl BackgroundTaskCoordinator { - fn push(&mut self, name: &'static str, handle: SchedulerHandle) { - self.handles.push((name, handle)); - } - - pub async fn shutdown(self) { - tracing::info!("正在关闭后台任务"); - for (name, handle) in self.handles { - handle.shutdown().await; - tracing::info!(task = name, "后台任务已停止"); - } - tracing::info!("所有后台任务已停止"); - } -} - -pub async fn start_background_tasks( - state: &AppState, - db_pool: SqlitePool, - config: &AppConfig, -) -> BackgroundTaskCoordinator { - let mut coordinator = BackgroundTaskCoordinator::default(); - - // 加载指纹：优先数据库 auto_update，否则使用默认 - let fingerprint_repo = FingerprintRepository::new(db_pool.clone()); - let fingerprint = match fingerprint_repo.load_latest_auto_updated().await { - Ok(Some(fp)) => { - tracing::info!( - version = %fp.app_version, - build = %fp.build_number, - "已从数据库加载 auto_update fingerprint" - ); - fp - } - Ok(None) => { - let fp = Fingerprint::default_codex_desktop(); - tracing::info!( - version = %fp.app_version, - build = %fp.build_number, - "使用默认 fingerprint" - ); - fp - } - Err(e) => { - tracing::warn!(error = %e, "从数据库加载 fingerprint 失败，使用默认值"); - Fingerprint::default_codex_desktop() - } - }; - - // 指纹自动更新（3 天轮询一次） - let update_checker = UpdateChecker::new( - Some(db_pool.clone()), - fingerprint.app_version.clone(), - fingerprint.build_number.clone(), - ); - let update_handle = update_checker.start_background_checker(); - coordinator.push( - "fingerprint_update", - SchedulerHandle::from_join_handle(update_handle), - ); - tracing::info!("fingerprint 更新检查器已启动"); - - let refresh_scheduler = RefreshScheduler::new(state.account_service(), config.clone()); - coordinator.push("refresh", refresh_scheduler.start().await); - tracing::info!("token 刷新调度器已启动"); - - let session_cleanup = SessionCleanupScheduler::new( - AdminAuthRepository::new(db_pool), - config.admin.session_cleanup_interval_secs, - ); - coordinator.push("session_cleanup", session_cleanup.start()); - tracing::info!("会话清理调度器已启动"); - - if let Some(cookie_repository) = state.account_service().cookie_cleanup_repository() { - let cookie_cleanup = CookieCleanupScheduler::new(cookie_repository); - coordinator.push("cookie_cleanup", cookie_cleanup.start()); - tracing::info!("cookie 清理调度器已启动"); - } else { - tracing::warn!("cookie repository 不可用，cookie 清理调度器未启动"); - } - - let quota_refresher = QuotaRefresher::new(state.account_service()); - coordinator.push("quota", quota_refresher.start()); - tracing::info!("quota 刷新器已启动"); - - let model_refresher = ModelRefresher::new(state.model_service()); - coordinator.push("model", model_refresher.start()); - tracing::info!("模型刷新器已启动"); - - coordinator -} diff --git a/src/runtime/tasks/mod.rs b/src/runtime/tasks/mod.rs deleted file mode 100644 index 8b5618a86..000000000 --- a/src/runtime/tasks/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub mod coordinator; -pub mod types; - -pub use coordinator::{start_background_tasks, BackgroundTaskCoordinator}; -pub use types::{SchedulerError, SchedulerHandle}; diff --git a/src/runtime/tasks/types.rs b/src/runtime/tasks/types.rs deleted file mode 100644 index 16e644db1..000000000 --- a/src/runtime/tasks/types.rs +++ /dev/null @@ -1,38 +0,0 @@ -use thiserror::Error; -use tokio::task::JoinHandle; - -#[derive(Debug, Error)] -pub enum SchedulerError { - #[error("scheduler already destroyed")] - Destroyed, - #[error("account {0} not found")] - AccountNotFound(String), - #[error("refresh lock acquisition failed")] - LockFailed, -} - -pub enum SchedulerHandle { - Channel(tokio::sync::mpsc::Sender<()>), - JoinHandle(JoinHandle<()>), -} - -impl SchedulerHandle { - pub fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>) -> Self { - Self::Channel(shutdown_tx) - } - - pub fn from_join_handle(handle: JoinHandle<()>) -> Self { - Self::JoinHandle(handle) - } - - pub async fn shutdown(self) { - match self { - Self::Channel(shutdown_tx) => { - let _ = shutdown_tx.send(()).await; - } - Self::JoinHandle(handle) => { - handle.abort(); - } - } - } -} diff --git a/src/utils/mod.rs b/src/utils/mod.rs deleted file mode 100644 index 24759b9b5..000000000 --- a/src/utils/mod.rs +++ /dev/null @@ -1,2 +0,0 @@ -pub mod json; -pub mod pagination; diff --git a/src/web/assets.rs b/src/web/assets.rs deleted file mode 100644 index 036d2f760..000000000 --- a/src/web/assets.rs +++ /dev/null @@ -1 +0,0 @@ -//! Asset loading placeholder for the future admin Web console. diff --git a/src/web/mod.rs b/src/web/mod.rs deleted file mode 100644 index 58558fec1..000000000 --- a/src/web/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -//! Backend mount point for the future admin Web console. - -pub mod assets; -pub mod router; -pub mod security; -pub mod shell; diff --git a/src/web/router.rs b/src/web/router.rs deleted file mode 100644 index 051cff208..000000000 --- a/src/web/router.rs +++ /dev/null @@ -1 +0,0 @@ -//! Router placeholder for the future admin Web console. diff --git a/src/web/security.rs b/src/web/security.rs deleted file mode 100644 index 78fe5f29a..000000000 --- a/src/web/security.rs +++ /dev/null @@ -1 +0,0 @@ -//! Security header placeholder for the future admin Web console. diff --git a/src/web/shell.rs b/src/web/shell.rs deleted file mode 100644 index 23f3e4a59..000000000 --- a/src/web/shell.rs +++ /dev/null @@ -1 +0,0 @@ -//! Shell rendering placeholder for the future admin Web console. diff --git a/tests/admin.rs b/tests/admin.rs deleted file mode 100644 index 3d838662d..000000000 --- a/tests/admin.rs +++ /dev/null @@ -1,21 +0,0 @@ -mod support; - -mod admin { - mod api_contract; - mod client_keys_route; - mod logs_route; - mod models_route; - mod session; - mod session_login_route; - mod session_repository; - mod settings_route; - mod usage_stats_route; - - mod accounts { - mod cookies_quota; - mod import_export; - mod lifecycle; - mod list; - mod oauth; - } -} diff --git a/tests/admin/accounts/cookies_quota.rs b/tests/admin/accounts/cookies_quota.rs deleted file mode 100644 index 79e09a11f..000000000 --- a/tests/admin/accounts/cookies_quota.rs +++ /dev/null @@ -1,662 +0,0 @@ -use std::sync::Arc; - -use async_trait::async_trait; -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use chrono::{Duration, Utc}; -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::gateway::{ - fingerprint::model::Fingerprint, - oauth::{RefreshFailure, TokenPair, TokenRefresher}, -}; - -use crate::support::{ - admin_accounts::{ - admin_accounts_test_app, admin_accounts_test_app_with_base_url_and_fingerprint, - admin_accounts_test_app_with_base_url_and_refresher, - admin_accounts_test_app_with_refresher, import_test_account, test_jwt, - 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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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 refreshed_access_token = test_jwt( - Some("chatgpt-account"), - Some("chatgpt-user"), - Some("refresh@example.com"), - Some("plus"), - 3600, - ); - let (app, state, pool, _dir) = admin_accounts_test_app_with_refresher( - "admin-account-refresh.sqlite", - 28, - StaticTokenRefresher { - result: Ok(TokenPair { - access_token: refreshed_access_token.clone(), - 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("/api/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(&refreshed_access_token)); - 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, refreshed_access_token); -} - -#[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("/api/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; - let window_reset_at = Utc::now() + Duration::minutes(5); - sqlx::query( - "insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds, last_used_at) values (?, 7, 11, 13, 17, 5, 19, 23, 29, ?, ?, 300, ?)", - ) - .bind("acct_reset_usage") - .bind("2026-06-12T12:30:00Z") - .bind(window_reset_at.to_rfc3339()) - .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("/api/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); - type ResetUsageRow = ( - i64, - i64, - i64, - i64, - i64, - i64, - i64, - i64, - Option, - Option, - ); - let usage: ResetUsageRow = sqlx::query_as( - "select request_count, input_tokens, output_tokens, cached_tokens, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_reset_at, 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, - 0, - 0, - 0, - 0, - Some(window_reset_at.to_rfc3339()), - None - ) - ); - let acquired = state - .services - .accounts - .acquire_runtime_account("gpt-5.5") - .await - .unwrap(); - assert_eq!(acquired.request_count, 1); - assert_eq!(acquired.window_reset_at, Some(window_reset_at)); - assert_ne!( - acquired.last_used_at.as_deref(), - Some("2026-06-12T12:00:00Z") - ); -} - -#[tokio::test] -async fn admin_account_quota_should_fetch_usage_store_quota_and_not_return_secrets() { - let server = MockServer::start().await; - let fingerprint = quota_test_fingerprint(); - let expected_user_agent = fingerprint.user_agent(); - let expected_sec_ch_ua = fingerprint.sec_ch_ua(); - 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_and_fingerprint( - "admin-account-quota.sqlite", - 31, - server.uri(), - fingerprint, - ) - .await; - import_test_account(&app, "session_1", "acct_quota").await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/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")); - let requests = server.received_requests().await.unwrap(); - assert_eq!( - requests[0] - .headers - .get("user-agent") - .and_then(|value| value.to_str().ok()), - Some(expected_user_agent.as_str()) - ); - assert_eq!( - requests[0] - .headers - .get("sec-ch-ua") - .and_then(|value| value.to_str().ok()), - Some(expected_sec_ch_ua.as_str()) - ); -} - -fn quota_test_fingerprint() -> Fingerprint { - Fingerprint { - app_version: "27.111.222".to_string(), - build_number: "9001".to_string(), - chromium_version: "155".to_string(), - ..Fingerprint::default_for_tests() - } -} - -#[tokio::test] -async fn admin_account_quota_warnings_should_require_admin_session_cookie() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-account-quota-warnings-auth.sqlite", 33).await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/accounts/quota-warnings") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); -} - -#[tokio::test] -async fn admin_account_quota_warnings_should_return_threshold_matches_from_cached_quota() { - let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-account-quota-warnings.sqlite", 34).await; - import_test_account(&app, "session_1", "acct_warn").await; - import_test_account(&app, "session_1", "acct_quiet").await; - - sqlx::query( - "update accounts set quota_json = ?, quota_fetched_at = ?, updated_at = ? where id = ?", - ) - .bind( - json!({ - "rate_limit": { - "used_percent": 85, - "reset_at": 1770000100 - }, - "secondary_rate_limit": { - "used_percent": 91, - "reset_at": 1770000200 - } - }) - .to_string(), - ) - .bind("2026-06-13T00:00:00Z") - .bind("2026-06-13T00:00:00Z") - .bind("acct_warn") - .execute(&pool) - .await - .unwrap(); - sqlx::query( - "update accounts set quota_json = ?, quota_fetched_at = ?, updated_at = ? where id = ?", - ) - .bind( - json!({ - "rate_limit": { - "used_percent": 25, - "reset_at": 1770000300 - }, - "secondary_rate_limit": null - }) - .to_string(), - ) - .bind("2026-06-13T01:00:00Z") - .bind("2026-06-13T01:00:00Z") - .bind("acct_quiet") - .execute(&pool) - .await - .unwrap(); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/accounts/quota-warnings") - .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"]["updatedAt"], "2026-06-13T00:00:00Z"); - let warnings = body["data"]["warnings"].as_array().unwrap(); - assert_eq!(warnings.len(), 2); - assert!(warnings - .iter() - .all(|warning| warning["accountId"] == "acct_warn")); - assert!(warnings.iter().all(|warning| { - warning["email"] == "acct_warn@example.com" && warning["usedPercent"].as_f64().is_some() - })); - assert!(warnings.iter().any(|warning| { - warning["window"] == "primary" - && warning["level"] == "warning" - && warning["usedPercent"] == 85.0 - && warning["resetAt"] == 1770000100 - })); - assert!(warnings.iter().any(|warning| { - warning["window"] == "secondary" - && warning["level"] == "critical" - && warning["usedPercent"] == 91.0 - && warning["resetAt"] == 1770000200 - })); -} - -#[tokio::test] -async fn admin_accounts_health_check_should_refresh_oauth_without_touching_codex_backend() { - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/api/codex/usage")) - .respond_with(ResponseTemplate::new(500)) - .expect(0) - .mount(&server) - .await; - let calls = Arc::new(Mutex::new(Vec::new())); - let (app, _state, pool, _dir) = admin_accounts_test_app_with_base_url_and_refresher( - "admin-account-health.sqlite", - 32, - server.uri(), - HealthCheckTokenRefresher { - calls: calls.clone(), - }, - ) - .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("/api/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 mut refresh_calls = calls.lock().await.clone(); - refresh_calls.sort(); - assert_eq!( - refresh_calls, - vec![ - "refresh-acct_health_alive".to_string(), - "refresh-acct_health_dead".to_string(), - ] - ); - 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"); -} - -#[tokio::test] -async fn admin_accounts_health_check_should_reject_unsupported_stagger_ms_field() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-account-health-unsupported-field.sqlite", 32).await; - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/accounts/health-check") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(r#"{"stagger_ms":1000}"#)) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_REQUEST); - let body = response_json(response).await; - assert_eq!(body["code"], 40001); -} - -#[derive(Clone)] -struct HealthCheckTokenRefresher { - calls: Arc>>, -} - -#[async_trait] -impl TokenRefresher for HealthCheckTokenRefresher { - async fn refresh(&self, refresh_token: &str) -> Result { - self.calls.lock().await.push(refresh_token.to_string()); - match refresh_token { - "refresh-acct_health_alive" => Ok(TokenPair { - access_token: "new-health-access".to_string(), - refresh_token: Some("new-health-refresh".to_string()), - }), - "refresh-acct_health_dead" => Err(RefreshFailure::InvalidGrant), - _ => Err(RefreshFailure::Transport), - } - } -} diff --git a/tests/admin/accounts/import_export.rs b/tests/admin/accounts/import_export.rs deleted file mode 100644 index b43d44dcc..000000000 --- a/tests/admin/accounts/import_export.rs +++ /dev/null @@ -1,863 +0,0 @@ -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::{ - codex::accounts::repository::AccountRepository, codex::gateway::oauth::TokenPair, - platform::crypto::SecretBox, platform::storage::db::connect_sqlite, runtime::build_router, - runtime::state::AppState, -}; - -use crate::support::{ - 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("/api/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("/api/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("/api/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_reject_non_native_export_shape() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-external-import.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": "external-data", - "version": 1, - "legacy": [], - "accounts": [ - { - "platform": "openai", - "type": "oauth", - "credentials": { - "access_token": "Bearer external-access-secret", - "refresh_token": "rt_external", - "email": "team@example.com", - "chatgpt_account_id": "chatgpt-account", - "chatgpt_user_id": "chatgpt-user", - "plan_type": "team" - } - } - ] - }); - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/accounts/import") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_external_import") - .body(Body::from(import_body.to_string())) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_REQUEST); - let body = response_json(response).await; - assert_eq!(body["message"], "No importable accounts found"); - let stored: (i64,) = sqlx::query_as("select count(*) from accounts") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(stored.0, 0); -} - -#[tokio::test] -async fn admin_accounts_import_should_reject_native_payload_with_unknown_account_fields() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-native-unknown-field.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_native_unknown", - "token": "native-access-secret", - "refreshToken": "native-refresh-secret", - "email": "native@example.com", - "accountId": "native-account", - "userId": "native-user", - "label": "Native Unknown", - "planType": "plus", - "legacyField": "ignored-secret", - "status": "active" - }] - }); - - let import_response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/accounts/import") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_native_unknown") - .body(Body::from(import_body.to_string())) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(import_response.status(), StatusCode::BAD_REQUEST); - let body = response_json(import_response).await; - assert_eq!(body["message"], "No importable accounts found"); -} - -#[tokio::test] -async fn admin_accounts_import_should_reject_native_payload_with_unknown_container_fields() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-native-container-field.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_native_extra_container", - "token": "native-access-secret", - "refreshToken": "native-refresh-secret", - "email": "native@example.com", - "accountId": "native-account", - "userId": "native-user", - "label": "Native", - "planType": "plus", - "status": "active" - }], - "legacyContainer": true - }); - - let import_response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/accounts/import") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_native_container") - .body(Body::from(import_body.to_string())) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(import_response.status(), StatusCode::BAD_REQUEST); - let body = response_json(import_response).await; - assert_eq!(body["message"], "No importable accounts found"); -} - -#[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("/api/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("legacyField").is_none()); - - let invalid = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/accounts/export?format=external") - .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_reject_unsupported_external_format() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-account-export-external.sqlite", 25).await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/accounts/export?format=external") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::BAD_REQUEST); -} - -#[tokio::test] -async fn admin_accounts_export_should_reject_unsupported_full_format() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-account-export-full-alias.sqlite", 25).await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/accounts/export?format=full") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::BAD_REQUEST); -} - -#[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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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/lifecycle.rs b/tests/admin/accounts/lifecycle.rs deleted file mode 100644 index cda699a0d..000000000 --- a/tests/admin/accounts/lifecycle.rs +++ /dev/null @@ -1,330 +0,0 @@ -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use serde_json::json; -use tower::ServiceExt; - -use crate::support::{ - 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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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/list.rs b/tests/admin/accounts/list.rs deleted file mode 100644 index 79d7febe0..000000000 --- a/tests/admin/accounts/list.rs +++ /dev/null @@ -1,57 +0,0 @@ -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use chrono::Utc; -use tower::ServiceExt; - -use codex_proxy_rs::{ - platform::crypto::SecretBox, platform::storage::db::connect_sqlite, runtime::build_router, - runtime::state::AppState, -}; - -use crate::support::{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("/api/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/oauth.rs b/tests/admin/accounts/oauth.rs deleted file mode 100644 index d30abb282..000000000 --- a/tests/admin/accounts/oauth.rs +++ /dev/null @@ -1,678 +0,0 @@ -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::gateway::oauth::{DeviceCode, OAuthError, TokenPair}, - platform::crypto::SecretBox, -}; - -use crate::support::{ - admin_accounts::{ - admin_accounts_test_app, admin_accounts_test_app_with_config, - admin_accounts_test_app_with_oauth_client, import_test_account, test_config, 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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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("/api/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_login_start_should_use_configured_oauth_authorize_endpoint() { - let dir = tempfile::tempdir().unwrap(); - let db = dir - .path() - .join("admin-auth-login-start-configured-oauth.sqlite"); - let mut config = test_config(format!("sqlite://{}", db.display())); - config.auth.oauth_client_id = "app_configured_client".to_string(); - config.auth.oauth_auth_endpoint = "https://auth.example.test/oauth/authorize".to_string(); - config.auth.oauth_token_endpoint = "https://auth.example.test/oauth/token".to_string(); - let (app, _state, _pool, _dir) = admin_accounts_test_app_with_config(dir, config, 42).await; - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/api/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_configured_oauth") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - let auth_url = body["data"]["authUrl"].as_str().unwrap(); - assert!(auth_url.starts_with("https://auth.example.test/oauth/authorize?")); - assert!(auth_url.contains("client_id=app_configured_client")); -} - -#[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("/api/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("/api/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("/api/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("/api/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_auth_callback_should_not_expose_unsupported_admin_callback_path() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-auth-callback-unsupported-admin.sqlite", 45).await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/auth/callback?code=callback-code&state=unsupported") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::NOT_FOUND); -} - -#[tokio::test] -async fn admin_auth_callback_should_not_expose_unsupported_openai_callback_path() { - let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-auth-callback-unsupported-openai.sqlite", 46).await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/auth/openai/callback?code=callback-code&state=unsupported") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::NOT_FOUND); -} diff --git a/tests/admin/session_repository.rs b/tests/admin/session_repository.rs deleted file mode 100644 index 7c41afcaf..000000000 --- a/tests/admin/session_repository.rs +++ /dev/null @@ -1,57 +0,0 @@ -use chrono::{Duration, Utc}; - -use codex_proxy_rs::{ - admin::session::repository::AdminAuthRepository, - platform::{identity::admin_session::hash_admin_password, storage::db::connect_sqlite}, -}; - -#[tokio::test] -async fn admin_auth_repository_should_create_and_load_default_admin_once() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-auth.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AdminAuthRepository::new(pool); - let password_hash = hash_admin_password("admin-password").unwrap(); - - let created = repo.ensure_default_admin(&password_hash).await.unwrap(); - let skipped = repo.ensure_default_admin("different-hash").await.unwrap(); - let admin = repo.load_first_admin().await.unwrap().unwrap(); - - assert!(created); - assert!(!skipped); - assert_eq!(admin.password_hash, password_hash); -} - -#[tokio::test] -async fn admin_auth_repository_should_create_validate_and_cleanup_sessions() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-sessions.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AdminAuthRepository::new(pool); - let password_hash = hash_admin_password("admin-password").unwrap(); - repo.ensure_default_admin(&password_hash).await.unwrap(); - let admin = repo.load_first_admin().await.unwrap().unwrap(); - let future_session = "sess_future"; - let expired_session = "sess_expired"; - - repo.create_session(future_session, &admin.id, Utc::now() + Duration::minutes(5)) - .await - .unwrap(); - repo.create_session( - expired_session, - &admin.id, - Utc::now() - Duration::minutes(5), - ) - .await - .unwrap(); - - assert!(repo.validate_session(future_session).await.unwrap()); - assert!(!repo.validate_session(expired_session).await.unwrap()); - - let deleted = repo.cleanup_expired_sessions(Utc::now()).await.unwrap(); - - assert_eq!(deleted, 1); - assert!(repo.validate_session(future_session).await.unwrap()); -} diff --git a/tests/admin/usage_stats_route.rs b/tests/admin/usage_stats_route.rs deleted file mode 100644 index fdda32189..000000000 --- a/tests/admin/usage_stats_route.rs +++ /dev/null @@ -1,246 +0,0 @@ -use std::collections::BTreeMap; - -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use chrono::Utc; -use tower::ServiceExt; - -use codex_proxy_rs::{ - config::{ - AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, - QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, - }, - platform::storage::db::connect_sqlite, - runtime::build_router, - runtime::state::AppState, -}; - -use crate::support::{response_json, seed_admin_session}; - -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(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - 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, - }, - ws_pool: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - capacity: 2_000, - capture_body: false, - }, - } -} - -#[tokio::test] -async fn admin_usage_stats_should_cursor_page_account_usage() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-usage-stats.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - seed_usage(&pool).await; - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let first_response = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage-stats?limit=1") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(first_response.status(), StatusCode::OK); - let first_body = response_json(first_response).await; - assert_eq!(first_body["code"], 200); - assert_eq!(first_body["requestId"], "req_usage"); - assert_eq!(first_body["data"].as_array().unwrap().len(), 1); - assert_eq!(first_body["data"][0]["accountId"], "acct_b"); - assert_eq!(first_body["data"][0]["requestCount"], 2); - assert_eq!(first_body["data"][0]["emptyResponseCount"], 1); - assert_eq!(first_body["data"][0]["inputTokens"], 7); - assert_eq!(first_body["page"]["limit"], 1); - let cursor = first_body["page"]["nextCursor"].as_str().unwrap(); - - let second_response = app - .oneshot( - Request::builder() - .method("GET") - .uri(format!("/api/admin/usage-stats?limit=1&cursor={cursor}")) - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(second_response.status(), StatusCode::OK); - let second_body = response_json(second_response).await; - assert_eq!(second_body["data"][0]["accountId"], "acct_a"); - assert!(second_body["page"]["nextCursor"].is_null()); -} - -#[tokio::test] -async fn admin_usage_stats_summary_should_return_usage_totals() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-usage-stats.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - seed_usage(&pool).await; - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage-stats/summary") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage_summary") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["code"], 200); - assert_eq!(body["data"]["accountCount"], 2); - assert_eq!(body["data"]["requestCount"], 5); - assert_eq!(body["data"]["emptyResponseCount"], 1); - assert_eq!(body["data"]["inputTokens"], 19); - assert_eq!(body["data"]["outputTokens"], 8); - assert_eq!(body["data"]["cachedTokens"], 3); -} - -#[tokio::test] -async fn admin_usage_stats_should_reject_missing_admin_session_cookie() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-usage-stats.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let app = build_router(AppState::with_pool(test_config(url), pool)); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage-stats") - .header("x-request-id", "req_usage") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - let body = response_json(response).await; - assert_eq!(body["code"], 40101); - assert_eq!(body["requestId"], "req_usage"); -} - -async fn seed_usage(pool: &sqlx::SqlitePool) { - let now = Utc::now().to_rfc3339(); - for (id, email, label, plan_type) in [ - ("acct_a", "a@example.com", "primary", "plus"), - ("acct_b", "b@example.com", "backup", "free"), - ] { - sqlx::query( - "insert into accounts (id, email, label, plan_type, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?, 'active', ?, ?)", - ) - .bind(id) - .bind(email) - .bind(label) - .bind(plan_type) - .bind("encrypted") - .bind(&now) - .bind(&now) - .execute(pool) - .await - .unwrap(); - } - sqlx::query( - "insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)", - ) - .bind("acct_a") - .bind(3_i64) - .bind(0_i64) - .bind(12_i64) - .bind(5_i64) - .bind(1_i64) - .bind("2026-06-11T00:00:00Z") - .execute(pool) - .await - .unwrap(); - sqlx::query( - "insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)", - ) - .bind("acct_b") - .bind(2_i64) - .bind(1_i64) - .bind(7_i64) - .bind(3_i64) - .bind(2_i64) - .bind("2026-06-11T00:01:00Z") - .execute(pool) - .await - .unwrap(); -} diff --git a/tests/architecture.rs b/tests/architecture.rs index 8400f2988..d1197d036 100644 --- a/tests/architecture.rs +++ b/tests/architecture.rs @@ -1,10 +1,14 @@ mod architecture { mod accounts_boundary; - mod admin_accounts_split; - mod admin_boundary; - mod codex_activity_boundary; + mod dependency_direction; + mod directory_shape; + mod forbidden_imports; + mod forbidden_legacy_paths; mod gateway_boundary; + mod placeholder_implementations; mod platform_boundary; mod runtime_tasks_boundary; - mod serving_boundary; + mod server_startup; + mod test_migration; + mod workspace_scaffold; } diff --git a/tests/architecture/accounts_boundary.rs b/tests/architecture/accounts_boundary.rs index 9ada0647b..be8fe372c 100644 --- a/tests/architecture/accounts_boundary.rs +++ b/tests/architecture/accounts_boundary.rs @@ -1,7 +1,15 @@ -use codex_proxy_rs::codex::{accounts::cookies::jar::CookieJar, models::catalog::ModelCatalog}; +use codex_proxy_core::{ + accounts::jwt::{jwt_expiry, JwtExpiry}, + auth::oauth::{RefreshPolicy, RefreshTrigger}, + gateway::conversation::{build_conversation_identity, ConversationIdentity}, +}; #[test] -fn accounts_exports_account_scoped_assets() { - let _cookie_jar_type = std::any::type_name::(); - let _model_catalog_type = std::any::type_name::(); +fn core_exports_account_domain_helpers() { + let _expiry_type = std::any::type_name::(); + let _identity_type = std::any::type_name::(); + let _refresh_policy_type = std::any::type_name::(); + let _refresh_trigger_type = std::any::type_name::(); + let _jwt_expiry_fn = jwt_expiry; + let _identity_fn = build_conversation_identity; } diff --git a/tests/architecture/admin_accounts_split.rs b/tests/architecture/admin_accounts_split.rs deleted file mode 100644 index 18adfeb97..000000000 --- a/tests/architecture/admin_accounts_split.rs +++ /dev/null @@ -1,51 +0,0 @@ -use codex_proxy_rs::admin::api::accounts::{ - cookies::{AccountCookiesData, SetAccountCookiesRequest}, - create::CreateAccountRequest, - delete::{BatchDeleteAccountsRequest, DeleteAccountData}, - export::AccountExportQuery, - health::{HealthCheckData, HealthCheckRequest}, - import::{AccountImportData, ImportCliAuthRequest}, - lifecycle::{ - BatchUpdateAccountStatusRequest, ResetAccountUsageData, UpdateAccountLabelRequest, - UpdateAccountStatusRequest, - }, - list::AccountsQuery, - oauth::{ - auth_callback, auth_code_relay, auth_device_login, auth_device_poll, auth_login_start, - AdminAuthCallbackQuery, AdminAuthCodeRelayData, AdminAuthCodeRelayRequest, - AdminAuthDeviceLoginData, AdminAuthDevicePollData, AdminAuthLoginStartData, - }, - quota::{AccountQuotaData, AccountQuotaWarningsData}, -}; - -#[test] -fn admin_accounts_handlers_are_split_by_workflow() { - let _list_query = std::any::type_name::(); - let _export_query = std::any::type_name::(); - let _create_request = std::any::type_name::(); - let _import_data = std::any::type_name::(); - let _cli_request = std::any::type_name::(); - let _label_request = std::any::type_name::(); - let _status_request = std::any::type_name::(); - let _reset_data = std::any::type_name::(); - let _batch_status = std::any::type_name::(); - let _delete_data = std::any::type_name::(); - let _batch_delete = std::any::type_name::(); - let _health_request = std::any::type_name::(); - let _health_data = std::any::type_name::(); - let _quota_data = std::any::type_name::(); - let _quota_warnings = std::any::type_name::(); - let _set_cookies = std::any::type_name::(); - let _cookies_data = std::any::type_name::(); - let _device_login = std::any::type_name::(); - let _device_poll = std::any::type_name::(); - let _login_start = std::any::type_name::(); - let _code_relay = std::any::type_name::(); - let _callback_query = std::any::type_name::(); - let _code_relay_data = std::any::type_name::(); - let _device_login_handler = auth_device_login; - let _device_poll_handler = auth_device_poll; - let _login_start_handler = auth_login_start; - let _code_relay_handler = auth_code_relay; - let _callback_handler = auth_callback; -} diff --git a/tests/architecture/admin_boundary.rs b/tests/architecture/admin_boundary.rs index f30f7ee74..8c5e00928 100644 --- a/tests/architecture/admin_boundary.rs +++ b/tests/architecture/admin_boundary.rs @@ -1,13 +1,12 @@ -use codex_proxy_rs::admin::{ - api::router::router, client_keys::service::ApiKeyService, session::service::AdminAuthService, - settings::SettingsService, tasks::session_cleanup::SessionCleanupScheduler, +use codex_proxy_core::admin::{ + auth::AdminAuthService, client_keys::ClientKeyService, settings::SettingsService, }; +use codex_proxy_server::admin_api::router::router; #[test] fn admin_exports_control_plane_modules() { let _admin_auth_type = std::any::type_name::(); - let _api_key_type = std::any::type_name::(); + let _api_key_type = std::any::type_name::(); let _settings_type = std::any::type_name::(); - let _session_cleanup_type = std::any::type_name::(); let _router_fn = router; } diff --git a/tests/architecture/codex_activity_boundary.rs b/tests/architecture/codex_activity_boundary.rs index 86105e590..487f7da8e 100644 --- a/tests/architecture/codex_activity_boundary.rs +++ b/tests/architecture/codex_activity_boundary.rs @@ -1,19 +1,13 @@ -use codex_proxy_rs::codex::{ - events::{event::EventLog, repository::EventLogRepository, service::LogService}, - tasks::{ - model_refresh::ModelRefresher, quota_refresh::QuotaRefresher, - token_refresh::RefreshScheduler, - }, +use codex_proxy_core::{ + events::{model::EventLog, service::EventLogService}, usage::service::UsageService, }; +use codex_proxy_runtime::tasks::model_refresh::ModelRefreshTask; #[test] fn codex_exports_activity_modules() { let _event_type = std::any::type_name::(); - let _event_repo_type = std::any::type_name::(); - let _log_service_type = std::any::type_name::(); + let _log_service_type = std::any::type_name::(); let _usage_service_type = std::any::type_name::(); - let _refresh_scheduler_type = std::any::type_name::(); - let _quota_refresher_type = std::any::type_name::(); - let _model_refresher_type = std::any::type_name::(); + let _model_refresher_type = std::any::type_name::(); } diff --git a/tests/architecture/dependency_direction.rs b/tests/architecture/dependency_direction.rs new file mode 100644 index 000000000..eef00f637 --- /dev/null +++ b/tests/architecture/dependency_direction.rs @@ -0,0 +1,93 @@ +use std::{collections::BTreeMap, fs, path::PathBuf}; + +fn repository_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) +} + +#[test] +fn workspace_dependencies_should_follow_architecture_direction() { + let manifests = workspace_dependency_map(); + + assert_no_workspace_deps(&manifests, "core"); + assert_no_workspace_deps(&manifests, "platform"); + assert_no_workspace_deps(&manifests, "assets"); + assert_forbidden_deps(&manifests, "adapters", &["runtime", "server"]); + assert_forbidden_deps(&manifests, "runtime", &["server"]); + assert_forbidden_deps(&manifests, "xtask", &["assets"]); +} + +fn workspace_dependency_map() -> BTreeMap> { + let root = repository_root(); + [ + "core", "adapters", "runtime", "server", "platform", "assets", "xtask", + ] + .into_iter() + .map(|crate_name| { + let manifest_path = root.join("crates").join(crate_name).join("Cargo.toml"); + let manifest = fs::read_to_string(&manifest_path).unwrap_or_else(|_| { + panic!("expected readable manifest at {}", manifest_path.display()) + }); + (crate_name.to_string(), workspace_deps(&manifest)) + }) + .collect() +} + +fn workspace_deps(manifest: &str) -> Vec { + let mut deps = Vec::new(); + let mut in_dependencies = false; + + for line in manifest.lines() { + let trimmed = line.trim(); + if trimmed.starts_with('[') { + in_dependencies = trimmed == "[dependencies]"; + continue; + } + if !in_dependencies { + continue; + } + for crate_name in [ + "codex-proxy-core", + "codex-proxy-adapters", + "codex-proxy-runtime", + "codex-proxy-server", + "codex-proxy-platform", + "codex-proxy-assets", + ] { + if trimmed.starts_with(crate_name) { + deps.push(crate_name.strip_prefix("codex-proxy-").unwrap().to_string()); + } + } + } + + deps +} + +fn assert_no_workspace_deps(manifests: &BTreeMap>, crate_name: &str) { + let deps = manifests + .get(crate_name) + .unwrap_or_else(|| panic!("expected manifest entry for {crate_name}")); + assert!( + deps.is_empty(), + "{crate_name} must not depend on workspace crates, found {deps:?}" + ); +} + +fn assert_forbidden_deps( + manifests: &BTreeMap>, + crate_name: &str, + forbidden: &[&str], +) { + let deps = manifests + .get(crate_name) + .unwrap_or_else(|| panic!("expected manifest entry for {crate_name}")); + let violations = deps + .iter() + .filter(|dep| forbidden.contains(&dep.as_str())) + .cloned() + .collect::>(); + + assert!( + violations.is_empty(), + "{crate_name} has forbidden workspace dependencies: {violations:?}" + ); +} diff --git a/tests/architecture/directory_shape.rs b/tests/architecture/directory_shape.rs new file mode 100644 index 000000000..4110e6c4b --- /dev/null +++ b/tests/architecture/directory_shape.rs @@ -0,0 +1,368 @@ +use std::{ + collections::BTreeSet, + fs, + path::{Path, PathBuf}, +}; + +fn repository_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) +} + +#[test] +fn rust_source_tree_should_match_architecture_whitelist() { + let root = repository_root(); + let expected = expected_rust_source_files(); + let actual = actual_rust_source_files(&root); + + assert_eq!( + actual, + expected, + "Rust source tree must match docs/architecture.md exactly.\nmissing:\n{}\nextra:\n{}", + format_missing(&expected, &actual), + format_extra(&expected, &actual) + ); +} + +#[test] +fn rust_source_directories_should_match_architecture_whitelist() { + let root = repository_root(); + let expected = expected_rust_source_directories(); + let actual = actual_rust_source_directories(&root); + + assert_eq!( + actual, + expected, + "Rust source directories must match docs/architecture.md exactly.\nmissing:\n{}\nextra:\n{}", + format_missing(&expected, &actual), + format_extra(&expected, &actual) + ); +} + +#[test] +fn workspace_members_should_match_architecture_crates() { + let cargo_toml = fs::read_to_string(repository_root().join("Cargo.toml")) + .expect("expected root Cargo.toml to be readable"); + + assert!( + cargo_toml.contains("members = [\"crates/*\"]"), + "root workspace should keep members = [\"crates/*\"]" + ); +} + +fn actual_rust_source_files(root: &Path) -> BTreeSet { + let mut files = BTreeSet::new(); + collect_rust_source_files(&root.join("src"), root, &mut files); + let crates_dir = root.join("crates"); + let entries = fs::read_dir(&crates_dir).expect("expected crates directory to be readable"); + for entry in entries.flatten() { + let path = entry.path().join("src"); + collect_rust_source_files(&path, root, &mut files); + } + files +} + +fn actual_rust_source_directories(root: &Path) -> BTreeSet { + let mut directories = BTreeSet::new(); + collect_rust_source_directories(&root.join("src"), root, &mut directories); + let crates_dir = root.join("crates"); + let entries = fs::read_dir(&crates_dir).expect("expected crates directory to be readable"); + for entry in entries.flatten() { + let path = entry.path().join("src"); + collect_rust_source_directories(&path, root, &mut directories); + } + directories +} + +fn collect_rust_source_files(dir: &Path, root: &Path, files: &mut BTreeSet) { + let Ok(entries) = fs::read_dir(dir) else { + return; + }; + + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + collect_rust_source_files(&path, root, files); + continue; + } + + let Some(extension) = path.extension().and_then(|extension| extension.to_str()) else { + continue; + }; + + if extension == "rs" || extension == "sql" { + let relative = path + .strip_prefix(root) + .expect("source file should be below repository root") + .to_string_lossy() + .replace('\\', "/"); + files.insert(relative); + } + } +} + +fn collect_rust_source_directories(dir: &Path, root: &Path, directories: &mut BTreeSet) { + if !dir.is_dir() { + return; + } + + let relative = dir + .strip_prefix(root) + .expect("source directory should be below repository root") + .to_string_lossy() + .replace('\\', "/"); + directories.insert(relative); + + let entries = fs::read_dir(dir).expect("expected source directory to be readable"); + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + collect_rust_source_directories(&path, root, directories); + } + } +} + +fn format_missing(expected: &BTreeSet, actual: &BTreeSet) -> String { + expected + .difference(actual) + .cloned() + .collect::>() + .join("\n") +} + +fn format_extra(expected: &BTreeSet, actual: &BTreeSet) -> String { + actual + .difference(expected) + .cloned() + .collect::>() + .join("\n") +} + +fn expected_rust_source_files() -> BTreeSet { + [ + "crates/core/src/lib.rs", + "crates/core/src/error.rs", + "crates/core/src/admin/mod.rs", + "crates/core/src/admin/ports.rs", + "crates/core/src/admin/auth.rs", + "crates/core/src/admin/client_keys.rs", + "crates/core/src/admin/settings.rs", + "crates/core/src/accounts/mod.rs", + "crates/core/src/accounts/model.rs", + "crates/core/src/accounts/ports.rs", + "crates/core/src/accounts/service.rs", + "crates/core/src/accounts/pool.rs", + "crates/core/src/accounts/lifecycle.rs", + "crates/core/src/accounts/cloudflare.rs", + "crates/core/src/accounts/cookies.rs", + "crates/core/src/accounts/jwt.rs", + "crates/core/src/accounts/usage.rs", + "crates/core/src/auth/mod.rs", + "crates/core/src/auth/ports.rs", + "crates/core/src/auth/oauth.rs", + "crates/core/src/auth/session.rs", + "crates/core/src/models/mod.rs", + "crates/core/src/models/model.rs", + "crates/core/src/models/ports.rs", + "crates/core/src/models/catalog.rs", + "crates/core/src/models/service.rs", + "crates/core/src/events/mod.rs", + "crates/core/src/events/model.rs", + "crates/core/src/events/ports.rs", + "crates/core/src/events/service.rs", + "crates/core/src/usage/mod.rs", + "crates/core/src/usage/model.rs", + "crates/core/src/usage/ports.rs", + "crates/core/src/usage/service.rs", + "crates/core/src/protocol/mod.rs", + "crates/core/src/protocol/openai/mod.rs", + "crates/core/src/protocol/openai/chat.rs", + "crates/core/src/protocol/openai/responses.rs", + "crates/core/src/protocol/openai/models.rs", + "crates/core/src/protocol/openai/errors.rs", + "crates/core/src/protocol/codex/mod.rs", + "crates/core/src/protocol/codex/chat.rs", + "crates/core/src/protocol/codex/responses.rs", + "crates/core/src/protocol/codex/events.rs", + "crates/core/src/protocol/codex/sse.rs", + "crates/core/src/protocol/codex/websocket.rs", + "crates/core/src/protocol/codex/schema.rs", + "crates/core/src/gateway/mod.rs", + "crates/core/src/gateway/ports.rs", + "crates/core/src/gateway/fingerprint.rs", + "crates/core/src/gateway/conversation.rs", + "crates/core/src/gateway/installation.rs", + "crates/core/src/serving/mod.rs", + "crates/core/src/serving/chat.rs", + "crates/core/src/serving/responses.rs", + "crates/core/src/serving/errors.rs", + "crates/core/src/serving/routing.rs", + "crates/core/src/serving/fallback.rs", + "crates/core/src/serving/affinity.rs", + "crates/core/src/serving/quota.rs", + "crates/core/src/serving/implicit_resume.rs", + "crates/core/src/serving/reasoning_replay.rs", + "crates/core/src/serving/stream.rs", + "crates/core/src/serving/recovery.rs", + "crates/core/src/serving/usage.rs", + "crates/adapters/src/lib.rs", + "crates/adapters/src/sqlite/mod.rs", + "crates/adapters/src/sqlite/accounts.rs", + "crates/adapters/src/sqlite/account_tokens.rs", + "crates/adapters/src/sqlite/account_usage.rs", + "crates/adapters/src/sqlite/refresh_leases.rs", + "crates/adapters/src/sqlite/cookies.rs", + "crates/adapters/src/sqlite/events.rs", + "crates/adapters/src/sqlite/models.rs", + "crates/adapters/src/sqlite/session_affinity.rs", + "crates/adapters/src/sqlite/admin_sessions.rs", + "crates/adapters/src/sqlite/client_keys.rs", + "crates/adapters/src/codex/mod.rs", + "crates/adapters/src/codex/client.rs", + "crates/adapters/src/codex/models.rs", + "crates/adapters/src/codex/fingerprint.rs", + "crates/adapters/src/codex/websocket/mod.rs", + "crates/adapters/src/codex/websocket/connect.rs", + "crates/adapters/src/codex/websocket/pool.rs", + "crates/adapters/src/codex/websocket/deflate.rs", + "crates/adapters/src/codex/websocket/opening.rs", + "crates/adapters/src/oauth/mod.rs", + "crates/adapters/src/oauth/openai.rs", + "crates/runtime/src/lib.rs", + "crates/runtime/src/bootstrap.rs", + "crates/runtime/src/state.rs", + "crates/runtime/src/services.rs", + "crates/runtime/src/repositories.rs", + "crates/runtime/src/upstream.rs", + "crates/runtime/src/config.rs", + "crates/runtime/src/tasks/mod.rs", + "crates/runtime/src/tasks/coordinator.rs", + "crates/runtime/src/tasks/token_refresh.rs", + "crates/runtime/src/tasks/quota_refresh.rs", + "crates/runtime/src/tasks/model_refresh.rs", + "crates/runtime/src/tasks/cookie_cleanup.rs", + "crates/runtime/src/tasks/session_cleanup.rs", + "crates/runtime/src/tasks/session_affinity_cleanup.rs", + "crates/runtime/src/tasks/fingerprint_update.rs", + "crates/server/src/main.rs", + "crates/server/src/lib.rs", + "crates/server/src/router.rs", + "crates/server/src/error/mod.rs", + "crates/server/src/error/admin.rs", + "crates/server/src/error/openai.rs", + "crates/server/src/middleware/mod.rs", + "crates/server/src/middleware/request_id.rs", + "crates/server/src/middleware/trace.rs", + "crates/server/src/middleware/auth.rs", + "crates/server/src/middleware/cors.rs", + "crates/server/src/admin_api/mod.rs", + "crates/server/src/admin_api/router.rs", + "crates/server/src/admin_api/response.rs", + "crates/server/src/admin_api/session.rs", + "crates/server/src/admin_api/settings.rs", + "crates/server/src/admin_api/diagnostics.rs", + "crates/server/src/admin_api/models.rs", + "crates/server/src/admin_api/usage.rs", + "crates/server/src/admin_api/accounts/mod.rs", + "crates/server/src/admin_api/accounts/list.rs", + "crates/server/src/admin_api/accounts/create.rs", + "crates/server/src/admin_api/accounts/import.rs", + "crates/server/src/admin_api/accounts/import_cli.rs", + "crates/server/src/admin_api/accounts/export.rs", + "crates/server/src/admin_api/accounts/lifecycle.rs", + "crates/server/src/admin_api/accounts/quota.rs", + "crates/server/src/admin_api/accounts/cookies.rs", + "crates/server/src/admin_api/accounts/oauth.rs", + "crates/server/src/admin_api/accounts/health.rs", + "crates/server/src/admin_api/client_keys/mod.rs", + "crates/server/src/admin_api/client_keys/list.rs", + "crates/server/src/admin_api/client_keys/create.rs", + "crates/server/src/admin_api/client_keys/import.rs", + "crates/server/src/admin_api/client_keys/export.rs", + "crates/server/src/admin_api/client_keys/lifecycle.rs", + "crates/server/src/admin_api/logs/mod.rs", + "crates/server/src/admin_api/logs/query.rs", + "crates/server/src/admin_api/logs/detail.rs", + "crates/server/src/admin_api/logs/state.rs", + "crates/server/src/openai_api/mod.rs", + "crates/server/src/openai_api/router.rs", + "crates/server/src/openai_api/auth.rs", + "crates/server/src/openai_api/chat.rs", + "crates/server/src/openai_api/responses.rs", + "crates/server/src/openai_api/models.rs", + "crates/server/src/openai_api/diagnostics.rs", + "crates/server/src/openai_api/error.rs", + "crates/server/src/openai_api/sse.rs", + "crates/platform/src/lib.rs", + "crates/platform/src/config/mod.rs", + "crates/platform/src/config/loader.rs", + "crates/platform/src/config/types.rs", + "crates/platform/src/crypto/mod.rs", + "crates/platform/src/crypto/secret_box.rs", + "crates/platform/src/crypto/hash.rs", + "crates/platform/src/identity/mod.rs", + "crates/platform/src/identity/admin_password.rs", + "crates/platform/src/identity/client_key.rs", + "crates/platform/src/storage/mod.rs", + "crates/platform/src/storage/sqlite.rs", + "crates/platform/src/storage/schema.sql", + "crates/platform/src/storage/paths.rs", + "crates/platform/src/logging/mod.rs", + "crates/platform/src/logging/rotation.rs", + "crates/platform/src/json/mod.rs", + "crates/platform/src/json/pagination.rs", + "crates/assets/src/lib.rs", + "crates/assets/src/router.rs", + "crates/assets/src/headers.rs", + "crates/xtask/src/main.rs", + "crates/xtask/src/build_web.rs", + "crates/xtask/src/check_architecture.rs", + "crates/xtask/src/release.rs", + ] + .into_iter() + .map(String::from) + .collect() +} + +fn expected_rust_source_directories() -> BTreeSet { + [ + "crates/core/src", + "crates/core/src/admin", + "crates/core/src/accounts", + "crates/core/src/auth", + "crates/core/src/models", + "crates/core/src/events", + "crates/core/src/usage", + "crates/core/src/protocol", + "crates/core/src/protocol/openai", + "crates/core/src/protocol/codex", + "crates/core/src/gateway", + "crates/core/src/serving", + "crates/adapters/src", + "crates/adapters/src/sqlite", + "crates/adapters/src/codex", + "crates/adapters/src/codex/websocket", + "crates/adapters/src/oauth", + "crates/runtime/src", + "crates/runtime/src/tasks", + "crates/server/src", + "crates/server/src/error", + "crates/server/src/middleware", + "crates/server/src/admin_api", + "crates/server/src/admin_api/accounts", + "crates/server/src/admin_api/client_keys", + "crates/server/src/admin_api/logs", + "crates/server/src/openai_api", + "crates/platform/src", + "crates/platform/src/config", + "crates/platform/src/crypto", + "crates/platform/src/identity", + "crates/platform/src/storage", + "crates/platform/src/logging", + "crates/platform/src/json", + "crates/assets/src", + "crates/xtask/src", + ] + .into_iter() + .map(String::from) + .collect() +} diff --git a/tests/architecture/forbidden_imports.rs b/tests/architecture/forbidden_imports.rs new file mode 100644 index 000000000..75e01324c --- /dev/null +++ b/tests/architecture/forbidden_imports.rs @@ -0,0 +1,123 @@ +use std::{fs, path::Path}; + +#[test] +fn source_imports_should_respect_layer_boundaries() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")); + let mut offenders = Vec::new(); + + collect_forbidden_imports(root, &mut offenders); + + assert!( + offenders.is_empty(), + "forbidden imports or source names found:\n{}", + offenders.join("\n") + ); +} + +fn collect_forbidden_imports(root: &Path, offenders: &mut Vec) { + for crate_name in ["core", "adapters", "runtime", "platform", "server"] { + let src = root.join("crates").join(crate_name).join("src"); + collect_crate_imports(crate_name, &src, offenders); + } +} + +fn collect_crate_imports(crate_name: &str, dir: &Path, offenders: &mut Vec) { + let Ok(entries) = fs::read_dir(dir) else { + return; + }; + + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + if path.file_name().and_then(|name| name.to_str()) == Some("tests") { + continue; + } + if crate_name == "core" + && path + .file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| matches!(name, "repository" | "transport")) + { + offenders.push(path.display().to_string()); + } + collect_crate_imports(crate_name, &path, offenders); + continue; + } + + let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else { + continue; + }; + if crate_name == "core" && file_name.ends_with("_repository.rs") { + offenders.push(path.display().to_string()); + } + if crate_name == "server" && file_name == "facade.rs" { + offenders.push(path.display().to_string()); + } + + let Some(extension) = path.extension().and_then(|extension| extension.to_str()) else { + continue; + }; + if extension != "rs" { + continue; + } + let Ok(content) = fs::read_to_string(&path) else { + continue; + }; + for forbidden in forbidden_patterns(crate_name) { + if content.contains(forbidden) { + offenders.push(format!("{} contains `{forbidden}`", path.display())); + } + } + } +} + +fn forbidden_patterns(crate_name: &str) -> &'static [&'static str] { + match crate_name { + "core" => &[ + "use axum::", + "axum::", + "use sqlx::", + "sqlx::", + "use reqwest::", + "reqwest::", + "use tokio_tungstenite::", + "tokio_tungstenite::", + "use tungstenite::", + "tungstenite::", + "use rustls::", + "rustls::", + "use tokio_rustls::", + "tokio_rustls::", + "std::fs", + "std::env", + "dirs::", + "std::path", + "PathBuf", + "Path::", + ], + "adapters" | "platform" => &["use axum::", "axum::"], + "runtime" => &[ + "use axum::", + "axum::", + "use sqlx::", + "sqlx::", + "use reqwest::", + "reqwest::", + "use tokio_tungstenite::", + "tokio_tungstenite::", + "use tungstenite::", + "tungstenite::", + ], + "server" => &[ + "use sqlx::", + "sqlx::", + "use reqwest::", + "reqwest::", + "use tokio_tungstenite::", + "tokio_tungstenite::", + "use tungstenite::", + "tungstenite::", + ], + _ => &[], + } +} diff --git a/tests/architecture/forbidden_legacy_paths.rs b/tests/architecture/forbidden_legacy_paths.rs new file mode 100644 index 000000000..76ece8770 --- /dev/null +++ b/tests/architecture/forbidden_legacy_paths.rs @@ -0,0 +1,54 @@ +use std::{fs, path::PathBuf}; + +fn repository_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) +} + +#[test] +fn source_tree_should_not_reference_legacy_gateway_oauth_module() { + let root = repository_root(); + let mut offenders = Vec::new(); + + for relative_dir in ["src", "tests"] { + let dir = root.join(relative_dir); + collect_legacy_oauth_references(&dir, &mut offenders); + } + + assert!( + offenders.is_empty(), + "legacy gateway::oauth references should be removed:\n{}", + offenders.join("\n") + ); +} + +fn collect_legacy_oauth_references(dir: &PathBuf, offenders: &mut Vec) { + let Ok(entries) = fs::read_dir(dir) else { + return; + }; + + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + collect_legacy_oauth_references(&path, offenders); + continue; + } + + let Some(extension) = path.extension().and_then(|ext| ext.to_str()) else { + continue; + }; + if extension != "rs" { + continue; + } + + let Ok(content) = fs::read_to_string(&path) else { + continue; + }; + if path.ends_with("tests/architecture/forbidden_legacy_paths.rs") { + continue; + } + + if content.contains("gateway::oauth") { + offenders.push(path.display().to_string()); + } + } +} diff --git a/tests/architecture/gateway_boundary.rs b/tests/architecture/gateway_boundary.rs index efb4f33f1..3fd822cf5 100644 --- a/tests/architecture/gateway_boundary.rs +++ b/tests/architecture/gateway_boundary.rs @@ -1,12 +1,37 @@ -use codex_proxy_rs::codex::gateway::{ - fingerprint::model::Fingerprint, oauth::TokenPair, protocol::schema::ResponseFormat, - transport::http_client::CodexBackendClient, +use codex_proxy_core::{ + gateway::{ + fingerprint::Fingerprint, + installation::{generate_installation_id, parse_installation_id}, + }, + protocol::{ + codex::{ + events::{ + cooldown_with_jitter, extract_usage, parse_rate_limit_headers, + retry_after_seconds_from_body, ParsedRateLimits, TokenUsage, + }, + responses::CodexResponsesRequest, + }, + openai::{ + chat::{translate_chat_to_codex, ChatCompletionRequest, ChatMessage}, + responses::ResponseFormat, + }, + }, }; #[test] -fn gateway_exports_chatgpt_integration_modules() { +fn gateway_exports_core_protocol_and_identity_modules() { let _fingerprint_type = std::any::type_name::(); - let _token_pair_type = std::any::type_name::(); let _response_format_type = std::any::type_name::(); - let _client_type = std::any::type_name::(); + let _request_type = std::any::type_name::(); + let _usage_type = std::any::type_name::(); + let _rate_limits_type = std::any::type_name::(); + let _translate_fn = translate_chat_to_codex; + let _extract_usage_fn = extract_usage; + let _parse_rate_limit_headers_fn = parse_rate_limit_headers; + let _retry_after_fn = retry_after_seconds_from_body; + let _cooldown_with_jitter_fn = cooldown_with_jitter; + let _chat_request_type = std::any::type_name::(); + let _chat_message_type = std::any::type_name::(); + let _generate_installation_id_fn = generate_installation_id; + let _parse_installation_id_fn = parse_installation_id; } diff --git a/tests/architecture/placeholder_implementations.rs b/tests/architecture/placeholder_implementations.rs new file mode 100644 index 000000000..0f0fa830c --- /dev/null +++ b/tests/architecture/placeholder_implementations.rs @@ -0,0 +1,128 @@ +use std::{fs, path::Path}; + +#[test] +fn rust_sources_should_not_keep_placeholder_markers() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")); + let mut offenders = Vec::new(); + + collect_placeholder_markers(&root.join("src"), &mut offenders); + collect_placeholder_markers(&root.join("crates"), &mut offenders); + + assert!( + offenders.is_empty(), + "placeholder implementation markers found:\n{}", + offenders.join("\n") + ); +} + +#[test] +fn serving_modules_should_not_keep_uncalled_toy_helpers() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")); + let mut offenders = Vec::new(); + + for (relative_path, marker) in [ + ( + "crates/core/src/serving/fallback.rs", + "pub fn should_fallback(", + ), + ( + "crates/core/src/serving/recovery.rs", + "pub fn is_recoverable(", + ), + ( + "crates/core/src/serving/responses.rs", + "pub fn prefers_websocket(", + ), + ] { + let path = root.join(relative_path); + let content = fs::read_to_string(&path).unwrap_or_default(); + if content.contains(marker) { + offenders.push(format!("{relative_path} keeps toy helper `{marker}`")); + } + } + + if production_call_count(root, "quota_reached(", "crates/core/src/serving/quota.rs") == 0 { + offenders.push( + "crates/core/src/serving/quota.rs exposes `quota_reached` without a production caller" + .to_string(), + ); + } + + assert!( + offenders.is_empty(), + "uncalled serving toy helpers found:\n{}", + offenders.join("\n") + ); +} + +fn collect_placeholder_markers(dir: &Path, offenders: &mut Vec) { + let Ok(entries) = fs::read_dir(dir) else { + return; + }; + + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + collect_placeholder_markers(&path, offenders); + continue; + } + + if path.extension().and_then(|extension| extension.to_str()) != Some("rs") { + continue; + } + + let Ok(content) = fs::read_to_string(&path) else { + continue; + }; + + for marker in [ + "fn main() {}", + "not wired yet", + "后续承载", + "pub struct CodexWebSocketConnection;", + "pub struct CodexWebSocketPool;", + ] { + if content.contains(marker) { + offenders.push(format!("{} contains `{marker}`", path.display())); + } + } + } +} + +fn production_call_count(root: &Path, needle: &str, defining_file: &str) -> usize { + let mut count = 0; + count_production_calls( + &root.join("crates"), + needle, + &root.join(defining_file), + &mut count, + ); + count +} + +fn count_production_calls(dir: &Path, needle: &str, defining_file: &Path, count: &mut usize) { + let Ok(entries) = fs::read_dir(dir) else { + return; + }; + + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + if path.file_name().and_then(|name| name.to_str()) == Some("tests") { + continue; + } + count_production_calls(&path, needle, defining_file, count); + continue; + } + if path == defining_file { + continue; + } + if path.extension().and_then(|extension| extension.to_str()) != Some("rs") { + continue; + } + let Ok(content) = fs::read_to_string(&path) else { + continue; + }; + *count += content.matches(needle).count(); + } +} diff --git a/tests/architecture/platform_boundary.rs b/tests/architecture/platform_boundary.rs index ce7fca743..a28c64df9 100644 --- a/tests/architecture/platform_boundary.rs +++ b/tests/architecture/platform_boundary.rs @@ -1,16 +1,19 @@ -use codex_proxy_rs::platform::{ - crypto::SecretBox, - http::auth::{admin_session_id, client_api_key}, - identity::{admin_session::hash_admin_password, client_key::ApiKeyHasher}, - storage::db::connect_sqlite, -}; +use std::{fs, path::PathBuf}; #[test] fn platform_exports_foundation_modules() { - let _secret_box_type = std::any::type_name::(); - let _hasher = ApiKeyHasher::new([7; 32]); - let _hash_fn = hash_admin_password; - let _client_key_fn = client_api_key; - let _admin_session_fn = admin_session_id; - let _connect_fn = connect_sqlite; + let platform_lib = fs::read_to_string( + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("crates/platform/src/lib.rs"), + ) + .expect("expected platform lib.rs to be readable"); + + assert!( + platform_lib.contains("pub mod config;") + && platform_lib.contains("pub mod crypto;") + && platform_lib.contains("pub mod identity;") + && platform_lib.contains("pub mod json;") + && platform_lib.contains("pub mod logging;") + && platform_lib.contains("pub mod storage;"), + "platform crate should expose foundation modules directly from its own crate" + ); } diff --git a/tests/architecture/runtime_tasks_boundary.rs b/tests/architecture/runtime_tasks_boundary.rs index d4f874fa1..530071b15 100644 --- a/tests/architecture/runtime_tasks_boundary.rs +++ b/tests/architecture/runtime_tasks_boundary.rs @@ -1,12 +1,10 @@ -use codex_proxy_rs::runtime::tasks::{ - coordinator::{start_background_tasks, BackgroundTaskCoordinator}, - types::{SchedulerError, SchedulerHandle}, +use codex_proxy_runtime::tasks::coordinator::{ + start_background_tasks, BackgroundTaskCoordinator, SchedulerHandle, }; #[test] fn runtime_exports_background_task_coordinator_modules() { let _coordinator_type = std::any::type_name::(); let _handle_type = std::any::type_name::(); - let _error_type = std::any::type_name::(); let _start_fn = start_background_tasks; } diff --git a/tests/architecture/server_startup.rs b/tests/architecture/server_startup.rs new file mode 100644 index 000000000..15e94e314 --- /dev/null +++ b/tests/architecture/server_startup.rs @@ -0,0 +1,64 @@ +use std::{fs, path::Path}; + +#[test] +fn server_entrypoint_should_own_background_task_lifecycle() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")); + let main_rs = root.join("crates/server/src/main.rs"); + let content = fs::read_to_string(&main_rs).expect("server main should be readable"); + + assert!( + content.contains("start_background_tasks(&state).await") + && content.contains("task_coordinator.shutdown().await"), + "server main must start runtime background tasks and shut them down after axum exits" + ); +} + +#[test] +fn server_entrypoint_should_inject_runtime_installation_id() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")); + let main_rs = root.join("crates/server/src/main.rs"); + let content = fs::read_to_string(&main_rs).expect("server main should be readable"); + + assert!( + content.contains("load_or_create_installation_id(Some(&data_dir))") + && content.contains("with_pool_secret_api_key_hasher_and_installation_id"), + "server main must load the platform installation id and inject it into runtime state" + ); +} + +#[test] +fn server_entrypoint_should_restore_session_affinity_before_serving() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")); + let main_rs = root.join("crates/server/src/main.rs"); + let content = fs::read_to_string(&main_rs).expect("server main should be readable"); + + assert!( + content.contains("restore_session_affinity_from_repository_now().await"), + "server main must restore session affinity mappings from SQLite during startup" + ); +} + +#[test] +fn server_entrypoint_should_restore_account_pool_before_serving() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")); + let main_rs = root.join("crates/server/src/main.rs"); + let content = fs::read_to_string(&main_rs).expect("server main should be readable"); + + assert!( + content.contains("restore_account_pool_from_repository().await"), + "server main must restore active runtime accounts from SQLite during startup" + ); +} + +#[test] +fn server_entrypoint_should_ensure_default_admin_before_serving() { + let root = Path::new(env!("CARGO_MANIFEST_DIR")); + let main_rs = root.join("crates/server/src/main.rs"); + let content = fs::read_to_string(&main_rs).expect("server main should be readable"); + + assert!( + content.contains("ensure_default_admin") + && content.contains("config.admin.default_password"), + "server main must initialize the configured default admin before serving admin routes" + ); +} diff --git a/tests/architecture/serving_boundary.rs b/tests/architecture/serving_boundary.rs deleted file mode 100644 index f0cbb94b5..000000000 --- a/tests/architecture/serving_boundary.rs +++ /dev/null @@ -1,47 +0,0 @@ -use codex_proxy_rs::codex::serving::{ - chat::ChatService, diagnostics::DiagnosticsService, dispatch::affinity::SessionAffinityMap, - http::router::router, responses::ResponsesService, -}; -use std::{fs, path::PathBuf}; - -#[test] -fn serving_exports_client_facing_proxy_modules() { - let _chat_type = std::any::type_name::(); - let _responses_type = std::any::type_name::(); - let _diagnostics_type = std::any::type_name::(); - let _affinity_type = std::any::type_name::(); - let _router_fn = router; -} - -#[test] -fn serving_dispatch_has_dedicated_recovery_transition_boundary() { - let manifest_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")); - let transition_path = manifest_dir.join("src/codex/serving/dispatch/transition.rs"); - let transition = fs::read_to_string(&transition_path).unwrap_or_else(|error| { - panic!( - "expected recovery transition boundary at {}: {error}", - transition_path.display() - ) - }); - let dispatch_mod = - fs::read_to_string(manifest_dir.join("src/codex/serving/dispatch/mod.rs")).unwrap(); - let fallback = - fs::read_to_string(manifest_dir.join("src/codex/serving/dispatch/fallback.rs")).unwrap(); - - assert!( - dispatch_mod.contains("mod transition;"), - "dispatch should declare a transition module" - ); - assert!( - transition.contains("execute_upstream_account_recovery_transition_with_deps"), - "account recovery side effects should live in the transition boundary" - ); - assert!( - transition.contains("execute_upstream_request_recovery_transition_with_deps"), - "request recovery side effects should live in the transition boundary" - ); - assert!( - !fallback.contains("acquire_with("), - "fallback classification should not acquire fallback accounts directly" - ); -} diff --git a/tests/architecture/test_migration.rs b/tests/architecture/test_migration.rs new file mode 100644 index 000000000..578fcb664 --- /dev/null +++ b/tests/architecture/test_migration.rs @@ -0,0 +1,449 @@ +use std::{ + collections::BTreeSet, + fs, + path::{Path, PathBuf}, +}; + +fn repository_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) +} + +#[test] +fn root_tests_should_only_keep_architecture_tests_and_shared_fixtures() { + let root = repository_root(); + let mut offenders = Vec::new(); + collect_root_test_rs_files(&root.join("tests"), &root, &mut offenders); + + assert!( + offenders.is_empty(), + "root behavior tests must be migrated into crates/*/tests:\n{}", + offenders.join("\n") + ); +} + +#[test] +fn baseline_behavior_test_suites_should_have_crate_local_migration_coverage() { + let root = repository_root(); + let current_tests = current_crate_test_files(&root); + let mut missing = Vec::new(); + + for suite in baseline_behavior_suites() { + for required in suite.required_current_files { + if !current_tests.contains(*required) { + missing.push(format!( + "{} should be covered by crate-local test file {required}", + suite.baseline_path + )); + } + } + } + + assert!( + missing.is_empty(), + "old behavior test suites lack crate-local migration evidence:\n{}", + missing.join("\n") + ); +} + +#[test] +fn core_service_modules_should_have_production_callers() { + let root = repository_root(); + let mut offenders = Vec::new(); + + for service in [ + CoreService { + path: "crates/core/src/accounts/service.rs", + symbol: "AccountService::", + }, + CoreService { + path: "crates/core/src/usage/service.rs", + symbol: "UsageService::", + }, + CoreService { + path: "crates/core/src/admin/settings.rs", + symbol: "SettingsService::", + }, + ] { + if production_call_count(&root, service.symbol, service.path) == 0 { + offenders.push(format!( + "{} exposes {} without a production caller", + service.path, service.symbol + )); + } + } + + assert!( + offenders.is_empty(), + "core service modules still look like architecture placeholders:\n{}", + offenders.join("\n") + ); +} + +fn collect_root_test_rs_files(dir: &Path, root: &Path, offenders: &mut Vec) { + let Ok(entries) = fs::read_dir(dir) else { + return; + }; + + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + let relative = relative_path(root, &path); + if relative == "tests/architecture" || relative.starts_with("tests/fixtures") { + continue; + } + collect_root_test_rs_files(&path, root, offenders); + continue; + } + + if path.extension().and_then(|extension| extension.to_str()) != Some("rs") { + continue; + } + + let relative = relative_path(root, &path); + if relative != "tests/architecture.rs" && !relative.starts_with("tests/architecture/") { + offenders.push(relative); + } + } +} + +fn current_crate_test_files(root: &Path) -> BTreeSet { + let mut files = BTreeSet::new(); + collect_crate_test_files(&root.join("crates"), root, &mut files); + files +} + +fn collect_crate_test_files(dir: &Path, root: &Path, files: &mut BTreeSet) { + let Ok(entries) = fs::read_dir(dir) else { + return; + }; + + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + collect_crate_test_files(&path, root, files); + continue; + } + if path.extension().and_then(|extension| extension.to_str()) == Some("rs") + && relative_path(root, &path).contains("/tests/") + { + files.insert(relative_path(root, &path)); + } + } +} + +fn production_call_count(root: &Path, needle: &str, defining_file: &str) -> usize { + let mut count = 0; + count_production_calls( + &root.join("crates"), + needle, + &root.join(defining_file), + &mut count, + ); + count +} + +fn count_production_calls(dir: &Path, needle: &str, defining_file: &Path, count: &mut usize) { + let Ok(entries) = fs::read_dir(dir) else { + return; + }; + + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + if path.file_name().and_then(|name| name.to_str()) == Some("tests") { + continue; + } + count_production_calls(&path, needle, defining_file, count); + continue; + } + if path == defining_file { + continue; + } + if path.extension().and_then(|extension| extension.to_str()) != Some("rs") { + continue; + } + let Ok(content) = fs::read_to_string(&path) else { + continue; + }; + *count += content.matches(needle).count(); + } +} + +fn relative_path(root: &Path, path: &Path) -> String { + path.strip_prefix(root) + .expect("path should be below repository root") + .to_string_lossy() + .replace('\\', "/") +} + +struct BaselineSuite { + baseline_path: &'static str, + required_current_files: &'static [&'static str], +} + +struct CoreService { + path: &'static str, + symbol: &'static str, +} + +fn baseline_behavior_suites() -> Vec { + vec![ + BaselineSuite { + baseline_path: "tests/admin/accounts/cookies_quota.rs", + required_current_files: &[ + "crates/server/tests/admin_accounts_routes.rs", + "crates/adapters/tests/cookies.rs", + "crates/runtime/tests/quota_refresh.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/admin/accounts/import_export.rs", + required_current_files: &["crates/server/tests/admin_accounts_routes.rs"], + }, + BaselineSuite { + baseline_path: "tests/admin/accounts/lifecycle.rs", + required_current_files: &["crates/server/tests/admin_accounts_routes.rs"], + }, + BaselineSuite { + baseline_path: "tests/admin/accounts/list.rs", + required_current_files: &["crates/server/tests/admin_accounts_routes.rs"], + }, + BaselineSuite { + baseline_path: "tests/admin/accounts/oauth.rs", + required_current_files: &[ + "crates/server/tests/admin_accounts_routes.rs", + "crates/adapters/tests/oauth.rs", + "crates/core/tests/auth.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/admin/api_contract.rs", + required_current_files: &["crates/server/tests/admin_api_contract.rs"], + }, + BaselineSuite { + baseline_path: "tests/admin/client_keys_route.rs", + required_current_files: &[ + "crates/server/tests/admin_client_keys_routes.rs", + "crates/adapters/tests/client_keys.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/admin/logs_route.rs", + required_current_files: &[ + "crates/server/tests/admin_logs_routes.rs", + "crates/adapters/tests/events.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/admin/models_route.rs", + required_current_files: &[ + "crates/server/tests/admin_models_routes.rs", + "crates/core/tests/models.rs", + "crates/adapters/tests/models.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/admin/session.rs", + required_current_files: &["crates/platform/tests/admin_password.rs"], + }, + BaselineSuite { + baseline_path: "tests/admin/session_login_route.rs", + required_current_files: &["crates/server/tests/admin_session_routes.rs"], + }, + BaselineSuite { + baseline_path: "tests/admin/session_repository.rs", + required_current_files: &["crates/adapters/tests/admin_sessions.rs"], + }, + BaselineSuite { + baseline_path: "tests/admin/settings_route.rs", + required_current_files: &[ + "crates/server/tests/admin_settings_routes.rs", + "crates/core/tests/admin_settings.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/admin/usage_stats_route.rs", + required_current_files: &["crates/server/tests/admin_accounts_routes.rs"], + }, + BaselineSuite { + baseline_path: "tests/codex_accounts/cookie_store.rs", + required_current_files: &[ + "crates/adapters/tests/cookies.rs", + "crates/runtime/tests/tasks.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_accounts/pool_scheduling.rs", + required_current_files: &["crates/core/tests/account_pool.rs"], + }, + BaselineSuite { + baseline_path: "tests/codex_accounts/refresh.rs", + required_current_files: &[ + "crates/core/tests/protocol.rs", + "crates/runtime/tests/token_refresh.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_accounts/refresh_scheduler.rs", + required_current_files: &["crates/core/tests/account_pool.rs"], + }, + BaselineSuite { + baseline_path: "tests/codex_accounts/repository.rs", + required_current_files: &[ + "crates/adapters/tests/account_repository.rs", + "crates/adapters/tests/refresh_leases.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_accounts/service_refresh.rs", + required_current_files: &["crates/runtime/tests/token_refresh.rs"], + }, + BaselineSuite { + baseline_path: "tests/codex_events/logs_pagination.rs", + required_current_files: &[ + "crates/adapters/tests/events.rs", + "crates/server/tests/admin_logs_routes.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_gateway/cli_auth_import.rs", + required_current_files: &["crates/server/tests/admin_accounts_routes.rs"], + }, + BaselineSuite { + baseline_path: "tests/codex_gateway/fingerprint_update.rs", + required_current_files: &[ + "crates/adapters/tests/codex.rs", + "crates/runtime/tests/fingerprint.rs", + "crates/runtime/tests/tasks.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_gateway/headers.rs", + required_current_files: &[ + "crates/adapters/tests/codex.rs", + "crates/core/tests/protocol.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_gateway/http_client.rs", + required_current_files: &["crates/adapters/tests/codex.rs"], + }, + BaselineSuite { + baseline_path: "tests/codex_gateway/oauth_refresh.rs", + required_current_files: &["crates/adapters/tests/oauth.rs"], + }, + BaselineSuite { + baseline_path: "tests/codex_gateway/usage_events.rs", + required_current_files: &["crates/core/tests/protocol.rs"], + }, + BaselineSuite { + baseline_path: "tests/codex_gateway/websocket.rs", + required_current_files: &[ + "crates/core/tests/protocol.rs", + "crates/adapters/tests/codex.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_gateway/websocket/pool.rs", + required_current_files: &[ + "crates/adapters/tests/codex.rs", + "crates/runtime/tests/upstream.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_models/catalog.rs", + required_current_files: &[ + "crates/core/tests/models.rs", + "crates/adapters/tests/models.rs", + "crates/server/tests/admin_models_routes.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_serving/chat_completions.rs", + required_current_files: &[ + "crates/server/tests/openai_chat_upstream.rs", + "crates/core/tests/protocol.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_serving/diagnostics_route.rs", + required_current_files: &[ + "crates/server/tests/openai_diagnostics_routes.rs", + "crates/server/tests/admin_settings_routes.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_serving/responses_http_sse.rs", + required_current_files: &["crates/server/tests/openai_chat_upstream.rs"], + }, + BaselineSuite { + baseline_path: "tests/codex_serving/responses_websocket.rs", + required_current_files: &[ + "crates/server/tests/openai_chat_upstream.rs", + "crates/adapters/tests/codex.rs", + "crates/runtime/tests/upstream.rs", + "crates/runtime/tests/session_affinity.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_serving/routes_chat.rs", + required_current_files: &["crates/core/tests/protocol.rs"], + }, + BaselineSuite { + baseline_path: "tests/codex_serving/routes_responses.rs", + required_current_files: &[ + "crates/server/tests/openai_responses_routes.rs", + "crates/server/tests/openai_models_auth.rs", + "crates/core/tests/models.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/codex_serving/upstream_errors.rs", + required_current_files: &["crates/server/tests/openai_chat_upstream.rs"], + }, + BaselineSuite { + baseline_path: "tests/codex_serving/upstream_fallback.rs", + required_current_files: &["crates/server/tests/openai_chat_upstream.rs"], + }, + BaselineSuite { + baseline_path: "tests/config.rs", + required_current_files: &["crates/platform/tests/config.rs"], + }, + BaselineSuite { + baseline_path: "tests/platform/client_key_auth.rs", + required_current_files: &["crates/platform/tests/client_key_auth.rs"], + }, + BaselineSuite { + baseline_path: "tests/platform/crypto.rs", + required_current_files: &["crates/platform/tests/crypto.rs"], + }, + BaselineSuite { + baseline_path: "tests/platform/http_auth.rs", + required_current_files: &[ + "crates/platform/tests/client_key_auth.rs", + "crates/server/tests/openai_models_auth.rs", + "crates/server/tests/admin_session_routes.rs", + ], + }, + BaselineSuite { + baseline_path: "tests/platform/log_rotation.rs", + required_current_files: &["crates/platform/tests/log_rotation.rs"], + }, + BaselineSuite { + baseline_path: "tests/platform/storage_schema.rs", + required_current_files: &["crates/platform/tests/storage_schema.rs"], + }, + BaselineSuite { + baseline_path: "tests/runtime/http_trace.rs", + required_current_files: &["crates/server/tests/http_trace_middleware.rs"], + }, + BaselineSuite { + baseline_path: "tests/runtime/startup.rs", + required_current_files: &[ + "crates/runtime/tests/account_pool_restore.rs", + "crates/runtime/tests/session_affinity.rs", + "crates/server/tests/openai_chat_upstream.rs", + ], + }, + ] +} diff --git a/tests/architecture/workspace_scaffold.rs b/tests/architecture/workspace_scaffold.rs new file mode 100644 index 000000000..716ef90ed --- /dev/null +++ b/tests/architecture/workspace_scaffold.rs @@ -0,0 +1,198 @@ +use std::{fs, path::PathBuf}; + +fn repository_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) +} + +#[test] +fn workspace_should_declare_crates_glob_member() { + let cargo_toml = fs::read_to_string(repository_root().join("Cargo.toml")) + .expect("expected root Cargo.toml to be readable"); + + assert!( + cargo_toml.contains("[workspace]") && cargo_toml.contains("members = [\"crates/*\"]"), + "root Cargo.toml should declare a workspace with members = [\"crates/*\"]" + ); +} + +#[test] +fn workspace_should_scaffold_required_member_crates() { + let root = repository_root(); + + for crate_name in [ + "core", "adapters", "runtime", "server", "platform", "assets", "xtask", + ] { + let crate_dir = root.join("crates").join(crate_name); + assert!( + crate_dir.join("Cargo.toml").is_file(), + "expected {crate_name} crate to define Cargo.toml at {}", + crate_dir.join("Cargo.toml").display() + ); + + let entry_paths = if crate_name == "server" { + vec![crate_dir.join("src/lib.rs"), crate_dir.join("src/main.rs")] + } else if crate_name == "xtask" { + vec![crate_dir.join("src/main.rs")] + } else { + vec![crate_dir.join("src/lib.rs")] + }; + + for entry_path in entry_paths { + assert!( + entry_path.is_file(), + "expected {crate_name} crate to define source entry at {}", + entry_path.display() + ); + } + } +} + +#[test] +fn workspace_should_scaffold_runtime_server_and_assets_shapes() { + let root = repository_root(); + + for relative_path in [ + "crates/core/src/admin/mod.rs", + "crates/core/src/admin/ports.rs", + "crates/core/src/admin/client_keys.rs", + "crates/core/src/accounts/model.rs", + "crates/core/src/auth/mod.rs", + "crates/core/src/auth/oauth.rs", + "crates/core/src/gateway/ports.rs", + "crates/core/src/models/mod.rs", + "crates/core/src/models/model.rs", + "crates/core/src/models/catalog.rs", + "crates/core/src/models/ports.rs", + "crates/core/src/models/service.rs", + "crates/core/src/protocol/codex/events.rs", + "crates/adapters/src/codex/mod.rs", + "crates/adapters/src/codex/client.rs", + "crates/adapters/src/codex/models.rs", + "crates/adapters/src/codex/fingerprint.rs", + "crates/adapters/src/codex/websocket/mod.rs", + "crates/adapters/src/codex/websocket/connect.rs", + "crates/adapters/src/codex/websocket/pool.rs", + "crates/adapters/src/codex/websocket/deflate.rs", + "crates/adapters/src/codex/websocket/opening.rs", + "crates/adapters/src/sqlite/mod.rs", + "crates/adapters/src/sqlite/client_keys.rs", + "crates/adapters/src/sqlite/models.rs", + "crates/assets/src/router.rs", + "crates/assets/src/headers.rs", + "crates/runtime/src/bootstrap.rs", + "crates/runtime/src/state.rs", + "crates/runtime/src/services.rs", + "crates/runtime/src/repositories.rs", + "crates/runtime/src/upstream.rs", + "crates/runtime/src/config.rs", + "crates/runtime/src/tasks/mod.rs", + "crates/runtime/src/tasks/coordinator.rs", + "crates/runtime/src/tasks/token_refresh.rs", + "crates/runtime/src/tasks/quota_refresh.rs", + "crates/runtime/src/tasks/model_refresh.rs", + "crates/runtime/src/tasks/cookie_cleanup.rs", + "crates/runtime/src/tasks/session_cleanup.rs", + "crates/runtime/src/tasks/session_affinity_cleanup.rs", + "crates/runtime/src/tasks/fingerprint_update.rs", + "crates/server/src/router.rs", + "crates/server/src/error/mod.rs", + "crates/server/src/error/admin.rs", + "crates/server/src/error/openai.rs", + "crates/server/src/middleware/mod.rs", + "crates/server/src/middleware/request_id.rs", + "crates/server/src/middleware/trace.rs", + "crates/server/src/middleware/auth.rs", + "crates/server/src/middleware/cors.rs", + "crates/server/src/admin_api/mod.rs", + "crates/server/src/admin_api/router.rs", + "crates/server/src/openai_api/mod.rs", + "crates/server/src/openai_api/auth.rs", + "crates/server/src/openai_api/chat.rs", + "crates/server/src/openai_api/diagnostics.rs", + "crates/server/src/openai_api/error.rs", + "crates/server/src/openai_api/models.rs", + "crates/server/src/openai_api/responses.rs", + "crates/server/src/openai_api/router.rs", + "crates/server/src/openai_api/sse.rs", + "crates/xtask/src/build_web.rs", + "crates/xtask/src/check_architecture.rs", + "crates/xtask/src/release.rs", + ] { + assert!( + root.join(relative_path).is_file(), + "expected architecture scaffold file at {}", + root.join(relative_path).display() + ); + } +} + +#[test] +fn migrated_modules_should_not_keep_legacy_duplicates() { + let root = repository_root(); + + for relative_path in [ + "src/main.rs", + "src/runtime/bootstrap.rs", + "src/runtime/mod.rs", + "src/runtime/router.rs", + "src/runtime/state.rs", + "src/runtime/tasks/coordinator.rs", + "src/runtime/tasks/mod.rs", + "src/runtime/tasks/types.rs", + "src/codex/accounts/jwt.rs", + "src/codex/gateway/conversation_identity.rs", + "src/codex/gateway/installation_id.rs", + "src/codex/gateway/fingerprint/mod.rs", + "src/codex/gateway/fingerprint/model.rs", + "src/codex/gateway/protocol/openai_to_codex.rs", + "src/codex/gateway/protocol/schema.rs", + "src/codex/gateway/protocol/tuple_schema.rs", + "src/codex/gateway/transport/types.rs", + "src/codex/gateway/transport/endpoints.rs", + "src/codex/gateway/transport/custom_ca.rs", + "src/codex/gateway/fingerprint/repository.rs", + "src/codex/gateway/fingerprint/updater.rs", + "src/codex/gateway/fingerprint/update_checker.rs", + "src/codex/gateway/transport/sse.rs", + "src/codex/gateway/protocol/codex_to_openai.rs", + "src/codex/gateway/transport/usage_events.rs", + "src/codex/gateway/transport/retry_after.rs", + "src/codex/gateway/transport/rate_limits.rs", + "src/codex/gateway/transport/headers.rs", + "src/codex/gateway/transport/http_client.rs", + "src/codex/gateway/oauth/mod.rs", + "src/codex/gateway/oauth/codex_cli.rs", + "src/codex/gateway/oauth/client.rs", + "src/codex/gateway/oauth/token.rs", + "src/codex/gateway/oauth/refresh.rs", + "src/codex/models/catalog.rs", + "src/codex/models/repository.rs", + "src/codex/models/service.rs", + "src/codex/serving/http/auth.rs", + "src/codex/serving/http/chat.rs", + "src/codex/serving/http/diagnostics.rs", + "src/codex/serving/http/errors.rs", + "src/codex/serving/http/mod.rs", + "src/codex/serving/http/models.rs", + "src/codex/serving/http/responses.rs", + "src/codex/serving/http/router.rs", + "tests/codex_gateway/http_client.rs", + "tests/codex_models/catalog.rs", + ] { + assert!( + !root.join(relative_path).exists(), + "migrated legacy module should be removed: {}", + root.join(relative_path).display() + ); + } +} + +#[test] +fn root_src_should_be_removed_after_workspace_migration() { + let root = repository_root(); + + assert!( + !root.join("src").exists(), + "root src should be removed after member crates expose architecture boundaries directly" + ); +} diff --git a/tests/codex_accounts.rs b/tests/codex_accounts.rs deleted file mode 100644 index e93833c16..000000000 --- a/tests/codex_accounts.rs +++ /dev/null @@ -1,10 +0,0 @@ -mod support; - -mod codex_accounts { - mod cookie_store; - mod pool_scheduling; - mod refresh; - mod refresh_scheduler; - mod repository; - mod service_refresh; -} diff --git a/tests/codex_accounts/cookie_store.rs b/tests/codex_accounts/cookie_store.rs deleted file mode 100644 index e7584ba2b..000000000 --- a/tests/codex_accounts/cookie_store.rs +++ /dev/null @@ -1,265 +0,0 @@ -use chrono::Utc; - -use codex_proxy_rs::{ - codex::accounts::cookies::{jar::CookieJar, repository::CookieRepository}, - codex::tasks::cookie_cleanup::CookieCleanupScheduler, - platform::crypto::SecretBox, - platform::storage::db::connect_sqlite, -}; - -#[test] -fn cookie_jar_captures_and_replays_account_scoped_cookies() { - let mut jar = CookieJar::default(); - jar.capture_set_cookie( - "acct_a", - "cf_clearance=abc; Domain=chatgpt.com; Path=/; HttpOnly", - ); - jar.capture_set_cookie( - "acct_b", - "cf_clearance=def; Domain=chatgpt.com; Path=/; HttpOnly", - ); - - assert_eq!( - jar.cookie_header("acct_a", "chatgpt.com"), - Some("cf_clearance=abc".to_string()) - ); - assert_eq!( - jar.cookie_header("acct_b", "chatgpt.com"), - Some("cf_clearance=def".to_string()) - ); -} - -#[tokio::test] -async fn cookie_repository_encrypts_and_replays_account_scoped_cookies() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("cookies.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_account(&pool, "acct_a").await; - seed_account(&pool, "acct_b").await; - let repo = CookieRepository::new(pool.clone(), SecretBox::new([9u8; 32])); - - repo.capture_set_cookie( - "acct_a", - "cf_clearance=secret-cookie; Domain=.chatgpt.com; Path=/", - ) - .await - .unwrap(); - - let stored: (String,) = sqlx::query_as( - "select value_cipher from account_cookies where account_id = ? and name = ?", - ) - .bind("acct_a") - .bind("cf_clearance") - .fetch_one(&pool) - .await - .unwrap(); - assert!(stored.0.starts_with("v1:")); - assert!(!stored.0.contains("secret-cookie")); - assert_eq!( - repo.cookie_header("acct_a", "chatgpt.com").await.unwrap(), - Some("cf_clearance=secret-cookie".to_string()) - ); - assert_eq!( - repo.cookie_header("acct_b", "chatgpt.com").await.unwrap(), - None - ); -} - -#[tokio::test] -async fn cookie_repository_should_scope_cookie_replay_by_domain_and_path() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("path-scoped-cookies.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_account(&pool, "acct_a").await; - let repo = CookieRepository::new(pool, SecretBox::new([11u8; 32])); - - repo.capture_set_cookie("acct_a", "cf_clearance=root; Domain=.chatgpt.com; Path=/") - .await - .unwrap(); - repo.capture_set_cookie("acct_a", "cf_clearance=api; Domain=.chatgpt.com; Path=/api") - .await - .unwrap(); - - assert_eq!( - repo.cookie_header_for_request("acct_a", "chatgpt.com", "/codex/responses") - .await - .unwrap(), - Some("cf_clearance=root".to_string()) - ); - assert_eq!( - repo.cookie_header_for_request("acct_a", "chatgpt.com", "/api/codex/usage") - .await - .unwrap(), - Some("cf_clearance=api; cf_clearance=root".to_string()) - ); -} - -#[tokio::test] -async fn cookie_repository_should_only_auto_capture_cf_clearance() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("capturable-cookies.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_account(&pool, "acct_a").await; - let repo = CookieRepository::new(pool.clone(), SecretBox::new([13u8; 32])); - - repo.capture_set_cookie("acct_a", "__cf_bm=bot-session; Domain=.chatgpt.com; Path=/") - .await - .unwrap(); - - let stored_count: (i64,) = - sqlx::query_as("select count(*) from account_cookies where account_id = ?") - .bind("acct_a") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(stored_count.0, 0); - assert_eq!( - repo.cookie_header("acct_a", "chatgpt.com").await.unwrap(), - None - ); -} - -#[tokio::test] -async fn cookie_repository_manual_cookie_header_should_allow_non_capturable_cookies() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("manual-cookies.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_account(&pool, "acct_a").await; - let repo = CookieRepository::new(pool, SecretBox::new([17u8; 32])); - - let stored = repo - .set_cookie_header("acct_a", "cf_clearance=clear; __cf_bm=bm") - .await - .unwrap(); - - assert_eq!(stored, 2); - assert_eq!( - repo.cookie_header("acct_a", "chatgpt.com").await.unwrap(), - Some("__cf_bm=bm; cf_clearance=clear".to_string()) - ); -} - -#[tokio::test] -async fn cookie_repository_should_not_replay_expired_cookies() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("expired-cookies.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_account(&pool, "acct_a").await; - let repo = CookieRepository::new(pool, SecretBox::new([19u8; 32])); - - repo.capture_set_cookie( - "acct_a", - "cf_clearance=expired; Domain=.chatgpt.com; Path=/; Expires=Wed, 21 Oct 2015 07:28:00 GMT", - ) - .await - .unwrap(); - - assert_eq!( - repo.cookie_header("acct_a", "chatgpt.com").await.unwrap(), - None - ); -} - -#[tokio::test] -async fn cookie_repository_cleanup_expired_cookies_should_delete_only_expired_rows() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("cleanup-expired-cookies.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_account(&pool, "acct_a").await; - let repo = CookieRepository::new(pool.clone(), SecretBox::new([29u8; 32])); - - repo.capture_set_cookie( - "acct_a", - "cf_clearance=expired; Domain=.chatgpt.com; Path=/; Expires=Wed, 21 Oct 2015 07:28:00 GMT", - ) - .await - .unwrap(); - repo.set_cookie_header("acct_a", "__cf_bm=active") - .await - .unwrap(); - - let deleted = repo.cleanup_expired(Utc::now()).await.unwrap(); - let stored_count: (i64,) = - sqlx::query_as("select count(*) from account_cookies where account_id = ?") - .bind("acct_a") - .fetch_one(&pool) - .await - .unwrap(); - - assert_eq!(deleted, 1); - assert_eq!(stored_count.0, 1); - assert_eq!( - repo.cookie_header("acct_a", "chatgpt.com").await.unwrap(), - Some("__cf_bm=active".to_string()) - ); -} - -#[tokio::test] -async fn cookie_cleanup_scheduler_should_delete_expired_cookie_rows() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("scheduled-cookie-cleanup.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_account(&pool, "acct_a").await; - let repo = CookieRepository::new(pool.clone(), SecretBox::new([31u8; 32])); - - repo.capture_set_cookie( - "acct_a", - "cf_clearance=expired; Domain=.chatgpt.com; Path=/; Max-Age=0", - ) - .await - .unwrap(); - - let scheduler = CookieCleanupScheduler::new(repo); - let deleted = scheduler.cleanup_once().await.unwrap(); - let stored_count: (i64,) = sqlx::query_as("select count(*) from account_cookies") - .fetch_one(&pool) - .await - .unwrap(); - - assert_eq!(deleted, 1); - assert_eq!(stored_count.0, 0); -} - -#[tokio::test] -async fn cookie_repository_should_treat_max_age_as_higher_priority_than_expires() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("max-age-cookies.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_account(&pool, "acct_a").await; - let repo = CookieRepository::new(pool, SecretBox::new([23u8; 32])); - - repo.capture_set_cookie( - "acct_a", - "cf_clearance=expired; Domain=.chatgpt.com; Path=/; Max-Age=0; Expires=Wed, 21 Oct 2999 07:28:00 GMT", - ) - .await - .unwrap(); - - assert_eq!( - repo.cookie_header("acct_a", "chatgpt.com").await.unwrap(), - None - ); -} - -async fn seed_account(pool: &sqlx::SqlitePool, id: &str) { - let now = Utc::now().to_rfc3339(); - sqlx::query( - "insert into accounts (id, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?)", - ) - .bind(id) - .bind("v1:test:test") - .bind("active") - .bind(&now) - .bind(now) - .execute(pool) - .await - .unwrap(); -} diff --git a/tests/codex_accounts/refresh.rs b/tests/codex_accounts/refresh.rs deleted file mode 100644 index 5f76bbaea..000000000 --- a/tests/codex_accounts/refresh.rs +++ /dev/null @@ -1,160 +0,0 @@ -use std::{ - sync::{ - atomic::{AtomicUsize, Ordering}, - Arc, - }, - time::Duration, -}; - -use async_trait::async_trait; -use chrono::{Duration as ChronoDuration, Utc}; -use tokio::time::sleep; - -use codex_proxy_rs::{ - codex::accounts::model::{Account, AccountStatus}, - codex::gateway::oauth::{ - RefreshFailure, RefreshPolicy, RefreshScheduler, RefreshTrigger, TokenPair, TokenRefresher, - }, -}; - -#[tokio::test] -async fn refresh_scheduler_refreshes_before_expiry_and_preserves_refresh_token() { - let now = Utc::now(); - let mut account = Account::test("acct_1", AccountStatus::Active); - account.access_token_expires_at = Some(now + ChronoDuration::seconds(60)); - account.refresh_token = Some("rt_keep".to_string()); - let scheduler = RefreshScheduler::new( - RefreshPolicy { - refresh_margin_seconds: 300, - refresh_concurrency: 2, - }, - StaticRefreshClient { - result: Ok(TokenPair { - access_token: "new-access".to_string(), - refresh_token: None, - }), - }, - ); - - let refreshed = scheduler - .refresh_account_at(&account, RefreshTrigger::BeforeExpiry, now) - .await - .unwrap(); - - assert_eq!(refreshed.access_token, "new-access"); - assert_eq!(refreshed.refresh_token.as_deref(), Some("rt_keep")); - assert_eq!(refreshed.status, AccountStatus::Active); -} - -#[tokio::test] -async fn refresh_scheduler_refreshes_immediately_after_unauthorized() { - let now = Utc::now(); - let mut account = Account::test("acct_1", AccountStatus::Active); - account.access_token_expires_at = Some(now + ChronoDuration::hours(6)); - let scheduler = RefreshScheduler::new( - RefreshPolicy { - refresh_margin_seconds: 300, - refresh_concurrency: 2, - }, - StaticRefreshClient { - result: Ok(TokenPair { - access_token: "new-access-after-401".to_string(), - refresh_token: Some("rt_rotated".to_string()), - }), - }, - ); - - let refreshed = scheduler - .refresh_account_at(&account, RefreshTrigger::Unauthorized, now) - .await - .unwrap(); - - assert_eq!(refreshed.access_token, "new-access-after-401"); - assert_eq!(refreshed.refresh_token.as_deref(), Some("rt_rotated")); -} - -#[tokio::test] -async fn refresh_scheduler_maps_refresh_failures_to_account_status() { - let now = Utc::now(); - let account = Account::test("acct_1", AccountStatus::Active); - let scheduler = RefreshScheduler::new( - RefreshPolicy { - refresh_margin_seconds: 300, - refresh_concurrency: 2, - }, - StaticRefreshClient { - result: Err(RefreshFailure::Banned), - }, - ); - - let refreshed = scheduler - .refresh_account_at(&account, RefreshTrigger::Unauthorized, now) - .await - .unwrap(); - - assert_eq!(refreshed.status, AccountStatus::Banned); -} - -#[tokio::test] -async fn refresh_scheduler_limits_refresh_concurrency() { - let now = Utc::now(); - let client = CountingRefreshClient::default(); - let scheduler = RefreshScheduler::new( - RefreshPolicy { - refresh_margin_seconds: 300, - refresh_concurrency: 2, - }, - client.clone(), - ); - - let mut handles = Vec::new(); - for index in 0..5 { - let mut account = Account::test(&format!("acct_{index}"), AccountStatus::Active); - account.access_token_expires_at = Some(now + ChronoDuration::seconds(1)); - let scheduler = scheduler.clone(); - handles.push(tokio::spawn(async move { - scheduler - .refresh_account_at(&account, RefreshTrigger::BeforeExpiry, now) - .await - .unwrap() - })); - } - - for handle in handles { - let refreshed = handle.await.unwrap(); - assert_eq!(refreshed.status, AccountStatus::Active); - } - assert_eq!(client.max_seen.load(Ordering::SeqCst), 2); -} - -#[derive(Clone)] -struct StaticRefreshClient { - result: Result, -} - -#[async_trait] -impl TokenRefresher for StaticRefreshClient { - async fn refresh(&self, _refresh_token: &str) -> Result { - self.result.clone() - } -} - -#[derive(Clone, Default)] -struct CountingRefreshClient { - active: Arc, - max_seen: Arc, -} - -#[async_trait] -impl TokenRefresher for CountingRefreshClient { - async fn refresh(&self, refresh_token: &str) -> Result { - let current = self.active.fetch_add(1, Ordering::SeqCst) + 1; - self.max_seen.fetch_max(current, Ordering::SeqCst); - sleep(Duration::from_millis(25)).await; - self.active.fetch_sub(1, Ordering::SeqCst); - Ok(TokenPair { - access_token: format!("access-{refresh_token}"), - refresh_token: None, - }) - } -} diff --git a/tests/codex_accounts/refresh_scheduler.rs b/tests/codex_accounts/refresh_scheduler.rs deleted file mode 100644 index a7a9c762d..000000000 --- a/tests/codex_accounts/refresh_scheduler.rs +++ /dev/null @@ -1,17 +0,0 @@ -use codex_proxy_rs::codex::accounts::{ - model::{Account, AccountStatus}, - pool::AccountPool, -}; - -#[test] -fn account_pool_skips_expired_disabled_banned_and_quota_exhausted_accounts() { - let mut pool = AccountPool::default(); - pool.insert(Account::test("active", AccountStatus::Active)); - pool.insert(Account::test("expired", AccountStatus::Expired)); - pool.insert(Account::test("disabled", AccountStatus::Disabled)); - pool.insert(Account::test("banned", AccountStatus::Banned)); - pool.insert(Account::test("quota", AccountStatus::QuotaExhausted)); - - let acquired = pool.acquire("gpt-5.5").unwrap(); - assert_eq!(acquired.id, "active"); -} diff --git a/tests/codex_accounts/repository.rs b/tests/codex_accounts/repository.rs deleted file mode 100644 index cfd61580c..000000000 --- a/tests/codex_accounts/repository.rs +++ /dev/null @@ -1,609 +0,0 @@ -use chrono::{Duration, Utc}; -use secrecy::{ExposeSecret, SecretString}; - -use codex_proxy_rs::{ - codex::accounts::{ - model::AccountStatus, - repository::{ - AccountClaimsUpdate, AccountRepository, AccountUsageRepository, NewAccount, - TokenUpdate, UsageDelta, - }, - }, - platform::crypto::SecretBox, - platform::storage::db::connect_sqlite, -}; - -#[tokio::test] -async fn account_repository_should_encrypt_tokens_and_load_decrypted_account() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([4u8; 32])); - let expires_at = Utc::now() + Duration::hours(1); - - repo.insert(NewAccount { - id: "acct_a".to_string(), - email: Some("user@example.com".to_string()), - account_id: Some("chatgpt-account".to_string()), - user_id: Some("chatgpt-user".to_string()), - label: Some("primary".to_string()), - plan_type: Some("plus".to_string()), - access_token: SecretString::new("access-secret".to_string().into()), - refresh_token: Some(SecretString::new("refresh-secret".to_string().into())), - access_token_expires_at: Some(expires_at), - status: AccountStatus::Active, - }) - .await - .unwrap(); - - let stored: (String, String) = sqlx::query_as( - "select access_token_cipher, refresh_token_cipher from accounts where id = ?", - ) - .bind("acct_a") - .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 loaded = repo.get("acct_a").await.unwrap().unwrap(); - assert_eq!(loaded.access_token.expose_secret(), "access-secret"); - assert_eq!( - loaded.refresh_token.unwrap().expose_secret(), - "refresh-secret" - ); - assert_eq!(loaded.email.as_deref(), Some("user@example.com")); - assert_eq!(loaded.status, AccountStatus::Active); -} - -#[tokio::test] -async fn account_repository_should_cursor_page_accounts_by_added_at_desc() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([5u8; 32]); - let repo = AccountRepository::new(pool.clone(), secret_box.clone()); - - seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await; - seed_repo_account(&pool, &secret_box, "acct_b", "2026-06-11T00:01:00Z").await; - seed_repo_account(&pool, &secret_box, "acct_c", "2026-06-11T00:02:00Z").await; - - let first_page = repo.list(None, 2).await.unwrap(); - assert_eq!( - first_page - .items - .iter() - .map(|account| account.id.as_str()) - .collect::>(), - ["acct_c", "acct_b"] - ); - - let second_page = repo.list(first_page.next_cursor, 2).await.unwrap(); - assert_eq!( - second_page - .items - .iter() - .map(|account| account.id.as_str()) - .collect::>(), - ["acct_a"] - ); -} - -#[tokio::test] -async fn account_repository_should_update_status_and_label_without_rewriting_tokens() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([6u8; 32])); - - repo.insert(NewAccount { - id: "acct_a".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-secret".to_string().into()), - refresh_token: None, - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - let before: (String,) = sqlx::query_as("select access_token_cipher from accounts where id = ?") - .bind("acct_a") - .fetch_one(&pool) - .await - .unwrap(); - - assert!(repo - .set_status("acct_a", AccountStatus::Disabled) - .await - .unwrap()); - assert!(repo - .set_label("acct_a", Some("work".to_string())) - .await - .unwrap()); - - let after: (String,) = sqlx::query_as("select access_token_cipher from accounts where id = ?") - .bind("acct_a") - .fetch_one(&pool) - .await - .unwrap(); - let loaded = repo.get("acct_a").await.unwrap().unwrap(); - assert_eq!(after.0, before.0); - assert_eq!(loaded.status, AccountStatus::Disabled); - assert_eq!(loaded.label.as_deref(), Some("work")); -} - -#[tokio::test] -async fn account_repository_should_accumulate_usage_counters() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([7u8; 32])); - - repo.insert(NewAccount { - id: "acct_a".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-secret".to_string().into()), - refresh_token: None, - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - - repo.record_usage( - "acct_a", - UsageDelta { - input_tokens: 10, - output_tokens: 4, - cached_tokens: 3, - empty_response_count: 0, - ..UsageDelta::default() - }, - ) - .await - .unwrap(); - repo.record_usage( - "acct_a", - UsageDelta { - input_tokens: 8, - output_tokens: 2, - cached_tokens: 1, - empty_response_count: 1, - ..UsageDelta::default() - }, - ) - .await - .unwrap(); - - let usage = repo.get_usage("acct_a").await.unwrap().unwrap(); - assert_eq!(usage.request_count, 2); - assert_eq!(usage.input_tokens, 18); - assert_eq!(usage.output_tokens, 6); - assert_eq!(usage.cached_tokens, 4); - assert_eq!(usage.empty_response_count, 1); - assert!(usage.last_used_at.is_some()); -} - -#[tokio::test] -async fn account_repository_should_accumulate_usage_window_counters() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([32u8; 32])); - - repo.insert(NewAccount { - id: "acct_a".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-secret".to_string().into()), - refresh_token: None, - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - - repo.record_usage( - "acct_a", - UsageDelta { - input_tokens: 10, - output_tokens: 4, - cached_tokens: 3, - empty_response_count: 0, - ..UsageDelta::default() - }, - ) - .await - .unwrap(); - repo.record_usage( - "acct_a", - UsageDelta { - input_tokens: 8, - output_tokens: 2, - cached_tokens: 1, - empty_response_count: 0, - ..UsageDelta::default() - }, - ) - .await - .unwrap(); - - let usage: (i64, i64, i64, i64) = sqlx::query_as( - "select window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens from account_usage where account_id = ?", - ) - .bind("acct_a") - .fetch_one(&pool) - .await - .unwrap(); - - assert_eq!(usage, (2, 18, 6, 4)); -} - -#[tokio::test] -async fn account_repository_should_accumulate_image_generation_usage_counters() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([35u8; 32])); - - repo.insert(NewAccount { - id: "acct_a".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-secret".to_string().into()), - refresh_token: None, - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - - repo.record_usage( - "acct_a", - UsageDelta { - image_input_tokens: 31, - image_output_tokens: 9, - image_request_count: 1, - ..UsageDelta::default() - }, - ) - .await - .unwrap(); - repo.record_usage( - "acct_a", - UsageDelta { - image_request_failed_count: 1, - ..UsageDelta::default() - }, - ) - .await - .unwrap(); - - let usage = repo.get_usage("acct_a").await.unwrap().unwrap(); - assert_eq!(usage.image_input_tokens, 31); - assert_eq!(usage.image_output_tokens, 9); - assert_eq!(usage.image_request_count, 1); - assert_eq!(usage.image_request_failed_count, 1); - assert_eq!(usage.window_image_input_tokens, 31); - assert_eq!(usage.window_image_output_tokens, 9); - assert_eq!(usage.window_image_request_count, 1); - assert_eq!(usage.window_image_request_failed_count, 1); -} - -#[tokio::test] -async fn account_repository_should_load_usage_count_into_runtime_pool_accounts() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([17u8; 32])); - - repo.insert(NewAccount { - id: "acct_a".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-secret".to_string().into()), - refresh_token: None, - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - repo.record_usage( - "acct_a", - UsageDelta { - input_tokens: 1, - output_tokens: 2, - cached_tokens: 3, - empty_response_count: 0, - ..UsageDelta::default() - }, - ) - .await - .unwrap(); - - let account = repo.list_pool_accounts().await.unwrap().remove(0); - - assert_eq!(account.request_count, 1); -} - -#[tokio::test] -async fn account_repository_should_restore_window_usage_into_runtime_pool_accounts() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([33u8; 32]); - let repo = AccountRepository::new(pool.clone(), secret_box.clone()); - let window_started_at = Utc::now() - Duration::minutes(2); - let window_reset_at = Utc::now() + Duration::minutes(3); - - seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await; - sqlx::query( - "insert into account_usage (account_id, request_count, image_input_tokens, image_output_tokens, image_request_count, image_request_failed_count, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_image_input_tokens, window_image_output_tokens, window_image_request_count, window_image_request_failed_count, window_started_at, window_reset_at, limit_window_seconds) values (?, 9, 31, 9, 2, 1, 7, 11, 13, 17, 19, 5, 1, 1, ?, ?, 300)", - ) - .bind("acct_a") - .bind(window_started_at.to_rfc3339()) - .bind(window_reset_at.to_rfc3339()) - .execute(&pool) - .await - .unwrap(); - - let account = repo.list_pool_accounts().await.unwrap().remove(0); - - assert_eq!(account.window_request_count, 7); - assert_eq!(account.image_input_tokens, 31); - assert_eq!(account.image_output_tokens, 9); - assert_eq!(account.image_request_count, 2); - assert_eq!(account.image_request_failed_count, 1); - assert_eq!(account.window_input_tokens, 11); - assert_eq!(account.window_output_tokens, 13); - assert_eq!(account.window_cached_tokens, 17); - assert_eq!(account.window_image_input_tokens, 19); - assert_eq!(account.window_image_output_tokens, 5); - assert_eq!(account.window_image_request_count, 1); - assert_eq!(account.window_image_request_failed_count, 1); - assert_eq!(account.window_started_at, Some(window_started_at)); - assert_eq!(account.window_reset_at, Some(window_reset_at)); - assert_eq!(account.limit_window_seconds, Some(300)); -} - -#[tokio::test] -async fn account_repository_reset_usage_should_preserve_rate_limit_window_metadata() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([35u8; 32]); - let repo = AccountRepository::new(pool.clone(), secret_box.clone()); - let window_started_at = Utc::now() - Duration::minutes(10); - let window_reset_at = Utc::now() + Duration::minutes(5); - - seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await; - sqlx::query( - "insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds, last_used_at) values (?, 9, 31, 9, 2, 7, 11, 13, 17, ?, ?, 300, ?)", - ) - .bind("acct_a") - .bind(window_started_at.to_rfc3339()) - .bind(window_reset_at.to_rfc3339()) - .bind(Utc::now().to_rfc3339()) - .execute(&pool) - .await - .unwrap(); - - AccountUsageRepository::new(pool.clone()) - .reset_account("acct_a") - .await - .unwrap(); - - let usage = repo.get_usage("acct_a").await.unwrap().unwrap(); - assert_eq!(usage.request_count, 0); - assert_eq!(usage.window_request_count, 0); - assert_eq!(usage.window_reset_at, Some(window_reset_at)); - assert_eq!(usage.limit_window_seconds, Some(300)); - assert!(usage - .window_started_at - .is_some_and(|reset_at| { reset_at > window_started_at && reset_at <= Utc::now() })); -} - -#[tokio::test] -async fn account_repository_should_restore_quota_verify_required_into_runtime_pool_accounts() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([34u8; 32]); - let repo = AccountRepository::new(pool.clone(), secret_box.clone()); - - seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await; - sqlx::query("update accounts set quota_verify_required = 1 where id = ?") - .bind("acct_a") - .execute(&pool) - .await - .unwrap(); - - let account = repo.list_pool_accounts().await.unwrap().remove(0); - - assert!(account.quota_verify_required); -} - -#[tokio::test] -async fn account_repository_should_exclude_refresh_lease_until_expiry() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([18u8; 32])); - let now = Utc::now(); - repo.insert(NewAccount { - id: "acct_a".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-secret".to_string().into()), - refresh_token: None, - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - - assert!(repo - .try_acquire_refresh_lease("acct_a", "owner-a", now + Duration::minutes(5)) - .await - .unwrap()); - assert!(!repo - .try_acquire_refresh_lease("acct_a", "owner-b", now + Duration::minutes(5)) - .await - .unwrap()); - sqlx::query("update account_refresh_leases set expires_at = ? where account_id = ?") - .bind((now - Duration::seconds(1)).to_rfc3339()) - .bind("acct_a") - .execute(&pool) - .await - .unwrap(); - assert!(repo - .try_acquire_refresh_lease("acct_a", "owner-b", now + Duration::minutes(5)) - .await - .unwrap()); -} - -#[tokio::test] -async fn account_repository_should_preserve_refresh_token_when_update_omits_one() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool, SecretBox::new([8u8; 32])); - - repo.insert(NewAccount { - id: "acct_a".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("old-access".to_string().into()), - refresh_token: Some(SecretString::new("old-refresh".to_string().into())), - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - - assert!(repo - .update_tokens( - "acct_a", - TokenUpdate { - access_token: SecretString::new("new-access".to_string().into()), - refresh_token: None, - access_token_expires_at: Some(Utc::now() + Duration::hours(2)), - }, - ) - .await - .unwrap()); - - let loaded = repo.get("acct_a").await.unwrap().unwrap(); - assert_eq!(loaded.access_token.expose_secret(), "new-access"); - assert_eq!(loaded.refresh_token.unwrap().expose_secret(), "old-refresh"); - assert_eq!(loaded.status, AccountStatus::Active); - assert!(loaded.access_token_expires_at.is_some()); -} - -#[tokio::test] -async fn account_repository_should_preserve_refresh_token_when_claims_update_omits_one() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool, SecretBox::new([34u8; 32])); - - repo.insert(NewAccount { - id: "acct_a".to_string(), - email: Some("old@example.com".to_string()), - account_id: Some("old-account".to_string()), - user_id: Some("old-user".to_string()), - label: Some("old-label".to_string()), - plan_type: Some("plus".to_string()), - access_token: SecretString::new("old-access".to_string().into()), - refresh_token: Some(SecretString::new("old-refresh".to_string().into())), - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - - assert!(repo - .update_from_claims( - "acct_a", - AccountClaimsUpdate { - email: Some("new@example.com".to_string()), - account_id: Some("new-account".to_string()), - user_id: Some("new-user".to_string()), - plan_type: Some("pro".to_string()), - access_token: SecretString::new("new-access".to_string().into()), - refresh_token: None, - access_token_expires_at: Some(Utc::now() + Duration::hours(2)), - status: AccountStatus::Active, - }, - ) - .await - .unwrap()); - - let loaded = repo.get("acct_a").await.unwrap().unwrap(); - assert_eq!(loaded.access_token.expose_secret(), "new-access"); - assert_eq!(loaded.refresh_token.unwrap().expose_secret(), "old-refresh"); - assert_eq!(loaded.email.as_deref(), Some("new@example.com")); - assert_eq!(loaded.account_id.as_deref(), Some("new-account")); - assert_eq!(loaded.user_id.as_deref(), Some("new-user")); - assert_eq!(loaded.plan_type.as_deref(), Some("pro")); - assert!(loaded.access_token_expires_at.is_some()); -} - -async fn seed_repo_account( - pool: &sqlx::SqlitePool, - secret_box: &SecretBox, - id: &str, - added_at: &str, -) { - let token = SecretString::new(format!("access-{id}").into()); - let cipher = secret_box.encrypt(&token).unwrap(); - sqlx::query( - "insert into accounts (id, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?)", - ) - .bind(id) - .bind(cipher) - .bind("active") - .bind(added_at) - .bind(added_at) - .execute(pool) - .await - .unwrap(); -} diff --git a/tests/codex_accounts/service_refresh.rs b/tests/codex_accounts/service_refresh.rs deleted file mode 100644 index c7fb1cd7c..000000000 --- a/tests/codex_accounts/service_refresh.rs +++ /dev/null @@ -1,171 +0,0 @@ -use std::{ - sync::{ - atomic::{AtomicUsize, Ordering}, - Arc, - }, - time::Duration as StdDuration, -}; - -use async_trait::async_trait; -use secrecy::SecretString; -use tokio::sync::Mutex; - -use codex_proxy_rs::{ - codex::{ - accounts::{ - model::AccountStatus, - pool::AccountPool, - repository::{AccountRepository, AccountUsageRepository, NewAccount, TokenUpdate}, - service::{AccountService, AccountServiceDependencies}, - }, - gateway::{ - fingerprint::model::Fingerprint, - oauth::{RefreshFailure, TokenPair, TokenRefresher}, - transport::websocket::CodexWebSocketPool, - }, - }, - platform::{crypto::SecretBox, storage::db::connect_sqlite}, -}; - -use crate::support::admin_accounts::test_config; - -#[tokio::test] -async fn account_service_refresh_should_retry_with_latest_disk_refresh_token_after_invalid_grant() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([19u8; 32])); - repo.insert(NewAccount { - id: "acct_stale".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-old".to_string().into()), - refresh_token: Some(SecretString::new("refresh-old".to_string().into())), - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - let calls = Arc::new(Mutex::new(Vec::new())); - let refresher = DiskRotatingRefresher { - repo: repo.clone(), - calls: calls.clone(), - first_call_seen: Arc::new(AtomicUsize::new(0)), - }; - let service = AccountService::new( - Arc::new(test_config(url)), - AccountServiceDependencies { - repository: Some(repo), - usage_repository: Some(AccountUsageRepository::new(pool)), - cookie_repository: None, - token_refresher: Some(Arc::new(refresher)), - account_pool: Arc::new(Mutex::new(AccountPool::default())), - websocket_pool: Arc::new(CodexWebSocketPool::with_default_max_age()), - fingerprint: Fingerprint::default_for_tests(), - }, - ); - - let result = service.refresh_account("acct_stale").await.unwrap(); - - assert_eq!(result.outcome.as_str(), "alive"); - assert_eq!( - calls.lock().await.as_slice(), - ["refresh-old", "refresh-latest"] - ); -} - -#[tokio::test] -async fn account_service_refresh_probe_should_not_persist_permanent_status_on_first_failure() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("accounts_scheduler.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = AccountRepository::new(pool.clone(), SecretBox::new([20u8; 32])); - repo.insert(NewAccount { - id: "acct_scheduler".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-old".to_string().into()), - refresh_token: Some(SecretString::new("refresh-old".to_string().into())), - access_token_expires_at: None, - status: AccountStatus::Active, - }) - .await - .unwrap(); - let service = AccountService::new( - Arc::new(test_config(url)), - AccountServiceDependencies { - repository: Some(repo.clone()), - usage_repository: Some(AccountUsageRepository::new(pool)), - cookie_repository: None, - token_refresher: Some(Arc::new(AlwaysFailingRefresher)), - account_pool: Arc::new(Mutex::new(AccountPool::default())), - websocket_pool: Arc::new(CodexWebSocketPool::with_default_max_age()), - fingerprint: Fingerprint::default_for_tests(), - }, - ); - - let result = service - .probe_account_refresh("acct_scheduler") - .await - .unwrap(); - - assert_eq!(result.outcome.as_str(), "dead"); - assert!(result.status.is_none()); - let stored = repo.get("acct_scheduler").await.unwrap().unwrap(); - assert_eq!(stored.status, AccountStatus::Active); -} - -#[derive(Clone)] -struct DiskRotatingRefresher { - repo: AccountRepository, - calls: Arc>>, - first_call_seen: Arc, -} - -#[derive(Clone)] -struct AlwaysFailingRefresher; - -#[async_trait] -impl TokenRefresher for AlwaysFailingRefresher { - async fn refresh(&self, _refresh_token: &str) -> Result { - Err(RefreshFailure::InvalidGrant) - } -} - -#[async_trait] -impl TokenRefresher for DiskRotatingRefresher { - async fn refresh(&self, refresh_token: &str) -> Result { - self.calls.lock().await.push(refresh_token.to_string()); - if refresh_token == "refresh-old" { - self.first_call_seen.fetch_add(1, Ordering::SeqCst); - self.repo - .update_tokens( - "acct_stale", - TokenUpdate { - access_token: SecretString::new("external-access".to_string().into()), - refresh_token: Some(SecretString::new("refresh-latest".to_string().into())), - access_token_expires_at: None, - }, - ) - .await - .unwrap(); - tokio::time::sleep(StdDuration::from_millis(5)).await; - return Err(RefreshFailure::InvalidGrant); - } - if refresh_token == "refresh-latest" { - return Ok(TokenPair { - access_token: "access-after-retry".to_string(), - refresh_token: Some("refresh-after-retry".to_string()), - }); - } - Err(RefreshFailure::InvalidGrant) - } -} diff --git a/tests/codex_events.rs b/tests/codex_events.rs deleted file mode 100644 index f4dcf6c3f..000000000 --- a/tests/codex_events.rs +++ /dev/null @@ -1,3 +0,0 @@ -mod codex_events { - mod logs_pagination; -} diff --git a/tests/codex_events/logs_pagination.rs b/tests/codex_events/logs_pagination.rs deleted file mode 100644 index 0baff98e6..000000000 --- a/tests/codex_events/logs_pagination.rs +++ /dev/null @@ -1,125 +0,0 @@ -use codex_proxy_rs::{ - codex::events::{ - event::{EventLevel, EventLog}, - repository::EventLogRepository, - service::LogService, - }, - config::LoggingConfig, - platform::storage::db::connect_sqlite, -}; -use serde_json::json; - -#[tokio::test] -async fn event_logs_are_cursor_paginated() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("logs.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let repo = EventLogRepository::new(pool); - - for idx in 0..3 { - repo.insert(EventLog::new( - "request", - EventLevel::Info, - format!("event {idx}"), - )) - .await - .unwrap(); - } - - let first = repo.list(None, 2).await.unwrap(); - assert_eq!(first.items.len(), 2); - assert!(first.next_cursor.is_some()); - - let second = repo.list(first.next_cursor, 2).await.unwrap(); - assert_eq!(second.items.len(), 1); - assert!(second.next_cursor.is_none()); -} - -#[tokio::test] -async fn log_service_should_skip_record_when_disabled() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("logs-disabled.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let service = LogService::new( - logging_config(false, 10, false), - Some(EventLogRepository::new(pool.clone())), - ); - - let recorded = service - .record(EventLog::new("request", EventLevel::Info, "disabled")) - .await - .unwrap(); - - let stored_count: (i64,) = sqlx::query_as("select count(*) from event_logs") - .fetch_one(&pool) - .await - .unwrap(); - assert!(!recorded); - assert_eq!(stored_count.0, 0); -} - -#[tokio::test] -async fn log_service_should_trim_to_capacity_after_record() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("logs-capacity.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let service = LogService::new( - logging_config(true, 1, false), - Some(EventLogRepository::new(pool.clone())), - ); - let mut first = EventLog::new("request", EventLevel::Info, "first"); - first.created_at = "2026-06-14T00:00:00Z".to_string(); - let mut second = EventLog::new("request", EventLevel::Info, "second"); - second.created_at = "2026-06-14T00:00:01Z".to_string(); - - service.record(first).await.unwrap(); - service.record(second).await.unwrap(); - - let messages: Vec<(String,)> = - sqlx::query_as("select message from event_logs order by created_at desc") - .fetch_all(&pool) - .await - .unwrap(); - assert_eq!(messages, vec![("second".to_string(),)]); -} - -#[tokio::test] -async fn log_service_should_remove_body_metadata_when_capture_body_disabled() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("logs-body-policy.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let service = LogService::new( - logging_config(true, 10, false), - Some(EventLogRepository::new(pool.clone())), - ); - let mut event = EventLog::new("request", EventLevel::Info, "body"); - event.metadata = json!({ - "body": "secret", - "requestBody": "secret", - "responseBody": "secret", - "safe": "visible" - }); - - service.record(event).await.unwrap(); - - let metadata_json: (String,) = sqlx::query_as("select metadata_json from event_logs") - .fetch_one(&pool) - .await - .unwrap(); - let metadata: serde_json::Value = serde_json::from_str(&metadata_json.0).unwrap(); - assert_eq!(metadata, json!({ "safe": "visible" })); -} - -fn logging_config(enabled: bool, capacity: u32, capture_body: bool) -> LoggingConfig { - LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled, - capacity, - capture_body, - } -} diff --git a/tests/codex_gateway.rs b/tests/codex_gateway.rs deleted file mode 100644 index bc3addc98..000000000 --- a/tests/codex_gateway.rs +++ /dev/null @@ -1,9 +0,0 @@ -mod codex_gateway { - mod cli_auth_import; - mod fingerprint_update; - mod headers; - mod http_client; - mod oauth_refresh; - mod usage_events; - mod websocket; -} diff --git a/tests/codex_gateway/cli_auth_import.rs b/tests/codex_gateway/cli_auth_import.rs deleted file mode 100644 index 540004717..000000000 --- a/tests/codex_gateway/cli_auth_import.rs +++ /dev/null @@ -1,44 +0,0 @@ -use std::fs; - -use codex_proxy_rs::codex::gateway::oauth::{parse_cli_auth_json, read_cli_auth_from_home}; - -#[test] -fn parse_cli_auth_json_should_accept_codex_cli_tokens() { - let auth = parse_cli_auth_json( - r#"{ - "access_token": "cli-access-secret", - "refresh_token": "cli-refresh-secret", - "id_token": "cli-id-token", - "expires_at": 1893456000 - }"#, - ) - .unwrap(); - - assert_eq!(auth.access_token(), "cli-access-secret"); - assert_eq!(auth.refresh_token(), Some("cli-refresh-secret")); -} - -#[test] -fn parse_cli_auth_json_should_reject_missing_access_token() { - let error = parse_cli_auth_json(r#"{"refresh_token":"cli-refresh-secret"}"#).unwrap_err(); - - assert_eq!( - error.to_string(), - "CLI auth.json does not contain access_token" - ); -} - -#[test] -fn read_cli_auth_from_home_should_load_auth_json_from_codex_home() { - let dir = tempfile::tempdir().unwrap(); - fs::write( - dir.path().join("auth.json"), - r#"{"access_token":"file-access-secret","refresh_token":"file-refresh-secret"}"#, - ) - .unwrap(); - - let auth = read_cli_auth_from_home(dir.path()).unwrap(); - - assert_eq!(auth.access_token(), "file-access-secret"); - assert_eq!(auth.refresh_token(), Some("file-refresh-secret")); -} diff --git a/tests/codex_gateway/fingerprint_update.rs b/tests/codex_gateway/fingerprint_update.rs deleted file mode 100644 index 52e22469e..000000000 --- a/tests/codex_gateway/fingerprint_update.rs +++ /dev/null @@ -1,72 +0,0 @@ -use codex_proxy_rs::{ - codex::gateway::fingerprint::{ - repository::FingerprintRepository, - updater::{FingerprintUpdater, CODEX_DESKTOP_UPDATE_SOURCE}, - }, - codex::gateway::transport::http_client::build_reqwest_client, - platform::storage::db::connect_sqlite, -}; -use wiremock::{ - matchers::{method, path}, - Mock, MockServer, ResponseTemplate, -}; - -#[tokio::test] -async fn fingerprint_updater_fetches_manifest_and_persists_history() { - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/desktop/update.json")) - .respond_with(ResponseTemplate::new(200).set_body_raw( - r#"{"version":"26.700.111","build_number":"5002"}"#, - "application/json", - )) - .mount(&server) - .await; - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("fingerprints.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .unwrap(); - let repo = FingerprintRepository::new(pool); - let updater = FingerprintUpdater::new( - build_reqwest_client(false).unwrap(), - repo.clone(), - format!("{}/desktop/update.json", server.uri()), - ); - - updater.poll_once().await.unwrap(); - - let latest = repo.latest().await.unwrap().unwrap(); - assert_eq!(latest.app_version, "26.700.111"); - assert_eq!(latest.build_number, "5002"); - assert_eq!(latest.source, CODEX_DESKTOP_UPDATE_SOURCE); -} - -#[tokio::test] -async fn fingerprint_repository_should_upsert_auto_update_record() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("fingerprints.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .unwrap(); - let repo = FingerprintRepository::new(pool.clone()); - - repo.upsert_auto_update("26.800.1", "6001", Some("147")) - .await - .unwrap(); - repo.upsert_auto_update("26.800.2", "6002", None) - .await - .unwrap(); - - let stored = repo.load_latest_auto_updated().await.unwrap().unwrap(); - let count: (i64,) = - sqlx::query_as("select count(*) from fingerprints where id = 'auto_updated'") - .fetch_one(&pool) - .await - .unwrap(); - - assert_eq!(count.0, 1); - assert_eq!(stored.app_version, "26.800.2"); - assert_eq!(stored.build_number, "6002"); - assert_eq!(stored.chromium_version, "146"); -} diff --git a/tests/codex_gateway/headers.rs b/tests/codex_gateway/headers.rs deleted file mode 100644 index 4cc4b20dd..000000000 --- a/tests/codex_gateway/headers.rs +++ /dev/null @@ -1,176 +0,0 @@ -use codex_proxy_rs::{ - codex::gateway::fingerprint::{model::Fingerprint, updater::parse_update_manifest}, - codex::gateway::transport::{ - headers::build_codex_headers, - http_client::{build_reqwest_client, CodexBackendClient, CodexRequestContext}, - types::CodexResponsesRequest, - }, -}; - -use wiremock::{ - matchers::{header, method, path}, - Mock, MockServer, ResponseTemplate, -}; - -#[test] -fn codex_headers_include_desktop_identity_and_turn_state() { - let fp = Fingerprint::default_for_tests(); - let headers = build_codex_headers( - &fp, - "access-token", - Some("acct_123"), - Some("turn-state"), - "rid_1", - ); - - assert_eq!(headers.get("originator").unwrap(), "Codex Desktop"); - assert!(headers.get("user-agent").unwrap().contains("Codex")); - assert_eq!(headers.get("authorization").unwrap(), "Bearer access-token"); - assert_eq!(headers.get("chatgpt-account-id").unwrap(), "acct_123"); - assert_eq!(headers.get("x-codex-turn-state").unwrap(), "turn-state"); - assert_eq!(headers.get("x-client-request-id").unwrap(), "rid_1"); -} - -#[test] -fn fingerprint_default_header_order_should_match_original_config() { - assert_eq!( - Fingerprint::default_header_order(), - vec![ - "authorization", - "chatgpt-account-id", - "originator", - "x-openai-internal-codex-residency", - "x-client-request-id", - "x-codex-installation-id", - "x-codex-turn-state", - "openai-beta", - "user-agent", - "sec-ch-ua", - "sec-ch-ua-mobile", - "sec-ch-ua-platform", - "accept-encoding", - "accept-language", - "sec-fetch-site", - "sec-fetch-mode", - "sec-fetch-dest", - "content-type", - "accept", - "cookie", - ] - ); -} - -#[test] -fn update_manifest_updates_app_version_and_build_number() { - let manifest = r#"{"version":"26.600.12345","build_number":"4001"}"#; - let update = parse_update_manifest(manifest).unwrap(); - assert_eq!(update.app_version, "26.600.12345"); - assert_eq!(update.build_number, "4001"); -} - -/// 验证使用数据库指纹构造的请求头是否正确 -#[tokio::test] -async fn database_fingerprint_headers_should_match_expected_format() { - let server = MockServer::start().await; - - // 模拟数据库中的最新指纹（auto_update） - let db_fingerprint = Fingerprint { - originator: "Codex Desktop".to_string(), - app_version: "26.609.41114".to_string(), - build_number: "3888".to_string(), - platform: "darwin".to_string(), - arch: "arm64".to_string(), - chromium_version: "146".to_string(), - user_agent_template: "Codex Desktop/{version} ({platform}; {arch})".to_string(), - default_headers: Fingerprint::default_headers(), - header_order: Fingerprint::default_header_order(), - }; - - let expected_user_agent = "Codex Desktop/26.609.41114 (darwin; arm64)"; - - // 验证关键请求头 - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("user-agent", expected_user_agent)) - .and(header("originator", "Codex Desktop")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string( - "event: response.completed\ndata: {\"response\":{\"id\":\"resp_1\"}}\n\n", - ), - ) - .expect(1) - .mount(&server) - .await; - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - server.uri(), - db_fingerprint, - ); - let mut request = CodexResponsesRequest::new_http_sse("gpt-4", "", Vec::new()); - request.force_http_sse = true; - - let response = client - .create_response( - &request, - CodexRequestContext { - access_token: "test-token", - account_id: None, - request_id: "req_test", - 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, - installation_id: None, - session_id: None, - }, - ) - .await; - - assert!( - response.is_ok(), - "Request should succeed with correct headers" - ); -} - -#[test] -fn fingerprint_user_agent_should_reflect_updated_version_fields() { - let baseline_fingerprint = Fingerprint::default_codex_desktop(); - let updated_fingerprint = Fingerprint { - originator: "Codex Desktop".to_string(), - app_version: "26.609.41114".to_string(), - build_number: "3888".to_string(), - platform: "darwin".to_string(), - arch: "arm64".to_string(), - chromium_version: "146".to_string(), - user_agent_template: "Codex Desktop/{version} ({platform}; {arch})".to_string(), - default_headers: Fingerprint::default_headers(), - header_order: Fingerprint::default_header_order(), - }; - - let baseline_ua = baseline_fingerprint.user_agent(); - let updated_ua = updated_fingerprint.user_agent(); - - assert!(baseline_ua.starts_with("Codex Desktop/")); - assert!(updated_ua.starts_with("Codex Desktop/")); - assert!(baseline_ua.contains("darwin")); - assert!(updated_ua.contains("darwin")); - - assert_ne!( - baseline_ua, updated_ua, - "fingerprint version fields should affect User-Agent" - ); - - assert!(updated_ua.contains("26.609.41114")); - - assert_eq!( - updated_fingerprint.sec_ch_ua(), - "\"Chromium\";v=\"146\", \"Not:A-Brand\";v=\"24\"" - ); -} diff --git a/tests/codex_gateway/http_client.rs b/tests/codex_gateway/http_client.rs deleted file mode 100644 index 03265e53f..000000000 --- a/tests/codex_gateway/http_client.rs +++ /dev/null @@ -1,554 +0,0 @@ -use std::time::Duration; - -use serde_json::json; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, -}; -use wiremock::{ - matchers::{body_json, header, method, path}, - Mock, MockServer, ResponseTemplate, -}; - -use codex_proxy_rs::{ - codex::gateway::fingerprint::model::Fingerprint, - codex::gateway::transport::{ - http_client::{ - build_reqwest_client, CodexBackendClient, CodexClientError, CodexRequestContext, - }, - types::{CodexCompactRequest, CodexResponsesRequest}, - usage_events::TokenUsage, - }, -}; - -#[test] -fn reqwest_http2_feature_should_be_enabled_for_fallback_parity() { - let builder = reqwest::Client::builder().http2_prior_knowledge(); - - drop(builder); -} - -#[tokio::test] -async fn codex_backend_client_should_send_desktop_headers_and_capture_response_metadata() { - let server = MockServer::start().await; - let sse_body = concat!( - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_1\",\"usage\":{\"input_tokens\":2,\"output_tokens\":3,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n", - "\n", - ); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-token")) - .and(header("chatgpt-account-id", "chatgpt-account")) - .and(header("originator", "Codex Desktop")) - .and(header("x-client-request-id", "req_1")) - .and(header("x-codex-turn-state", "turn_1")) - .and(header("cookie", "cf_clearance=old")) - .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=/", - ) - .insert_header("x-codex-turn-state", "turn_2") - .set_body_string(sse_body), - ) - .mount(&server) - .await; - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - server.uri(), - Fingerprint::default_for_tests(), - ); - let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new()); - request.force_http_sse = true; - - let response = client - .create_response( - &request, - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_1", - turn_state: Some("turn_1"), - turn_metadata: None, - beta_features: None, - include_timing_metrics: None, - version: None, - codex_window_id: None, - parent_thread_id: None, - cookie_header: Some("cf_clearance=old"), - installation_id: None, - session_id: None, - }, - ) - .await - .unwrap(); - - assert_eq!( - response.usage, - Some(TokenUsage { - input_tokens: 2, - output_tokens: 3, - cached_tokens: 1, - image_input_tokens: 0, - image_output_tokens: 0, - total_tokens: 5, - }) - ); - assert_eq!(response.turn_state.as_deref(), Some("turn_2")); - assert_eq!( - response.set_cookie_headers, - vec!["cf_clearance=new; Domain=.chatgpt.com; Path=/".to_string()] - ); -} - -#[tokio::test] -async fn codex_backend_client_usage_should_use_original_auxiliary_headers() { - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/api/codex/usage")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "rate_limit": { - "limit_reached": false - } - }))) - .mount(&server) - .await; - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - server.uri(), - Fingerprint::default_for_tests(), - ); - - let usage = client - .fetch_usage(CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_aux", - turn_state: Some("turn-state"), - turn_metadata: Some("turn-meta"), - beta_features: Some("feature-a"), - include_timing_metrics: Some("true"), - version: Some("26.318.11754"), - codex_window_id: Some("cw_1"), - parent_thread_id: Some("parent-1"), - cookie_header: Some("cf_clearance=old"), - installation_id: Some("install-1"), - session_id: Some("session-1"), - }) - .await - .unwrap(); - - assert_eq!(usage["rate_limit"]["limit_reached"], false); - let requests = server.received_requests().await.unwrap(); - let headers = &requests[0].headers; - assert_eq!( - headers.get("accept").and_then(|value| value.to_str().ok()), - Some("application/json") - ); - assert_eq!( - headers - .get("accept-encoding") - .and_then(|value| value.to_str().ok()), - Some("gzip, deflate") - ); - assert!(headers.get("content-type").is_none()); - assert!(headers.get("openai-beta").is_none()); - assert!(headers.get("x-openai-internal-codex-residency").is_none()); - assert!(headers.get("x-client-request-id").is_none()); - assert!(headers.get("x-codex-installation-id").is_none()); - assert!(headers.get("session_id").is_none()); -} - -#[tokio::test] -async fn codex_backend_client_models_should_use_original_auxiliary_headers() { - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/codex/models")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "models": [ - {"slug": "gpt-5.5", "title": "GPT 5.5"} - ] - }))) - .mount(&server) - .await; - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - server.uri(), - Fingerprint::default_for_tests(), - ); - - let models = client - .fetch_models(CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_models", - turn_state: Some("turn-state"), - turn_metadata: None, - beta_features: None, - include_timing_metrics: None, - version: None, - codex_window_id: None, - parent_thread_id: None, - cookie_header: None, - installation_id: Some("install-1"), - session_id: Some("session-1"), - }) - .await - .unwrap(); - - assert_eq!(models.len(), 1); - let requests = server.received_requests().await.unwrap(); - let models_request = requests - .iter() - .find(|request| request.url.path() == "/codex/models") - .unwrap(); - let headers = &models_request.headers; - assert_eq!( - headers.get("accept").and_then(|value| value.to_str().ok()), - Some("application/json") - ); - assert!(headers.get("content-type").is_none()); - assert!(headers.get("openai-beta").is_none()); - assert!(headers.get("x-openai-internal-codex-residency").is_none()); - assert!(headers.get("x-client-request-id").is_none()); - assert!(headers.get("x-codex-installation-id").is_none()); - assert!(headers.get("session_id").is_none()); -} - -#[tokio::test] -async fn codex_backend_client_should_cap_non_success_error_body_at_one_mib() { - let server = MockServer::start().await; - let large_error_body = "x".repeat(1024 * 1024 + 17); - Mock::given(method("POST")) - .and(path("/codex/responses")) - .respond_with(ResponseTemplate::new(500).set_body_string(large_error_body)) - .mount(&server) - .await; - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - server.uri(), - Fingerprint::default_for_tests(), - ); - let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new()); - request.force_http_sse = true; - - let result = client - .create_response( - &request, - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_large_error", - 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, - installation_id: None, - session_id: None, - }, - ) - .await; - - let Err(CodexClientError::Upstream { status, body, .. }) = result else { - panic!("expected upstream error"); - }; - assert_eq!(status, reqwest::StatusCode::INTERNAL_SERVER_ERROR); - assert_eq!(body.len(), 1024 * 1024); -} - -#[tokio::test] -async fn codex_backend_client_should_send_http_sse_headers_in_fingerprint_order() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.unwrap(); - let request = read_http_request(&mut stream).await; - write_completed_sse_response(&mut stream).await; - request - }); - - let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new()); - request.force_http_sse = true; - request.turn_metadata = Some("turn-meta".to_string()); - request.beta_features = Some("beta-a".to_string()); - request.include_timing_metrics = Some("true".to_string()); - request.version = Some("26.519.81530".to_string()); - request.codex_window_id = Some("cw_1".to_string()); - request.parent_thread_id = Some("parent-1".to_string()); - let client = CodexBackendClient::new( - build_reqwest_client(true).unwrap(), - format!("http://{addr}"), - Fingerprint::default_for_tests(), - ); - - client - .create_response( - &request, - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_order", - turn_state: Some("turn-state"), - 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: Some("cf_clearance=old"), - installation_id: Some("install-1"), - session_id: Some("session-1"), - }, - ) - .await - .unwrap(); - - let raw_request = server.await.unwrap(); - let header_names = header_names(&raw_request); - assert_header_subsequence( - &header_names, - &[ - "authorization", - "chatgpt-account-id", - "originator", - "user-agent", - "sec-ch-ua", - "sec-ch-ua-mobile", - "sec-ch-ua-platform", - "accept-encoding", - "accept-language", - "sec-fetch-site", - "sec-fetch-mode", - "sec-fetch-dest", - "content-type", - "cookie", - "accept", - "openai-beta", - "x-openai-internal-codex-residency", - "x-client-request-id", - "x-codex-installation-id", - "session_id", - "x-codex-window-id", - "x-codex-turn-state", - "x-codex-turn-metadata", - "x-codex-beta-features", - "x-responsesapi-include-timing-metrics", - "version", - "x-codex-parent-thread-id", - ], - ); -} - -#[tokio::test] -async fn codex_backend_client_should_send_compact_headers_in_fingerprint_order() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.unwrap(); - let request = read_http_request(&mut stream).await; - write_compact_json_response(&mut stream).await; - request - }); - let client = CodexBackendClient::new( - build_reqwest_client(true).unwrap(), - format!("http://{addr}"), - Fingerprint::default_for_tests(), - ); - - client - .create_compact_response( - &CodexCompactRequest { - model: "gpt-5.5".to_string(), - input: Vec::new(), - instructions: String::new(), - tools: None, - parallel_tool_calls: None, - reasoning: None, - text: None, - }, - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_compact", - 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: Some("cf_clearance=old"), - installation_id: Some("install-1"), - session_id: None, - }, - ) - .await - .unwrap(); - - let raw_request = server.await.unwrap(); - let header_names = header_names(&raw_request); - assert_header_subsequence( - &header_names, - &[ - "authorization", - "chatgpt-account-id", - "originator", - "user-agent", - "sec-ch-ua", - "sec-ch-ua-mobile", - "sec-ch-ua-platform", - "accept-encoding", - "accept-language", - "sec-fetch-site", - "sec-fetch-mode", - "sec-fetch-dest", - "content-type", - "cookie", - "openai-beta", - "x-openai-internal-codex-residency", - "x-client-request-id", - "x-codex-installation-id", - ], - ); -} - -#[tokio::test] -async fn build_reqwest_client_should_reuse_cached_connection_pool() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut first_stream, _) = listener.accept().await.unwrap(); - read_http_request(&mut first_stream).await; - write_empty_http_response(&mut first_stream).await; - - tokio::select! { - request = read_http_request(&mut first_stream) => { - write_empty_http_response(&mut first_stream).await; - !request.is_empty() - } - accepted = listener.accept() => { - let (mut second_stream, _) = accepted.unwrap(); - read_http_request(&mut second_stream).await; - write_empty_http_response(&mut second_stream).await; - false - } - _ = tokio::time::sleep(Duration::from_millis(500)) => false, - } - }); - - let url = format!("http://{addr}/reuse"); - let first_client = build_reqwest_client(false).unwrap(); - first_client - .get(&url) - .send() - .await - .unwrap() - .text() - .await - .unwrap(); - let second_client = build_reqwest_client(false).unwrap(); - second_client - .get(&url) - .send() - .await - .unwrap() - .text() - .await - .unwrap(); - - assert!(server.await.unwrap()); -} - -fn header_names(request: &str) -> Vec { - request - .lines() - .skip(1) - .take_while(|line| !line.is_empty()) - .filter_map(|line| { - line.split_once(':') - .map(|(name, _)| name.to_ascii_lowercase()) - }) - .collect() -} - -fn assert_header_subsequence(actual: &[String], expected: &[&str]) { - let mut offset = 0; - for expected_name in expected { - let Some(position) = actual[offset..] - .iter() - .position(|actual_name| actual_name == expected_name) - else { - panic!("missing header {expected_name}; actual order: {actual:?}"); - }; - offset += position + 1; - } -} - -async fn read_http_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 write_empty_http_response(stream: &mut TcpStream) { - stream - .write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n") - .await - .unwrap(); -} - -async fn write_completed_sse_response(stream: &mut TcpStream) { - let body = concat!( - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_order\",\"status\":\"completed\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n", - "\n", - ); - stream - .write_all( - format!( - "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n{body}", - body.len() - ) - .as_bytes(), - ) - .await - .unwrap(); -} - -async fn write_compact_json_response(stream: &mut TcpStream) { - let body = r#"{"output":[]}"#; - stream - .write_all( - format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\n\r\n{body}", - body.len() - ) - .as_bytes(), - ) - .await - .unwrap(); -} diff --git a/tests/codex_gateway/usage_events.rs b/tests/codex_gateway/usage_events.rs deleted file mode 100644 index a10435693..000000000 --- a/tests/codex_gateway/usage_events.rs +++ /dev/null @@ -1,245 +0,0 @@ -use serde_json::json; - -use codex_proxy_rs::codex::gateway::transport::{ - sse::{encode_sse_event, parse_sse_events, SseEvent, MAX_SSE_EVENT_BUFFER_BYTES}, - usage_events::{extract_sse_usage, extract_usage, TokenUsage}, -}; - -const SSE_STANDARD_PARSE_GOLDEN: &str = - include_str!("../fixtures/responses/golden/sse_standard_events.json"); -const SSE_PRETTY_JSON_PARSE_GOLDEN: &str = - include_str!("../fixtures/responses/golden/sse_pretty_json_events.json"); -const NON_SSE_BODY_PARSE_GOLDEN: &str = - include_str!("../fixtures/responses/golden/non_sse_body_event.json"); - -#[test] -fn parse_sse_events_combines_multiline_data_and_metadata() { - let input = concat!( - ": keepalive\n", - "id: evt_1\n", - "event: response.output_text.delta\n", - "data: {\"delta\":\"hel\"}\n", - "data: {\"delta\":\"lo\"}\n", - "retry: 5000\n", - "\n", - ); - - let events = parse_sse_events(input).unwrap(); - - assert_eq!(events.len(), 1); - assert_eq!(events[0].id.as_deref(), Some("evt_1")); - assert_eq!( - events[0].event.as_deref(), - Some("response.output_text.delta") - ); - assert_eq!(events[0].data, "{\"delta\":\"hel\"}\n{\"delta\":\"lo\"}"); - assert_eq!(events[0].retry, Some(5000)); - assert_eq!(sse_events_json(&events), SSE_STANDARD_PARSE_GOLDEN); -} - -#[test] -fn parse_sse_events_should_keep_non_prefixed_json_continuation_lines() { - let input = concat!( - "event: error\n", - "data: {\n", - " \"error\": {\n", - " \"message\": \"bad upstream\"\n", - " }\n", - "}\n", - "\n", - ); - - let events = parse_sse_events(input).unwrap(); - - assert_eq!(events.len(), 1); - assert_eq!( - events[0].data, - "{\n \"error\": {\n \"message\": \"bad upstream\"\n }\n}" - ); - assert_eq!(sse_events_json(&events), SSE_PRETTY_JSON_PARSE_GOLDEN); -} - -#[test] -fn parse_sse_events_should_ignore_done_sentinel_without_event_type() { - let events = parse_sse_events("data: [DONE]\r\n\r\n").unwrap(); - - assert!(events.is_empty()); -} - -#[test] -fn parse_sse_events_should_convert_non_sse_json_body_to_error_event() { - let events = parse_sse_events(r#"{"error":{"message":"not an SSE stream"}}"#).unwrap(); - - assert_eq!(events.len(), 1); - assert_eq!(events[0].event.as_deref(), Some("error")); - let data: serde_json::Value = serde_json::from_str(&events[0].data).unwrap(); - assert_eq!(data["error"]["code"], "non_sse_response"); - assert_eq!(data["error"]["message"], "not an SSE stream"); - assert_eq!(sse_events_json(&events), NON_SSE_BODY_PARSE_GOLDEN); -} - -#[test] -fn parse_sse_events_should_reject_single_event_buffer_above_original_limit() { - let mut input = String::from("event: response.output_text.delta\ndata: "); - input.extend(std::iter::repeat_n('x', MAX_SSE_EVENT_BUFFER_BYTES)); - - let Err(error) = parse_sse_events(&input) else { - panic!("oversized SSE event was accepted"); - }; - - assert_eq!( - error.to_string(), - "SSE buffer exceeded 67108864 bytes — aborting stream" - ); -} - -#[test] -fn encode_sse_event_emits_event_and_data_frame() { - let encoded = encode_sse_event("error", r#"{"error":{"code":"upstream_error"}}"#); - - assert_eq!( - encoded, - "event: error\ndata: {\"error\":{\"code\":\"upstream_error\"}}\n\n" - ); -} - -#[test] -fn extract_usage_reads_codex_token_usage_shape() { - let body = json!({ - "usage": { - "input_tokens": 12, - "output_tokens": 5, - "input_tokens_details": { - "cached_tokens": 3 - } - } - }); - - let usage = extract_usage(&body).unwrap(); - - assert_eq!( - usage, - TokenUsage { - input_tokens: 12, - output_tokens: 5, - cached_tokens: 3, - image_input_tokens: 0, - image_output_tokens: 0, - total_tokens: 17, - } - ); -} - -#[test] -fn extract_usage_reads_image_generation_tool_usage_separately() { - let body = json!({ - "usage": { - "input_tokens": 12, - "output_tokens": 5, - "input_tokens_details": { - "cached_tokens": 3 - } - }, - "tool_usage": { - "image_gen": { - "input_tokens": 31, - "output_tokens": 9 - } - } - }); - - let usage = extract_usage(&body).unwrap(); - let serialized = serde_json::to_value(usage).unwrap(); - - assert_eq!(serialized["imageInputTokens"], 31); - assert_eq!(serialized["imageOutputTokens"], 9); - assert_eq!(serialized["inputTokens"], 12); - assert_eq!(serialized["outputTokens"], 5); -} - -#[test] -fn extract_usage_reads_openai_token_usage_shape_and_merges() { - let first = extract_usage(&json!({ - "usage": { - "prompt_tokens": 8, - "completion_tokens": 4, - "prompt_tokens_details": { - "cached_tokens": 2 - } - } - })) - .unwrap(); - let second = TokenUsage { - input_tokens: 1, - output_tokens: 2, - cached_tokens: 0, - image_input_tokens: 0, - image_output_tokens: 0, - total_tokens: 3, - }; - - assert_eq!( - first.merged(second), - TokenUsage { - input_tokens: 9, - output_tokens: 6, - cached_tokens: 2, - image_input_tokens: 0, - image_output_tokens: 0, - total_tokens: 15, - } - ); -} - -#[test] -fn extract_sse_usage_should_use_completed_response_usage_without_merging_earlier_usage() { - let body = concat!( - "event: response.created\n", - "data: {\"response\":{\"id\":\"resp_1\",\"usage\":{\"input_tokens\":3,\"output_tokens\":5,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n", - "event: response.completed\n", - "data: {\"response\":{\"id\":\"resp_1\",\"usage\":{\"input_tokens\":3,\"output_tokens\":5,\"input_tokens_details\":{\"cached_tokens\":1}}}}\n\n", - ); - - let usage = extract_sse_usage(body).unwrap().unwrap(); - - assert_eq!( - usage, - TokenUsage { - input_tokens: 3, - output_tokens: 5, - cached_tokens: 1, - image_input_tokens: 0, - image_output_tokens: 0, - total_tokens: 8, - } - ); -} - -#[test] -fn extract_sse_usage_should_read_completed_image_generation_tool_usage() { - let body = include_str!("../fixtures/responses/http_sse/completed_image_usage.sse"); - - let usage = extract_sse_usage(body).unwrap().unwrap(); - let serialized = serde_json::to_value(usage).unwrap(); - - assert_eq!(serialized["imageInputTokens"], 31); - assert_eq!(serialized["imageOutputTokens"], 9); -} - -fn sse_events_json(events: &[SseEvent]) -> String { - let value = serde_json::Value::Array( - events - .iter() - .map(|event| { - json!({ - "event": event.event, - "data": event.data, - "id": event.id, - "retry": event.retry, - }) - }) - .collect(), - ); - - format!("{}\n", serde_json::to_string_pretty(&value).unwrap()) -} diff --git a/tests/codex_gateway/websocket.rs b/tests/codex_gateway/websocket.rs deleted file mode 100644 index b9f265bfc..000000000 --- a/tests/codex_gateway/websocket.rs +++ /dev/null @@ -1,4983 +0,0 @@ -use std::{ - path::{Path, PathBuf}, - sync::Arc, - time::Duration, -}; - -use flate2::{Compress, Compression, FlushCompress}; -use futures::{SinkExt, StreamExt}; -use serde_json::{json, Value}; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::TcpListener, - sync::{oneshot, Mutex}, - time::sleep, -}; -use tokio_tungstenite::{ - accept_async, accept_hdr_async, - tungstenite::{ - handshake::derive_accept_key, - handshake::server::{Request as WsRequest, Response as WsResponse}, - Message, - }, -}; - -use codex_proxy_rs::codex::gateway::transport::{ - http_client::{ - build_reqwest_client, CodexBackendClient, CodexClientError, CodexRequestContext, - }, - types::CodexResponsesRequest, - websocket::{ - http_sse_fallback_allowed, transport_for_request, websocket_opening_audit_snapshot, - websocket_parity_diff, websocket_payload_audit_snapshot, - write_websocket_audit_artifact_for_dir, CodexTransport, CodexWebSocketError, - CodexWebSocketPool, CodexWebSocketPoolConfig, OpeningAuditHeader, OpeningAuditSnapshot, - PayloadAuditSnapshot, WebSocketAuditArtifact, WebSocketAuditErrorSnapshot, - }, -}; - -mod pool; - -const WS_COMPLETED_SSE_GOLDEN: &str = - include_str!("../fixtures/responses/golden/websocket_completed.sse"); - -#[test] -fn transport_for_request_should_default_to_websocket_without_history() { - let request = base_request(); - - assert_eq!( - transport_for_request(&request), - CodexTransport::WebSocketPreferred - ); - assert!(http_sse_fallback_allowed(&request)); -} - -#[test] -fn transport_for_request_should_require_websocket_without_fallback_for_previous_response_id() { - let mut request = base_request(); - request.previous_response_id = Some("resp_123".to_string()); - - assert_eq!( - transport_for_request(&request), - CodexTransport::WebSocketRequired - ); - assert!(!http_sse_fallback_allowed(&request)); -} - -#[test] -fn transport_for_request_should_allow_forced_http_sse() { - let mut request = base_request(); - request.force_http_sse = true; - - assert_eq!(transport_for_request(&request), CodexTransport::HttpSse); - assert!(http_sse_fallback_allowed(&request)); -} - -#[test] -fn transport_for_request_should_prefer_websocket_with_fallback_for_explicit_websocket_without_history( -) { - let mut request = base_request(); - request.use_websocket = true; - - assert_eq!( - transport_for_request(&request), - CodexTransport::WebSocketPreferred - ); - assert!(http_sse_fallback_allowed(&request)); -} - -#[test] -fn use_websocket_should_not_serialize_to_upstream_json() { - let mut request = base_request(); - request.use_websocket = true; - - let body = serde_json::to_value(&request).unwrap(); - - assert!(body.get("use_websocket").is_none()); - assert!(body.get("useWebSocket").is_none()); -} - -#[test] -fn websocket_opening_audit_snapshot_should_redact_sensitive_headers() { - let request = tokio_tungstenite::tungstenite::http::Request::builder() - .uri("ws://example.test:8080/codex/responses?source=audit") - .header("Sec-WebSocket-Key", "test-websocket-key") - .header("Authorization", "Bearer access-token") - .header("ChatGPT-Account-Id", "acct-secret") - .header("User-Agent", "Codex Desktop/26.609.71450") - .header("originator", "Codex Desktop") - .header("OpenAI-Beta", "responses_websockets=2026-02-06") - .header( - "x-codex-beta-features", - "terminal_resize_reflow,memories,network_proxy,prevent_idle_sleep,remote_compaction_v2", - ) - .header("x-client-request-id", "req_secret") - .header("session_id", "thread-secret") - .header("x-codex-window-id", "window-secret") - .header("x-codex-turn-metadata", r#"{"request_kind":"prewarm"}"#) - .body(()) - .unwrap(); - - let snapshot = websocket_opening_audit_snapshot(&request).unwrap(); - - assert_eq!( - snapshot.request_line, - "GET /codex/responses?source=audit HTTP/1.1" - ); - assert_eq!( - snapshot - .headers - .iter() - .map(|header| (header.name.as_str(), header.value.as_str())) - .collect::>(), - vec![ - ("Host", "example.test:8080"), - ("Connection", "Upgrade"), - ("Upgrade", "websocket"), - ("Sec-WebSocket-Version", "13"), - ("Sec-WebSocket-Key", "test-websocket-key"), - ("chatgpt-account-id", ""), - ("authorization", ""), - ("user-agent", "Codex Desktop/26.609.71450"), - ("originator", "Codex Desktop"), - ("openai-beta", "responses_websockets=2026-02-06"), - ( - "x-codex-beta-features", - "terminal_resize_reflow,memories,network_proxy,prevent_idle_sleep,remote_compaction_v2", - ), - ("x-client-request-id", ""), - ("session-id", ""), - ("thread-id", ""), - ("x-codex-window-id", ""), - ("x-codex-turn-metadata", ""), - ( - "sec-websocket-extensions", - "permessage-deflate; client_max_window_bits" - ), - ] - ); -} - -#[test] -fn websocket_payload_audit_snapshot_should_redact_user_content() { - let mut request = CodexResponsesRequest::new_http_sse( - "gpt-5.5", - "private instructions", - vec![json!({ - "role": "user", - "content": "private prompt", - })], - ); - request.previous_response_id = Some("resp_secret".to_string()); - request.reasoning = Some(json!({"effort": "medium"})); - request.tools = Some(vec![json!({ - "type": "function", - "name": "private_tool", - "description": "private schema", - })]); - request.text = Some(json!({"format": {"type": "text"}})); - request.service_tier = Some("flex".to_string()); - request.prompt_cache_key = Some("cache-secret".to_string()); - request.generate = Some(false); - request.include = Some(vec!["reasoning.encrypted_content".to_string()]); - request.client_metadata = Some(json!({ - "thread_id": "thread-secret", - "safe": "value", - })); - - let snapshot = websocket_payload_audit_snapshot(&request); - - assert_eq!( - snapshot.top_level_keys, - vec![ - "type", - "model", - "instructions", - "previous_response_id", - "input", - "tools", - "tool_choice", - "parallel_tool_calls", - "reasoning", - "store", - "stream", - "include", - "service_tier", - "prompt_cache_key", - "text", - "generate", - "client_metadata", - ] - ); - assert_eq!(snapshot.body["type"], "response.create"); - assert_eq!(snapshot.body["model"], "gpt-5.5"); - assert_eq!(snapshot.body["stream"], true); - assert_eq!(snapshot.body["instructions"], ""); - assert_eq!(snapshot.body["input"], ""); - assert_eq!(snapshot.body["previous_response_id"], ""); - assert_eq!(snapshot.body["prompt_cache_key"], ""); - assert_eq!(snapshot.body["client_metadata"], ""); - assert_eq!(snapshot.body["tools"], ""); -} - -#[test] -fn websocket_parity_diff_should_report_header_order_changes() { - let current = audit_artifact_with_header_names(&["Host", "Connection", "Authorization"]); - let reference = audit_artifact_with_header_names(&["Host", "Authorization", "Connection"]); - - let diff = websocket_parity_diff(¤t, &reference); - - assert!(diff.differences.iter().any(|difference| { - difference.path == "opening.header_order" - && difference.current == json!(["Host", "Connection", "Authorization"]) - && difference.reference == json!(["Host", "Authorization", "Connection"]) - })); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn previous_response_id_should_use_websocket_transport() { - let received_headers = Arc::new(Mutex::new(None)); - let received_request = Arc::new(Mutex::new(None)); - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let headers_for_task = Arc::clone(&received_headers); - let request_for_task = Arc::clone(&received_request); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = - accept_hdr_async(stream, |request: &WsRequest, response: WsResponse| { - let mut headers = Vec::new(); - for (name, value) in request.headers() { - let value = value.to_str().unwrap_or_default().to_string(); - headers.push((name.as_str().to_string(), value)); - } - let headers_for_callback = Arc::clone(&headers_for_task); - tokio::spawn(async move { - *headers_for_callback.lock().await = Some(headers); - }); - Ok(response) - }) - .await - .unwrap(); - let message = websocket.next().await.unwrap().unwrap(); - let text = message.into_text().unwrap(); - *request_for_task.lock().await = Some(serde_json::from_str::(&text).unwrap()); - websocket - .send(Message::Text( - json!({ - "type": "response.completed", - "response": { - "id": "resp_ws", - "object": "response", - "usage": { - "input_tokens": 2, - "output_tokens": 1, - "total_tokens": 3 - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "be brief", Vec::new()); - request.previous_response_id = Some("resp_prev".to_string()); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &request, - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_ws", - 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, - installation_id: None, - session_id: None, - }, - ) - .await - .unwrap(); - - server.await.unwrap(); - let request = received_request.lock().await.clone().unwrap(); - assert_eq!(request["type"], "response.create"); - assert_eq!(request["model"], "gpt-5.5"); - assert_eq!(request["instructions"], "be brief"); - assert_eq!(request["previous_response_id"], "resp_prev"); - assert_eq!(request["stream"], true); - assert!(response.body.contains("event: response.completed")); - assert!(response.body.contains("\"id\":\"resp_ws\"")); - assert_eq!(response.usage.unwrap().input_tokens, 2); - let headers = received_headers.lock().await.clone().unwrap(); - assert!(headers - .iter() - .any(|(name, value)| { name == "authorization" && value == "Bearer access-token" })); -} - -#[tokio::test] -async fn websocket_audit_artifact_should_require_explicit_gate() { - let request = tokio_tungstenite::tungstenite::http::Request::builder() - .uri("ws://example.test/codex/responses") - .header("Sec-WebSocket-Key", "test-websocket-key") - .header("Authorization", "Bearer access-token") - .header("ChatGPT-Account-Id", "chatgpt-account") - .header("Cookie", "session=secret") - .body(()) - .unwrap(); - let mut payload_request = CodexResponsesRequest::new_http_sse( - "gpt-5.5", - "private instructions", - vec![json!({"role": "user", "content": "private prompt"})], - ); - payload_request.prompt_cache_key = Some("cache-secret".to_string()); - payload_request.client_metadata = Some(json!({"thread_id": "thread-secret"})); - let artifact = WebSocketAuditArtifact { - transport_mode: "websocket_preferred".to_string(), - fallback_allowed: true, - opening: Some(websocket_opening_audit_snapshot(&request).unwrap()), - payload: Some(websocket_payload_audit_snapshot(&payload_request)), - error: Some(WebSocketAuditErrorSnapshot { - classification: "opening_failed".to_string(), - message: "connection refused".to_string(), - }), - }; - let dir = tempfile::tempdir().unwrap(); - - let disabled = write_websocket_audit_artifact_for_dir(None, &artifact) - .await - .unwrap(); - - assert!(disabled.is_none()); - assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 0); - - let written = write_websocket_audit_artifact_for_dir(Some(dir.path()), &artifact) - .await - .unwrap() - .unwrap(); - let body = tokio::fs::read_to_string(written).await.unwrap(); - assert!(body.contains("websocket_preferred")); - assert!(body.contains("response.create")); - assert!(!body.contains("access-token")); - assert!(!body.contains("session=secret")); - assert!(!body.contains("acct-secret")); - assert!(!body.contains("private prompt")); - assert!(!body.contains("private instructions")); - assert!(!body.contains("cache-secret")); - assert!(!body.contains("thread-secret")); -} - -#[tokio::test] -async fn websocket_capture_harness_should_record_opening_bytes() { - let artifact_path = PathBuf::from(".codex-ws-audit/rs-current-capture.json"); - let _ = std::fs::remove_file(&artifact_path); - let mut request = CodexResponsesRequest::new_http_sse( - "gpt-5.5", - "private capture instructions", - vec![json!({ - "role": "user", - "content": "private capture prompt", - })], - ); - request.prompt_cache_key = Some("capture-cache-secret".to_string()); - request.client_metadata = Some(json!({ - "thread_id": "capture-thread-secret", - "safe": "capture", - })); - request.generate = Some(false); - - let context = CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_ws_capture", - turn_state: None, - turn_metadata: Some( - r#"{"installation_id":"install-1","session_id":"session-1","thread_id":"session-1","turn_id":"","window_id":"session-1:0","request_kind":"prewarm","sandbox":"seccomp"}"#, - ), - beta_features: Some( - "terminal_resize_reflow,memories,network_proxy,prevent_idle_sleep,remote_compaction_v2", - ), - include_timing_metrics: None, - version: None, - codex_window_id: Some("session-1:0"), - parent_thread_id: None, - cookie_header: None, - installation_id: Some("install-1"), - session_id: Some("session-1"), - }; - - let capture = capture_codex_websocket_exchange(&request, context, &artifact_path).await; - - assert!(capture - .opening_bytes - .starts_with("GET /codex/responses HTTP/1.1\r\n")); - assert_headers_appear_in_order( - &capture.opening_bytes, - &[ - "Host: ", - "Connection: Upgrade\r\n", - "Upgrade: websocket\r\n", - "Sec-WebSocket-Version: 13\r\n", - "Sec-WebSocket-Key: ", - "chatgpt-account-id: chatgpt-account\r\n", - "authorization: Bearer access-token\r\n", - "user-agent: Codex Desktop/26.519.81530 (darwin; arm64)\r\n", - "originator: Codex Desktop\r\n", - "openai-beta: responses_websockets=2026-02-06\r\n", - "x-codex-beta-features: terminal_resize_reflow,memories,network_proxy,prevent_idle_sleep,remote_compaction_v2\r\n", - "x-client-request-id: session-1\r\n", - "session-id: session-1\r\n", - "thread-id: session-1\r\n", - "x-codex-window-id: session-1:0\r\n", - "x-codex-turn-metadata: {\"installation_id\":\"install-1\",\"session_id\":\"session-1\",\"thread_id\":\"session-1\",\"turn_id\":\"\",\"window_id\":\"session-1:0\",\"request_kind\":\"prewarm\",\"sandbox\":\"seccomp\"}\r\n", - "sec-websocket-extensions: permessage-deflate; client_max_window_bits\r\n", - ], - ); - assert!( - capture.sec_websocket_key_offset < capture.authorization_header_offset, - "Sec-WebSocket-Key should appear before business headers in:\n{}", - capture.opening_bytes - ); - let payload = serde_json::from_str::(&capture.first_frame_text).unwrap(); - assert_eq!(payload["type"], "response.create"); - assert_eq!(payload["instructions"], "private capture instructions"); - assert_eq!(payload["input"][0]["content"], "private capture prompt"); - assert_eq!(payload["tools"], json!([])); - assert_eq!(payload["reasoning"], Value::Null); - assert_eq!(payload["include"], json!([])); - assert_eq!(payload["generate"], false); - assert_substrings_appear_in_order( - &capture.first_frame_text, - &[ - "\"type\":\"response.create\"", - "\"model\":\"gpt-5.5\"", - "\"instructions\":\"private capture instructions\"", - "\"input\":", - "\"tools\":[]", - "\"tool_choice\":\"auto\"", - "\"parallel_tool_calls\":true", - "\"reasoning\":null", - "\"store\":false", - "\"stream\":true", - "\"include\":[]", - "\"prompt_cache_key\":\"session-1\"", - "\"generate\":false", - "\"client_metadata\":", - ], - ); - - let artifact = tokio::fs::read_to_string(&artifact_path).await.unwrap(); - assert!(artifact.contains("websocket_preferred")); - assert!(artifact.contains("response.create")); - assert!(!artifact.contains("private capture instructions")); - assert!(!artifact.contains("private capture prompt")); - assert!(!artifact.contains("capture-cache-secret")); - assert!(!artifact.contains("capture-thread-secret")); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn ordinary_response_should_use_websocket_transport_by_default() { - let received_request = Arc::new(Mutex::new(None)); - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let request_for_task = Arc::clone(&received_request); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = - accept_hdr_async(stream, |_request: &WsRequest, response: WsResponse| { - Ok(response) - }) - .await - .unwrap(); - let message = websocket.next().await.unwrap().unwrap(); - let text = message.into_text().unwrap(); - *request_for_task.lock().await = Some(serde_json::from_str::(&text).unwrap()); - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_default", 4, 2).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let request = base_request(); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &request, - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_ws_default", - 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, - installation_id: None, - session_id: None, - }, - ) - .await - .unwrap(); - - server.await.unwrap(); - let request = received_request.lock().await.clone().unwrap(); - assert_eq!(request["type"], "response.create"); - assert_eq!(request["model"], "gpt-5.5"); - assert!(request.get("previous_response_id").is_none()); - assert_eq!( - response.body, - with_sse_terminal_separator(WS_COMPLETED_SSE_GOLDEN) - ); - assert!(response.body.contains("event: response.completed")); - assert!(response.body.contains("\"id\":\"resp_ws_default\"")); - assert_eq!(response.usage.unwrap().input_tokens, 4); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_request_should_forward_security_chain_body_fields() { - let received_headers = Arc::new(Mutex::new(None)); - let received_request = Arc::new(Mutex::new(None)); - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let headers_for_task = Arc::clone(&received_headers); - let request_for_task = Arc::clone(&received_request); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = - accept_hdr_async(stream, |request: &WsRequest, response: WsResponse| { - let mut headers = Vec::new(); - for (name, value) in request.headers() { - let value = value.to_str().unwrap_or_default().to_string(); - headers.push((name.as_str().to_string(), value)); - } - let headers_for_callback = Arc::clone(&headers_for_task); - tokio::spawn(async move { - *headers_for_callback.lock().await = Some(headers); - }); - Ok(response) - }) - .await - .unwrap(); - let message = websocket.next().await.unwrap().unwrap(); - let text = message.into_text().unwrap(); - *request_for_task.lock().await = Some(serde_json::from_str::(&text).unwrap()); - websocket - .send(Message::Text( - json!({ - "type": "response.completed", - "response": { - "id": "resp_ws_security", - "object": "response" - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let mut request = base_request(); - request.use_websocket = true; - request.prompt_cache_key = Some("client-thread".to_string()); - request.client_metadata = Some(json!({ - "safe": "yes", - "x-openai-subagent": "review" - })); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - client - .create_response( - &request, - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_ws_security", - turn_state: Some("turn-state"), - turn_metadata: Some("{\"thread_source\":\"subagent\"}"), - beta_features: Some("feature-a"), - include_timing_metrics: Some("true"), - version: Some("26.318.11754"), - codex_window_id: Some("cw_derived"), - parent_thread_id: Some("parent-456"), - cookie_header: None, - installation_id: Some("install-123"), - session_id: Some("cp_derived"), - }, - ) - .await - .unwrap(); - - server.await.unwrap(); - let request = received_request.lock().await.clone().unwrap(); - assert_eq!(request["prompt_cache_key"], "cp_derived"); - assert_eq!( - request["client_metadata"], - json!({ - "safe": "yes", - "x-openai-subagent": "review", - "x-codex-installation-id": "install-123", - "x-codex-window-id": "cw_derived", - "x-codex-turn-metadata": "{\"thread_source\":\"subagent\"}", - "x-codex-parent-thread-id": "parent-456" - }) - ); - let headers = received_headers.lock().await.clone().unwrap(); - assert!(headers - .iter() - .any(|(name, value)| { name == "x-client-request-id" && value == "cp_derived" })); - assert!(headers - .iter() - .any(|(name, value)| { name == "x-openai-subagent" && value == "review" })); - assert!(headers.iter().all(|(name, _)| name != "content-type")); - assert!(headers.iter().all(|(name, _)| name != "accept")); -} - -#[tokio::test] -async fn websocket_handshake_429_should_surface_as_upstream_error_before_body_is_sent() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - 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("authorization: Bearer access-token")); - let body = r#"{"error":{"message":"rate limited"}}"#; - let response = format!( - "HTTP/1.1 429 Too Many Requests\r\nretry-after: 33\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", - body.len() - ); - stream.write_all(response.as_bytes()).await.unwrap(); - }); - - let mut request = base_request(); - request.use_websocket = true; - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &request, - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_ws_429", - 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, - installation_id: None, - session_id: None, - }, - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { - status, - retry_after_seconds, - body, - } = error - else { - panic!("expected upstream error, found {error:?}"); - }; - assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS); - assert_eq!(retry_after_seconds, Some(33)); - assert!(body.contains("rate limited")); -} - -#[tokio::test] -async fn websocket_handshake_should_offer_original_permessage_deflate_extension() { - 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 (mut stream, _) = listener.accept().await.unwrap(); - let request = read_http_upgrade_request(&mut stream).await; - request_tx.send(request).unwrap(); - stream - .write_all( - b"HTTP/1.1 400 Bad Request\r\ncontent-length: 0\r\nconnection: close\r\n\r\n", - ) - .await - .unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - let result = client - .websocket_stream_response(&base_request(), request_context("req_ws_extensions", None)) - .await; - assert!(result.is_err()); - - server.await.unwrap(); - let raw_request = request_rx.await.unwrap(); - assert_headers_appear_in_order( - &raw_request, - &[ - "Host: ", - "Connection: Upgrade\r\n", - "Upgrade: websocket\r\n", - "Sec-WebSocket-Version: 13\r\n", - "Sec-WebSocket-Key: ", - "chatgpt-account-id: chatgpt-account\r\n", - "authorization: Bearer access-token\r\n", - "user-agent: Codex Desktop/26.519.81530 (darwin; arm64)\r\n", - "originator: Codex Desktop\r\n", - "openai-beta: responses_websockets=2026-02-06\r\n", - "x-client-request-id: req_ws_extensions\r\n", - "sec-websocket-extensions: permessage-deflate; client_max_window_bits\r\n", - ], - ); -} - -#[tokio::test] -async fn websocket_should_decode_permessage_deflate_response_frame_when_server_accepts_extension() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.unwrap(); - let request = read_http_upgrade_request(&mut stream).await; - let accept_key = websocket_accept_key(&request); - let response = format!( - "HTTP/1.1 101 Switching Protocols\r\n\ - Connection: Upgrade\r\n\ - Upgrade: websocket\r\n\ - Sec-WebSocket-Accept: {accept_key}\r\n\ - Sec-WebSocket-Extensions: permessage-deflate\r\n\ - \r\n" - ); - stream.write_all(response.as_bytes()).await.unwrap(); - let _client_request = read_client_websocket_frame(&mut stream).await; - let payload = websocket_completed_response("resp_ws_deflate", 5, 2); - let compressed_frame = compressed_server_text_frame(&payload); - stream.write_all(&compressed_frame).await.unwrap(); - sleep(Duration::from_millis(50)).await; - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - let mut request = base_request(); - request.previous_response_id = Some("resp_deflate_prev".to_string()); - - let response = client - .create_response(&request, request_context("req_ws_deflate", None)) - .await - .unwrap(); - - server.await.unwrap(); - assert!(response.body.contains("event: response.completed")); - assert!(response.body.contains("\"id\":\"resp_ws_deflate\"")); - assert_eq!(response.usage.unwrap().input_tokens, 5); -} - -#[tokio::test] -async fn websocket_should_reply_to_server_ping_before_completed_event() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.unwrap(); - let request = read_http_upgrade_request(&mut stream).await; - let accept_key = websocket_accept_key(&request); - let response = format!( - "HTTP/1.1 101 Switching Protocols\r\n\ - Connection: Upgrade\r\n\ - Upgrade: websocket\r\n\ - Sec-WebSocket-Accept: {accept_key}\r\n\ - \r\n" - ); - stream.write_all(response.as_bytes()).await.unwrap(); - let _client_request = read_client_websocket_frame(&mut stream).await; - - stream - .write_all(&server_ping_frame(b"codex-ping")) - .await - .unwrap(); - let (opcode, payload) = tokio::time::timeout( - Duration::from_secs(1), - read_client_websocket_frame_with_opcode(&mut stream), - ) - .await - .expect("client should reply to server ping before terminal response"); - assert_eq!(opcode, 0x0a); - assert_eq!(payload, b"codex-ping"); - - let payload = websocket_completed_response("resp_ws_pong", 6, 3); - stream - .write_all(&server_text_frame(payload.as_bytes())) - .await - .unwrap(); - sleep(Duration::from_millis(50)).await; - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response(&base_request(), request_context("req_ws_ping", None)) - .await - .unwrap(); - - server.await.unwrap(); - assert!(response.body.contains("event: response.completed")); - assert!(response.body.contains("\"id\":\"resp_ws_pong\"")); - assert_eq!(response.usage.unwrap().input_tokens, 6); -} - -#[tokio::test] -async fn websocket_invalid_text_frame_should_be_ignored_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.unwrap(); - let request = read_http_upgrade_request(&mut stream).await; - let accept_key = websocket_accept_key(&request); - let response = format!( - "HTTP/1.1 101 Switching Protocols\r\n\ - Connection: Upgrade\r\n\ - Upgrade: websocket\r\n\ - Sec-WebSocket-Accept: {accept_key}\r\n\ - \r\n" - ); - stream.write_all(response.as_bytes()).await.unwrap(); - let _client_request = read_client_websocket_frame(&mut stream).await; - stream - .write_all(&server_text_frame(b"not-json-from-upstream")) - .await - .unwrap(); - let payload = websocket_completed_response("resp_ws_after_invalid", 7, 4); - stream - .write_all(&server_text_frame(payload.as_bytes())) - .await - .unwrap(); - sleep(Duration::from_millis(50)).await; - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_invalid_text", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("not-json-from-upstream")); - assert!(response.body.contains("event: response.completed")); - assert!(response.body.contains("\"id\":\"resp_ws_after_invalid\"")); - assert_eq!(response.usage.unwrap().input_tokens, 7); -} - -#[tokio::test(start_paused = true)] -async fn websocket_silent_upstream_should_timeout_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let (request_seen_tx, request_seen_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(); - request_seen_tx.send(()).unwrap(); - futures::future::pending::<()>().await; - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - let response_task = tokio::spawn(async move { - client - .create_response( - &base_request(), - request_context("req_ws_silent_timeout", None), - ) - .await - }); - - request_seen_rx.await.unwrap(); - tokio::task::yield_now().await; - tokio::time::advance(Duration::from_secs(20)).await; - tokio::task::yield_now().await; - - let error = response_task.await.unwrap().unwrap_err(); - server.abort(); - let CodexClientError::WebSocket(CodexWebSocketError::ReceiveIdleTimeout { timeout }) = error - else { - panic!("expected receive idle timeout, found {error:?}"); - }; - assert_eq!(timeout, Duration::from_secs(20)); -} - -#[tokio::test] -async fn websocket_binary_response_frame_should_error_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.unwrap(); - let request = read_http_upgrade_request(&mut stream).await; - let accept_key = websocket_accept_key(&request); - let response = format!( - "HTTP/1.1 101 Switching Protocols\r\n\ - Connection: Upgrade\r\n\ - Upgrade: websocket\r\n\ - Sec-WebSocket-Accept: {accept_key}\r\n\ - \r\n" - ); - stream.write_all(response.as_bytes()).await.unwrap(); - let _client_request = read_client_websocket_frame(&mut stream).await; - let payload = websocket_completed_response("resp_ws_binary", 6, 3); - stream - .write_all(&server_binary_frame(payload.as_bytes())) - .await - .unwrap(); - sleep(Duration::from_millis(50)).await; - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response(&base_request(), request_context("req_ws_binary", None)) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::WebSocket(CodexWebSocketError::UnexpectedBinaryEvent) = error else { - panic!("expected unexpected binary websocket event, found {error:?}"); - }; -} - -#[tokio::test] -async fn websocket_close_frame_before_first_event_should_error_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.unwrap(); - let request = read_http_upgrade_request(&mut stream).await; - let accept_key = websocket_accept_key(&request); - let response = format!( - "HTTP/1.1 101 Switching Protocols\r\n\ - Connection: Upgrade\r\n\ - Upgrade: websocket\r\n\ - Sec-WebSocket-Accept: {accept_key}\r\n\ - \r\n" - ); - stream.write_all(response.as_bytes()).await.unwrap(); - let _client_request = read_client_websocket_frame(&mut stream).await; - stream.write_all(&server_close_frame()).await.unwrap(); - sleep(Duration::from_millis(50)).await; - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response(&base_request(), request_context("req_ws_close", None)) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::WebSocket(CodexWebSocketError::ClosedByServerBeforeCompleted) = error - else { - panic!("expected server close before response.completed, found {error:?}"); - }; -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_first_error_frame_should_surface_as_upstream_error_without_http_fallback() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = - accept_hdr_async(stream, |request: &WsRequest, response: WsResponse| { - assert_eq!( - request - .headers() - .get("authorization") - .and_then(|value| value.to_str().ok()), - Some("Bearer access-token") - ); - Ok(response) - }) - .await - .unwrap(); - let _message = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - json!({ - "type": "response.failed", - "response": { - "id": "resp_limit", - "error": { - "code": "usage_limit_reached", - "message": "weekly limit reached", - "resets_in_seconds": 45 - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let request = base_request(); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &request, - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_ws_error_frame", - 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, - installation_id: None, - session_id: None, - }, - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { - status, - retry_after_seconds, - body, - } = error - else { - panic!("expected upstream error, found {error:?}"); - }; - assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS); - assert_eq!(retry_after_seconds, Some(45)); - assert!(body.contains("usage_limit_reached")); -} - -#[tokio::test] -async fn websocket_malformed_completed_response_should_error_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.completed", - "response": { - "object": "response", - "usage": { - "input_tokens": 1, - "output_tokens": 1 - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_malformed_completed", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - assert!( - error - .to_string() - .contains("failed to parse ResponseCompleted"), - "unexpected error: {error:?}" - ); -} - -#[tokio::test] -async fn websocket_completed_response_with_incomplete_usage_should_error_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.completed", - "response": { - "id": "resp_incomplete_usage", - "object": "response", - "usage": { - "input_tokens": 1, - "output_tokens": 1 - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_incomplete_completed_usage", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - assert!( - error - .to_string() - .contains("failed to parse ResponseCompleted"), - "unexpected error: {error:?}" - ); -} - -#[tokio::test] -async fn websocket_completed_without_response_should_not_finish_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.completed" - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_completed_without_response", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::WebSocket(CodexWebSocketError::ClosedByServerBeforeCompleted) = error - else { - panic!( - "expected response-less completed frame to be ignored before close, found {error:?}" - ); - }; -} - -#[tokio::test] -async fn websocket_created_without_response_should_be_ignored_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.created" - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_created_without_response", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_created_without_response", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.created")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_output_text_delta_without_delta_should_be_ignored_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.output_text.delta" - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_output_text_delta_without_delta", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_output_text_delta_without_delta", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_text.delta")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_delta_events_missing_required_fields_should_be_ignored_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.custom_tool_call_input.delta", - "delta": "partial input" - }), - json!({ - "type": "response.reasoning_summary_text.delta", - "delta": "summary" - }), - json!({ - "type": "response.reasoning_text.delta", - "delta": "reasoning" - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_delta_missing_required_fields", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_delta_missing_required_fields", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response - .body - .contains("event: response.custom_tool_call_input.delta")); - assert!(!response - .body - .contains("event: response.reasoning_summary_text.delta")); - assert!(!response - .body - .contains("event: response.reasoning_text.delta")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_output_item_events_without_item_should_be_ignored_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done" - }), - json!({ - "type": "response.output_item.added" - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_output_item_without_item", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_output_item_without_item", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_output_item_events_with_non_object_item_should_be_ignored_like_official_client() -{ - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": 123 - }), - json!({ - "type": "response.output_item.added", - "item": 123 - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_output_item_non_object_item", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_output_item_non_object_item", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_output_item_events_with_invalid_item_type_tag_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "role": "assistant", - "content": [] - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": 123, - "role": "assistant", - "content": [] - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_output_item_invalid_type_tag", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_output_item_invalid_type_tag", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_output_item_events_with_invalid_metadata_should_be_ignored_like_official_client() -{ - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "message", - "role": "assistant", - "content": [], - "metadata": [] - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "web_search_call", - "metadata": { - "turn_id": {} - } - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_output_item_invalid_metadata", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_output_item_invalid_metadata", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_message_output_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "message", - "content": [] - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "message", - "role": "assistant", - "content": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_message_output_item_invalid_shape", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_message_output_item_invalid_shape", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_message_output_item_events_with_invalid_optional_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "message", - "id": {}, - "role": "assistant", - "content": [] - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "message", - "role": "assistant", - "content": [], - "phase": "draft" - } - }), - json!({ - "type": "response.output_item.done", - "item": { - "type": "message", - "role": "assistant", - "content": [], - "phase": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_message_output_item_invalid_optional", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_message_output_item_invalid_optional", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_message_output_item_events_with_invalid_content_items_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "message", - "role": "assistant", - "content": [{"type": "output_text"}] - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "message", - "role": "assistant", - "content": [{ - "type": "input_image", - "image_url": {}, - "detail": "thumbnail" - }] - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_message_content_item_invalid", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_message_content_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_agent_message_output_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "agent_message", - "recipient": "user", - "content": [] - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "agent_message", - "author": "assistant", - "recipient": {}, - "content": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_agent_message_item_invalid", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_agent_message_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_agent_message_output_item_events_with_invalid_content_items_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "agent_message", - "author": "assistant", - "recipient": "user", - "content": [{"type": "input_text"}] - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "agent_message", - "author": "assistant", - "recipient": "user", - "content": [{ - "type": "encrypted_content", - "encrypted_content": {} - }] - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_agent_message_content_item_invalid", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_agent_message_content_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_reasoning_output_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "reasoning", - "encrypted_content": "enc" - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "reasoning", - "id": {}, - "summary": [], - "content": {}, - "encrypted_content": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_reasoning_item_invalid", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_reasoning_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_reasoning_output_item_events_with_invalid_nested_items_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "reasoning", - "summary": [{"type": "summary_text"}] - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "reasoning", - "summary": [], - "content": [{ - "type": "reasoning_text", - "text": {} - }] - } - }), - json!({ - "type": "response.output_item.done", - "item": { - "type": "reasoning", - "summary": [], - "content": [{"type": "unsupported"}] - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_reasoning_nested_item_invalid", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_reasoning_nested_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_function_call_output_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "function_call", - "arguments": "{}", - "call_id": "call_ws_invalid_function" - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "function_call", - "name": "lookup", - "arguments": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_function_call_output_item_invalid", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_function_call_output_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_function_call_output_item_events_with_invalid_optional_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "function_call", - "id": {}, - "name": "lookup", - "arguments": "{}", - "call_id": "call_ws_invalid_function_optional" - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "function_call", - "name": "lookup", - "namespace": {}, - "arguments": "{}", - "call_id": "call_ws_invalid_function_optional" - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response( - "resp_ws_function_call_output_item_invalid_optional", - 2, - 1, - ) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_function_call_output_item_invalid_optional", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_tool_search_call_output_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "tool_search_call", - "arguments": { - "query": "calendar create" - } - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "tool_search_call", - "execution": "client" - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_tool_search_call_output_item_invalid", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_tool_search_call_output_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_tool_search_call_output_item_events_with_invalid_optional_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "tool_search_call", - "id": {}, - "execution": "client", - "arguments": { - "query": "calendar create" - } - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "tool_search_call", - "call_id": {}, - "status": {}, - "execution": "client", - "arguments": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response( - "resp_ws_tool_search_call_output_item_invalid_optional", - 2, - 1, - ) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_tool_search_call_output_item_invalid_optional", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_function_call_output_payload_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "function_call_output", - "output": "ok" - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "function_call_output", - "call_id": "call_ws_invalid_function_output", - "output": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_function_call_output_item_invalid", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_function_call_output_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_custom_tool_call_output_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "custom_tool_call", - "name": "render", - "input": "{}" - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "custom_tool_call", - "call_id": "call_ws_invalid_custom_tool", - "name": "render", - "input": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_custom_tool_call_item_invalid", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_custom_tool_call_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_custom_tool_call_output_item_events_with_invalid_optional_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "custom_tool_call", - "id": {}, - "call_id": "call_ws_invalid_custom_tool_optional", - "name": "render", - "input": "{}" - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "custom_tool_call", - "status": {}, - "call_id": "call_ws_invalid_custom_tool_optional", - "name": "render", - "input": "{}" - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response( - "resp_ws_custom_tool_call_item_invalid_optional", - 2, - 1, - ) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_custom_tool_call_item_invalid_optional", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_custom_tool_call_output_payload_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "custom_tool_call_output", - "output": "ok" - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "custom_tool_call_output", - "call_id": "call_ws_invalid_custom_tool_output", - "output": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_custom_tool_call_output_item_invalid", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_custom_tool_call_output_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_custom_tool_call_output_result_item_events_with_invalid_optional_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "custom_tool_call_output", - "call_id": "call_ws_invalid_custom_tool_output_optional", - "name": {}, - "output": "ok" - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "custom_tool_call_output", - "call_id": "call_ws_invalid_custom_tool_output_optional", - "name": [], - "output": "ok" - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response( - "resp_ws_custom_tool_call_output_item_invalid_optional", - 2, - 1, - ) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_custom_tool_call_output_item_invalid_optional", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_function_output_payload_item_events_with_invalid_content_items_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "function_call_output", - "call_id": "call_ws_invalid_function_output_content", - "output": [{"type": "input_text"}] - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "custom_tool_call_output", - "call_id": "call_ws_invalid_custom_tool_output_content", - "output": [{ - "type": "input_image", - "image_url": "data:image/png;base64,aaa", - "detail": "full" - }] - } - }), - json!({ - "type": "response.output_item.done", - "item": { - "type": "custom_tool_call_output", - "call_id": "call_ws_invalid_custom_tool_output_encrypted", - "output": [{ - "type": "encrypted_content", - "encrypted_content": {} - }] - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_function_output_content_item_invalid", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_function_output_content_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_tool_search_output_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "tool_search_output", - "execution": "client", - "tools": [] - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "tool_search_output", - "status": "completed", - "execution": "client", - "tools": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_tool_search_output_item_invalid", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_tool_search_output_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_tool_search_output_item_events_with_invalid_optional_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.output_item.done", - "item": { - "type": "tool_search_output", - "call_id": {}, - "status": "completed", - "execution": "server", - "tools": [] - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket - .send(Message::Text( - websocket_completed_response( - "resp_ws_tool_search_output_item_optional_invalid", - 2, - 1, - ) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_tool_search_output_item_optional_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_local_shell_call_output_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "local_shell_call", - "action": { - "type": "exec", - "command": ["pwd"] - } - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "local_shell_call", - "status": "completed", - "action": { - "type": "exec", - "command": "pwd" - } - } - }), - json!({ - "type": "response.output_item.done", - "item": { - "type": "local_shell_call", - "call_id": {}, - "status": "done", - "action": { - "type": "exec", - "command": ["pwd"], - "timeout_ms": "1000", - "env": { - "PATH": {} - } - } - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_local_shell_call_item_invalid", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_local_shell_call_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_web_search_call_output_item_events_with_invalid_optional_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "web_search_call", - "status": {}, - "action": { - "type": "search", - "query": "weather" - } - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "web_search_call", - "status": "completed", - "action": { - "type": "search", - "query": {}, - "queries": ["weather", {}] - } - } - }), - json!({ - "type": "response.output_item.done", - "item": { - "type": "web_search_call", - "id": {}, - "action": { - "type": "find_in_page", - "url": "https://example.com", - "pattern": {} - } - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_web_search_call_item_invalid", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_web_search_call_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_image_generation_call_output_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "image_generation_call", - "status": "completed", - "result": "Zm9v" - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "image_generation_call", - "id": "ig_ws_invalid", - "status": "completed", - "result": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_image_generation_call_item_invalid", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_image_generation_call_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_image_generation_call_output_item_events_with_invalid_optional_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.output_item.done", - "item": { - "type": "image_generation_call", - "id": "ig_ws_optional_invalid", - "status": "completed", - "revised_prompt": {}, - "result": "Zm9v" - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket - .send(Message::Text( - websocket_completed_response( - "resp_ws_image_generation_call_item_optional_invalid", - 2, - 1, - ) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_image_generation_call_item_optional_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_compaction_output_item_events_with_invalid_required_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "compaction" - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "compaction_summary", - "encrypted_content": {} - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_compaction_item_invalid", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_compaction_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_context_compaction_output_item_events_with_invalid_optional_fields_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.output_item.done", - "item": { - "type": "context_compaction", - "encrypted_content": {} - } - }), - json!({ - "type": "response.output_item.added", - "item": { - "type": "context_compaction", - "encrypted_content": [] - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_context_compaction_item_invalid", 2, 1) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_context_compaction_item_invalid", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_reasoning_summary_part_added_without_summary_index_should_be_ignored_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.reasoning_summary_part.added" - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket - .send(Message::Text( - websocket_completed_response( - "resp_ws_reasoning_summary_part_added_without_summary_index", - 2, - 1, - ) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context( - "req_ws_reasoning_summary_part_added_without_summary_index", - None, - ), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response - .body - .contains("event: response.reasoning_summary_part.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_null_option_fields_should_be_ignored_like_missing_fields_in_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - for event in [ - json!({ - "type": "response.created", - "response": null - }), - json!({ - "type": "response.output_text.delta", - "delta": null - }), - json!({ - "type": "response.output_item.done", - "item": null - }), - json!({ - "type": "response.output_item.added", - "item": null - }), - json!({ - "type": "response.completed", - "response": null - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .unwrap(); - } - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_null_option_fields", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_null_options", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.created")); - assert!(!response.body.contains("event: response.output_text.delta")); - assert!(!response.body.contains("event: response.output_item.done")); - assert!(!response.body.contains("event: response.output_item.added")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_frames_with_invalid_official_event_shape_should_be_ignored_like_official_client() -{ - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.output_text.delta", - "delta": 123 - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_invalid_event_shape", 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &base_request(), - request_context("req_ws_invalid_event_shape", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("event: response.output_text.delta")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -async fn websocket_response_failed_rate_limit_message_should_surface_retry_after_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.failed", - "response": { - "id": "resp_rate_limit_delay", - "error": { - "code": "rate_limit_exceeded", - "message": "Rate limit reached. Please try again in 11.054s." - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_rate_limit_message_delay", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { - status, - retry_after_seconds, - body, - } = error - else { - panic!("expected rate-limit response.failed to surface as upstream error, found {error:?}"); - }; - assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS); - assert_eq!(retry_after_seconds, Some(12)); - assert!(body.contains("Please try again in 11.054s")); -} - -#[tokio::test] -async fn websocket_response_failed_retry_after_message_should_require_rate_limit_code_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.failed", - "response": { - "id": "resp_non_rate_limit_delay", - "error": { - "code": "upstream_transient_error", - "message": "Try again in 35 seconds." - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_non_rate_limit_message_delay", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { - status, - retry_after_seconds, - body, - } = error - else { - panic!( - "expected non-rate-limit response.failed to surface as upstream error, found {error:?}" - ); - }; - assert_eq!(status, reqwest::StatusCode::SERVICE_UNAVAILABLE); - assert_eq!(retry_after_seconds, None); - assert!(body.contains("Try again in 35 seconds")); -} - -#[tokio::test] -async fn websocket_response_failed_server_overloaded_should_surface_as_503_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.failed", - "response": { - "id": "resp_overloaded", - "error": { - "code": "server_is_overloaded", - "message": "The server is overloaded" - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_server_overloaded", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { status, body, .. } = error else { - panic!("expected server overload response.failed to surface as upstream error, found {error:?}"); - }; - assert_eq!(status, reqwest::StatusCode::SERVICE_UNAVAILABLE); - assert!(body.contains("server_is_overloaded")); - assert!(body.contains("The server is overloaded")); -} - -#[tokio::test] -async fn websocket_unknown_response_failed_should_surface_as_503_like_official_retryable() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.failed", - "response": { - "id": "resp_unknown_failure", - "error": { - "code": "upstream_transient_error", - "message": "Try again shortly" - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_unknown_failed", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { status, body, .. } = error else { - panic!("expected unknown response.failed to surface as upstream error, found {error:?}"); - }; - assert_eq!(status, reqwest::StatusCode::SERVICE_UNAVAILABLE); - assert!(body.contains("upstream_transient_error")); - assert!(body.contains("Try again shortly")); -} - -#[tokio::test] -async fn websocket_response_failed_special_codes_should_use_official_status_classes() { - let cases = [ - ("insufficient_quota", reqwest::StatusCode::PAYMENT_REQUIRED), - ("quota_exceeded", reqwest::StatusCode::PAYMENT_REQUIRED), - ("context_length_exceeded", reqwest::StatusCode::BAD_REQUEST), - ("invalid_prompt", reqwest::StatusCode::BAD_REQUEST), - ("cyber_policy", reqwest::StatusCode::BAD_REQUEST), - ("usage_not_included", reqwest::StatusCode::TOO_MANY_REQUESTS), - ]; - - for (code, expected_status) in cases { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.failed", - "response": { - "id": format!("resp_{code}"), - "error": { - "code": code, - "message": format!("special failure {code}") - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - let error = client - .create_response( - &base_request(), - request_context("req_ws_special_failed", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { status, body, .. } = error else { - panic!("expected {code} to surface as upstream error, found {error:?}"); - }; - assert_eq!(status, expected_status, "unexpected status for {code}"); - assert!(body.contains(code)); - } -} - -#[tokio::test] -async fn websocket_wrapped_error_status_should_surface_as_upstream_error_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "error", - "status": 400, - "error": { - "type": "invalid_request_error", - "message": "Model does not support this request" - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_wrapped_error", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { - status, - retry_after_seconds, - body, - } = error - else { - panic!("expected wrapped websocket error to surface as upstream error, found {error:?}"); - }; - assert_eq!(status, reqwest::StatusCode::BAD_REQUEST); - assert_eq!(retry_after_seconds, None); - assert!(body.contains("invalid_request_error")); - assert!(body.contains("Model does not support this request")); -} - -#[tokio::test] -async fn websocket_unmapped_success_status_error_should_not_return_successful_error_event_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "error", - "status": 200, - "error": { - "type": "non_terminal_notice", - "message": "This frame should not complete the response" - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_unmapped_success_error", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::WebSocket(CodexWebSocketError::ClosedByServerBeforeCompleted) = error - else { - panic!("expected unmapped success-status error frame to be ignored until close, found {error:?}"); - }; -} - -#[tokio::test] -async fn websocket_wrapped_error_retry_after_header_should_be_preserved_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "error", - "status": 429, - "error": { - "type": "rate_limit_exceeded", - "message": "Too many requests" - }, - "headers": { - "retry-after": "37" - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_wrapped_error_retry_after", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { - status, - retry_after_seconds, - body, - } = error - else { - panic!( - "expected wrapped websocket error retry-after to surface as upstream error, found {error:?}" - ); - }; - assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS); - assert_eq!(retry_after_seconds, Some(37)); - assert!(body.contains("rate_limit_exceeded")); - assert!(body.contains("retry-after")); -} - -#[tokio::test] -async fn websocket_wrapped_connection_limit_should_use_retryable_503_precedence_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "error", - "status": 400, - "error": { - "code": "websocket_connection_limit_reached", - "type": "invalid_request_error", - "message": "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue." - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_wrapped_connection_limit", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { status, body, .. } = error else { - panic!("expected wrapped websocket connection limit to surface as upstream error, found {error:?}"); - }; - assert_eq!(status, reqwest::StatusCode::SERVICE_UNAVAILABLE); - assert!(body.contains("websocket_connection_limit_reached")); -} - -#[tokio::test] -async fn websocket_midstream_wrapped_error_status_should_surface_as_upstream_error_like_official_client( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.created", - "response": { - "id": "resp_midstream_wrapped_error" - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket - .send(Message::Text( - json!({ - "type": "error", - "status": 400, - "error": { - "type": "invalid_request_error", - "message": "Mid-stream wrapped websocket error" - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &base_request(), - request_context("req_ws_midstream_wrapped_error", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { - status, - retry_after_seconds, - body, - } = error - else { - panic!("expected mid-stream wrapped websocket error to surface as upstream error, found {error:?}"); - }; - assert_eq!(status, reqwest::StatusCode::BAD_REQUEST); - assert_eq!(retry_after_seconds, None); - assert!(body.contains("invalid_request_error")); - assert!(body.contains("Mid-stream wrapped websocket error")); -} - -#[tokio::test] -async fn websocket_incomplete_event_should_error_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.incomplete", - "response": { - "id": "resp_incomplete", - "incomplete_details": { - "reason": "max_output_tokens" - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response(&base_request(), request_context("req_ws_incomplete", None)) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::WebSocket(CodexWebSocketError::IncompleteResponse { reason }) = error - else { - panic!("expected response.incomplete websocket error, found {error:?}"); - }; - assert_eq!(reason, "max_output_tokens"); -} - -#[tokio::test] -async fn websocket_connection_limit_failed_frame_should_surface_as_503_like_official_retryable() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.failed", - "response": { - "id": "resp_ws_connection_limit", - "error": { - "code": "websocket_connection_limit_reached", - "message": "connection limit reached" - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response(&base_request(), request_context("req_ws_limit", None)) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { status, body, .. } = error else { - panic!("expected upstream error, found {error:?}"); - }; - assert_eq!(status, reqwest::StatusCode::SERVICE_UNAVAILABLE); - assert!(body.contains("websocket_connection_limit_reached")); -} - -#[tokio::test] -async fn websocket_pooled_connection_limit_frame_should_surface_as_503() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text( - json!({ - "type": "response.failed", - "response": { - "id": "resp_pooled_ws_connection_limit", - "error": { - "code": "websocket_connection_limit_reached", - "message": "connection limit reached" - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let pool = Arc::new(CodexWebSocketPool::with_config(manual_pool_config( - Duration::from_secs(60), - 8, - ))); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ) - .with_websocket_pool(pool, "chatgpt-account"); - let mut request = base_request(); - request.prompt_cache_key = Some("chatgpt-account:conversation".to_string()); - - let error = client - .create_response(&request, request_context("req_pooled_ws_limit", None)) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::Upstream { status, body, .. } = error else { - panic!("expected upstream error, found {error:?}"); - }; - assert_eq!(status, reqwest::StatusCode::SERVICE_UNAVAILABLE); - assert!(body.contains("websocket_connection_limit_reached")); -} - -#[tokio::test] -async fn websocket_pooled_upstream_error_should_discard_connection_like_official_client() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (first_stream, _) = listener.accept().await.unwrap(); - let mut first_websocket = accept_async(first_stream).await.unwrap(); - let _first_message = first_websocket.next().await.unwrap().unwrap(); - first_websocket - .send(Message::Text( - json!({ - "type": "response.failed", - "response": { - "id": "resp_pooled_rate_limit", - "error": { - "code": "rate_limit_exceeded", - "message": "Rate limit reached. Please try again in 1s." - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - - let (second_stream, _) = listener.accept().await.unwrap(); - let mut second_websocket = accept_async(second_stream).await.unwrap(); - let _second_message = second_websocket.next().await.unwrap().unwrap(); - second_websocket - .send(Message::Text( - websocket_completed_response("resp_after_pooled_error", 8, 2).into(), - )) - .await - .unwrap(); - second_websocket.close(None).await.unwrap(); - }); - - let pool = Arc::new(CodexWebSocketPool::with_config(manual_pool_config( - Duration::from_secs(60), - 8, - ))); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ) - .with_websocket_pool(pool, "chatgpt-account"); - let mut request = base_request(); - request.prompt_cache_key = Some("chatgpt-account:conversation".to_string()); - - let error = client - .create_response(&request, request_context("req_pooled_ws_error", None)) - .await - .unwrap_err(); - let CodexClientError::Upstream { status, body, .. } = error else { - panic!("expected first pooled request to surface upstream error, found {error:?}"); - }; - assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS); - assert!(body.contains("rate_limit_exceeded")); - - let response = client - .create_response(&request, request_context("req_pooled_ws_after_error", None)) - .await - .unwrap(); - - server.await.unwrap(); - assert!(response.body.contains("\"id\":\"resp_after_pooled_error\"")); - assert_eq!(response.usage.unwrap().input_tokens, 8); -} - -#[tokio::test] -async fn ordinary_response_should_not_fallback_to_http_sse_when_websocket_transport_fails() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut websocket_stream, _) = listener.accept().await.unwrap(); - let websocket_request = read_http_upgrade_request(&mut websocket_stream).await; - assert!(websocket_request.starts_with("GET /codex/responses HTTP/1.1")); - drop(websocket_stream); - }); - - let request = base_request(); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let error = client - .create_response( - &request, - request_context("req_ws_transport_no_fallback", None), - ) - .await - .unwrap_err(); - - server.await.unwrap(); - let CodexClientError::WebSocket(CodexWebSocketError::Transport(_)) = error else { - panic!("expected websocket transport error without HTTP SSE fallback, found {error:?}"); - }; -} - -#[tokio::test] -async fn ordinary_response_should_fallback_to_http_sse_when_websocket_upgrade_required() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut websocket_stream, _) = listener.accept().await.unwrap(); - let websocket_request = read_http_upgrade_request(&mut websocket_stream).await; - assert!(websocket_request.starts_with("GET /codex/responses HTTP/1.1")); - websocket_stream - .write_all( - b"HTTP/1.1 426 Upgrade Required\r\ncontent-length: 0\r\nconnection: close\r\n\r\n", - ) - .await - .unwrap(); - - let (mut http_stream, _) = listener.accept().await.unwrap(); - let http_request = read_http_upgrade_request(&mut http_stream).await; - assert!(http_request.starts_with("POST /codex/responses HTTP/1.1")); - let body = "event: response.completed\ndata: {\"response\":{\"id\":\"resp_http_fallback\",\"usage\":{\"input_tokens\":3,\"output_tokens\":1}}}\n\n"; - let response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", - body.len() - ); - http_stream.write_all(response.as_bytes()).await.unwrap(); - }); - - let request = base_request(); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &request, - request_context("req_ws_upgrade_required_fallback", None), - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(response.body.contains("event: response.completed")); - assert!(response.body.contains("\"id\":\"resp_http_fallback\"")); - assert_eq!(response.usage.unwrap().input_tokens, 3); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_success_should_capture_handshake_headers_and_rate_limit_headers() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = - accept_hdr_async(stream, |_request: &WsRequest, mut response: WsResponse| { - response - .headers_mut() - .insert("x-codex-turn-state", "turn-ws".parse().unwrap()); - response.headers_mut().insert( - "set-cookie", - "cf_clearance=ws; Domain=.chatgpt.com; Path=/" - .parse() - .unwrap(), - ); - response - .headers_mut() - .insert("x-ratelimit-remaining-requests", "41".parse().unwrap()); - Ok(response) - }) - .await - .unwrap(); - let _message = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - json!({ - "type": "response.completed", - "response": { - "id": "resp_ws_headers", - "object": "response", - "usage": { - "input_tokens": 1, - "output_tokens": 1, - "total_tokens": 2 - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let mut request = base_request(); - request.use_websocket = true; - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &request, - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_ws_headers", - 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, - installation_id: None, - session_id: None, - }, - ) - .await - .unwrap(); - - server.await.unwrap(); - assert_eq!(response.turn_state.as_deref(), Some("turn-ws")); - assert_eq!( - response.set_cookie_headers, - vec!["cf_clearance=ws; Domain=.chatgpt.com; Path=/".to_string()] - ); - assert!(response - .rate_limit_headers - .iter() - .any(|(name, value)| { name == "x-ratelimit-remaining-requests" && value == "41" })); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_success_should_capture_internal_rate_limit_events_without_forwarding_them() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _message = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - json!({ - "type": "codex.rate_limits", - "rate_limits": { - "primary": { - "used_percent": 100, - "window_minutes": 5, - "reset_at": 1893456300 - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_rate_limits", 1, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let mut request = base_request(); - request.use_websocket = true; - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response( - &request, - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_ws_rate_limits", - 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, - installation_id: None, - session_id: None, - }, - ) - .await - .unwrap(); - - server.await.unwrap(); - assert!(!response.body.contains("codex.rate_limits")); - assert!(response - .rate_limit_headers - .iter() - .any(|(name, value)| { name == "x-codex-primary-used-percent" && value == "100" })); - assert!(response - .rate_limit_headers - .iter() - .any(|(name, value)| { name == "x-codex-primary-reset-at" && value == "1893456300" })); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_metadata_event_should_update_turn_state_without_forwarding() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _message = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - json!({ - "type": "response.metadata", - "headers": { - "x-codex-turn-state": "turn-from-metadata" - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_ws_metadata", 1, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let mut request = base_request(); - request.use_websocket = true; - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - - let response = client - .create_response(&request, request_context("req_ws_metadata", None)) - .await - .unwrap(); - - server.await.unwrap(); - assert_eq!(response.turn_state.as_deref(), Some("turn-from-metadata")); - assert!(!response.body.contains("response.metadata")); - assert!(response.body.contains("event: response.completed")); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_stream_should_error_when_connection_closes_before_terminal_frame() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _request = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - json!({ - "type": "response.output_text.delta", - "delta": "partial" - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - let mut stream = client - .websocket_stream_response(&base_request(), request_context("req_mid_close", None)) - .await - .unwrap() - .body_stream; - - let first_chunk = stream.next().await.unwrap().unwrap(); - assert!(first_chunk.contains("partial")); - let error = stream.next().await.unwrap().unwrap_err(); - assert!(error - .to_string() - .contains("websocket closed by server before response.completed")); - - server.await.unwrap(); -} - -fn base_request() -> CodexResponsesRequest { - CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new()) -} - -fn manual_pool_config(max_age: Duration, max_per_account: usize) -> CodexWebSocketPoolConfig { - CodexWebSocketPoolConfig { - enabled: true, - max_age, - max_per_account, - maintenance_interval: None, - ping_interval: None, - ping_timeout: Duration::from_millis(50), - liveness_timeout: None, - } -} - -fn keepalive_pool_config(ping_timeout: Duration) -> CodexWebSocketPoolConfig { - CodexWebSocketPoolConfig { - ping_interval: Some(Duration::from_millis(1)), - ping_timeout, - liveness_timeout: Some(Duration::from_secs(60)), - ..manual_pool_config(Duration::from_secs(60), 8) - } -} - -fn request_context<'a>( - request_id: &'a str, - session_id: Option<&'a str>, -) -> CodexRequestContext<'a> { - CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - 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, - installation_id: None, - session_id, - } -} - -fn audit_artifact_with_header_names(header_names: &[&str]) -> WebSocketAuditArtifact { - WebSocketAuditArtifact { - transport_mode: "websocket_preferred".to_string(), - fallback_allowed: true, - opening: Some(OpeningAuditSnapshot { - request_line: "GET /codex/responses HTTP/1.1".to_string(), - headers: header_names - .iter() - .map(|name| OpeningAuditHeader { - name: (*name).to_string(), - value: "value".to_string(), - }) - .collect(), - }), - payload: Some(PayloadAuditSnapshot { - top_level_keys: vec!["type".to_string(), "model".to_string()], - body: json!({ - "type": "response.create", - "model": "gpt-5.5", - }), - }), - error: None, - } -} - -#[derive(Debug)] -struct CapturedWebSocketExchange { - opening_bytes: String, - first_frame_text: String, - sec_websocket_key_offset: usize, - authorization_header_offset: usize, -} - -async fn capture_codex_websocket_exchange( - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - artifact_path: &Path, -) -> CapturedWebSocketExchange { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.unwrap(); - let opening_bytes = read_http_upgrade_request(&mut stream).await; - let accept_key = websocket_accept_key(&opening_bytes); - let response = format!( - "HTTP/1.1 101 Switching Protocols\r\n\ - Connection: Upgrade\r\n\ - Upgrade: websocket\r\n\ - Sec-WebSocket-Accept: {accept_key}\r\n\ - \r\n" - ); - stream.write_all(response.as_bytes()).await.unwrap(); - let first_frame = read_client_websocket_frame(&mut stream).await; - let first_frame_text = String::from_utf8(first_frame).unwrap(); - let payload = websocket_completed_response("resp_ws_capture", 3, 1); - stream - .write_all(&encode_server_frame(0x1, false, payload.as_bytes())) - .await - .unwrap(); - let lower_opening_bytes = opening_bytes.to_ascii_lowercase(); - CapturedWebSocketExchange { - sec_websocket_key_offset: lower_opening_bytes - .find("sec-websocket-key: ") - .expect("capture should include Sec-WebSocket-Key"), - authorization_header_offset: lower_opening_bytes - .find("authorization: ") - .expect("capture should include Authorization"), - opening_bytes, - first_frame_text, - } - }); - - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ); - client.create_response(request, context).await.unwrap(); - - let capture = server.await.unwrap(); - write_capture_artifact(artifact_path, request, &capture).await; - capture -} - -async fn write_capture_artifact( - artifact_path: &Path, - request: &CodexResponsesRequest, - capture: &CapturedWebSocketExchange, -) { - if let Some(parent) = artifact_path.parent() { - tokio::fs::create_dir_all(parent).await.unwrap(); - } - let artifact = json!({ - "transport_mode": "websocket_preferred", - "fallback_allowed": true, - "opening": { - "request_line": capture.opening_bytes.lines().next().unwrap_or_default(), - "headers": redacted_headers_from_raw_opening(&capture.opening_bytes), - "sec_websocket_key_offset": capture.sec_websocket_key_offset, - "authorization_header_offset": capture.authorization_header_offset, - }, - "first_frame": websocket_payload_audit_snapshot(request), - "tls": { - "captured": false, - "reason": "local WSS capture is not used in this test because the production client validates against native system roots" - } - }); - let body = serde_json::to_vec_pretty(&artifact).unwrap(); - tokio::fs::write(artifact_path, body).await.unwrap(); -} - -fn redacted_headers_from_raw_opening(raw_request: &str) -> Vec { - raw_request - .split("\r\n") - .skip(1) - .take_while(|line| !line.is_empty()) - .filter_map(|line| { - let (name, value) = line.split_once(": ")?; - Some(json!({ - "name": name, - "value": redacted_capture_header_value(name, value), - })) - }) - .collect() -} - -fn redacted_capture_header_value(name: &str, value: &str) -> String { - match name.to_ascii_lowercase().as_str() { - "authorization" - | "chatgpt-account-id" - | "cookie" - | "x-client-request-id" - | "x-codex-installation-id" - | "session_id" - | "session-id" - | "thread-id" - | "x-codex-window-id" - | "x-codex-turn-state" - | "x-codex-turn-metadata" - | "x-codex-parent-thread-id" => "".to_string(), - "sec-websocket-key" => "".to_string(), - _ => value.to_string(), - } -} - -fn websocket_completed_response( - response_id: &str, - input_tokens: i64, - output_tokens: i64, -) -> String { - json!({ - "type": "response.completed", - "response": { - "id": response_id, - "object": "response", - "usage": { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "total_tokens": input_tokens + output_tokens - } - } - }) - .to_string() -} - -fn with_sse_terminal_separator(body: &str) -> String { - if body.ends_with("\n\n") { - body.to_string() - } else { - format!("{body}\n") - } -} - -async fn read_http_upgrade_request(stream: &mut tokio::net::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 read_client_websocket_frame(stream: &mut tokio::net::TcpStream) -> Vec { - read_client_websocket_frame_with_opcode(stream).await.1 -} - -async fn read_client_websocket_frame_with_opcode( - stream: &mut tokio::net::TcpStream, -) -> (u8, Vec) { - let mut header = [0_u8; 2]; - stream.read_exact(&mut header).await.unwrap(); - let opcode = header[0] & 0x0f; - let mut payload_len = u64::from(header[1] & 0x7f); - if payload_len == 126 { - let mut extended = [0_u8; 2]; - stream.read_exact(&mut extended).await.unwrap(); - payload_len = u64::from(u16::from_be_bytes(extended)); - } else if payload_len == 127 { - let mut extended = [0_u8; 8]; - stream.read_exact(&mut extended).await.unwrap(); - payload_len = u64::from_be_bytes(extended); - } - let mut mask = [0_u8; 4]; - if header[1] & 0x80 != 0 { - stream.read_exact(&mut mask).await.unwrap(); - } - let mut payload = vec![0_u8; payload_len as usize]; - stream.read_exact(&mut payload).await.unwrap(); - if header[1] & 0x80 != 0 { - for (index, byte) in payload.iter_mut().enumerate() { - *byte ^= mask[index % mask.len()]; - } - } - (opcode, payload) -} - -fn websocket_accept_key(raw_request: &str) -> String { - let key = raw_request - .lines() - .find_map(|line| line.strip_prefix("Sec-WebSocket-Key: ")) - .expect("raw websocket request should contain Sec-WebSocket-Key"); - derive_accept_key(key.as_bytes()) -} - -fn compressed_server_text_frame(text: &str) -> Vec { - let mut compressor = Compress::new(Compression::fast(), false); - let mut payload = Vec::with_capacity(text.len() + 32); - compressor - .compress_vec(text.as_bytes(), &mut payload, FlushCompress::Sync) - .unwrap(); - if payload.ends_with(&[0x00, 0x00, 0xff, 0xff]) { - payload.truncate(payload.len() - 4); - } - encode_server_frame(0x1, true, &payload) -} - -fn server_ping_frame(payload: &[u8]) -> Vec { - server_frame(0x89, payload) -} - -fn server_text_frame(payload: &[u8]) -> Vec { - server_frame(0x81, payload) -} - -fn server_binary_frame(payload: &[u8]) -> Vec { - server_frame(0x82, payload) -} - -fn server_close_frame() -> Vec { - server_frame(0x88, &1000_u16.to_be_bytes()) -} - -fn server_frame(first_byte: u8, payload: &[u8]) -> Vec { - let mut frame = vec![first_byte]; - if payload.len() < 126 { - frame.push(payload.len() as u8); - } else if u16::try_from(payload.len()).is_ok() { - frame.push(126); - frame.extend_from_slice(&(payload.len() as u16).to_be_bytes()); - } else { - frame.push(127); - frame.extend_from_slice(&(payload.len() as u64).to_be_bytes()); - } - frame.extend_from_slice(payload); - frame -} - -fn encode_server_frame(opcode: u8, rsv1: bool, payload: &[u8]) -> Vec { - let mut frame = Vec::new(); - let first = 0x80 | if rsv1 { 0x40 } else { 0 } | opcode; - frame.push(first); - match payload.len() { - len @ 0..=125 => frame.push(len as u8), - len @ 126..=65_535 => { - frame.push(126); - frame.extend_from_slice(&(len as u16).to_be_bytes()); - } - len => { - frame.push(127); - frame.extend_from_slice(&(len as u64).to_be_bytes()); - } - } - frame.extend_from_slice(payload); - frame -} - -fn assert_headers_appear_in_order(raw_request: &str, expected_headers: &[&str]) { - assert_substrings_appear_in_order(raw_request, expected_headers); -} - -fn assert_substrings_appear_in_order(haystack: &str, expected_values: &[&str]) { - let mut offset = 0; - for expected in expected_values { - let Some(index) = haystack[offset..].find(expected) else { - panic!("missing expected substring `{expected}` in:\n{haystack}"); - }; - offset += index + expected.len(); - } -} diff --git a/tests/codex_gateway/websocket/pool.rs b/tests/codex_gateway/websocket/pool.rs deleted file mode 100644 index 0daae2ad4..000000000 --- a/tests/codex_gateway/websocket/pool.rs +++ /dev/null @@ -1,559 +0,0 @@ -use super::*; - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let (release_first_tx, release_first_rx) = oneshot::channel(); - let server = tokio::spawn(async move { - let (first_stream, _) = listener.accept().await.unwrap(); - let mut first_ws = - accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _first_request = first_ws.next().await.unwrap().unwrap(); - first_ws - .send(Message::Text( - json!({ - "type": "response.output_text.delta", - "delta": "first connection is still busy" - }) - .to_string() - .into(), - )) - .await - .unwrap(); - - let (second_stream, _) = listener.accept().await.unwrap(); - let mut second_ws = - accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _second_request = second_ws.next().await.unwrap().unwrap(); - second_ws - .send(Message::Text( - websocket_completed_response("resp_busy_second", 2, 1).into(), - )) - .await - .unwrap(); - - let third_on_new_connection = tokio::select! { - message = second_ws.next() => { - if matches!(message, Some(Ok(message)) if message.is_text()) { - false - } else { - let (third_stream, _) = listener.accept().await.unwrap(); - let mut third_ws = accept_hdr_async( - third_stream, - |_request: &WsRequest, response| Ok(response), - ) - .await - .unwrap(); - let _third_request = third_ws.next().await.unwrap().unwrap(); - third_ws - .send(Message::Text( - websocket_completed_response("resp_busy_third", 2, 1).into(), - )) - .await - .unwrap(); - third_ws.close(None).await.unwrap(); - true - } - } - accepted = listener.accept() => { - let (third_stream, _) = accepted.unwrap(); - let mut third_ws = accept_hdr_async( - third_stream, - |_request: &WsRequest, response| Ok(response), - ) - .await - .unwrap(); - let _third_request = third_ws.next().await.unwrap().unwrap(); - third_ws - .send(Message::Text( - websocket_completed_response("resp_busy_third", 2, 1).into(), - )) - .await - .unwrap(); - third_ws.close(None).await.unwrap(); - true - } - }; - - release_first_rx.await.unwrap(); - first_ws - .send(Message::Text( - websocket_completed_response("resp_busy_first", 2, 1).into(), - )) - .await - .unwrap(); - first_ws.close(None).await.unwrap(); - third_on_new_connection - }); - - let pool = Arc::new(CodexWebSocketPool::with_default_max_age()); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ) - .with_websocket_pool(pool, "entry-a"); - let mut request = base_request(); - request.prompt_cache_key = Some("conversation-a".to_string()); - - let mut first_stream = client - .websocket_stream_response(&request, request_context("req_busy_first", Some("conv-a"))) - .await - .unwrap() - .body_stream; - let first_chunk = first_stream.next().await.unwrap().unwrap(); - assert!(first_chunk.contains("first connection is still busy")); - - let second = client - .create_response(&request, request_context("req_busy_second", Some("conv-a"))) - .await - .unwrap(); - assert!(second.body.contains("\"id\":\"resp_busy_second\"")); - - let third = client - .create_response(&request, request_context("req_busy_third", Some("conv-a"))) - .await - .unwrap(); - assert!(third.body.contains("\"id\":\"resp_busy_third\"")); - - release_first_tx.send(()).unwrap(); - while first_stream.next().await.transpose().unwrap().is_some() {} - - assert!( - server.await.unwrap(), - "busy pool key reused a second pooled websocket instead of one-shot bypass" - ); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_pool_should_bypass_new_keys_after_account_cap() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (first_stream, _) = listener.accept().await.unwrap(); - let mut first_ws = - accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _first_request = first_ws.next().await.unwrap().unwrap(); - first_ws - .send(Message::Text( - websocket_completed_response("resp_cap_first", 2, 1).into(), - )) - .await - .unwrap(); - - let (second_stream, _) = listener.accept().await.unwrap(); - let mut second_ws = - accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _second_request = second_ws.next().await.unwrap().unwrap(); - second_ws - .send(Message::Text( - websocket_completed_response("resp_cap_second", 2, 1).into(), - )) - .await - .unwrap(); - - let third_on_new_connection = tokio::select! { - message = second_ws.next() => { - if matches!(message, Some(Ok(message)) if message.is_text()) { - false - } else { - let (third_stream, _) = listener.accept().await.unwrap(); - let mut third_ws = accept_hdr_async( - third_stream, - |_request: &WsRequest, response| Ok(response), - ) - .await - .unwrap(); - let _third_request = third_ws.next().await.unwrap().unwrap(); - third_ws - .send(Message::Text( - websocket_completed_response("resp_cap_third", 2, 1).into(), - )) - .await - .unwrap(); - third_ws.close(None).await.unwrap(); - true - } - } - accepted = listener.accept() => { - let (third_stream, _) = accepted.unwrap(); - let mut third_ws = accept_hdr_async( - third_stream, - |_request: &WsRequest, response| Ok(response), - ) - .await - .unwrap(); - let _third_request = third_ws.next().await.unwrap().unwrap(); - third_ws - .send(Message::Text( - websocket_completed_response("resp_cap_third", 2, 1).into(), - )) - .await - .unwrap(); - third_ws.close(None).await.unwrap(); - true - } - }; - - first_ws.close(None).await.unwrap(); - third_on_new_connection - }); - - let pool = Arc::new(CodexWebSocketPool::with_limits(Duration::from_secs(60), 1)); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ) - .with_websocket_pool(pool, "entry-cap"); - - let mut first_request = base_request(); - first_request.prompt_cache_key = Some("conversation-one".to_string()); - let first = client - .create_response( - &first_request, - request_context("req_cap_first", Some("conv-one")), - ) - .await - .unwrap(); - assert!(first.body.contains("\"id\":\"resp_cap_first\"")); - - let mut second_request = base_request(); - second_request.prompt_cache_key = Some("conversation-two".to_string()); - let second = client - .create_response( - &second_request, - request_context("req_cap_second", Some("conv-two")), - ) - .await - .unwrap(); - assert!(second.body.contains("\"id\":\"resp_cap_second\"")); - - let third = client - .create_response( - &second_request, - request_context("req_cap_third", Some("conv-two")), - ) - .await - .unwrap(); - assert!(third.body.contains("\"id\":\"resp_cap_third\"")); - - assert!( - server.await.unwrap(), - "account cap bypass still allowed a capped conversation to be pooled" - ); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_pool_keepalive_should_keep_idle_connection_when_pong_is_received() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _first_request = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_keepalive_first", 2, 1).into(), - )) - .await - .unwrap(); - - let ping = websocket.next().await.unwrap().unwrap(); - assert!( - ping.is_ping(), - "expected pool maintenance ping, got {ping:?}" - ); - websocket - .send(Message::Pong(Vec::new().into())) - .await - .unwrap(); - - let reused = tokio::select! { - message = websocket.next() => { - let message = message.unwrap().unwrap(); - if !message.is_text() { - return false; - } - websocket - .send(Message::Text( - websocket_completed_response("resp_keepalive_second", 2, 1).into(), - )) - .await - .unwrap(); - true - } - accepted = listener.accept() => { - let (stream, _) = accepted.unwrap(); - let mut websocket = accept_hdr_async( - stream, - |_request: &WsRequest, response| Ok(response), - ) - .await - .unwrap(); - let _request = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_keepalive_second", 2, 1).into(), - )) - .await - .unwrap(); - false - } - }; - - websocket.close(None).await.unwrap(); - reused - }); - - let pool = Arc::new(CodexWebSocketPool::with_config(keepalive_pool_config( - Duration::from_millis(50), - ))); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ) - .with_websocket_pool(Arc::clone(&pool), "entry-keepalive"); - let mut request = base_request(); - request.prompt_cache_key = Some("conversation-keepalive".to_string()); - - let first = client - .create_response( - &request, - request_context("req_keepalive_first", Some("conv-keepalive")), - ) - .await - .unwrap(); - assert!(first.body.contains("\"id\":\"resp_keepalive_first\"")); - - pool.maintain_idle_connections().await; - - let second = client - .create_response( - &request, - request_context("req_keepalive_second", Some("conv-keepalive")), - ) - .await - .unwrap(); - assert!(second.body.contains("\"id\":\"resp_keepalive_second\"")); - - assert!( - server.await.unwrap(), - "keepalive pong should leave the pooled websocket reusable" - ); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_pool_keepalive_should_evict_idle_connection_without_pong() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (first_stream, _) = listener.accept().await.unwrap(); - let mut first_ws = - accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _first_request = first_ws.next().await.unwrap().unwrap(); - first_ws - .send(Message::Text( - websocket_completed_response("resp_no_pong_first", 2, 1).into(), - )) - .await - .unwrap(); - - let (second_stream, _) = listener.accept().await.unwrap(); - let mut second_ws = - accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _second_request = second_ws.next().await.unwrap().unwrap(); - second_ws - .send(Message::Text( - websocket_completed_response("resp_no_pong_second", 2, 1).into(), - )) - .await - .unwrap(); - second_ws.close(None).await.unwrap(); - }); - - let pool = Arc::new(CodexWebSocketPool::with_config(keepalive_pool_config( - Duration::from_millis(20), - ))); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ) - .with_websocket_pool(Arc::clone(&pool), "entry-no-pong"); - let mut request = base_request(); - request.prompt_cache_key = Some("conversation-no-pong".to_string()); - - let first = client - .create_response( - &request, - request_context("req_no_pong_first", Some("conv-no-pong")), - ) - .await - .unwrap(); - assert!(first.body.contains("\"id\":\"resp_no_pong_first\"")); - - pool.maintain_idle_connections().await; - - let second = client - .create_response( - &request, - request_context("req_no_pong_second", Some("conv-no-pong")), - ) - .await - .unwrap(); - assert!(second.body.contains("\"id\":\"resp_no_pong_second\"")); - - server.await.unwrap(); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_pool_gc_sweep_should_close_expired_idle_connections() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - for response_id in ["resp_gc_first", "resp_gc_second"] { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = - accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _request = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - websocket_completed_response(response_id, 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - } - }); - - let pool = Arc::new(CodexWebSocketPool::with_config(manual_pool_config( - Duration::from_millis(5), - 8, - ))); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ) - .with_websocket_pool(Arc::clone(&pool), "entry-gc"); - let mut request = base_request(); - request.prompt_cache_key = Some("conversation-gc".to_string()); - - let first = client - .create_response(&request, request_context("req_gc_first", Some("conv-gc"))) - .await - .unwrap(); - assert!(first.body.contains("\"id\":\"resp_gc_first\"")); - - sleep(Duration::from_millis(15)).await; - pool.gc_sweep().await; - - let second = client - .create_response(&request, request_context("req_gc_second", Some("conv-gc"))) - .await - .unwrap(); - assert!(second.body.contains("\"id\":\"resp_gc_second\"")); - - server.await.unwrap(); -} - -#[tokio::test] -#[expect( - clippy::result_large_err, - reason = "tokio-tungstenite handshake callbacks use a large error response type" -)] -async fn websocket_pool_shutdown_should_close_idle_connections_and_disable_pooling() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - for response_id in ["resp_shutdown_first", "resp_shutdown_second"] { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = - accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response)) - .await - .unwrap(); - let _request = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - websocket_completed_response(response_id, 2, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - } - }); - - let pool = Arc::new(CodexWebSocketPool::with_config(manual_pool_config( - Duration::from_secs(60), - 8, - ))); - let client = CodexBackendClient::new( - build_reqwest_client(false).unwrap(), - format!("http://{addr}"), - codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(), - ) - .with_websocket_pool(Arc::clone(&pool), "entry-shutdown"); - let mut request = base_request(); - request.prompt_cache_key = Some("conversation-shutdown".to_string()); - - let first = client - .create_response( - &request, - request_context("req_shutdown_first", Some("conv-shutdown")), - ) - .await - .unwrap(); - assert!(first.body.contains("\"id\":\"resp_shutdown_first\"")); - - pool.shutdown().await; - - let second = client - .create_response( - &request, - request_context("req_shutdown_second", Some("conv-shutdown")), - ) - .await - .unwrap(); - assert!(second.body.contains("\"id\":\"resp_shutdown_second\"")); - - server.await.unwrap(); -} diff --git a/tests/codex_models.rs b/tests/codex_models.rs deleted file mode 100644 index 9faed2216..000000000 --- a/tests/codex_models.rs +++ /dev/null @@ -1,3 +0,0 @@ -mod codex_models { - mod catalog; -} diff --git a/tests/codex_models/catalog.rs b/tests/codex_models/catalog.rs deleted file mode 100644 index b070595ab..000000000 --- a/tests/codex_models/catalog.rs +++ /dev/null @@ -1,145 +0,0 @@ -use std::collections::BTreeMap; - -use codex_proxy_rs::{ - codex::models::{ - catalog::{ - BackendModelEntry, BackendReasoningEffort, BackendTruncationPolicy, ModelCatalog, - ModelPlanSnapshot, ParsedModelName, - }, - repository::ModelSnapshotRepository, - }, - config::ModelConfig, - platform::storage::db::connect_sqlite, -}; - -#[test] -fn model_catalog_should_parse_alias_reasoning_and_service_tier_suffixes() { - let mut aliases = BTreeMap::new(); - aliases.insert("codex-fast".to_string(), "gpt-5.5".to_string()); - let catalog = ModelCatalog::from_config(&ModelConfig { - default_model: "gpt-5.5".to_string(), - default_reasoning_effort: None, - service_tier: None, - aliases, - }); - - let parsed = catalog.parse_model_name("codex-fast-high-flex"); - - assert_eq!( - parsed, - ParsedModelName { - model_id: "gpt-5.5".to_string(), - reasoning_effort: Some("high".to_string()), - service_tier: Some("flex".to_string()) - } - ); -} - -#[test] -fn model_catalog_should_fall_back_to_default_for_unknown_model_names() { - let catalog = ModelCatalog::from_config(&ModelConfig { - default_model: "gpt-5.5".to_string(), - default_reasoning_effort: None, - service_tier: None, - aliases: BTreeMap::new(), - }); - - let parsed = catalog.parse_model_name("unknown-medium"); - - assert_eq!(parsed.model_id, "gpt-5.5"); - assert_eq!(parsed.reasoning_effort, Some("medium".to_string())); -} - -#[test] -fn model_catalog_should_parse_none_and_minimal_reasoning_suffixes() { - let catalog = ModelCatalog::from_config(&ModelConfig { - default_model: "gpt-5.5".to_string(), - default_reasoning_effort: None, - service_tier: None, - aliases: BTreeMap::new(), - }); - - let none = catalog.parse_model_name("gpt-5.5-none"); - let minimal = catalog.parse_model_name("gpt-5.5-minimal-fast"); - - assert_eq!(none.reasoning_effort, Some("none".to_string())); - assert_eq!(minimal.reasoning_effort, Some("minimal".to_string())); - assert_eq!(minimal.service_tier, Some("fast".to_string())); -} - -#[test] -fn model_catalog_should_merge_backend_snapshots_and_build_plan_allowlist() { - let snapshot = ModelPlanSnapshot::from_backend_entries( - "plus", - vec![BackendModelEntry { - slug: Some("gpt-6".to_string()), - display_name: Some("GPT-6".to_string()), - description: Some("Backend model".to_string()), - is_default: Some(true), - default_reasoning_level: Some("minimal".to_string()), - supported_reasoning_levels: vec![BackendReasoningEffort { - effort: Some("minimal".to_string()), - description: Some("Minimal".to_string()), - ..BackendReasoningEffort::default() - }], - input_modalities: Some(vec!["text".to_string(), "image".to_string()]), - output_modalities: Some(vec!["text".to_string()]), - supports_personality: Some(true), - context_window: Some(200_000), - max_output_tokens: Some(16_384), - truncation_policy: Some(BackendTruncationPolicy { - limit: Some(131_072), - }), - ..BackendModelEntry::default() - }], - ); - let catalog = ModelCatalog::from_config_and_snapshots( - &ModelConfig { - default_model: "gpt-5.5".to_string(), - default_reasoning_effort: None, - service_tier: None, - aliases: BTreeMap::new(), - }, - &[snapshot], - ); - - let info = catalog.model_info("gpt-6").unwrap(); - assert_eq!(info.display_name, "GPT-6"); - assert_eq!(info.default_reasoning_effort, "minimal"); - assert_eq!( - info.supported_reasoning_efforts[0].reasoning_effort, - "minimal" - ); - assert_eq!(info.context_window, Some(200_000)); - assert_eq!(info.truncation_policy_limit, Some(131_072)); - assert_eq!( - catalog.model_plan_allowlist().get("gpt-6").unwrap(), - &vec!["plus".to_string()] - ); -} - -#[tokio::test] -async fn model_snapshot_repository_should_replace_and_load_plan_snapshots() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("models.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .unwrap(); - let repo = ModelSnapshotRepository::new(pool); - let snapshot = ModelPlanSnapshot::from_backend_entries( - "team", - vec![BackendModelEntry { - id: Some("gpt-team".to_string()), - name: Some("GPT Team".to_string()), - ..BackendModelEntry::default() - }], - ); - - repo.replace_plan_snapshot(&snapshot).await.unwrap(); - let loaded = repo.list_plan_snapshots().await.unwrap(); - - assert_eq!(loaded.len(), 1); - assert_eq!(loaded[0].plan_type, "team"); - assert_eq!(loaded[0].models[0].id, "gpt-team"); - assert_eq!(loaded[0].models[0].source, "backend"); -} diff --git a/tests/codex_serving.rs b/tests/codex_serving.rs deleted file mode 100644 index 1421a7ddb..000000000 --- a/tests/codex_serving.rs +++ /dev/null @@ -1,12 +0,0 @@ -mod support; - -mod codex_serving { - mod chat_completions; - mod diagnostics_route; - mod responses_http_sse; - mod responses_websocket; - mod routes_chat; - mod routes_responses; - mod upstream_errors; - mod upstream_fallback; -} diff --git a/tests/codex_serving/chat_completions.rs b/tests/codex_serving/chat_completions.rs deleted file mode 100644 index 67650f2fb..000000000 --- a/tests/codex_serving/chat_completions.rs +++ /dev/null @@ -1,910 +0,0 @@ -use std::collections::BTreeMap; - -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use serde_json::{json, Value}; -use tower::ServiceExt; -use wiremock::{ - matchers::{header, method, path}, - Mock, MockServer, ResponseTemplate, -}; - -use codex_proxy_rs::{ - codex::gateway::conversation_identity::build_conversation_identity, - config::{ - AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, - QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, - }, - platform::crypto::SecretBox, - platform::identity::{client_key::ApiKeyHasher, client_key_repository::ClientApiKeyRepository}, - platform::storage::db::connect_sqlite, - runtime::build_router, - runtime::state::AppState, -}; - -use crate::support::{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"); -const CHAT_TUPLE_OBJECT_SSE: &str = include_str!("../fixtures/chat/tuple_object.sse"); -const CHAT_NO_AVAILABLE_ACCOUNTS_GOLDEN: &str = - include_str!("../fixtures/chat/golden/no_available_accounts.json"); -const CHAT_RATE_LIMIT_FALLBACK_EXHAUSTED_GOLDEN: &str = - include_str!("../fixtures/chat/golden/rate_limit_fallback_exhausted.json"); -const CHAT_QUOTA_EXHAUSTED_FALLBACK_EXHAUSTED_GOLDEN: &str = - include_str!("../fixtures/chat/golden/quota_exhausted_fallback_exhausted.json"); - -struct ImportedApp { - app: axum::Router, - client_api_key: String, - _tempdir: tempfile::TempDir, -} - -struct ImportAccount { - id: &'static str, - account_id: &'static str, - token: &'static str, - refresh_token: &'static str, -} - -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(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - 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, - }, - ws_pool: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, - }, - logging: LoggingConfig { - directory: "logs".to_string(), - 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_chat", - 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("chat-completions.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([61u8; 32]); - let hasher = ApiKeyHasher::new([62u8; 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, - secret_box, - hasher, - )); - if !accounts.is_empty() { - let accounts = accounts - .iter() - .map(|account| { - json!({ - "id": account.id, - "accountId": account.account_id, - "token": account.token, - "refreshToken": account.refresh_token, - "accessTokenExpiresAt": "2999-01-01T00:00:00Z", - "status": "active" - }) - }) - .collect::>(); - - let import_response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/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, - client_api_key: generated.plaintext, - _tempdir: tempdir, - } -} - -async fn single_upstream_post_body(server: &MockServer) -> Value { - let requests = server.received_requests().await.unwrap(); - let post_requests = requests - .iter() - .filter(|request| request.method.as_str() == "POST") - .collect::>(); - assert_eq!(post_requests.len(), 1); - assert_eq!(requests.len(), 1); - serde_json::from_slice(&post_requests[0].body).unwrap() -} - -#[tokio::test] -async fn chat_completions_should_return_openai_error_when_no_accounts_are_available() { - let server = MockServer::start().await; - let imported = build_imported_app_with_accounts(server.uri(), &[]).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "messages": [ - {"role": "user", "content": "Say ok"} - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); - let body = response_json(response).await; - let expected: Value = serde_json::from_str(CHAT_NO_AVAILABLE_ACCOUNTS_GOLDEN) - .expect("chat no-accounts golden should parse"); - assert_eq!(body, expected); -} - -#[tokio::test] -async fn chat_completions_should_return_openai_error_when_429_fallback_is_exhausted() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-secret")) - .respond_with(ResponseTemplate::new(429).set_body_json(json!({ - "error": {"message": "rate limited"} - }))) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "messages": [ - {"role": "user", "content": "Say ok"} - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); - let body = response_json(response).await; - let expected: Value = serde_json::from_str(CHAT_RATE_LIMIT_FALLBACK_EXHAUSTED_GOLDEN) - .expect("chat rate-limit golden should parse"); - assert_eq!(body, expected); -} - -#[tokio::test] -async fn chat_completions_should_return_openai_error_when_402_fallback_is_exhausted() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-secret")) - .respond_with(ResponseTemplate::new(402).set_body_json(json!({ - "error": {"message": "quota reached"} - }))) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "messages": [ - {"role": "user", "content": "Say ok"} - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::PAYMENT_REQUIRED); - let body = response_json(response).await; - let expected: Value = serde_json::from_str(CHAT_QUOTA_EXHAUSTED_FALLBACK_EXHAUSTED_GOLDEN) - .expect("chat quota-exhausted golden should parse"); - assert_eq!(body, expected); -} - -#[tokio::test] -async fn chat_completions_should_use_user_as_client_conversation_id() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-secret")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(CHAT_NON_STREAM_TEXT_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "user": "chat-client-session", - "messages": [ - {"role": "user", "content": "Say ok"} - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let upstream_body = single_upstream_post_body(&server).await; - let identity = build_conversation_identity(Some("chat-client-session"), None, "acct_chat"); - assert_eq!( - upstream_body["prompt_cache_key"], - identity.conversation_id.unwrap() - ); -} - -#[tokio::test] -async fn chat_completions_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(CHAT_RETRY_AFTER_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/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "stream": true, - "messages": [ - {"role": "user", "content": "Say ok"} - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let frames = response_sse_data(response).await; - assert_eq!(frames.last().unwrap(), "[DONE]"); - let chunks = frames - .iter() - .filter(|frame| frame.as_str() != "[DONE]") - .map(|frame| serde_json::from_str::(frame).unwrap()) - .collect::>(); - assert_eq!(chunks[1]["choices"][0]["delta"]["content"], "ok"); - assert_eq!(chunks[2]["usage"]["prompt_tokens"], 4); -} - -#[tokio::test] -async fn chat_completions_translates_messages_and_returns_openai_response() { - 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")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(CHAT_NON_STREAM_TEXT_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "messages": [ - {"role": "system", "content": "You are concise."}, - {"role": "user", "content": "Say hello"} - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert!(body["id"].as_str().unwrap().starts_with("chatcmpl-")); - assert_eq!(body["object"], "chat.completion"); - assert_eq!(body["model"], "gpt-5.5"); - assert_eq!(body["choices"][0]["message"]["role"], "assistant"); - assert_eq!(body["choices"][0]["message"]["content"], "hello"); - assert_eq!(body["choices"][0]["finish_reason"], "stop"); - assert_eq!(body["usage"]["prompt_tokens"], 9); - assert_eq!(body["usage"]["completion_tokens"], 3); - assert_eq!(body["usage"]["total_tokens"], 12); - let upstream_body = single_upstream_post_body(&server).await; - assert_eq!(upstream_body["model"], "gpt-5.5"); - assert_eq!(upstream_body["instructions"], "You are concise."); - assert_eq!(upstream_body["input"][0]["role"], "user"); - assert_eq!(upstream_body["input"][0]["content"], "Say hello"); - assert!(upstream_body["client_metadata"]["x-codex-installation-id"].is_string()); -} - -#[tokio::test] -async fn chat_completions_should_convert_and_reconvert_tuple_schema() { - 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(CHAT_TUPLE_OBJECT_SSE), - ) - .expect(1) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "messages": [{"role": "user", "content": "Return JSON"}], - "response_format": { - "type": "json_schema", - "json_schema": { - "name": "TupleAnswer", - "schema": { - "type": "object", - "properties": { - "point": { - "type": "array", - "prefixItems": [ - {"type": "number"}, - {"type": "number"} - ] - } - } - }, - "strict": true - } - } - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!( - body["choices"][0]["message"]["content"], - "{\"point\":[1,2]}" - ); - let upstream_body = single_upstream_post_body(&server).await; - let schema = &upstream_body["text"]["format"]["schema"]; - assert!(schema["properties"]["point"].get("prefixItems").is_none()); - assert_eq!(schema["properties"]["point"]["type"], "object"); -} - -#[tokio::test] -async fn chat_completions_streams_openai_chunks_from_codex_sse() { - 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")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(CHAT_STREAM_TEXT_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "stream": true, - "messages": [ - {"role": "user", "content": "Say hello"} - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let frames = response_sse_data(response).await; - assert_eq!(frames.last().unwrap(), "[DONE]"); - let chunks = frames - .iter() - .filter(|frame| frame.as_str() != "[DONE]") - .map(|frame| serde_json::from_str::(frame).unwrap()) - .collect::>(); - - assert_eq!(chunks[0]["object"], "chat.completion.chunk"); - assert_eq!(chunks[0]["model"], "gpt-5.5"); - assert_eq!(chunks[0]["choices"][0]["delta"]["role"], "assistant"); - assert_eq!(chunks[1]["choices"][0]["delta"]["content"], "hel"); - assert_eq!(chunks[2]["choices"][0]["delta"]["content"], "lo"); - assert_eq!(chunks[3]["choices"][0]["finish_reason"], "stop"); - assert_eq!(chunks[3]["usage"]["prompt_tokens"], 9); - assert_eq!(chunks[3]["usage"]["completion_tokens"], 3); - assert_eq!(chunks[3]["usage"]["total_tokens"], 12); - assert_eq!( - chunks[3]["usage"]["prompt_tokens_details"]["cached_tokens"], - 4 - ); -} - -#[tokio::test] -async fn chat_completions_stream_should_reconvert_tuple_schema() { - 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(CHAT_TUPLE_OBJECT_SSE), - ) - .expect(1) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "stream": true, - "messages": [{"role": "user", "content": "Return JSON"}], - "response_format": { - "type": "json_schema", - "json_schema": { - "name": "TupleAnswer", - "schema": { - "type": "object", - "properties": { - "point": { - "type": "array", - "prefixItems": [ - {"type": "number"}, - {"type": "number"} - ] - } - } - }, - "strict": true - } - } - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let frames = response_sse_data(response).await; - let chunks = frames - .iter() - .filter(|frame| frame.as_str() != "[DONE]") - .map(|frame| serde_json::from_str::(frame).unwrap()) - .collect::>(); - - assert_eq!( - chunks[1]["choices"][0]["delta"]["content"], - "{\"point\":[1,2]}" - ); -} - -#[tokio::test] -async fn chat_completions_forwards_parallel_tools_and_model_suffix_options() { - 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")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(CHAT_PARALLEL_TOOLS_SUCCESS_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5-high-fast", - "parallel_tool_calls": true, - "messages": [ - {"role": "user", "content": "Use tools if needed"} - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["model"], "gpt-5.5-high-fast"); - let upstream_body = single_upstream_post_body(&server).await; - assert_eq!(upstream_body["reasoning"]["effort"], "high"); - assert_eq!(upstream_body["reasoning"]["summary"], "auto"); - assert_eq!(upstream_body["service_tier"], "priority"); - assert_eq!(upstream_body["parallel_tool_calls"], true); -} - -#[tokio::test] -async fn chat_completions_collects_tool_calls_and_reasoning_content() { - 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")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(CHAT_TOOL_REASONING_COMPLETE_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "reasoning_effort": "high", - "messages": [ - {"role": "user", "content": "Find it"} - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - let message = &body["choices"][0]["message"]; - assert_eq!(message["content"], Value::Null); - assert_eq!(message["reasoning_content"], "thinking "); - assert_eq!(message["tool_calls"][0]["id"], "call_abc"); - assert_eq!(message["tool_calls"][0]["type"], "function"); - assert_eq!(message["tool_calls"][0]["function"]["name"], "lookup"); - assert_eq!( - message["tool_calls"][0]["function"]["arguments"], - "{\"city\":\"Paris\"}" - ); - assert_eq!(body["choices"][0]["finish_reason"], "tool_calls"); - assert_eq!( - body["usage"]["completion_tokens_details"]["reasoning_tokens"], - 7 - ); -} - -#[tokio::test] -async fn chat_completions_streams_tool_call_chunks_from_codex_sse() { - 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")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(CHAT_TOOL_STREAM_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/chat/completions") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "stream": true, - "messages": [ - {"role": "user", "content": "Find it"} - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let frames = response_sse_data(response).await; - let chunks = frames - .iter() - .filter(|frame| frame.as_str() != "[DONE]") - .map(|frame| serde_json::from_str::(frame).unwrap()) - .collect::>(); - - let start_call = &chunks[1]["choices"][0]["delta"]["tool_calls"][0]; - assert_eq!(start_call["index"], 0); - assert_eq!(start_call["id"], "call_abc"); - assert_eq!(start_call["type"], "function"); - assert_eq!(start_call["function"]["name"], "lookup"); - assert_eq!(start_call["function"]["arguments"], ""); - assert_eq!( - chunks[2]["choices"][0]["delta"]["tool_calls"][0]["function"]["arguments"], - "{\"city\"" - ); - assert_eq!( - chunks[3]["choices"][0]["delta"]["tool_calls"][0]["function"]["arguments"], - ":\"Paris\"}" - ); - assert_eq!(chunks[4]["choices"][0]["finish_reason"], "tool_calls"); -} diff --git a/tests/codex_serving/diagnostics_route.rs b/tests/codex_serving/diagnostics_route.rs deleted file mode 100644 index aad3735e4..000000000 --- a/tests/codex_serving/diagnostics_route.rs +++ /dev/null @@ -1,364 +0,0 @@ -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use tower::ServiceExt; -use wiremock::{ - matchers::{header, method, path}, - Mock, MockServer, ResponseTemplate, -}; - -use codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint; - -use crate::support::{ - response_json, - upstream::{build_imported_app, build_imported_app_with_fingerprint}, -}; - -fn diagnostics_fingerprint() -> Fingerprint { - Fingerprint { - originator: "Codex Desktop".to_string(), - app_version: "27.100.200".to_string(), - build_number: "9001".to_string(), - platform: "linux".to_string(), - arch: "x64".to_string(), - chromium_version: "147".to_string(), - user_agent_template: "Codex Desktop/{version} ({platform}; {arch})".to_string(), - default_headers: Fingerprint::default_headers(), - header_order: Fingerprint::default_header_order(), - } -} - -#[tokio::test] -async fn debug_diagnostics_should_return_local_runtime_pool_transport_and_path_summary() { - let imported = build_imported_app("https://chatgpt.test/backend-api".to_string()).await; - - let response = imported - .app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/debug/diagnostics") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["status"], "ok"); - assert_eq!(body["runtime"]["packageVersion"], env!("CARGO_PKG_VERSION")); - assert_eq!(body["paths"]["config"], "config.yaml"); - assert_eq!(body["paths"]["localConfig"], "local.yaml"); - assert!(body["paths"]["databaseUrl"] - .as_str() - .unwrap() - .starts_with("sqlite://")); - assert_eq!( - body["transport"]["backendBaseUrl"], - "https://chatgpt.test/backend-api" - ); - assert_eq!(body["transport"]["tls"]["forceHttp11"], false); - assert_eq!( - body["transport"]["fingerprint"]["originator"], - "Codex Desktop" - ); - assert_eq!(body["accounts"]["repositoryAvailable"], true); - assert_eq!(body["accounts"]["pool"]["total"], 1); - assert_eq!(body["accounts"]["pool"]["active"], 1); - assert_eq!(body["accounts"]["capacity"]["maxConcurrentPerAccount"], 3); - assert_eq!(body["accounts"]["capacity"]["totalSlots"], 3); - assert_eq!(body["accounts"]["capacity"]["availableSlots"], 3); - assert_eq!(body["settings"]["defaultModel"], "gpt-5.5"); - - let serialized = serde_json::to_string(&body).unwrap(); - assert!(!serialized.contains("access-secret")); - assert!(!serialized.contains("refresh-secret")); -} - -#[tokio::test] -async fn debug_diagnostics_should_report_runtime_fingerprint() { - let imported = build_imported_app_with_fingerprint( - "https://chatgpt.test/backend-api".to_string(), - diagnostics_fingerprint(), - ) - .await; - - let response = imported - .app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/debug/diagnostics") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["transport"]["fingerprint"]["appVersion"], "27.100.200"); - assert_eq!(body["transport"]["fingerprint"]["buildNumber"], "9001"); - assert_eq!(body["transport"]["fingerprint"]["chromiumVersion"], "147"); - assert_eq!( - body["transport"]["fingerprint"]["userAgent"], - "Codex Desktop/27.100.200 (linux; x64)" - ); -} - -#[tokio::test] -async fn debug_diagnostics_should_reject_forwarded_remote_requests() { - let imported = build_imported_app("https://chatgpt.test/backend-api".to_string()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("GET") - .uri("/debug/diagnostics") - .header("x-forwarded-for", "203.0.113.10") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::FORBIDDEN); -} - -#[tokio::test] -async fn debug_fingerprint_should_return_runtime_fingerprint_summary() { - let imported = build_imported_app("https://chatgpt.test/backend-api".to_string()).await; - - let response = imported - .app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/debug/fingerprint") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["source"], "runtime"); - assert_eq!(body["originator"], "Codex Desktop"); - assert_eq!(body["appVersion"], "26.519.81530"); - assert_eq!(body["buildNumber"], "3178"); - assert_eq!( - body["userAgent"], - "Codex Desktop/26.519.81530 (darwin; arm64)" - ); -} - -#[tokio::test] -async fn debug_fingerprint_should_reject_forwarded_remote_requests() { - let imported = build_imported_app("https://chatgpt.test/backend-api".to_string()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("GET") - .uri("/debug/fingerprint") - .header("x-real-ip", "203.0.113.20") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::FORBIDDEN); -} - -#[tokio::test] -async fn debug_upstream_should_probe_with_runtime_fingerprint() { - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/codex/models")) - .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({ - "error": { - "message": "missing or invalid token" - } - }))) - .expect(1) - .mount(&server) - .await; - let imported = - build_imported_app_with_fingerprint(server.uri(), diagnostics_fingerprint()).await; - - let response = imported - .app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/debug/upstream") - .header("x-request-id", "req_debug_probe_fingerprint") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - let requests = server.received_requests().await.unwrap(); - assert_eq!(body["statusCode"], 401); - assert_eq!(requests.len(), 1); - let headers = &requests[0].headers; - assert_eq!( - headers - .get("user-agent") - .and_then(|value| value.to_str().ok()), - Some("Codex Desktop/27.100.200 (linux; x64)") - ); - assert_eq!( - headers - .get("sec-ch-ua") - .and_then(|value| value.to_str().ok()), - Some("\"Chromium\";v=\"147\", \"Not:A-Brand\";v=\"24\"") - ); -} - -#[tokio::test] -async fn debug_upstream_should_probe_codex_models_endpoint_without_returning_secrets() { - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/codex/models")) - .and(header("originator", "Codex Desktop")) - .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({ - "error": { - "message": "missing or invalid token" - } - }))) - .expect(1) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/debug/upstream") - .header("x-request-id", "req_debug_probe") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["target"], "codexModels"); - assert_eq!(body["backendBaseUrl"], server.uri()); - assert_eq!(body["reachable"], true); - assert_eq!(body["statusCode"], 401); - assert_eq!(body["authorization"], "rejected"); - assert!(body["endpoint"].as_str().unwrap().contains("/codex/models")); - - let serialized = serde_json::to_string(&body).unwrap(); - assert!(!serialized.contains("access-secret")); - assert!(!serialized.contains("refresh-secret")); - - let requests = server.received_requests().await.unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].headers.get("x-client-request-id").is_none()); -} - -#[tokio::test] -async fn debug_upstream_should_reject_forwarded_remote_requests_without_probe() { - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/codex/models")) - .respond_with(ResponseTemplate::new(200)) - .expect(0) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("GET") - .uri("/debug/upstream") - .header("x-forwarded-for", "203.0.113.50") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::FORBIDDEN); -} - -#[tokio::test] -async fn admin_diagnostics_should_require_admin_session_cookie() { - let imported = build_imported_app("https://chatgpt.test/backend-api".to_string()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/diagnostics") - .header("x-request-id", "req_admin_diagnostics") - .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_admin_diagnostics"); -} - -#[tokio::test] -async fn admin_diagnostics_should_return_admin_enveloped_runtime_summary_without_secrets() { - let imported = build_imported_app("https://chatgpt.test/backend-api".to_string()).await; - - let response = imported - .app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/diagnostics") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_admin_diagnostics") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["code"], 200); - assert_eq!(body["requestId"], "req_admin_diagnostics"); - assert_eq!(body["data"]["status"], "ok"); - assert_eq!( - body["data"]["transport"]["backendBaseUrl"], - "https://chatgpt.test/backend-api" - ); - assert_eq!(body["data"]["accounts"]["pool"]["active"], 1); - assert_eq!(body["data"]["accounts"]["capacity"]["totalSlots"], 3); - - let serialized = serde_json::to_string(&body).unwrap(); - assert!(!serialized.contains("access-secret")); - assert!(!serialized.contains("refresh-secret")); -} diff --git a/tests/codex_serving/responses_http_sse.rs b/tests/codex_serving/responses_http_sse.rs deleted file mode 100644 index 4e06143c5..000000000 --- a/tests/codex_serving/responses_http_sse.rs +++ /dev/null @@ -1,1766 +0,0 @@ -use std::{ - sync::Arc, - time::{Duration, Instant}, -}; - -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use chrono::{DateTime, Utc}; -use futures::StreamExt; -use serde_json::{json, Value}; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, - sync::{oneshot, Mutex}, -}; -use tower::ServiceExt; -use wiremock::{ - matchers::{header, method, path}, - Mock, MockServer, ResponseTemplate, -}; - -use codex_proxy_rs::{ - codex::accounts::cookies::repository::CookieRepository, - codex::gateway::conversation_identity::build_conversation_identity, -}; - -use crate::support::{ - response_json, response_text, - upstream::{ - build_imported_app, build_imported_app_with_accounts_and_config, enable_runtime_logging, - fetch_v1_event_log, ImportAccount, - }, -}; - -const COMPLETED_USAGE_SSE: &str = - include_str!("../fixtures/responses/http_sse/completed_usage.sse"); -const COMPLETED_IMAGE_USAGE_SSE: &str = - include_str!("../fixtures/responses/http_sse/completed_image_usage.sse"); -const COMPLETED_FIELDS_SSE: &str = - include_str!("../fixtures/responses/http_sse/completed_fields.sse"); -const COMPLETED_REASONING_INCLUDE_SSE: &str = - include_str!("../fixtures/responses/http_sse/completed_reasoning_include.sse"); -const TEXT_DELTAS_COMPLETED_SSE: &str = - include_str!("../fixtures/responses/http_sse/text_deltas_completed.sse"); -const DONE_ITEM_COMPLETED_SSE: &str = - include_str!("../fixtures/responses/http_sse/done_item_completed.sse"); -const STREAM_USAGE_SSE: &str = include_str!("../fixtures/responses/http_sse/stream_usage.sse"); -const DEFAULT_STREAM_SSE: &str = include_str!("../fixtures/responses/http_sse/default_stream.sse"); -const EMPTY_COMPLETED_SSE: &str = - include_str!("../fixtures/responses/http_sse/empty_completed.sse"); -const COMPLETED_TUPLE_OBJECT_SSE: &str = - include_str!("../fixtures/responses/http_sse/completed_tuple_object.sse"); -const TUPLE_RECONVERT_RESPONSE_GOLDEN: &str = - include_str!("../fixtures/responses/golden/tuple_reconvert_response.json"); -const TUPLE_RECONVERT_STREAM_GOLDEN: &str = concat!( - include_str!("../fixtures/responses/golden/tuple_reconvert_stream.sse"), - "\n" -); -const COMPACT_RATE_LIMIT_FALLBACK_EXHAUSTED_GOLDEN: &str = - include_str!("../fixtures/responses/golden/compact_rate_limit_fallback_exhausted.json"); -const COMPACT_NO_AVAILABLE_ACCOUNTS_GOLDEN: &str = - include_str!("../fixtures/responses/golden/compact_no_available_accounts.json"); - -#[tokio::test] -async fn v1_responses_should_reject_invalid_json_without_upstream_request() { - let server = MockServer::start().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("{")) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_REQUEST); - let body = response_json(response).await; - assert_eq!(body["error"]["code"], "invalid_request"); - let requests = server.received_requests().await.unwrap(); - assert!(requests.is_empty()); -} - -#[tokio::test] -async fn v1_responses_should_reject_non_object_json_without_upstream_request() { - let server = MockServer::start().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("[]")) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_REQUEST); - let body = response_json(response).await; - assert_eq!(body["error"]["code"], "invalid_request"); - let requests = server.received_requests().await.unwrap(); - assert!(requests.is_empty()); -} - -#[tokio::test] -async fn v1_responses_should_return_responses_error_when_no_accounts_are_available() { - let server = MockServer::start().await; - let imported = build_imported_app_with_accounts_and_config(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( - json!({ - "model": "gpt-5.5", - "input": [], - "stream": false - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); - let body = response_json(response).await; - assert_eq!(body["type"], "error"); - assert_eq!(body["error"]["type"], "server_error"); - assert_eq!(body["error"]["code"], "no_available_accounts"); -} - -#[tokio::test] -async fn v1_responses_should_skip_event_log_when_logging_disabled() { - 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(COMPLETED_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_logs_disabled") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let stored_count: (i64,) = sqlx::query_as( - "select count(*) from event_logs where request_id = ? and kind = 'v1.response'", - ) - .bind("req_logs_disabled") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(stored_count.0, 0); -} - -#[tokio::test] -async fn v1_responses_should_honor_explicit_http_sse_transport() { - 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(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":[],"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_text(response).await; - assert!(body.contains("event: response.output_text.delta")); - let requests = server.received_requests().await.unwrap(); - assert_eq!(requests.len(), 1); - assert_eq!(requests[0].method.as_str(), "POST"); - let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert_eq!(upstream_body["model"], "gpt-5.5"); - assert!(upstream_body.get("use_websocket").is_none()); -} - -#[tokio::test] -async fn v1_responses_should_stagger_same_account_requests_before_sending_upstream() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let request_times = Arc::new(Mutex::new(Vec::new())); - let request_times_for_server = Arc::clone(&request_times); - let (first_seen_tx, first_seen_rx) = oneshot::channel(); - let (release_first_tx, release_first_rx) = oneshot::channel(); - let server = tokio::spawn(async move { - let (mut first_stream, _) = listener.accept().await.unwrap(); - request_times_for_server.lock().await.push(Instant::now()); - first_seen_tx.send(()).unwrap(); - let _first_request = read_http_request(&mut first_stream).await; - - let (mut second_stream, _) = listener.accept().await.unwrap(); - request_times_for_server.lock().await.push(Instant::now()); - let _second_request = read_http_request(&mut second_stream).await; - write_http_sse_response(&mut second_stream, COMPLETED_USAGE_SSE).await; - - release_first_rx.await.unwrap(); - write_http_sse_response(&mut first_stream, COMPLETED_USAGE_SSE).await; - }); - let imported = build_imported_app_with_accounts_and_config( - format!("http://{addr}"), - &[ImportAccount { - id: "acct_stagger", - account_id: "chatgpt-stagger", - token: "access-stagger", - refresh_token: "refresh-stagger", - }], - |config| { - config.auth.max_concurrent_per_account = 2; - config.auth.request_interval_ms = 300; - }, - ) - .await; - - let first_app = imported.app.clone(); - let first_api_key = imported.client_api_key.clone(); - let first_response = tokio::spawn(async move { - first_app - .oneshot(v1_response_request(&first_api_key, "req_stagger_first")) - .await - .unwrap() - }); - first_seen_rx.await.unwrap(); - - let second = imported - .app - .clone() - .oneshot(v1_response_request( - &imported.client_api_key, - "req_stagger_second", - )) - .await - .unwrap(); - assert_eq!(second.status(), StatusCode::OK); - release_first_tx.send(()).unwrap(); - let first = first_response.await.unwrap(); - assert_eq!(first.status(), StatusCode::OK); - server.await.unwrap(); - - let times = request_times.lock().await; - assert_eq!(times.len(), 2); - let elapsed = times[1].duration_since(times[0]); - assert!( - elapsed >= Duration::from_millis(180), - "second upstream request was sent too early: {elapsed:?}" - ); -} - -#[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")) - .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; - enable_runtime_logging(&imported.app).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,"use_websocket":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_record_image_generation_usage_when_tool_succeeds() { - 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(COMPLETED_IMAGE_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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false,"use_websocket":false,"tools":[{"type":"image_generation"}]}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let usage: (i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64) = - sqlx::query_as( - "select request_count, input_tokens, output_tokens, cached_tokens, image_input_tokens, image_output_tokens, image_request_count, image_request_failed_count, window_image_input_tokens, window_image_output_tokens, window_image_request_count, window_image_request_failed_count from account_usage where account_id = ?", - ) - .bind("acct_imported") - .fetch_one(&imported.pool) - .await - .unwrap(); - - assert_eq!(usage, (1, 7, 4, 2, 31, 9, 1, 0, 31, 9, 1, 0)); -} - -#[tokio::test] -async fn v1_responses_should_record_failed_image_generation_attempt_when_tool_has_no_output() { - 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(COMPLETED_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") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false,"use_websocket":false,"tools":[{"type":"image_generation"}]}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let usage: (i64, i64, i64, i64) = sqlx::query_as( - "select image_request_count, image_request_failed_count, window_image_request_count, window_image_request_failed_count from account_usage where account_id = ?", - ) - .bind("acct_imported") - .fetch_one(&imported.pool) - .await - .unwrap(); - - assert_eq!(usage, (0, 1, 0, 1)); -} - -#[tokio::test] -async fn v1_responses_should_scope_upstream_cookie_by_codex_response_path() { - 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(COMPLETED_USAGE_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - let cookie_domain = reqwest::Url::parse(&server.uri()) - .unwrap() - .host_str() - .unwrap() - .to_string(); - let cookie_repo = CookieRepository::new(imported.pool.clone(), imported.secret_box.clone()); - cookie_repo - .capture_set_cookie( - "acct_imported", - &format!("cf_clearance=root; Domain={cookie_domain}; Path=/"), - ) - .await - .unwrap(); - cookie_repo - .capture_set_cookie( - "acct_imported", - &format!("cf_clearance=codex; Domain={cookie_domain}; Path=/codex"), - ) - .await - .unwrap(); - - 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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let requests = server.received_requests().await.unwrap(); - let cookie_header = requests - .iter() - .find(|request| request.url.path() == "/codex/responses") - .and_then(|request| request.headers.get("cookie")) - .and_then(|value| value.to_str().ok()); - assert_eq!(cookie_header, Some("cf_clearance=codex; cf_clearance=root")); -} - -#[tokio::test] -async fn v1_responses_should_passively_cache_rate_limit_headers() { - let server = MockServer::start().await; - let reset_at = Utc::now().timestamp() + 300; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .insert_header("x-codex-primary-used-percent", "100") - .insert_header("x-codex-primary-window-minutes", "5") - .insert_header("x-codex-primary-reset-at", reset_at.to_string()) - .set_body_string(COMPLETED_USAGE_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - sqlx::query( - r#"update accounts set quota_json = '{"credits":{"has_credits":true,"unlimited":false,"balance":12}}', quota_verify_required = 1 where id = ?"#, - ) - .bind("acct_imported") - .execute(&imported.pool) - .await - .unwrap(); - - 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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let stored: (String, i64, Option, i64) = sqlx::query_as( - "select quota_json, quota_limit_reached, quota_cooldown_until, quota_verify_required from accounts where id = ?", - ) - .bind("acct_imported") - .fetch_one(&imported.pool) - .await - .unwrap(); - let quota: Value = serde_json::from_str(&stored.0).unwrap(); - assert_eq!(quota["rate_limit"]["limit_reached"], true); - assert_eq!(quota["rate_limit"]["reset_at"], reset_at); - assert_eq!(quota["credits"]["balance"], 12); - assert_eq!(stored.1, 1); - assert!(stored.2.is_some()); - assert_eq!(stored.3, 0); - let window: (i64, i64, i64, i64, Option, Option, Option) = - sqlx::query_as( - "select window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds from account_usage where account_id = ?", - ) - .bind("acct_imported") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(window.0, 1); - assert_eq!(window.1, 7); - assert_eq!(window.2, 4); - assert_eq!(window.3, 2); - assert!(window.4.is_some()); - assert_eq!( - DateTime::parse_from_rfc3339(window.5.as_deref().unwrap()) - .unwrap() - .timestamp(), - reset_at - ); - assert_eq!(window.6, Some(300)); -} - -#[tokio::test] -async fn v1_responses_should_record_empty_response_attempts() { - 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(EMPTY_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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_GATEWAY); - let usage: (i64, i64, i64, i64) = sqlx::query_as( - "select request_count, empty_response_count, input_tokens, output_tokens from account_usage where account_id = ?", - ) - .bind("acct_imported") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(usage, (3, 3, 0, 0)); -} - -#[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-body")) - .and(header("x-codex-beta-features", "beta-body")) - .and(header("x-responsesapi-include-timing-metrics", "true")) - .and(header("version", "2026-06-12")) - .and(header("x-codex-parent-thread-id", "parent-body")) - .and(header("x-openai-subagent", "review")) - .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-codex-window-id", "window-header") - .header("x-codex-parent-thread-id", "parent-header") - .header("x-openai-subagent", "review") - .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", - "turnMetadata": "meta-body", - "betaFeatures": "beta-body", - "includeTimingMetrics": "true", - "version": "2026-06-12", - "codexWindowId": "window-body", - "parentThreadId": "parent-body", - "use_websocket": false - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_fields"); - let requests = server.received_requests().await.unwrap(); - assert_eq!(requests.len(), 1); - let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - let identity = - build_conversation_identity(Some("pcache"), Some("window-body"), "acct_imported"); - assert_eq!( - upstream_body["prompt_cache_key"], - identity.conversation_id.unwrap() - ); - assert_eq!(upstream_body["model"], "gpt-5.5"); - assert_eq!(upstream_body["service_tier"], "priority"); - assert_eq!(upstream_body["client_metadata"]["safe"], "yes"); - assert_eq!( - upstream_body["client_metadata"]["x-openai-subagent"], - "review" - ); - assert_eq!( - upstream_body["client_metadata"]["x-codex-window-id"], - identity.window_id.unwrap() - ); - assert_eq!( - upstream_body["client_metadata"]["x-codex-turn-metadata"], - "meta-body" - ); - assert_eq!( - upstream_body["client_metadata"]["x-codex-parent-thread-id"], - "parent-body" - ); - let installation_id = upstream_body["client_metadata"]["x-codex-installation-id"] - .as_str() - .unwrap(); - assert!(uuid::Uuid::parse_str(installation_id).is_ok()); - for local_field in [ - "turnState", - "turnMetadata", - "betaFeatures", - "includeTimingMetrics", - "version", - "codexWindowId", - "parentThreadId", - "use_websocket", - ] { - assert!(upstream_body.get(local_field).is_none()); - } -} - -#[tokio::test] -async fn v1_responses_should_convert_tuple_schema_before_upstream() { - 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(COMPLETED_FIELDS_SSE), - ) - .expect(1) - .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( - json!({ - "model": "gpt-5.5", - "stream": false, - "use_websocket": false, - "input": [], - "text": { - "format": { - "type": "json_schema", - "name": "TupleAnswer", - "schema": { - "type": "object", - "properties": { - "point": { - "type": "array", - "prefixItems": [ - {"type": "number"}, - {"type": "number"} - ], - "items": false - } - }, - "required": ["point"] - }, - "strict": true - } - } - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let requests = server.received_requests().await.unwrap(); - let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - let schema = &upstream_body["text"]["format"]["schema"]; - assert!(schema["properties"]["point"].get("prefixItems").is_none()); - assert!(schema["properties"]["point"].get("items").is_none()); - assert_eq!(schema["additionalProperties"], false); - assert_eq!( - schema["properties"]["point"], - json!({ - "type": "object", - "properties": { - "0": {"type": "number"}, - "1": {"type": "number"} - }, - "required": ["0", "1"], - "additionalProperties": false - }) - ); -} - -#[tokio::test] -async fn v1_responses_should_reconvert_tuple_schema_output_for_client() { - 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(COMPLETED_TUPLE_OBJECT_SSE), - ) - .expect(1) - .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( - json!({ - "model": "gpt-5.5", - "stream": false, - "use_websocket": false, - "input": [], - "text": { - "format": { - "type": "json_schema", - "name": "TupleAnswer", - "schema": { - "type": "object", - "properties": { - "point": { - "type": "array", - "prefixItems": [ - {"type": "number"}, - {"type": "number"} - ] - } - } - }, - "strict": true - } - } - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - let expected: Value = serde_json::from_str(TUPLE_RECONVERT_RESPONSE_GOLDEN).unwrap(); - assert_eq!(body, expected); -} - -#[tokio::test] -async fn v1_responses_stream_should_reconvert_tuple_schema_output_for_client() { - 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(COMPLETED_TUPLE_OBJECT_SSE), - ) - .expect(1) - .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( - json!({ - "model": "gpt-5.5", - "stream": true, - "use_websocket": false, - "input": [], - "text": { - "format": { - "type": "json_schema", - "name": "TupleAnswer", - "schema": { - "type": "object", - "properties": { - "point": { - "type": "array", - "prefixItems": [ - {"type": "number"}, - {"type": "number"} - ] - } - } - }, - "strict": true - } - } - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_text(response).await; - assert_eq!(body, TUPLE_RECONVERT_STREAM_GOLDEN); -} - -#[tokio::test] -async fn v1_responses_review_route_should_force_review_subagent_upstream() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("x-openai-subagent", "review")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(COMPLETED_FIELDS_SSE), - ) - .expect(1) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/responses/review") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "stream": false, - "use_websocket": false, - "input": [] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let requests = server.received_requests().await.unwrap(); - let post_request = requests - .iter() - .find(|request| request.method.as_str() == "POST" && !request.body.is_empty()) - .unwrap(); - let upstream_body: Value = serde_json::from_slice(&post_request.body).unwrap(); - assert_eq!( - upstream_body["client_metadata"]["x-openai-subagent"], - "review" - ); -} - -#[tokio::test] -async fn v1_responses_compact_should_post_json_to_codex_compact_upstream() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses/compact")) - .and(header("authorization", "Bearer access-secret")) - .and(header("chatgpt-account-id", "chatgpt-account")) - .and(header("content-type", "application/json")) - .and(header("openai-beta", "responses_websockets=2026-02-06")) - .and(header("x-openai-internal-codex-residency", "us")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "output": [ - { - "type": "message", - "role": "assistant", - "content": [{"type": "output_text", "text": "compacted"}] - } - ] - }))) - .expect(1) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/responses/compact") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5-fast", - "instructions": "compress the session", - "input": [ - {"role": "user", "content": "hello"}, - { - "type": "reasoning", - "id": "rs_1", - "status": "completed", - "summary": [{"type": "summary_text", "text": "kept"}], - "ignored": "drop" - }, - {"type": "compaction", "encrypted_content": "enc_compact"}, - {"type": "compaction", "id": "drop_missing_encrypted"} - ], - "tools": [{"type": "function", "name": "lookup"}], - "parallel_tool_calls": false, - "reasoning": {"effort": "high", "summary": "auto", "extra": "drop"}, - "text": { - "format": { - "type": "json_schema", - "name": "Compact", - "schema": {"type": "object"}, - "strict": true - } - }, - "stream": true, - "store": true, - "prompt_cache_key": "must_not_forward" - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["output"][0]["content"][0]["text"], "compacted"); - - let requests = server.received_requests().await.unwrap(); - let compact_request = requests - .iter() - .find(|request| request.url.path() == "/codex/responses/compact") - .unwrap(); - assert_ne!( - compact_request - .headers - .get("accept") - .and_then(|value| value.to_str().ok()), - Some("text/event-stream") - ); - assert!(compact_request - .headers - .get("x-codex-installation-id") - .is_some()); - let upstream_body: Value = serde_json::from_slice(&compact_request.body).unwrap(); - assert_eq!(upstream_body["model"], "gpt-5.5"); - assert_eq!(upstream_body["instructions"], "compress the session"); - assert_eq!(upstream_body["parallel_tool_calls"], false); - assert_eq!( - upstream_body["reasoning"], - json!({"effort": "high", "summary": "auto"}) - ); - assert_eq!( - upstream_body["tools"], - json!([{"type": "function", "name": "lookup"}]) - ); - assert_eq!(upstream_body["text"]["format"]["type"], "json_schema"); - assert!(upstream_body.get("stream").is_none()); - assert!(upstream_body.get("store").is_none()); - assert!(upstream_body.get("prompt_cache_key").is_none()); - assert_eq!(upstream_body["input"].as_array().unwrap().len(), 3); - assert!(upstream_body["input"][1].get("ignored").is_none()); - assert_eq!( - upstream_body["input"][2]["encrypted_content"], - "enc_compact" - ); -} - -#[tokio::test] -async fn v1_responses_compact_should_return_rate_limit_error_when_fallback_is_exhausted() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses/compact")) - .respond_with(ResponseTemplate::new(429).set_body_json(json!({ - "error": { - "code": "rate_limit_exceeded", - "message": "compact quota reached" - } - }))) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/responses/compact") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "input": [{"role": "user", "content": "hello"}] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); - let body = response_json(response).await; - let expected: Value = serde_json::from_str(COMPACT_RATE_LIMIT_FALLBACK_EXHAUSTED_GOLDEN) - .expect("compact rate-limit golden should parse"); - assert_eq!(body, expected); -} - -#[tokio::test] -async fn v1_responses_compact_should_return_responses_error_when_no_accounts_are_available() { - let server = MockServer::start().await; - let imported = build_imported_app_with_accounts_and_config(server.uri(), &[], |_| {}).await; - - let response = imported - .app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/responses/compact") - .header( - "authorization", - format!("Bearer {}", imported.client_api_key), - ) - .header("content-type", "application/json") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "input": [{"role": "user", "content": "hello"}] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); - let body = response_json(response).await; - let expected: Value = serde_json::from_str(COMPACT_NO_AVAILABLE_ACCOUNTS_GOLDEN) - .expect("compact no-accounts golden should parse"); - assert_eq!(body, expected); -} - -#[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")) - .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,"use_websocket":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"); - let requests = server.received_requests().await.unwrap(); - let post_request = requests - .iter() - .find(|request| request.method.as_str() == "POST") - .unwrap(); - let upstream_body: Value = serde_json::from_slice(&post_request.body).unwrap(); - assert_eq!( - upstream_body["include"], - json!(["reasoning.encrypted_content"]) - ); - assert_eq!(upstream_body["reasoning"]["summary"], "auto"); -} - -#[tokio::test] -async fn v1_responses_should_not_add_reasoning_include_when_client_include_is_non_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(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( - json!({ - "model": "gpt-5.5", - "stream": false, - "use_websocket": false, - "input": [], - "reasoning": {"effort": "high"}, - "include": ["file_search_call.results"] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let requests = server.received_requests().await.unwrap(); - let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert_eq!( - upstream_body["include"], - json!(["file_search_call.results"]) - ); -} - -#[tokio::test] -async fn v1_responses_should_sanitize_reasoning_and_compaction_input_before_upstream() { - 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(COMPLETED_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") - .body(Body::from( - json!({ - "model": "gpt-5.5", - "stream": false, - "use_websocket": false, - "input": [ - { - "type": "reasoning", - "id": "rs_1", - "status": "completed", - "summary": [ - {"type": "summary_text", "text": "valid summary"}, - {"type": "ignored", "text": "drop"} - ], - "encrypted_content": "enc_reasoning", - "content": [ - {"type": "reasoning_text", "text": "valid reasoning"}, - {"type": "ignored", "text": "drop"} - ], - "extra": "drop" - }, - { - "type": "reasoning", - "id": "", - "summary": [{"type": "summary_text", "text": "drop"}] - }, - { - "type": "compaction", - "id": "cmp_1", - "encrypted_content": "enc_compaction", - "extra": "drop" - }, - {"type": "compaction", "id": "cmp_drop"}, - {"type": "message", "role": "user", "content": "keep me", "extra": 42} - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let requests = server.received_requests().await.unwrap(); - let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap(); - assert_eq!( - upstream_body["input"], - json!([ - { - "type": "reasoning", - "id": "rs_1", - "summary": [{"type": "summary_text", "text": "valid summary"}], - "status": "completed", - "encrypted_content": "enc_reasoning", - "content": [{"type": "reasoning_text", "text": "valid reasoning"}] - }, - { - "type": "compaction", - "encrypted_content": "enc_compaction", - "id": "cmp_1" - }, - {"type": "message", "role": "user", "content": "keep me", "extra": 42} - ]) - ); -} - -#[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,"use_websocket":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,"use_websocket":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")) - .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; - enable_runtime_logging(&imported.app).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,"use_websocket":false}"#, - )) - .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_close_http_sse_upstream_when_client_disconnects() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.unwrap(); - let _request = read_http_request(&mut stream).await; - write_chunked_http_sse_headers(&mut stream).await; - write_http_chunk( - &mut stream, - b"event: response.output_text.delta\ndata: {\"delta\":\"hello\"}\n\n", - ) - .await; - - tokio::time::timeout( - Duration::from_secs(2), - wait_for_http_sse_upstream_disconnect(&mut stream), - ) - .await - .is_ok() - }); - 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":[],"stream":true,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let mut body = response.into_body().into_data_stream(); - let first_chunk = tokio::time::timeout(Duration::from_secs(1), body.next()) - .await - .expect("first SSE chunk should arrive before disconnect") - .expect("stream should yield a first chunk") - .expect("chunk should be readable"); - let first_sse = String::from_utf8(first_chunk.to_vec()).unwrap(); - assert!(first_sse.contains("event: response.output_text.delta")); - - drop(body); - assert!( - server.await.unwrap(), - "dropping the downstream stream should close the HTTP SSE upstream socket" - ); -} - -#[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")) - .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":[],"use_websocket":false}"#, - )) - .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")); -} - -fn v1_response_request(api_key: &str, request_id: &str) -> Request { - Request::builder() - .method("POST") - .uri("/v1/responses") - .header("authorization", format!("Bearer {api_key}")) - .header("content-type", "application/json") - .header("x-request-id", request_id) - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false,"use_websocket":false}"#, - )) - .unwrap() -} - -async fn read_http_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 write_http_sse_response(stream: &mut TcpStream, body: &str) { - let response = format!( - "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", - body.len() - ); - stream.write_all(response.as_bytes()).await.unwrap(); -} - -async fn write_chunked_http_sse_headers(stream: &mut TcpStream) { - stream - .write_all( - b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n", - ) - .await - .unwrap(); -} - -async fn write_http_chunk(stream: &mut TcpStream, chunk: &[u8]) { - stream - .write_all(format!("{:x}\r\n", chunk.len()).as_bytes()) - .await - .unwrap(); - stream.write_all(chunk).await.unwrap(); - stream.write_all(b"\r\n").await.unwrap(); - stream.flush().await.unwrap(); -} - -async fn wait_for_http_sse_upstream_disconnect(stream: &mut TcpStream) { - let mut buffer = [0u8; 1024]; - loop { - match stream.read(&mut buffer).await { - Ok(0) | Err(_) => return, - Ok(_) => {} - } - } -} diff --git a/tests/codex_serving/responses_websocket.rs b/tests/codex_serving/responses_websocket.rs deleted file mode 100644 index 07a7b9eeb..000000000 --- a/tests/codex_serving/responses_websocket.rs +++ /dev/null @@ -1,2863 +0,0 @@ -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use futures::{SinkExt, StreamExt}; -use serde_json::{json, Value}; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, - sync::oneshot, - time::{timeout, Duration}, -}; -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::runtime::{build_router, state::AppState}; - -use crate::support::{ - assert_response_failed_stream, response_json, response_text, - upstream::{ - build_imported_app, build_imported_app_with_accounts, - build_imported_app_with_accounts_and_config, ImportAccount, - }, -}; - -const WEBSOCKET_COMPLETED_RESPONSE: &str = - include_str!("../fixtures/responses/websocket/completed.json"); -const WEBSOCKET_HISTORY_RATE_LIMITED: &str = - include_str!("../fixtures/responses/websocket/history_rate_limited.json"); -const WEBSOCKET_RATE_LIMITED: &str = - include_str!("../fixtures/responses/websocket/rate_limited.json"); -const WEBSOCKET_TOKEN_REVOKED: &str = - include_str!("../fixtures/responses/websocket/token_revoked.json"); -const WEBSOCKET_FIRST_ACCOUNT_LIMITED: &str = - include_str!("../fixtures/responses/websocket/first_account_limited.json"); -const WEBSOCKET_SECOND_ACCOUNT_LIMITED: &str = - include_str!("../fixtures/responses/websocket/second_account_limited.json"); -const WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND: &str = - include_str!("../fixtures/responses/websocket/previous_response_not_found.json"); -const WEBSOCKET_UNANSWERED_FUNCTION_CALL: &str = - include_str!("../fixtures/responses/websocket/unanswered_function_call.json"); -const WEBSOCKET_PARTIAL_DELTA: &str = - include_str!("../fixtures/responses/websocket/partial_delta.json"); -const WEBSOCKET_INVALID_ENCRYPTED_CONTENT: &str = - include_str!("../fixtures/responses/websocket/invalid_encrypted_content.json"); -const WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY: &str = - include_str!("../fixtures/responses/websocket/completed_with_reasoning_replay.json"); -const REASONING_REPLAY_REQUEST_GOLDEN: &str = - include_str!("../fixtures/responses/golden/reasoning_replay_request.json"); - -#[expect( - clippy::result_large_err, - reason = "tungstenite 的 header callback API 固定使用较大的 handshake error response" -)] -fn assert_access_secret_header( - request: &WsRequest, - response: WsResponse, -) -> Result { - assert_eq!( - request - .headers() - .get("authorization") - .and_then(|value| value.to_str().ok()), - Some("Bearer access-secret") - ); - Ok(response) -} - -#[expect( - clippy::result_large_err, - reason = "tungstenite 的 header callback API 固定使用较大的 handshake error response" -)] -fn assert_access_a_header( - request: &WsRequest, - response: WsResponse, -) -> Result { - assert_eq!( - request - .headers() - .get("authorization") - .and_then(|value| value.to_str().ok()), - Some("Bearer access-a") - ); - Ok(response) -} - -#[expect( - clippy::result_large_err, - reason = "tungstenite 的 header callback API 固定使用较大的 handshake error response" -)] -fn assert_access_b_header( - request: &WsRequest, - response: WsResponse, -) -> Result { - assert_eq!( - request - .headers() - .get("authorization") - .and_then(|value| value.to_str().ok()), - Some("Bearer access-b") - ); - Ok(response) -} - -#[expect( - clippy::result_large_err, - reason = "tungstenite 的 header callback API 固定使用较大的 handshake error response" -)] -fn assert_metadata_turn_state_header( - request: &WsRequest, - response: WsResponse, -) -> Result { - assert_eq!( - request - .headers() - .get("x-codex-turn-state") - .and_then(|value| value.to_str().ok()), - Some("turn-from-metadata") - ); - Ok(response) -} - -#[tokio::test] -async fn v1_responses_should_use_websocket_upstream_by_default_while_serving_sse() { - 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_default", 6, 2).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":[],"generate":false}"#, - )) - .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_default\"")); - let request = request_rx.await.unwrap(); - assert_eq!(request["type"], "response.create"); - assert_eq!(request["model"], "gpt-5.5"); - assert_eq!(request["generate"], false); - assert!(request.get("previous_response_id").is_none()); - assert!(request["prompt_cache_key"] - .as_str() - .is_some_and(|value| value.starts_with("cp_"))); - assert!(request["client_metadata"]["x-codex-installation-id"].is_string()); - assert!(request["client_metadata"]["x-codex-window-id"] - .as_str() - .is_some_and(|value| value.starts_with("cp_") && value.ends_with(":0"))); - server.await.unwrap(); -} - -#[tokio::test] -async fn v1_responses_should_ignore_camel_case_use_websocket_field() { - 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_camel_case_transport", 6, 2).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":[],"useWebSocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let request = request_rx.await.unwrap(); - assert_eq!(request["type"], "response.create"); - assert!(request.get("useWebSocket").is_none()); - server.await.unwrap(); -} - -#[tokio::test] -async fn v1_responses_websocket_should_stream_first_frame_before_terminal_event() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let (first_frame_tx, first_frame_rx) = oneshot::channel(); - let (terminal_tx, terminal_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(); - websocket - .send(Message::Text( - json!({ - "type": "response.output_text.delta", - "delta": "first" - }) - .to_string() - .into(), - )) - .await - .unwrap(); - first_frame_tx.send(()).unwrap(); - terminal_rx.await.unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_route_ws_streaming", 4, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - }); - let imported = build_imported_app(format!("http://{addr}")).await; - let app = imported.app.clone(); - let request = 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(); - let mut response_task = tokio::spawn(async move { app.oneshot(request).await.unwrap() }); - - first_frame_rx.await.unwrap(); - let response = match timeout(Duration::from_millis(250), &mut response_task).await { - Ok(response) => response.unwrap(), - Err(_) => { - let _ = terminal_tx.send(()); - let _ = timeout(Duration::from_secs(1), response_task).await; - panic!("websocket response should be returned after the first non-error frame"); - } - }; - - assert_eq!(response.status(), StatusCode::OK); - let mut body = response.into_body().into_data_stream(); - let first_chunk = timeout(Duration::from_secs(1), body.next()) - .await - .expect("first SSE chunk should arrive before terminal frame") - .expect("stream should yield one chunk") - .expect("chunk should be readable"); - let first_sse = String::from_utf8(first_chunk.to_vec()).unwrap(); - assert!(first_sse.contains("event: response.output_text.delta")); - - terminal_tx.send(()).unwrap(); - while let Some(chunk) = timeout(Duration::from_secs(1), body.next()).await.unwrap() { - chunk.unwrap(); - } - server.await.unwrap(); -} - -#[tokio::test] -async fn v1_responses_websocket_stream_should_synthesize_response_failed_when_closed_before_terminal( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - 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(); - websocket - .send(Message::Text(WEBSOCKET_PARTIAL_DELTA.trim().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":[],"stream":true,"use_websocket":true}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_text(response).await; - assert!(body.contains("event: response.output_text.delta")); - assert!(body.contains("event: response.failed")); - assert!(body.contains("stream_disconnected")); - server.await.unwrap(); -} - -#[tokio::test] -async fn v1_responses_websocket_should_reuse_connection_for_recorded_conversation() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_async(stream).await.unwrap(); - let first_message = websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_pool_first", 4, 1).into(), - )) - .await - .unwrap(); - - loop { - tokio::select! { - message = websocket.next() => { - match message { - Some(Ok(message)) if message.is_text() => { - let second_request = serde_json::from_str::( - &message.into_text().unwrap(), - ) - .unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_pool_second", 3, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - break (true, first_request, second_request); - } - Some(_) => continue, - None => { - let second_request = accept_successful_websocket_response( - &listener, - "Bearer access-secret", - "resp_pool_second", - ) - .await; - break (false, first_request, second_request); - } - } - } - accepted = listener.accept() => { - let (stream, _) = accepted.unwrap(); - let mut second_websocket = accept_async(stream).await.unwrap(); - let second_message = second_websocket.next().await.unwrap().unwrap(); - let second_request = serde_json::from_str::( - &second_message.into_text().unwrap(), - ) - .unwrap(); - second_websocket - .send(Message::Text( - websocket_completed_response("resp_pool_second", 3, 1).into(), - )) - .await - .unwrap(); - second_websocket.close(None).await.unwrap(); - break (false, first_request, second_request); - } - } - } - }); - let imported = build_imported_app(format!("http://{addr}")).await; - - let first_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( - json!({ - "model": "gpt-5.5", - "input": [{ - "role": "user", - "content": "reuse this upstream websocket" - }] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(first_response.status(), StatusCode::OK); - let first_body = response_text(first_response).await; - assert!(first_body.contains("\"id\":\"resp_pool_first\"")); - - let second_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_pool_first"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::OK); - let second_body = response_text(second_response).await; - assert!(second_body.contains("\"id\":\"resp_pool_second\"")); - - let (reused_connection, first_request, second_request) = server.await.unwrap(); - assert!(reused_connection, "second request opened a new websocket"); - assert_eq!( - second_request["prompt_cache_key"], first_request["prompt_cache_key"], - "pooled websocket reuse should stay on the recorded conversation key" - ); - assert_eq!(second_request["previous_response_id"], "resp_pool_first"); -} - -#[tokio::test] -async fn v1_responses_websocket_should_not_reuse_connection_when_pool_is_disabled() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_async(stream).await.unwrap(); - let first_message = websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_disabled_pool_first", 4, 1).into(), - )) - .await - .unwrap(); - - loop { - tokio::select! { - message = websocket.next() => { - match message { - Some(Ok(message)) if message.is_text() => { - let second_request = serde_json::from_str::( - &message.into_text().unwrap(), - ) - .unwrap(); - websocket - .send(Message::Text( - websocket_completed_response( - "resp_disabled_pool_second", - 3, - 1, - ) - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - break (true, first_request, second_request); - } - Some(_) => continue, - None => { - let second_request = accept_successful_websocket_response( - &listener, - "Bearer access-secret", - "resp_disabled_pool_second", - ) - .await; - break (false, first_request, second_request); - } - } - } - accepted = listener.accept() => { - let (stream, _) = accepted.unwrap(); - let mut second_websocket = accept_async(stream).await.unwrap(); - let second_message = second_websocket.next().await.unwrap().unwrap(); - let second_request = serde_json::from_str::( - &second_message.into_text().unwrap(), - ) - .unwrap(); - second_websocket - .send(Message::Text( - websocket_completed_response("resp_disabled_pool_second", 3, 1).into(), - )) - .await - .unwrap(); - second_websocket.close(None).await.unwrap(); - break (false, first_request, second_request); - } - } - } - }); - let imported = build_imported_app_with_accounts_and_config( - format!("http://{addr}"), - &[ImportAccount { - id: "acct_imported", - account_id: "chatgpt-account", - token: "access-secret", - refresh_token: "refresh-secret", - }], - |config| { - config.ws_pool.enabled = false; - }, - ) - .await; - - let first_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( - json!({ - "model": "gpt-5.5", - "input": [{ - "role": "user", - "content": "do not reuse this upstream websocket" - }] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(first_response.status(), StatusCode::OK); - let first_body = response_text(first_response).await; - assert!(first_body.contains("\"id\":\"resp_disabled_pool_first\"")); - - let second_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_disabled_pool_first"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::OK); - let second_body = response_text(second_response).await; - assert!(second_body.contains("\"id\":\"resp_disabled_pool_second\"")); - - let (reused_connection, first_request, second_request) = server.await.unwrap(); - assert!( - !reused_connection, - "disabled pool reused the upstream websocket" - ); - assert_eq!( - second_request["prompt_cache_key"], first_request["prompt_cache_key"], - "disabling the pool must not change the recorded conversation key" - ); - assert_eq!( - second_request["previous_response_id"], - "resp_disabled_pool_first" - ); -} - -#[tokio::test] -async fn v1_responses_websocket_stream_should_record_metadata_turn_state_for_continuation() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (first_stream, _) = listener.accept().await.unwrap(); - let mut first_websocket = accept_async(first_stream).await.unwrap(); - let first_message = first_websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - first_websocket - .send(Message::Text( - json!({ - "type": "response.metadata", - "headers": { - "x-codex-turn-state": "turn-from-metadata" - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - first_websocket - .send(Message::Text( - websocket_completed_response("resp_metadata_turn_state", 4, 1).into(), - )) - .await - .unwrap(); - first_websocket.close(None).await.unwrap(); - - let (second_stream, _) = listener.accept().await.unwrap(); - let mut second_websocket = - accept_hdr_async(second_stream, assert_metadata_turn_state_header) - .await - .unwrap(); - let second_message = second_websocket.next().await.unwrap().unwrap(); - let second_request = - serde_json::from_str::(&second_message.into_text().unwrap()).unwrap(); - second_websocket - .send(Message::Text( - websocket_completed_response("resp_metadata_turn_state_next", 3, 1).into(), - )) - .await - .unwrap(); - second_websocket.close(None).await.unwrap(); - (first_request, second_request) - }); - let imported = build_imported_app_with_accounts_and_config( - format!("http://{addr}"), - &[ImportAccount { - id: "acct_imported", - account_id: "chatgpt-account", - token: "access-secret", - refresh_token: "refresh-secret", - }], - |config| { - config.ws_pool.enabled = false; - }, - ) - .await; - - let first_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!(first_response.status(), StatusCode::OK); - let first_body = response_text(first_response).await; - assert!(first_body.contains("\"id\":\"resp_metadata_turn_state\"")); - assert!(!first_body.contains("response.metadata")); - - let second_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_metadata_turn_state"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::OK); - let second_body = response_text(second_response).await; - assert!(second_body.contains("\"id\":\"resp_metadata_turn_state_next\"")); - - let (first_request, second_request) = server.await.unwrap(); - assert!(first_request.get("previous_response_id").is_none()); - assert_eq!( - second_request["previous_response_id"], - "resp_metadata_turn_state" - ); -} - -#[tokio::test] -async fn v1_responses_websocket_should_implicitly_resume_full_history_continuation_with_reasoning_replay( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_async(stream).await.unwrap(); - let first_message = websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY.trim().into(), - )) - .await - .unwrap(); - - let second_message = websocket.next().await.unwrap().unwrap(); - let second_request = - serde_json::from_str::(&second_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_implicit_resume_second", 4, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - (first_request, second_request) - }); - let imported = build_imported_app(format!("http://{addr}")).await; - - let first_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( - json!({ - "model": "gpt-5.5", - "instructions": "answer briefly", - "stream": false, - "input": [{ - "role": "user", - "content": "remember this" - }] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(first_response.status(), StatusCode::OK); - - let second_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( - json!({ - "model": "gpt-5.5", - "instructions": "answer briefly", - "stream": false, - "input": [ - { - "role": "user", - "content": "remember this" - }, - { - "role": "assistant", - "content": "cached answer" - }, - { - "role": "user", - "content": "continue" - } - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::OK); - - let (first_request, second_request) = server.await.unwrap(); - assert!(first_request["prompt_cache_key"].as_str().is_some()); - assert_eq!( - second_request["previous_response_id"], - "resp_implicit_resume_first" - ); - assert_eq!( - second_request["prompt_cache_key"], first_request["prompt_cache_key"], - "implicit resume should reuse the recorded conversation identity" - ); - assert_eq!( - second_request["input"][0]["encrypted_content"], - "enc_reasoning_replay" - ); - assert_eq!(second_request["input"][1]["content"], "continue"); - assert_eq!(second_request["input"].as_array().unwrap().len(), 2); - - let mut replay_request = second_request; - replay_request - .as_object_mut() - .unwrap() - .remove("prompt_cache_key"); - replay_request - .as_object_mut() - .unwrap() - .remove("client_metadata"); - let expected_replay_request: Value = - serde_json::from_str(REASONING_REPLAY_REQUEST_GOLDEN).unwrap(); - assert_eq!(replay_request, expected_replay_request); -} - -#[tokio::test] -async fn v1_responses_websocket_should_not_implicitly_resume_unmatched_function_call_output() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_async(stream).await.unwrap(); - let first_message = websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - websocket_completed_function_call_response("resp_call_first", "call_expected") - .into(), - )) - .await - .unwrap(); - - let second_message = websocket.next().await.unwrap().unwrap(); - let second_request = - serde_json::from_str::(&second_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_call_mismatch_second", 4, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - (first_request, second_request) - }); - let imported = build_imported_app(format!("http://{addr}")).await; - - let first_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( - json!({ - "model": "gpt-5.5", - "instructions": "use the lookup tool", - "stream": false, - "input": [{ - "role": "user", - "content": "call the lookup tool" - }] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(first_response.status(), StatusCode::OK); - - let second_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( - json!({ - "model": "gpt-5.5", - "instructions": "use the lookup tool", - "stream": false, - "input": [ - { - "role": "user", - "content": "call the lookup tool" - }, - { - "type": "function_call", - "call_id": "call_expected", - "name": "lookup", - "arguments": "{}" - }, - { - "type": "function_call_output", - "call_id": "call_missing", - "output": "tool output" - } - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::OK); - - let (first_request, second_request) = server.await.unwrap(); - assert!(first_request["prompt_cache_key"].as_str().is_some()); - assert!(second_request.get("previous_response_id").is_none()); - assert_eq!(second_request["input"].as_array().unwrap().len(), 3); - assert_eq!(second_request["input"][2]["call_id"], "call_missing"); -} - -#[tokio::test] -async fn v1_responses_websocket_should_not_implicitly_resume_self_contained_function_call_replay() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_async(stream).await.unwrap(); - let first_message = websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - websocket_completed_function_call_response( - "resp_self_contained_first", - "call_self", - ) - .into(), - )) - .await - .unwrap(); - - let second_message = websocket.next().await.unwrap().unwrap(); - let second_request = - serde_json::from_str::(&second_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_self_contained_second", 4, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - (first_request, second_request) - }); - let imported = build_imported_app(format!("http://{addr}")).await; - - let first_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( - json!({ - "model": "gpt-5.5", - "instructions": "use the lookup tool", - "stream": false, - "input": [{ - "role": "user", - "content": "call the lookup tool" - }] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(first_response.status(), StatusCode::OK); - - let second_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( - json!({ - "model": "gpt-5.5", - "instructions": "use the lookup tool", - "stream": false, - "input": [ - { - "role": "user", - "content": "call the lookup tool" - }, - { - "type": "function_call", - "call_id": "call_self", - "name": "lookup", - "arguments": "{}" - }, - { - "type": "function_call_output", - "call_id": "call_self", - "output": "tool output" - } - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::OK); - - let (first_request, second_request) = server.await.unwrap(); - assert!(first_request["prompt_cache_key"].as_str().is_some()); - assert!(second_request.get("previous_response_id").is_none()); - assert_eq!(second_request["input"].as_array().unwrap().len(), 3); - assert_eq!(second_request["input"][2]["call_id"], "call_self"); -} - -#[tokio::test] -async fn v1_responses_websocket_should_implicitly_resume_after_sqlite_affinity_restore() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (first_stream, _) = listener.accept().await.unwrap(); - let mut first_websocket = accept_async(first_stream).await.unwrap(); - let first_message = first_websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - first_websocket - .send(Message::Text( - WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY.trim().into(), - )) - .await - .unwrap(); - first_websocket.close(None).await.unwrap(); - - let (second_stream, _) = listener.accept().await.unwrap(); - let mut second_websocket = accept_async(second_stream).await.unwrap(); - let second_message = second_websocket.next().await.unwrap().unwrap(); - let second_request = - serde_json::from_str::(&second_message.into_text().unwrap()).unwrap(); - second_websocket - .send(Message::Text( - websocket_completed_response("resp_restored_implicit_resume", 4, 1).into(), - )) - .await - .unwrap(); - second_websocket.close(None).await.unwrap(); - (first_request, second_request) - }); - let imported = build_imported_app(format!("http://{addr}")).await; - - let first_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( - json!({ - "model": "gpt-5.5", - "instructions": "answer briefly", - "stream": false, - "input": [{ - "role": "user", - "content": "remember this" - }] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(first_response.status(), StatusCode::OK); - - let restored_state = AppState::with_pool_secret_and_api_key_hasher( - imported.config.clone(), - imported.pool.clone(), - imported.secret_box.clone(), - imported.api_key_hasher.clone(), - ); - assert_eq!( - restored_state - .reload_account_pool_from_repository() - .await - .unwrap(), - 1 - ); - assert_eq!( - restored_state - .reload_session_affinity_from_repository() - .await - .unwrap(), - 1 - ); - let restored_app = build_router(restored_state); - - let second_response = restored_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( - json!({ - "model": "gpt-5.5", - "instructions": "answer briefly", - "stream": false, - "input": [ - { - "role": "user", - "content": "remember this" - }, - { - "role": "assistant", - "content": "cached answer" - }, - { - "role": "user", - "content": "continue" - } - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::OK); - - let (first_request, second_request) = server.await.unwrap(); - assert_eq!( - second_request["prompt_cache_key"], first_request["prompt_cache_key"], - "SQLite-restored affinity should keep the recorded conversation identity" - ); - assert_eq!( - second_request["previous_response_id"], - "resp_implicit_resume_first" - ); - assert_eq!(second_request["input"].as_array().unwrap().len(), 1); - assert_eq!(second_request["input"][0]["content"], "continue"); -} - -#[tokio::test] -async fn v1_responses_websocket_should_not_implicitly_resume_across_codex_windows() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_async(stream).await.unwrap(); - let first_message = websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY.trim().into(), - )) - .await - .unwrap(); - - let second_message = websocket.next().await.unwrap().unwrap(); - let second_request = - serde_json::from_str::(&second_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_window_b", 8, 2).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - (first_request, second_request) - }); - let imported = build_imported_app(format!("http://{addr}")).await; - - let first_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( - json!({ - "model": "gpt-5.5", - "instructions": "answer briefly", - "prompt_cache_key": "shared-variant-session", - "codexWindowId": "window-a", - "stream": false, - "input": [{ - "role": "user", - "content": "remember this" - }] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(first_response.status(), StatusCode::OK); - - let second_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( - json!({ - "model": "gpt-5.5", - "instructions": "answer briefly", - "prompt_cache_key": "shared-variant-session", - "codexWindowId": "window-b", - "stream": false, - "input": [ - { - "role": "user", - "content": "remember this" - }, - { - "role": "assistant", - "content": "cached answer" - }, - { - "role": "user", - "content": "continue in another window" - } - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::OK); - - let (first_request, second_request) = server.await.unwrap(); - assert_eq!( - second_request["prompt_cache_key"], first_request["prompt_cache_key"], - "both windows still share the same upstream conversation identity" - ); - assert!(second_request.get("previous_response_id").is_none()); - assert_eq!(second_request["input"].as_array().unwrap().len(), 3); -} - -#[tokio::test] -async fn v1_responses_websocket_should_evict_reasoning_replay_after_invalid_encrypted_content() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_async(stream).await.unwrap(); - let first_message = websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY.trim().into(), - )) - .await - .unwrap(); - - let invalid_message = websocket.next().await.unwrap().unwrap(); - let invalid_request = - serde_json::from_str::(&invalid_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - WEBSOCKET_INVALID_ENCRYPTED_CONTENT.trim().into(), - )) - .await - .unwrap(); - - let (retried_stream, _) = listener.accept().await.unwrap(); - let mut retried_websocket = accept_async(retried_stream).await.unwrap(); - let retried_message = retried_websocket.next().await.unwrap().unwrap(); - let retried_request = - serde_json::from_str::(&retried_message.into_text().unwrap()).unwrap(); - retried_websocket - .send(Message::Text( - websocket_completed_response("resp_after_replay_eviction", 4, 1).into(), - )) - .await - .unwrap(); - retried_websocket.close(None).await.unwrap(); - (first_request, invalid_request, retried_request) - }); - let imported = build_imported_app(format!("http://{addr}")).await; - - let first_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( - json!({ - "model": "gpt-5.5", - "instructions": "answer briefly", - "stream": false, - "input": [{ - "role": "user", - "content": "remember this" - }] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(first_response.status(), StatusCode::OK); - - let continuation_body = json!({ - "model": "gpt-5.5", - "instructions": "answer briefly", - "stream": false, - "input": [ - { - "role": "user", - "content": "remember this" - }, - { - "role": "assistant", - "content": "cached answer" - }, - { - "role": "user", - "content": "continue" - } - ] - }) - .to_string(); - - let retried_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(continuation_body)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(retried_response.status(), StatusCode::OK); - - let (first_request, invalid_request, retried_request) = server.await.unwrap(); - assert_eq!( - invalid_request["previous_response_id"], - "resp_implicit_resume_first" - ); - assert_eq!( - invalid_request["input"][0]["encrypted_content"], - "enc_reasoning_replay" - ); - assert_eq!( - retried_request["prompt_cache_key"], first_request["prompt_cache_key"], - "the replay eviction must not change the conversation identity" - ); - assert!(retried_request.get("previous_response_id").is_none()); - let retried_input = retried_request["input"].as_array().unwrap(); - assert!( - retried_input - .iter() - .all(|item| item.get("encrypted_content").is_none()), - "invalid reasoning replay should be evicted before retry" - ); - assert_eq!(retried_input.last().unwrap()["content"], "continue"); -} - -#[tokio::test] -async fn v1_responses_websocket_should_restore_full_history_when_implicit_resume_previous_response_is_missing( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_async(stream).await.unwrap(); - let first_message = websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text( - WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY.trim().into(), - )) - .await - .unwrap(); - - let implicit_message = websocket.next().await.unwrap().unwrap(); - let implicit_request = - serde_json::from_str::(&implicit_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text(WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.into())) - .await - .unwrap(); - - let (restored_stream, _) = listener.accept().await.unwrap(); - let mut restored_websocket = accept_async(restored_stream).await.unwrap(); - let restored_message = restored_websocket.next().await.unwrap().unwrap(); - let restored_request = - serde_json::from_str::(&restored_message.into_text().unwrap()).unwrap(); - restored_websocket - .send(Message::Text( - websocket_completed_response("resp_implicit_resume_restored", 10, 2).into(), - )) - .await - .unwrap(); - restored_websocket.close(None).await.unwrap(); - (first_request, implicit_request, restored_request) - }); - let imported = build_imported_app(format!("http://{addr}")).await; - - let first_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( - json!({ - "model": "gpt-5.5", - "instructions": "answer briefly", - "stream": false, - "input": [{ - "role": "user", - "content": "remember this" - }] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(first_response.status(), StatusCode::OK); - - let second_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( - json!({ - "model": "gpt-5.5", - "instructions": "answer briefly", - "stream": false, - "input": [ - { - "role": "user", - "content": "remember this" - }, - { - "role": "assistant", - "content": "cached answer" - }, - { - "role": "user", - "content": "continue" - } - ] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::OK); - let body = response_json(second_response).await; - assert_eq!(body["id"], "resp_implicit_resume_restored"); - - let (_first_request, implicit_request, restored_request) = server.await.unwrap(); - assert_eq!( - implicit_request["previous_response_id"], - "resp_implicit_resume_first" - ); - assert_eq!(implicit_request["input"].as_array().unwrap().len(), 2); - assert!(restored_request.get("previous_response_id").is_none()); - assert_eq!(restored_request["input"].as_array().unwrap().len(), 3); - assert_eq!(restored_request["input"][0]["role"], "user"); - assert_eq!(restored_request["input"][1]["role"], "assistant"); - assert_eq!(restored_request["input"][2]["content"], "continue"); -} - -#[tokio::test] -async fn v1_responses_websocket_pool_should_be_evicted_after_admin_account_status_cycle() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_async(stream).await.unwrap(); - let _first_message = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - websocket_completed_response("resp_pool_status_first", 4, 1).into(), - )) - .await - .unwrap(); - - loop { - tokio::select! { - message = websocket.next() => { - match message { - Some(Ok(message)) if message.is_text() => { - websocket - .send(Message::Text( - websocket_completed_response("resp_pool_status_second", 3, 1).into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - break true; - } - Some(_) => continue, - None => { - accept_successful_websocket_response( - &listener, - "Bearer access-secret", - "resp_pool_status_second", - ) - .await; - break false; - } - } - } - accepted = listener.accept() => { - let (stream, _) = accepted.unwrap(); - let mut second_websocket = accept_async(stream).await.unwrap(); - let _second_message = second_websocket.next().await.unwrap().unwrap(); - second_websocket - .send(Message::Text( - websocket_completed_response("resp_pool_status_second", 3, 1).into(), - )) - .await - .unwrap(); - second_websocket.close(None).await.unwrap(); - break false; - } - } - } - }); - let imported = build_imported_app(format!("http://{addr}")).await; - - let first_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":[],"prompt_cache_key":"status-cycle"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(first_response.status(), StatusCode::OK); - let first_body = response_text(first_response).await; - assert!(first_body.contains("\"id\":\"resp_pool_status_first\"")); - - update_admin_account_status(&imported.app, "acct_imported", "disabled").await; - update_admin_account_status(&imported.app, "acct_imported", "active").await; - - let second_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_pool_status_first"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::OK); - let second_body = response_text(second_response).await; - assert!(second_body.contains("\"id\":\"resp_pool_status_second\"")); - - let reused_connection = server.await.unwrap(); - assert!( - !reused_connection, - "admin status lifecycle should evict the old pooled websocket" - ); -} - -#[tokio::test] -async fn v1_responses_should_route_previous_response_id_to_recorded_account() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let first = accept_successful_websocket_response( - &listener, - "Bearer access-a", - "resp_affinity_first", - ) - .await; - let second = accept_successful_websocket_response( - &listener, - "Bearer access-a", - "resp_affinity_second", - ) - .await; - (first, second) - }); - 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 first_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( - json!({ - "model": "gpt-5.5", - "instructions": "answer briefly", - "input": [{ - "role": "user", - "content": "keep this conversation on the same account" - }] - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(first_response.status(), StatusCode::OK); - let first_body = response_text(first_response).await; - assert!(first_body.contains("\"id\":\"resp_affinity_first\"")); - let stored_affinity: (String, String, String, Option, String) = sqlx::query_as( - "select account_id, conversation_id, function_call_ids_json, input_tokens, expires_at from session_affinities where response_id = ?", - ) - .bind("resp_affinity_first") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(stored_affinity.0, "acct_a"); - assert!(!stored_affinity.1.is_empty()); - assert_eq!(stored_affinity.2, "[]"); - assert_eq!(stored_affinity.3, Some(3)); - assert!(!stored_affinity.4.is_empty()); - - let second_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_affinity_first"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second_response.status(), StatusCode::OK); - let second_body = response_text(second_response).await; - assert!(second_body.contains("\"id\":\"resp_affinity_second\"")); - - let (first_request, second_request) = server.await.unwrap(); - assert_ne!( - first_request["prompt_cache_key"], - Value::Null, - "first request should establish an upstream prompt cache key" - ); - assert_eq!( - second_request["previous_response_id"], - "resp_affinity_first" - ); - assert_eq!( - second_request["prompt_cache_key"], first_request["prompt_cache_key"], - "previous_response_id should inherit the recorded conversation identity" - ); -} - -#[tokio::test] -async fn v1_responses_non_stream_previous_response_not_found_should_strip_history_and_retry_same_account( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_hdr_async(stream, assert_access_secret_header) - .await - .unwrap(); - let first_message = websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text(WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.into())) - .await - .unwrap(); - let (second_stream, _) = listener.accept().await.unwrap(); - let mut second_websocket = accept_hdr_async(second_stream, assert_access_secret_header) - .await - .unwrap(); - let second_message = second_websocket.next().await.unwrap().unwrap(); - let second_request = - serde_json::from_str::(&second_message.into_text().unwrap()).unwrap(); - second_websocket - .send(Message::Text( - websocket_completed_response("resp_after_history_strip", 3, 1).into(), - )) - .await - .unwrap(); - second_websocket.close(None).await.unwrap(); - (first_request, second_request) - }); - 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":[],"stream":false,"previous_response_id":"resp_missing"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_history_strip"); - let (first_request, second_request) = server.await.unwrap(); - assert_eq!(first_request["previous_response_id"], "resp_missing"); - assert!(second_request.get("previous_response_id").is_none()); -} - -#[tokio::test] -async fn v1_responses_stream_previous_response_not_found_should_strip_history_and_retry_same_account( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_hdr_async(stream, assert_access_secret_header) - .await - .unwrap(); - let first_message = websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text(WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.into())) - .await - .unwrap(); - let (second_stream, _) = listener.accept().await.unwrap(); - let mut second_websocket = accept_hdr_async(second_stream, assert_access_secret_header) - .await - .unwrap(); - let second_message = second_websocket.next().await.unwrap().unwrap(); - let second_request = - serde_json::from_str::(&second_message.into_text().unwrap()).unwrap(); - second_websocket - .send(Message::Text( - websocket_completed_response("resp_stream_after_history_strip", 3, 1).into(), - )) - .await - .unwrap(); - second_websocket.close(None).await.unwrap(); - (first_request, second_request) - }); - 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":[],"stream":true,"previous_response_id":"resp_missing"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_text(response).await; - assert!(body.contains("\"id\":\"resp_stream_after_history_strip\"")); - let (first_request, second_request) = server.await.unwrap(); - assert_eq!(first_request["previous_response_id"], "resp_missing"); - assert!(second_request.get("previous_response_id").is_none()); -} - -#[tokio::test] -async fn v1_responses_non_stream_unanswered_function_call_should_strip_history_and_retry_same_account( -) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_hdr_async(stream, assert_access_secret_header) - .await - .unwrap(); - let first_message = websocket.next().await.unwrap().unwrap(); - let first_request = - serde_json::from_str::(&first_message.into_text().unwrap()).unwrap(); - websocket - .send(Message::Text(WEBSOCKET_UNANSWERED_FUNCTION_CALL.into())) - .await - .unwrap(); - let (second_stream, _) = listener.accept().await.unwrap(); - let mut second_websocket = accept_hdr_async(second_stream, assert_access_secret_header) - .await - .unwrap(); - let second_message = second_websocket.next().await.unwrap().unwrap(); - let second_request = - serde_json::from_str::(&second_message.into_text().unwrap()).unwrap(); - second_websocket - .send(Message::Text( - websocket_completed_response("resp_after_function_call_strip", 3, 1).into(), - )) - .await - .unwrap(); - second_websocket.close(None).await.unwrap(); - (first_request, second_request) - }); - 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":[],"stream":false,"previous_response_id":"resp_with_call"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_function_call_strip"); - let (first_request, second_request) = server.await.unwrap(); - assert_eq!(first_request["previous_response_id"], "resp_with_call"); - assert!(second_request.get("previous_response_id").is_none()); -} - -#[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_retry_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(77), - WEBSOCKET_HISTORY_RATE_LIMITED, - ) - .await; - accept_successful_websocket_response(&listener, "Bearer access-b", "resp_history_fallback") - .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::OK); - let body = response_text(response).await; - assert!(body.contains("\"id\":\"resp_history_fallback\"")); - let fallback_request = server.await.unwrap(); - assert_eq!(fallback_request["previous_response_id"], "resp_prev"); - 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, 1); -} - -#[tokio::test] -async fn v1_responses_non_stream_previous_response_id_websocket_429_should_retry_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(77), - WEBSOCKET_HISTORY_RATE_LIMITED, - ) - .await; - accept_successful_websocket_response( - &listener, - "Bearer access-b", - "resp_history_fallback_non_stream", - ) - .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::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_history_fallback_non_stream"); - let fallback_request = server.await.unwrap(); - assert_eq!(fallback_request["previous_response_id"], "resp_prev"); - 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, 1); -} - -#[tokio::test] -async fn v1_responses_websocket_without_history_should_mark_expired_after_fallback_401() { - 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 - }); - 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 - .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 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.failed")); - assert!(body.contains("\"type\":\"invalid_request_error\"")); - assert!(body.contains("\"code\":\"authentication_error\"")); - assert!(body.contains("All accounts exhausted")); - assert!(body.contains("token_revoked")); - server.await.unwrap(); - let account_b_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_b") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_b_status.0, "expired"); - 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)); -} - -#[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::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_response_failed_stream( - &body, - "rate_limit_error", - "rate_limit_exceeded", - &[ - "All accounts exhausted (2 rate-limited)", - "second account limited", - ], - ); - 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_websocket_response_failed_quota_should_retry_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 { - let (stream, _) = listener.accept().await.unwrap(); - let mut first_websocket = accept_hdr_async(stream, assert_access_a_header) - .await - .unwrap(); - let _first_message = first_websocket.next().await.unwrap().unwrap(); - first_websocket - .send(Message::Text( - json!({ - "type": "response.failed", - "response": { - "id": "resp_ws_quota_failed", - "error": { - "code": "insufficient_quota", - "message": "quota exhausted" - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - - let (stream, _) = listener.accept().await.unwrap(); - let mut second_websocket = accept_hdr_async(stream, assert_access_b_header) - .await - .unwrap(); - let _second_message = second_websocket.next().await.unwrap().unwrap(); - second_websocket - .send(Message::Text( - websocket_completed_response("resp_after_ws_quota", 3, 1).into(), - )) - .await - .unwrap(); - second_websocket.close(None).await.unwrap(); - }); - let imported = build_imported_app_with_accounts( - format!("http://{addr}"), - &[ - ImportAccount { - id: "acct_ws_response_failed_quota_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_ws_response_failed_quota_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,"use_websocket":true}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_ws_quota"); - server.await.unwrap(); - let first_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_ws_response_failed_quota_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(first_status.0, "quota_exhausted"); -} - -#[tokio::test] -async fn v1_responses_websocket_without_history_should_return_stream_error_when_402_has_no_fallback( -) { - 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", - 402, - "Payment Required", - None, - r#"{"error":{"message":"quota reached"}}"#, - ) - .await; - }); - let imported = build_imported_app_with_accounts( - format!("http://{addr}"), - &[ImportAccount { - id: "acct_ws_402_single", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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::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_response_failed_stream( - &body, - "invalid_request_error", - "codex_api_error", - &[ - "All accounts exhausted (1 quota-exhausted)", - "quota reached", - ], - ); - server.await.unwrap(); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_ws_402_single") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "quota_exhausted"); -} - -#[tokio::test] -async fn v1_responses_websocket_without_history_should_return_stream_error_when_model_unsupported_has_no_fallback( -) { - 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", - 400, - "Bad Request", - None, - r#"{"error":{"code":"model_not_available","message":"Model gpt-5.5 is not available on this account plan"}}"#, - ) - .await; - }); - let imported = build_imported_app_with_accounts( - format!("http://{addr}"), - &[ImportAccount { - id: "acct_ws_model_single", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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::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_response_failed_stream( - &body, - "invalid_request_error", - "codex_api_error", - &[ - "No accounts available", - "model_not_available", - "not available", - ], - ); - server.await.unwrap(); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_ws_model_single") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "active"); -} - -#[tokio::test] -async fn v1_responses_websocket_with_history_should_return_stream_error_when_path_block_has_no_fallback( -) { - 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", 404, "Not Found", None, "") - .await; - }); - let imported = build_imported_app_with_accounts( - format!("http://{addr}"), - &[ImportAccount { - id: "acct_ws_path_block_single", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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,"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_response_failed_stream( - &body, - "server_error", - "codex_api_error", - &["No accounts available", "Cloudflare path-block"], - ); - server.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")); - let lower_request = request.to_ascii_lowercase(); - assert!( - lower_request.contains(&format!( - "authorization: {}", - expected_authorization.to_ascii_lowercase() - )), - "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["response"]["usage"]["total_tokens"] = json!(input_tokens + output_tokens); - value.to_string() -} - -fn websocket_completed_function_call_response(response_id: &str, call_id: &str) -> String { - let mut value: Value = serde_json::from_str(WEBSOCKET_COMPLETED_RESPONSE).unwrap(); - value["response"]["id"] = Value::String(response_id.to_string()); - value["response"]["output"] = json!([{ - "type": "function_call", - "id": format!("fc_{call_id}"), - "call_id": call_id, - "name": "lookup", - "arguments": "{}" - }]); - value["response"]["usage"]["input_tokens"] = json!(6); - value["response"]["usage"]["output_tokens"] = json!(1); - value["response"]["usage"]["total_tokens"] = json!(7); - 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() -} - -async fn update_admin_account_status(app: &axum::Router, account_id: &str, status: &str) { - let response = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri(format!("/api/admin/accounts/{account_id}/status")) - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(json!({ "status": status }).to_string())) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); -} diff --git a/tests/codex_serving/routes_chat.rs b/tests/codex_serving/routes_chat.rs deleted file mode 100644 index 809844b95..000000000 --- a/tests/codex_serving/routes_chat.rs +++ /dev/null @@ -1,35 +0,0 @@ -use serde_json::json; - -use codex_proxy_rs::codex::gateway::protocol::openai_to_codex::{ - translate_chat_to_codex, ChatCompletionRequest, ChatMessage, -}; - -#[test] -fn chat_completion_translates_to_codex_response_request() { - let req = ChatCompletionRequest { - model: "gpt-5.5".to_string(), - stream: true, - messages: vec![ChatMessage { - role: "user".to_string(), - content: Some(json!("hello")), - name: None, - tool_calls: None, - tool_call_id: None, - function_call: None, - }], - reasoning_effort: None, - service_tier: None, - tools: None, - tool_choice: None, - parallel_tool_calls: None, - functions: None, - response_format: None, - user: None, - }; - - let codex = translate_chat_to_codex(req).unwrap(); - assert_eq!(codex.model, "gpt-5.5"); - assert!(codex.stream); - assert!(codex.force_http_sse); - assert_eq!(codex.input[0]["role"], "user"); -} diff --git a/tests/codex_serving/routes_responses.rs b/tests/codex_serving/routes_responses.rs deleted file mode 100644 index 8e950f834..000000000 --- a/tests/codex_serving/routes_responses.rs +++ /dev/null @@ -1,372 +0,0 @@ -use std::collections::BTreeMap; - -use axum::{ - body::{to_bytes, Body}, - http::{Request, StatusCode}, - Router, -}; -use serde_json::{json, Value}; -use tower::ServiceExt; - -use codex_proxy_rs::{ - codex::models::{ - catalog::{BackendModelEntry, ModelPlanSnapshot}, - repository::ModelSnapshotRepository, - }, - config::{ - AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, - QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, - }, - platform::crypto::SecretBox, - platform::identity::{client_key::ApiKeyHasher, client_key_repository::ClientApiKeyRepository}, - platform::storage::db::connect_sqlite, - runtime::build_router, - runtime::state::AppState, -}; - -fn test_config() -> AppConfig { - let mut aliases = BTreeMap::new(); - aliases.insert("codex-fast".to_string(), "gpt-5.5".to_string()); - 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, - }, - 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(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - 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: "sqlite://:memory:".to_string(), - }, - 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, - }, - ws_pool: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - capacity: 2_000, - capture_body: false, - }, - } -} - -fn test_config_with_database(database_url: String) -> AppConfig { - let mut config = test_config(); - config.database.url = database_url; - config -} - -async fn test_app_with_client_api_key() -> (Router, String, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("routes-responses.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let hasher = ApiKeyHasher::new([52u8; 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_with_database(url), - pool, - SecretBox::new([53u8; 32]), - hasher, - )); - (app, generated.plaintext, dir) -} - -async fn test_app_with_backend_model_snapshot() -> (Router, String, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("routes-models.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let hasher = ApiKeyHasher::new([54u8; 32]); - let generated = hasher.generate_client_api_key("test"); - ClientApiKeyRepository::new(pool.clone()) - .insert_generated("test", &generated) - .await - .unwrap(); - let snapshot = ModelPlanSnapshot::from_backend_entries( - "plus", - vec![BackendModelEntry { - id: Some("gpt-6".to_string()), - name: Some("GPT-6".to_string()), - ..BackendModelEntry::default() - }], - ); - ModelSnapshotRepository::new(pool.clone()) - .replace_plan_snapshot(&snapshot) - .await - .unwrap(); - let app = build_router(AppState::with_pool_secret_and_api_key_hasher( - test_config_with_database(url), - pool, - SecretBox::new([55u8; 32]), - hasher, - )); - (app, generated.plaintext, dir) -} - -#[tokio::test] -async fn v1_requires_client_api_key_not_admin_cookie() { - let app = build_router(AppState::new(test_config())); - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/responses") - .body(Body::from("{}")) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); -} - -#[tokio::test] -async fn responses_route_rejects_unknown_models() { - let (app, api_key, _dir) = test_app_with_client_api_key().await; - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/responses") - .header("authorization", format!("Bearer {api_key}")) - .header("content-type", "application/json") - .body(Body::from(r#"{"model":"unknown-model","input":[]}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::NOT_FOUND); -} - -#[tokio::test] -async fn v1_response_has_request_id_header_without_admin_body_field() { - let (app, api_key, _dir) = test_app_with_client_api_key().await; - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/v1/responses") - .header("authorization", format!("Bearer {api_key}")) - .header("x-request-id", "req_client") - .body(Body::from("{}")) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!( - response.headers().get("x-request-id").unwrap(), - "req_client" - ); - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - let body: Value = serde_json::from_slice(&bytes).unwrap(); - assert!(body.get("requestId").is_none()); -} - -#[tokio::test] -async fn models_route_returns_openai_compatible_codex_model_list() { - let (app, api_key, _dir) = test_app_with_client_api_key().await; - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models") - .header("authorization", format!("Bearer {api_key}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - let body: Value = serde_json::from_slice(&bytes).unwrap(); - assert_eq!(body["object"], "list"); - assert_eq!(body["data"][0]["object"], "model"); - assert_eq!(body["data"][0]["id"], "gpt-5.5"); - assert_eq!(body["data"].as_array().unwrap().len(), 1); -} - -#[tokio::test] -async fn model_catalog_route_returns_codex_metadata_without_alias_entries() { - let (app, api_key, _dir) = test_app_with_client_api_key().await; - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models/catalog") - .header("authorization", format!("Bearer {api_key}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - let body: Value = serde_json::from_slice(&bytes).unwrap(); - assert_eq!(body.as_array().unwrap().len(), 1); - assert_eq!(body[0]["id"], "gpt-5.5"); - assert_eq!(body[0]["isDefault"], true); - assert_eq!(body[0]["defaultReasoningEffort"], "medium"); - assert_eq!( - body[0]["supportedReasoningEfforts"][0]["reasoningEffort"], - "low" - ); -} - -#[tokio::test] -async fn model_catalog_route_returns_cached_backend_models() { - let (app, api_key, _dir) = test_app_with_backend_model_snapshot().await; - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models/catalog") - .header("authorization", format!("Bearer {api_key}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - let body: Value = serde_json::from_slice(&bytes).unwrap(); - assert_eq!(body.as_array().unwrap().len(), 1); - assert_eq!(body[0]["id"], "gpt-6"); - assert_eq!(body[0]["source"], "backend"); -} - -#[tokio::test] -async fn model_detail_route_returns_openai_model_for_known_model() { - let (app, api_key, _dir) = test_app_with_client_api_key().await; - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models/gpt-5.5") - .header("authorization", format!("Bearer {api_key}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - let body: Value = serde_json::from_slice(&bytes).unwrap(); - assert_eq!(body["id"], "gpt-5.5"); - assert_eq!(body["object"], "model"); -} - -#[tokio::test] -async fn model_detail_route_rejects_unknown_model_with_openai_error() { - let (app, api_key, _dir) = test_app_with_client_api_key().await; - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models/not-a-codex-model") - .header("authorization", format!("Bearer {api_key}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::NOT_FOUND); - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - let body: Value = serde_json::from_slice(&bytes).unwrap(); - assert_eq!(body["error"]["code"], "model_not_found"); -} - -#[tokio::test] -async fn model_info_route_returns_extended_catalog_entry() { - let (app, api_key, _dir) = test_app_with_client_api_key().await; - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/models/gpt-5.5/info") - .header("authorization", format!("Bearer {api_key}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - let body: Value = serde_json::from_slice(&bytes).unwrap(); - assert_eq!(body["id"], "gpt-5.5"); - assert_eq!(body["outputModalities"], json!(["text"])); -} - -#[tokio::test] -async fn debug_models_route_returns_model_store_summary() { - let (app, api_key, _dir) = test_app_with_client_api_key().await; - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/debug/models") - .header("authorization", format!("Bearer {api_key}")) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - let body: Value = serde_json::from_slice(&bytes).unwrap(); - assert_eq!(body["totalModels"], 1); - assert_eq!(body["aliasCount"], 1); -} diff --git a/tests/codex_serving/upstream_errors.rs b/tests/codex_serving/upstream_errors.rs deleted file mode 100644 index 7b91fa842..000000000 --- a/tests/codex_serving/upstream_errors.rs +++ /dev/null @@ -1,278 +0,0 @@ -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use tower::ServiceExt; -use wiremock::{ - matchers::{method, path}, - Mock, MockServer, ResponseTemplate, -}; - -use crate::support::{ - response_json, response_text, - upstream::{build_imported_app, enable_runtime_logging, fetch_v1_event_log}, -}; - -const ERROR_EVENT_SSE: &str = include_str!("../fixtures/responses/http_sse/error_event.sse"); -const FAILED_EVENT_SSE: &str = include_str!("../fixtures/responses/http_sse/failed_event.sse"); -const STREAM_ERROR_EVENT_SSE: &str = - include_str!("../fixtures/responses/http_sse/stream_error_event.sse"); -const STREAM_FAILED_EVENT_SSE: &str = - include_str!("../fixtures/responses/http_sse/stream_failed_event.sse"); -const STREAM_PREMATURE_CLOSE_SSE: &str = - include_str!("../fixtures/responses/http_sse/stream_premature_close.sse"); -const STREAM_PREMATURE_CLOSE_GOLDEN: &str = - include_str!("../fixtures/responses/golden/stream_premature_close.sse"); - -#[tokio::test] -async fn v1_responses_non_stream_should_return_codex_api_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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_GATEWAY); - let body = response_json(response).await; - assert_eq!(body["type"], "error"); - assert_eq!(body["error"]["code"], "codex_api_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_rate_limit_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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); - let body = response_json(response).await; - assert_eq!(body["type"], "error"); - assert_eq!(body["error"]["code"], "rate_limit_exceeded"); - assert_eq!(body["error"]["type"], "rate_limit_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; - enable_runtime_logging(&imported.app).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,"use_websocket":false}"#, - )) - .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; - enable_runtime_logging(&imported.app).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,"use_websocket":false}"#, - )) - .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_stream_should_synthesize_response_failed_when_http_sse_closes_before_terminal( -) { - 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_PREMATURE_CLOSE_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app(server.uri()).await; - enable_runtime_logging(&imported.app).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_premature_http_sse") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":true,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_text(response).await; - assert_eq!( - redact_proxy_response_id(&body), - with_sse_terminal_separator(STREAM_PREMATURE_CLOSE_GOLDEN) - ); - assert!(body.contains("event: response.output_text.delta")); - assert!(body.contains("event: response.failed")); - assert!(body.contains("stream_disconnected")); - let event = fetch_v1_event_log(&imported.pool, "req_stream_premature_http_sse").await; - assert_eq!(event.3, 502); - assert_eq!(event.4["failureEvent"], "response.failed"); - assert_eq!(event.4["upstreamCode"], "stream_disconnected"); -} - -fn redact_proxy_response_id(body: &str) -> String { - const PREFIX: &str = r#""id":"resp_proxy_"#; - let mut redacted = String::with_capacity(body.len()); - let mut rest = body; - while let Some(index) = rest.find(PREFIX) { - redacted.push_str(&rest[..index]); - redacted.push_str(r#""id":"resp_proxy_redacted"#); - let suffix = &rest[index + PREFIX.len()..]; - let Some(end) = suffix.find('"') else { - rest = suffix; - break; - }; - rest = &suffix[end..]; - } - redacted.push_str(rest); - redacted -} - -fn with_sse_terminal_separator(body: &str) -> String { - if body.ends_with("\n\n") { - body.to_string() - } else { - format!("{body}\n") - } -} diff --git a/tests/codex_serving/upstream_fallback.rs b/tests/codex_serving/upstream_fallback.rs deleted file mode 100644 index 9123f7931..000000000 --- a/tests/codex_serving/upstream_fallback.rs +++ /dev/null @@ -1,2068 +0,0 @@ -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use chrono::{DateTime, Utc}; -use serde_json::json; -use std::net::TcpListener; -use tower::ServiceExt; -use wiremock::{ - matchers::{header, method, path}, - Mock, MockServer, ResponseTemplate, -}; - -use codex_proxy_rs::codex::accounts::cookies::repository::CookieRepository; - -use crate::support::{ - assert_response_failed_stream, response_json, response_text, - upstream::{build_imported_app_with_accounts, ImportAccount}, -}; - -const AFTER_429_SSE: &str = include_str!("../fixtures/responses/http_sse/after_429.sse"); -const STREAM_AFTER_429_SSE: &str = - include_str!("../fixtures/responses/http_sse/stream_after_429.sse"); -const AFTER_402_SSE: &str = include_str!("../fixtures/responses/http_sse/after_402.sse"); -const AFTER_403_SSE: &str = include_str!("../fixtures/responses/http_sse/after_403.sse"); -const AFTER_CLOUDFLARE_403_SSE: &str = - include_str!("../fixtures/responses/http_sse/after_cloudflare_403.sse"); -const AFTER_401_FALLBACK_SSE: &str = - include_str!("../fixtures/responses/http_sse/after_401_fallback.sse"); -const STREAM_AFTER_401_FALLBACK_SSE: &str = - include_str!("../fixtures/responses/http_sse/stream_after_401_fallback.sse"); -const SSE_RESPONSE_FAILED_QUOTA: &str = - include_str!("../fixtures/responses/http_sse/response_failed_quota.sse"); -const SSE_RESPONSE_FAILED_TOKEN_INVALID: &str = - include_str!("../fixtures/responses/http_sse/response_failed_token_invalid.sse"); -const SSE_RESPONSE_FAILED_MODEL_UNSUPPORTED: &str = - include_str!("../fixtures/responses/http_sse/response_failed_model_unsupported.sse"); -const SSE_PREVIOUS_RESPONSE_NOT_FOUND: &str = - include_str!("../fixtures/responses/http_sse/previous_response_not_found.sse"); -const SSE_UNANSWERED_FUNCTION_CALL: &str = - include_str!("../fixtures/responses/http_sse/unanswered_function_call.sse"); -const AFTER_5XX_RETRY_SSE: &str = - include_str!("../fixtures/responses/http_sse/after_5xx_retry.sse"); -const STREAM_AFTER_5XX_RETRY_SSE: &str = - include_str!("../fixtures/responses/http_sse/stream_after_5xx_retry.sse"); -const AFTER_MODEL_UNSUPPORTED_RETRY_SSE: &str = - include_str!("../fixtures/responses/http_sse/after_model_unsupported_retry.sse"); -const RESPONSE_FAILED_AUTH_RECOVERY_GOLDEN: &str = - include_str!("../fixtures/responses/golden/response_failed_auth_recovery.json"); -const CLOUDFLARE_CHALLENGE_NO_FALLBACK_GOLDEN: &str = - include_str!("../fixtures/responses/golden/cloudflare_challenge_no_fallback.json"); -const TOKEN_INVALID_NO_FALLBACK_GOLDEN: &str = - include_str!("../fixtures/responses/golden/token_invalid_no_fallback.json"); - -#[tokio::test] -async fn v1_responses_should_retry_next_account_after_429_retry_after() { - let started_at = Utc::now(); - 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,"use_websocket":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)); - let cooldown: (i64, Option) = sqlx::query_as( - "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", - ) - .bind("acct_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(cooldown.0, 1); - let cooldown_until = DateTime::parse_from_rfc3339(cooldown.1.as_deref().unwrap()) - .unwrap() - .with_timezone(&Utc); - assert!(cooldown_until > started_at); -} - -#[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,"use_websocket":false}"#, - )) - .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_retry_same_account_after_5xx_before_fallback() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(500).set_body_json(json!({ - "error": {"message": "temporary upstream failure"} - }))) - .up_to_n_times(2) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(AFTER_5XX_RETRY_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_5xx_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_5xx_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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_5xx_retry"); - - let requests = server.received_requests().await.unwrap(); - let access_a_posts = requests - .iter() - .filter(|request| { - request.method.as_str() == "POST" - && request - .headers - .get("authorization") - .and_then(|value| value.to_str().ok()) - == Some("Bearer access-a") - }) - .count(); - let access_b_posts = requests - .iter() - .filter(|request| { - request.method.as_str() == "POST" - && request - .headers - .get("authorization") - .and_then(|value| value.to_str().ok()) - == Some("Bearer access-b") - }) - .count(); - assert_eq!(access_a_posts, 3); - assert_eq!(access_b_posts, 0); -} - -#[tokio::test] -async fn v1_responses_should_record_request_count_when_5xx_retries_are_exhausted() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(500).set_body_json(json!({ - "error": {"message": "temporary upstream failure"} - }))) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ImportAccount { - id: "acct_5xx_exhausted", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); - let failed_usage: Option<(i64,)> = - sqlx::query_as("select request_count from account_usage where account_id = ?") - .bind("acct_5xx_exhausted") - .fetch_optional(&imported.pool) - .await - .unwrap(); - assert_eq!(failed_usage.map(|row| row.0).unwrap_or_default(), 1); - - let requests = server.received_requests().await.unwrap(); - let post_count = requests - .iter() - .filter(|request| request.method.as_str() == "POST") - .count(); - assert_eq!(post_count, 3); -} - -#[tokio::test] -async fn v1_responses_should_record_request_count_when_http_transport_fails() { - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let upstream_url = format!("http://{}", listener.local_addr().unwrap()); - drop(listener); - let imported = build_imported_app_with_accounts( - upstream_url, - &[ImportAccount { - id: "acct_transport_failed", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_GATEWAY); - let failed_usage: Option<(i64,)> = - sqlx::query_as("select request_count from account_usage where account_id = ?") - .bind("acct_transport_failed") - .fetch_optional(&imported.pool) - .await - .unwrap(); - assert_eq!(failed_usage.map(|row| row.0).unwrap_or_default(), 1); -} - -#[tokio::test] -async fn v1_responses_stream_should_retry_same_account_after_5xx_before_fallback() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(500).set_body_json(json!({ - "error": {"message": "temporary upstream failure"} - }))) - .up_to_n_times(2) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(STREAM_AFTER_5XX_RETRY_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_5xx_stream_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_5xx_stream_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":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_text(response).await; - assert!(body.contains("resp_stream_after_5xx_retry")); - - let requests = server.received_requests().await.unwrap(); - let access_a_posts = requests - .iter() - .filter(|request| { - request.method.as_str() == "POST" - && request - .headers - .get("authorization") - .and_then(|value| value.to_str().ok()) - == Some("Bearer access-a") - }) - .count(); - let access_b_posts = requests - .iter() - .filter(|request| { - request.method.as_str() == "POST" - && request - .headers - .get("authorization") - .and_then(|value| value.to_str().ok()) - == Some("Bearer access-b") - }) - .count(); - assert_eq!(access_a_posts, 3); - assert_eq!(access_b_posts, 0); -} - -#[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,"use_websocket":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"); - let failed_usage: Option<(i64,)> = - sqlx::query_as("select request_count from account_usage where account_id = ?") - .bind("acct_a") - .fetch_optional(&imported.pool) - .await - .unwrap(); - assert_eq!(failed_usage.map(|row| row.0).unwrap_or_default(), 1); -} - -#[tokio::test] -async fn v1_responses_should_classify_non_stream_sse_failure_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(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(SSE_RESPONSE_FAILED_QUOTA), - ) - .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_sse_failed_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_sse_failed_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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_402"); - let first_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_sse_failed_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(first_status.0, "quota_exhausted"); -} - -#[tokio::test] -async fn v1_responses_should_strip_history_after_http_sse_previous_response_not_found() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(SSE_PREVIOUS_RESPONSE_NOT_FOUND), - ) - .up_to_n_times(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .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_sse_previous_missing", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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,"use_websocket":false,"turnState":"turn-stale"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_402"); - let requests = server.received_requests().await.unwrap(); - let access_a_posts = requests - .iter() - .filter(|request| { - request.method.as_str() == "POST" - && request - .headers - .get("authorization") - .and_then(|value| value.to_str().ok()) - == Some("Bearer access-a") - }) - .collect::>(); - assert_eq!(access_a_posts.len(), 2); - assert_eq!( - access_a_posts[0] - .headers - .get("x-codex-turn-state") - .and_then(|value| value.to_str().ok()), - Some("turn-stale") - ); - assert!(access_a_posts[1] - .headers - .get("x-codex-turn-state") - .is_none()); -} - -#[tokio::test] -async fn v1_responses_should_strip_history_after_http_sse_unanswered_function_call() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(SSE_UNANSWERED_FUNCTION_CALL), - ) - .up_to_n_times(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .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_sse_unanswered_call", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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,"use_websocket":false,"turnState":"turn-stale"}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_402"); - let requests = server.received_requests().await.unwrap(); - let access_a_posts = requests - .iter() - .filter(|request| { - request.method.as_str() == "POST" - && request - .headers - .get("authorization") - .and_then(|value| value.to_str().ok()) - == Some("Bearer access-a") - }) - .collect::>(); - assert_eq!(access_a_posts.len(), 2); - assert_eq!( - access_a_posts[0] - .headers - .get("x-codex-turn-state") - .and_then(|value| value.to_str().ok()), - Some("turn-stale") - ); - assert!(access_a_posts[1] - .headers - .get("x-codex-turn-state") - .is_none()); -} - -#[tokio::test] -async fn v1_responses_should_retry_next_account_when_model_not_supported() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(400).set_body_json(json!({ - "error": { - "message": "Model gpt-5.5 is not supported on this account plan", - "code": "model_not_supported" - } - }))) - .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_MODEL_UNSUPPORTED_RETRY_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_model_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_model_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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_model_unsupported_retry"); - - let requests = server.received_requests().await.unwrap(); - let access_a_posts = requests - .iter() - .filter(|request| { - request.method.as_str() == "POST" - && request - .headers - .get("authorization") - .and_then(|value| value.to_str().ok()) - == Some("Bearer access-a") - }) - .count(); - let access_b_posts = requests - .iter() - .filter(|request| { - request.method.as_str() == "POST" - && request - .headers - .get("authorization") - .and_then(|value| value.to_str().ok()) - == Some("Bearer access-b") - }) - .count(); - assert_eq!(access_a_posts, 1); - assert_eq!(access_b_posts, 1); -} - -#[tokio::test] -async fn v1_responses_should_classify_sse_model_not_supported_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(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(SSE_RESPONSE_FAILED_MODEL_UNSUPPORTED), - ) - .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_MODEL_UNSUPPORTED_RETRY_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_sse_model_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_sse_model_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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_model_unsupported_retry"); -} - -#[tokio::test] -async fn v1_responses_should_retry_model_not_supported_only_once() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .respond_with(ResponseTemplate::new(400).set_body_json(json!({ - "error": { - "message": "Model gpt-5.5 is not available on this account plan", - "code": "model_not_available" - } - }))) - .up_to_n_times(2) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(AFTER_MODEL_UNSUPPORTED_RETRY_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_model_once_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_model_once_b", - account_id: "chatgpt-b", - token: "access-b", - refresh_token: "refresh-b", - }, - ImportAccount { - id: "acct_model_once_c", - account_id: "chatgpt-c", - token: "access-c", - refresh_token: "refresh-c", - }, - ], - ) - .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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_REQUEST); - let body = response_json(response).await; - assert_eq!( - body["error"]["message"], - "Codex upstream error: {\"error\":{\"code\":\"model_not_available\",\"message\":\"Model gpt-5.5 is not available on this account plan\"}}" - ); - - let requests = server.received_requests().await.unwrap(); - let post_count = requests - .iter() - .filter(|request| request.method.as_str() == "POST") - .count(); - assert_eq!(post_count, 2); -} - -#[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,"use_websocket":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_return_502_when_cloudflare_challenge_has_no_fallback() { - 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", - )) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ImportAccount { - id: "acct_cf_single", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::BAD_GATEWAY); - let body = response_json(response).await; - let expected: serde_json::Value = - serde_json::from_str(CLOUDFLARE_CHALLENGE_NO_FALLBACK_GOLDEN).unwrap(); - assert_eq!(body, expected); -} - -#[tokio::test] -async fn v1_responses_should_cool_down_cloudflare_403_and_retry_next_account() { - let started_at = Utc::now(); - 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 cookie_repo = CookieRepository::new(imported.pool.clone(), imported.secret_box.clone()); - cookie_repo - .set_cookie_header("acct_a", "cf_clearance=old") - .await - .unwrap(); - cookie_repo - .set_cookie_header("acct_b", "cf_clearance=keep") - .await - .unwrap(); - assert_eq!( - cookie_repo - .cookie_header("acct_a", "chatgpt.com") - .await - .unwrap() - .as_deref(), - Some("cf_clearance=old") - ); - - 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,"use_websocket":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"); - let (cooldown_until,): (Option,) = - sqlx::query_as("select cloudflare_cooldown_until from accounts where id = ?") - .bind("acct_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - let cooldown_until = DateTime::parse_from_rfc3339(cooldown_until.as_deref().unwrap()) - .unwrap() - .with_timezone(&Utc); - assert!(cooldown_until > started_at); - assert_eq!( - cookie_repo - .cookie_header("acct_a", "chatgpt.com") - .await - .unwrap(), - None - ); - assert_eq!( - cookie_repo - .cookie_header("acct_b", "chatgpt.com") - .await - .unwrap() - .as_deref(), - Some("cf_clearance=keep") - ); -} - -#[tokio::test] -async fn v1_responses_should_clear_cookies_and_retry_next_account_after_cloudflare_path_block_404() -{ - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(404)) - .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_cf_path_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_cf_path_b", - account_id: "chatgpt-b", - token: "access-b", - refresh_token: "refresh-b", - }, - ], - ) - .await; - let cookie_repo = CookieRepository::new(imported.pool.clone(), imported.secret_box.clone()); - cookie_repo - .set_cookie_header("acct_cf_path_a", "cf_clearance=old") - .await - .unwrap(); - cookie_repo - .set_cookie_header("acct_cf_path_b", "cf_clearance=keep") - .await - .unwrap(); - - 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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_cf"); - assert_eq!( - cookie_repo - .cookie_header("acct_cf_path_a", "chatgpt.com") - .await - .unwrap(), - None - ); - assert_eq!( - cookie_repo - .cookie_header("acct_cf_path_b", "chatgpt.com") - .await - .unwrap() - .as_deref(), - Some("cf_clearance=keep") - ); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_cf_path_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "active"); -} - -#[tokio::test] -async fn v1_responses_should_disable_account_after_three_cloudflare_path_block_404s() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(404)) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ImportAccount { - id: "acct_cf_path_disabled", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .await; - let cookie_repo = CookieRepository::new(imported.pool.clone(), imported.secret_box.clone()); - cookie_repo - .set_cookie_header("acct_cf_path_disabled", "cf_clearance=old") - .await - .unwrap(); - - for _ in 0..3 { - 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":false,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::BAD_GATEWAY); - } - - assert_eq!( - cookie_repo - .cookie_header("acct_cf_path_disabled", "chatgpt.com") - .await - .unwrap(), - None - ); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_cf_path_disabled") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "disabled"); -} - -#[tokio::test] -async fn v1_responses_should_mark_expired_after_401_and_retry_next_account_non_stream() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(401).set_body_json(json!({ - "error": {"code": "token_revoked", "message": "token revoked"} - }))) - .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_401_FALLBACK_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_401_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_401_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_401_non_stream") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["id"], "resp_after_401_fallback"); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_401_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "expired"); -} - -#[tokio::test] -async fn v1_responses_should_recover_auth_failure_from_response_failed_event() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(SSE_RESPONSE_FAILED_TOKEN_INVALID), - ) - .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_401_FALLBACK_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_response_failed_auth_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_response_failed_auth_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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - let expected: serde_json::Value = - serde_json::from_str(RESPONSE_FAILED_AUTH_RECOVERY_GOLDEN).unwrap(); - assert_eq!(body, expected); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_response_failed_auth_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "expired"); -} - -#[tokio::test] -async fn v1_responses_should_mark_expired_after_401_and_retry_next_account_stream() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(401).set_body_json(json!({ - "error": {"code": "token_revoked", "message": "token revoked"} - }))) - .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_401_FALLBACK_SSE), - ) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ - ImportAccount { - id: "acct_401_stream_a", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }, - ImportAccount { - id: "acct_401_stream_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_401_stream") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":true,"use_websocket":false}"#, - )) - .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("resp_stream_after_401_fallback")); - 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_401_stream_b") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(usage, (1, 4, 2, 0)); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_401_stream_a") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "expired"); -} - -#[tokio::test] -async fn v1_responses_should_mark_expired_and_return_401_when_401_has_no_fallback() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(401).set_body_json(json!({ - "error": {"code": "token_revoked", "message": "token revoked"} - }))) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ImportAccount { - id: "acct_401_single", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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_401_no_fallback") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_text(response).await; - assert_response_failed_stream( - &body, - "invalid_request_error", - "authentication_error", - &["All accounts exhausted (1 expired)", "token_revoked"], - ); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_401_single") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "expired"); -} - -#[tokio::test] -async fn v1_responses_should_mark_quota_exhausted_and_return_stream_error_when_402_has_no_fallback() -{ - 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 reached"} - }))) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ImportAccount { - id: "acct_402_single", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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":false}"#, - )) - .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_response_failed_stream( - &body, - "invalid_request_error", - "codex_api_error", - &[ - "All accounts exhausted (1 quota-exhausted)", - "quota reached", - ], - ); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_402_single") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "quota_exhausted"); -} - -#[tokio::test] -async fn v1_responses_should_return_stream_error_when_model_unsupported_has_no_fallback() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(400).set_body_json(json!({ - "error": { - "message": "Model gpt-5.5 is not available on this account plan", - "code": "model_not_available" - } - }))) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ImportAccount { - id: "acct_model_single", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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":false}"#, - )) - .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_response_failed_stream( - &body, - "invalid_request_error", - "codex_api_error", - &[ - "No accounts available", - "model_not_available", - "not available", - ], - ); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_model_single") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "active"); -} - -#[tokio::test] -async fn v1_responses_should_return_stream_error_when_cloudflare_path_block_has_no_fallback() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(404)) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ImportAccount { - id: "acct_path_block_single", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .await; - let cookie_repo = CookieRepository::new(imported.pool.clone(), imported.secret_box.clone()); - cookie_repo - .set_cookie_header("acct_path_block_single", "cf_clearance=old") - .await - .unwrap(); - - 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":false}"#, - )) - .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_response_failed_stream( - &body, - "server_error", - "codex_api_error", - &["No accounts available", "Cloudflare path-block"], - ); - assert_eq!( - cookie_repo - .cookie_header("acct_path_block_single", "chatgpt.com") - .await - .unwrap(), - None - ); -} - -#[tokio::test] -async fn v1_responses_non_stream_should_mark_expired_and_return_401_when_401_has_no_fallback() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(401).set_body_json(json!({ - "error": {"code": "token_revoked", "message": "token revoked"} - }))) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ImportAccount { - id: "acct_401_single_non_stream", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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_401_no_fallback_non_stream") - .body(Body::from( - r#"{"model":"gpt-5.5","input":[],"stream":false,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - let body = response_json(response).await; - let expected: serde_json::Value = - serde_json::from_str(TOKEN_INVALID_NO_FALLBACK_GOLDEN).unwrap(); - assert_eq!(body, expected); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_401_single_non_stream") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "expired"); -} - -#[tokio::test] -async fn v1_responses_should_mark_banned_when_401_says_account_deactivated() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/codex/responses")) - .and(header("authorization", "Bearer access-a")) - .respond_with(ResponseTemplate::new(401).set_body_json(json!({ - "error": {"message": "account deactivated"} - }))) - .mount(&server) - .await; - let imported = build_imported_app_with_accounts( - server.uri(), - &[ImportAccount { - id: "acct_401_deactivated", - account_id: "chatgpt-a", - token: "access-a", - refresh_token: "refresh-a", - }], - ) - .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,"use_websocket":false}"#, - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") - .bind("acct_401_deactivated") - .fetch_one(&imported.pool) - .await - .unwrap(); - assert_eq!(account_status.0, "banned"); -} diff --git a/tests/platform.rs b/tests/platform.rs deleted file mode 100644 index f1233b736..000000000 --- a/tests/platform.rs +++ /dev/null @@ -1,7 +0,0 @@ -mod platform { - mod client_key_auth; - mod crypto; - mod http_auth; - mod log_rotation; - mod storage_schema; -} diff --git a/tests/platform/http_auth.rs b/tests/platform/http_auth.rs deleted file mode 100644 index 51d7ea879..000000000 --- a/tests/platform/http_auth.rs +++ /dev/null @@ -1,38 +0,0 @@ -use axum::http::{HeaderMap, HeaderValue}; - -use codex_proxy_rs::platform::http::auth::{admin_session_id, client_api_key}; - -#[test] -fn client_api_key_should_accept_only_cpr_bearer_tokens() { - let mut headers = HeaderMap::new(); - headers.insert( - "authorization", - HeaderValue::from_static("Bearer cpr_local_client_key"), - ); - assert_eq!( - client_api_key(&headers).unwrap().as_str(), - "cpr_local_client_key" - ); - - headers.insert( - "authorization", - HeaderValue::from_static("Bearer sk-official-openai-key"), - ); - assert!(client_api_key(&headers).is_none()); -} - -#[test] -fn admin_session_id_should_ignore_client_api_key_and_extract_only_cookie() { - let mut headers = HeaderMap::new(); - headers.insert( - "authorization", - HeaderValue::from_static("Bearer cpr_local_client_key"), - ); - assert!(admin_session_id(&headers).is_none()); - - headers.insert( - "cookie", - HeaderValue::from_static("theme=dark; cpr_admin_session=sess_123; other=1"), - ); - assert_eq!(admin_session_id(&headers), Some("sess_123")); -} diff --git a/tests/runtime.rs b/tests/runtime.rs deleted file mode 100644 index fbf8392c4..000000000 --- a/tests/runtime.rs +++ /dev/null @@ -1,6 +0,0 @@ -mod support; - -mod runtime { - mod http_trace; - mod startup; -} diff --git a/tests/runtime/startup.rs b/tests/runtime/startup.rs deleted file mode 100644 index faaae7672..000000000 --- a/tests/runtime/startup.rs +++ /dev/null @@ -1,374 +0,0 @@ -use std::time::Instant; - -use axum::{ - body::Bytes, - http::{HeaderMap, StatusCode}, -}; -use chrono::{Duration, Utc}; -use futures::{SinkExt, StreamExt}; -use secrecy::SecretString; -use serde_json::{json, Value}; -use tokio::net::TcpListener; -use tokio_tungstenite::{ - accept_hdr_async, - tungstenite::{ - handshake::server::{Request as WsRequest, Response as WsResponse}, - Message, - }, -}; - -use codex_proxy_rs::{ - codex::accounts::{ - model::AccountStatus, - repository::{NewAccount, UsageDelta}, - }, - config::{ - AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, - QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, - }, - platform::crypto::SecretBox, - platform::storage::db::connect_sqlite, - runtime::state::AppState, -}; - -use crate::support::response_text; - -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: Default::default(), - }, - 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(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - 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, - }, - ws_pool: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - capacity: 2_000, - capture_body: false, - }, - } -} - -#[tokio::test] -async fn app_state_should_restore_account_pool_from_sqlite_accounts() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("startup-pool.sqlite"); - 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::codex::accounts::repository::AccountRepository::new( - pool.clone(), - secret_box.clone(), - ); - repo.insert(NewAccount { - id: "acct_restored".to_string(), - email: Some("user@example.com".to_string()), - account_id: Some("chatgpt-account".to_string()), - user_id: None, - label: Some("primary".to_string()), - plan_type: Some("plus".to_string()), - access_token: SecretString::new("access-secret".to_string().into()), - refresh_token: Some(SecretString::new("refresh-secret".to_string().into())), - access_token_expires_at: Some(Utc::now() + chrono::Duration::hours(1)), - status: AccountStatus::Active, - }) - .await - .unwrap(); - repo.record_usage( - "acct_restored", - UsageDelta { - input_tokens: 10, - output_tokens: 3, - cached_tokens: 2, - empty_response_count: 0, - ..UsageDelta::default() - }, - ) - .await - .unwrap(); - let state = AppState::with_pool_and_secret_box(test_config(url), pool, secret_box); - - let restored = state.reload_account_pool_from_repository().await.unwrap(); - - assert_eq!(restored, 1); - let acquired = state - .services - .accounts - .acquire_runtime_account("gpt-5.5") - .await - .unwrap(); - assert_eq!(acquired.id, "acct_restored"); - assert_eq!(acquired.access_token, "access-secret"); - assert!(acquired.last_used_at.is_some()); -} - -#[tokio::test] -async fn app_state_should_restore_persisted_cooldowns_into_runtime_pool() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("startup-pool-cooldown.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([32u8; 32]); - let repo = codex_proxy_rs::codex::accounts::repository::AccountRepository::new( - pool.clone(), - secret_box.clone(), - ); - repo.insert(NewAccount { - id: "acct_cooling".to_string(), - email: Some("cooling@example.com".to_string()), - account_id: Some("chatgpt-cooling".to_string()), - user_id: None, - label: None, - plan_type: Some("plus".to_string()), - access_token: SecretString::new("access-cooling".to_string().into()), - refresh_token: Some(SecretString::new("refresh-cooling".to_string().into())), - access_token_expires_at: Some(Utc::now() + Duration::hours(1)), - status: AccountStatus::Active, - }) - .await - .unwrap(); - let cooldown_until = (Utc::now() + Duration::minutes(10)).to_rfc3339(); - sqlx::query( - "update accounts set quota_limit_reached = 1, quota_cooldown_until = ?, cloudflare_cooldown_until = ? where id = ?", - ) - .bind(&cooldown_until) - .bind(&cooldown_until) - .bind("acct_cooling") - .execute(&pool) - .await - .unwrap(); - let state = AppState::with_pool_and_secret_box(test_config(url), pool, secret_box); - - let restored = state.reload_account_pool_from_repository().await.unwrap(); - - assert_eq!(restored, 1); - assert!(state - .services - .accounts - .acquire_runtime_account("gpt-5.5") - .await - .is_none()); -} - -#[tokio::test] -async fn app_state_should_restore_session_affinity_from_sqlite() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - accept_successful_websocket_response( - &listener, - "Bearer access-a", - Some("turn_restore"), - "resp_restore_second", - ) - .await - }); - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("startup-affinity.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([33u8; 32]); - let repo = codex_proxy_rs::codex::accounts::repository::AccountRepository::new( - pool.clone(), - secret_box.clone(), - ); - repo.insert(NewAccount { - id: "acct_a".to_string(), - email: Some("a@example.com".to_string()), - account_id: Some("chatgpt-a".to_string()), - user_id: None, - label: None, - plan_type: Some("plus".to_string()), - access_token: SecretString::new("access-a".to_string().into()), - refresh_token: Some(SecretString::new("refresh-a".to_string().into())), - access_token_expires_at: Some(Utc::now() + Duration::hours(1)), - status: AccountStatus::Active, - }) - .await - .unwrap(); - repo.insert(NewAccount { - id: "acct_b".to_string(), - email: Some("b@example.com".to_string()), - account_id: Some("chatgpt-b".to_string()), - user_id: None, - label: None, - plan_type: Some("plus".to_string()), - access_token: SecretString::new("access-b".to_string().into()), - refresh_token: Some(SecretString::new("refresh-b".to_string().into())), - access_token_expires_at: Some(Utc::now() + Duration::hours(1)), - status: AccountStatus::Active, - }) - .await - .unwrap(); - repo.record_usage( - "acct_a", - UsageDelta { - input_tokens: 1, - output_tokens: 0, - cached_tokens: 0, - empty_response_count: 0, - ..UsageDelta::default() - }, - ) - .await - .unwrap(); - let now = Utc::now(); - sqlx::query( - "insert into session_affinities (response_id, account_id, conversation_id, turn_state, instructions_hash, input_tokens, function_call_ids_json, variant_hash, expires_at, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind("resp_restore_first") - .bind("acct_a") - .bind("conv_restore") - .bind("turn_restore") - .bind(Option::::None) - .bind(3_i64) - .bind(r#"["call_restore"]"#) - .bind(Option::::None) - .bind((now + Duration::hours(3)).to_rfc3339()) - .bind(now.to_rfc3339()) - .execute(&pool) - .await - .unwrap(); - let mut config = test_config(url); - config.api.base_url = format!("http://{addr}"); - let state = AppState::with_pool_and_secret_box(config, pool, secret_box); - state.reload_account_pool_from_repository().await.unwrap(); - assert_eq!( - state - .reload_session_affinity_from_repository() - .await - .unwrap(), - 1 - ); - - let response = state - .services - .responses - .handle( - "req_restore_affinity", - HeaderMap::new(), - Bytes::from( - r#"{"model":"gpt-5.5","input":[],"previous_response_id":"resp_restore_first"}"#, - ), - Instant::now(), - ) - .await; - - assert_eq!(response.status(), StatusCode::OK); - let body = response_text(response).await; - assert!(body.contains("\"id\":\"resp_restore_second\"")); - let upstream_request = server.await.unwrap(); - assert_eq!( - upstream_request["previous_response_id"], - "resp_restore_first" - ); -} - -#[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, - expected_turn_state: Option<&str>, - response_id: &str, -) -> Value { - let (stream, _) = listener.accept().await.unwrap(); - let expected_authorization = expected_authorization.to_string(); - let expected_turn_state = expected_turn_state.map(ToString::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()) - ); - assert_eq!( - request - .headers() - .get("x-codex-turn-state") - .and_then(|value| value.to_str().ok()), - expected_turn_state.as_deref() - ); - 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", - "status": "completed", - "output": [], - "usage": { - "input_tokens": 1, - "output_tokens": 1, - "total_tokens": 2, - "input_tokens_details": {"cached_tokens": 0} - } - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - websocket.close(None).await.unwrap(); - request -} diff --git a/tests/support/admin_accounts.rs b/tests/support/admin_accounts.rs deleted file mode 100644 index 033721da6..000000000 --- a/tests/support/admin_accounts.rs +++ /dev/null @@ -1,395 +0,0 @@ -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::{ - codex::gateway::{ - fingerprint::model::Fingerprint, - oauth::{DeviceCode, OAuthClient, OAuthError, RefreshFailure, TokenPair, TokenRefresher}, - }, - config::{ - AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, - QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, - }, - platform::crypto::SecretBox, - platform::identity::client_key::ApiKeyHasher, - platform::storage::db::connect_sqlite, - runtime::build_router, - runtime::state::AppState, -}; - -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(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - 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, - }, - ws_pool: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, - }, - logging: LoggingConfig { - directory: "logs".to_string(), - 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 config = test_config_with_base_url(url, base_url); - admin_accounts_test_app_with_config(dir, config, key_byte).await -} - -pub async fn admin_accounts_test_app_with_config( - dir: tempfile::TempDir, - config: AppConfig, - key_byte: u8, -) -> (Router, AppState, sqlx::SqlitePool, tempfile::TempDir) { - let pool = connect_sqlite(&config.database.url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let state = - AppState::with_pool_and_secret_box(config, 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_base_url_and_fingerprint( - db_name: &str, - key_byte: u8, - base_url: String, - fingerprint: Fingerprint, -) -> (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_secret_api_key_hasher_and_fingerprint( - test_config_with_base_url(url, base_url), - pool.clone(), - SecretBox::new([key_byte; 32]), - ApiKeyHasher::new([key_byte; 32]), - fingerprint, - ); - 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_base_url_and_refresher( - db_name: &str, - key_byte: u8, - base_url: String, - 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_with_base_url(url, base_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("/api/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}"), - "accessTokenExpiresAt": "2999-01-01T00:00:00Z", - "status": "active" - }] - }); - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/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/support/mod.rs b/tests/support/mod.rs deleted file mode 100644 index 30e5c6b0b..000000000 --- a/tests/support/mod.rs +++ /dev/null @@ -1,98 +0,0 @@ -#![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() -} - -pub fn assert_response_failed_stream( - body: &str, - expected_type: &str, - expected_code: &str, - expected_message_parts: &[&str], -) { - assert!( - body.ends_with("\n\n"), - "stream should end with an SSE boundary" - ); - let mut event = None; - let mut data = None; - for line in body.lines() { - if let Some(value) = line.strip_prefix("event: ") { - event = Some(value); - } else if let Some(value) = line.strip_prefix("data: ") { - data = Some(value); - break; - } - } - - assert_eq!(event, Some("response.failed")); - let value: Value = - serde_json::from_str(data.expect("response.failed stream should include data")).unwrap(); - assert_eq!(value["type"], "response.failed"); - let response = &value["response"]; - assert!(response["id"] - .as_str() - .is_some_and(|id| id.starts_with("resp_proxy_"))); - assert_eq!(response["status"], "failed"); - assert_eq!(response["error"], value["error"]); - - let error = &value["error"]; - assert_eq!(error["type"], expected_type); - assert_eq!(error["code"], expected_code); - let message = error["message"] - .as_str() - .expect("response.failed error should include message"); - for expected in expected_message_parts { - assert!( - message.contains(expected), - "expected stream error message to contain {expected:?}, got {message:?}" - ); - } -} diff --git a/tests/support/upstream.rs b/tests/support/upstream.rs deleted file mode 100644 index 60b95dad9..000000000 --- a/tests/support/upstream.rs +++ /dev/null @@ -1,399 +0,0 @@ -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::{ - codex::gateway::fingerprint::model::Fingerprint, - codex::gateway::oauth::{RefreshFailure, TokenPair, TokenRefresher}, - config::{ - AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, ModelConfig, - QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, - UsageStatsConfig, - }, - platform::crypto::SecretBox, - platform::identity::{client_key::ApiKeyHasher, client_key_repository::ClientApiKeyRepository}, - platform::storage::db::connect_sqlite, - runtime::build_router, - runtime::state::AppState, -}; - -use super::seed_admin_session; - -pub struct ImportedApp { - pub app: axum::Router, - pub pool: sqlx::SqlitePool, - pub secret_box: SecretBox, - pub api_key_hasher: ApiKeyHasher, - pub config: AppConfig, - 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(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - 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, - }, - ws_pool: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - session_cleanup_interval_secs: 3600, - }, - logging: LoggingConfig { - directory: "logs".to_string(), - 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 { - build_imported_app_with_accounts_and_config(base_url, accounts, |_| {}).await -} - -pub async fn build_imported_app_with_accounts_and_config( - base_url: String, - accounts: &[ImportAccount], - configure: impl FnOnce(&mut AppConfig), -) -> ImportedApp { - build_imported_app_with_accounts_config_and_fingerprint(base_url, accounts, configure, None) - .await -} - -pub async fn build_imported_app_with_fingerprint( - base_url: String, - fingerprint: Fingerprint, -) -> ImportedApp { - build_imported_app_with_accounts_config_and_fingerprint( - base_url, - &[ImportAccount { - id: "acct_imported", - account_id: "chatgpt-account", - token: "access-secret", - refresh_token: "refresh-secret", - }], - |_| {}, - Some(fingerprint), - ) - .await -} - -async fn build_imported_app_with_accounts_config_and_fingerprint( - base_url: String, - accounts: &[ImportAccount], - configure: impl FnOnce(&mut AppConfig), - fingerprint: Option, -) -> 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 mut config = test_config(url, base_url); - configure(&mut config); - let state_config = config.clone(); - let state = match fingerprint { - Some(fingerprint) => AppState::with_pool_secret_api_key_hasher_and_fingerprint( - state_config, - pool.clone(), - secret_box.clone(), - hasher.clone(), - fingerprint, - ), - None => AppState::with_pool_secret_and_api_key_hasher( - state_config, - pool.clone(), - secret_box.clone(), - hasher.clone(), - ), - }; - let app = build_router(state); - let accounts = accounts - .iter() - .map(|account| { - json!({ - "id": account.id, - "accountId": account.account_id, - "token": account.token, - "refreshToken": account.refresh_token, - "accessTokenExpiresAt": "2999-01-01T00:00:00Z", - "status": "active" - }) - }) - .collect::>(); - - if !accounts.is_empty() { - let import_response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/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, - api_key_hasher: hasher, - config, - 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 config = test_config(url, base_url); - let app = build_router( - AppState::with_pool_secret_api_key_hasher_and_token_refresher( - config.clone(), - pool.clone(), - secret_box.clone(), - hasher.clone(), - token_refresher, - ), - ); - let accounts = accounts - .iter() - .map(|account| { - json!({ - "id": account.id, - "accountId": account.account_id, - "token": account.token, - "refreshToken": account.refresh_token, - "accessTokenExpiresAt": "2999-01-01T00:00:00Z", - "status": "active" - }) - }) - .collect::>(); - - let import_response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/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, - api_key_hasher: hasher, - config, - 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(), - ) -} - -pub async fn enable_runtime_logging(app: &axum::Router) { - let response = app - .clone() - .oneshot( - Request::builder() - .method("PATCH") - .uri("/api/admin/logs/state") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(json!({ "enabled": true }).to_string())) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); -}