mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 16:06:08 +00:00
fix: reuse cookies for usage refresh
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -85,6 +85,17 @@
|
||||
|
||||
这些是二进制字符串证据,不等价于完整源码控制流;但足以证明官方包中存在这些 transport、请求字段、恢复状态和完成边界。
|
||||
|
||||
追加取证到的更细字段:
|
||||
|
||||
- 官方二进制包含源码路径字符串:`codex-api/src/endpoint/responses_websocket.rs`、`codex-api/src/endpoint/responses.rs`、`codex-api/src/endpoint/compact.rs`、`codex-api/src/endpoint/models.rs`、`codex-client/src/request.rs`、`codex-client/src/default_client.rs`、`core/src/client.rs`。
|
||||
- WebSocket 相关符号包含 `ResponsesWebsocketClient`、`ResponsesWebsocketConnection`、`merge_request_headers`、`serialize_websocket_request`、`send_websocket_request`、`run_websocket_response_stream`、`map_wrapped_websocket_error_event`。
|
||||
- 官方连接日志包含 `responses_websocket.connect`、`responses_websocket.stream_request`、`responses.stream_request`、`transport`、`pre_compression_bytes`、`post_compression_bytes`、`compression_duration_ms`。
|
||||
- 官方协议能力包含 `requestAttestation`,注释为允许 client 接收 `attestation/generate` 请求并为上游生成 `x-oai-attestation`。
|
||||
- 官方 WebSocket 头相关字段包含 `ws_request_header_traceparent`、`ws_request_header_tracestate`、`ws_request_header_x_openai_internal_codex_responses_lite`。
|
||||
- 官方二进制出现 `x-openai-internal-codex-responses-lite`、`x-codex-ws-stream-request-start-ms`、`x-codex-beta-features`。这些字段当前只能确认存在,未确认每个字段的生成条件。
|
||||
- 官方错误文案包含 `Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue.`,说明 WebSocket 连接有 60 分钟生命周期或服务端连接限制。
|
||||
- 官方安全拦截文案包含 `This request has been flagged for possible cybersecurity risk.`,属于请求内容/工具调用侧风控信号,不是普通额度不足。
|
||||
|
||||
### Webview 额度逻辑
|
||||
|
||||
相关文件:
|
||||
@@ -176,6 +187,11 @@
|
||||
- `src/upstream/transport/client.rs` 复用同一套请求头构造 WebSocket opening headers。
|
||||
- WebSocket 请求会写 audit artifact,包含 opening snapshot 和 payload snapshot。
|
||||
- 如果启用了 pool,会按账号/session 复用连接。
|
||||
- opening request 固定补齐 `Host`、`Connection: Upgrade`、`Upgrade: websocket`、`Sec-WebSocket-Version: 13`、`Sec-WebSocket-Key`,业务头随后追加,最后追加 `sec-websocket-extensions: permessage-deflate; client_max_window_bits`。
|
||||
- 实际握手由 `tokio-tungstenite`/`tungstenite` 执行,TLS 走 rustls;HTTP/SSE 走 `reqwest` + rustls。
|
||||
- 首个文本帧是 `response.create`,payload 由 `CodexResponsesRequest` 序列化而来。
|
||||
- WebSocket opening 会移除 `content-type` 和 `accept`,因为这两个是 HTTP/SSE 请求头,不属于 WS opening 业务头。
|
||||
- 连接池 key 为 `base_url + account_id + conversation_id`,conversation id 优先取 `prompt_cache_key`,其次 `client_conversation_id`,再其次 `previous_response_id`。
|
||||
|
||||
当前 WebSocket fallback:
|
||||
|
||||
@@ -191,37 +207,43 @@
|
||||
|
||||
### 请求头
|
||||
|
||||
当前实现:`src/upstream/transport/client.rs`
|
||||
当前实现集中在 `src/upstream/transport/client.rs` 和 `src/upstream/transport/headers.rs`。不同请求类型不是同一组头:
|
||||
|
||||
基础 header 来自 fingerprint:
|
||||
| 请求类型 | 路径 | 传输 | 当前头部规则 |
|
||||
| --- | --- | --- | --- |
|
||||
| Responses HTTP SSE | `/codex/responses` | `reqwest` POST | 指纹基础头 + `content-type: application/json` + 可选 `cookie` + `accept: text/event-stream` + `openai-beta: responses_websockets=2026-02-06` + `x-openai-internal-codex-residency: us` + `x-client-request-id` + Codex 上下文头 |
|
||||
| Responses WebSocket | `/codex/responses` | `tokio-tungstenite` GET upgrade + 首帧 `response.create` | 先按 Responses HTTP 头构造,再移除 `content-type` 和 `accept`,追加标准 WebSocket upgrade 头和 `sec-websocket-extensions` |
|
||||
| Compact | `/codex/responses/compact` | `reqwest` POST | 指纹基础头 + `content-type` + 可选 `cookie` + `openai-beta` + `x-openai-internal-codex-residency` + 新生成的 `x-client-request-id` + 可选 `x-codex-installation-id` |
|
||||
| Models/探测 | `/codex/models`、`/models`、`/sentinel/chat-requirements` | `reqwest` GET | 指纹基础头 + 可选 `cookie` + `accept: application/json` + `accept-encoding: gzip, deflate` + 可选 `x-codex-installation-id` |
|
||||
| Usage/配额 | `/api/codex/usage`、`/wham/usage`、`/codex/usage` | `reqwest` GET | `user-agent` + `authorization` + `originator` + 可选 `chatgpt-account-id` + 可选 `cookie` + `accept: application/json` |
|
||||
|
||||
指纹基础头:
|
||||
|
||||
- `authorization`
|
||||
- `chatgpt-account-id`
|
||||
- `originator`
|
||||
- `user-agent`
|
||||
- `sec-ch-ua`
|
||||
- fingerprint 默认 header 和 header order
|
||||
- fingerprint 默认 header 和 header order,例如 `accept-encoding`、`accept-language`、`sec-ch-ua-mobile`、`sec-ch-ua-platform`、`sec-fetch-*`
|
||||
|
||||
Responses 请求额外写入:
|
||||
Responses/WS/Compact 的 Codex 上下文头:
|
||||
|
||||
- `content-type: application/json`
|
||||
- `cookie`
|
||||
- `accept: text/event-stream`
|
||||
- `openai-beta: responses_websockets=2026-02-06`
|
||||
- `x-openai-internal-codex-residency: us`
|
||||
- `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`
|
||||
- 条件写入 `x-openai-subagent`
|
||||
- `x-client-request-id`:主 Responses 请求优先使用 `session_id`,没有 session 时使用本次 `request_id`;Compact 每次新生成 UUID。
|
||||
- `x-codex-installation-id`:有 installation id 时写入。
|
||||
- `session_id`:仅主 Responses 请求有 conversation identity 时写入。
|
||||
- `x-codex-window-id`:由请求 `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`:客户端 body/header 透传后写入。
|
||||
- `x-openai-subagent`:只接受 `review`、`compact`、`memory_consolidation`、`collab_spawn`。
|
||||
|
||||
官方二进制还出现 `x-oai-attestation`。当前项目没有构造该头;仅凭静态字符串无法确认其生成条件、签名来源和是否每次请求都必需,不能用固定值或伪造值补齐。
|
||||
官方二进制还出现 `x-oai-attestation`、`requestAttestation`、`ws_request_header_traceparent`、`ws_request_header_tracestate`、`x-openai-internal-codex-responses-lite`、`x-codex-ws-stream-request-start-ms`。当前项目没有构造这些头;仅凭静态字符串无法确认生成条件、签名来源和失败行为,不能用固定值或伪造值补齐。
|
||||
|
||||
当前需要继续取证的头部差距:
|
||||
|
||||
- `x-oai-attestation` 应该来自 app-server 协议的 `attestation/generate` 能力,可能需要 Desktop host 配合签名。
|
||||
- `traceparent` / `tracestate` 是 OTEL 链路追踪头,官方只确认 WebSocket request header 字段存在,未确认是否每个请求都带。
|
||||
- `x-openai-internal-codex-responses-lite` 与 `ws_request_header_x_openai_internal_codex_responses_lite` 只确认官方存在,未确认和模型/feature 的关系。
|
||||
- `x-codex-ws-stream-request-start-ms` 看起来是 WebSocket stream 开始时间戳,未确认是否影响上游接受请求。
|
||||
|
||||
### 请求携带的数据
|
||||
|
||||
@@ -300,6 +322,8 @@ WebSocket 事件:
|
||||
- 403 body 命中 `cf-mitigated`、`cf-chl-bypass`、`_cf_chl`、`cf_chl`、`attention required`、`just a moment` 时进入 Cloudflare challenge cooldown。
|
||||
- 404 空 body 作为 path-block,立即删除账号 cookie,多次后禁用账号。
|
||||
- 调度层会在 Cloudflare challenge/path-block 后换号或返回明确错误。
|
||||
- 官方二进制出现 `This request has been flagged for possible cybersecurity risk.`。当前项目没有单独识别这个文案,若真实链路遇到,应作为独立风控/安全策略错误,而不是归入 quota 或 Cloudflare。
|
||||
- 官方错误枚举还包含 `usage_limit_exceeded`、`cyber_policy`、`response_stream_connection_failed`、`response_stream_disconnected` 等 CodexErrorInfo 类别;当前项目主要按 HTTP status、SSE failure code/message 和 WS wrapped error 分类。
|
||||
|
||||
### 额度和限流
|
||||
|
||||
@@ -320,6 +344,29 @@ WebSocket 事件:
|
||||
- additional/model 限额不覆盖 core 限额,作为独立 window 行保留;如果同一个 additional limit 同时有 5h/7d,也会拆成两行。
|
||||
- WebSocket `codex.rate_limits` 被动刷新只携带短周期窗口时,输出同样的 `snapshots` 结构,并保留已有 `monthly_limit`、`credits`、`spend_control` 和 additional snapshots,避免把 `/wham/usage` 拿到的月限额或 additional 限额冲掉。
|
||||
|
||||
#### 配额刷新触发
|
||||
|
||||
当前项目里配额刷新不是单一入口:
|
||||
|
||||
- 管理端主动刷新:`GET /api/admin/accounts/quota?id=...` 会调用 `fetch_usage`,拉取实时 usage 后归一化并写回 `quota_json`。
|
||||
- 管理端健康检查:`health_check_accounts` 也会调用 `fetch_usage`,但只判断 alive/dead,不写 quota。
|
||||
- 账号导入:导入时如果缺 plan/quota 或导入数据要求补全,会调用 `fetch_usage` 补齐。
|
||||
- 业务请求前验证:如果账号 `quota_verify_required = true`,调度层在真正发 `/codex/responses` 前会先调用 `fetch_usage`。如果仍然 `limit_reached`,会释放该账号、换号,最多尝试 5 个账号。
|
||||
- 业务响应后被动刷新:成功响应的 HTTP/WS rate-limit header,以及 WS 内部 `codex.rate_limits` 事件,会同步到账号池和持久化 quota。
|
||||
- 业务失败触发锁定:上游 HTTP 429 会写入 `quota_cooldown_until`;HTTP 402 或 SSE `quota_exceeded` / `insufficient_quota` 会把账号置为 `QuotaExhausted`。
|
||||
- 后台定时刷新:`QuotaRefreshTask` 默认按配置间隔扫描 active 且 `quota_limit_reached || quota_verify_required` 的账号;单个账号至少间隔 30 分钟,同批请求之间默认错峰 3 秒。
|
||||
|
||||
Usage 请求路径按 base URL 选择:
|
||||
|
||||
- base URL 带 `/backend-api` 时:优先 `/wham/usage`,再 `/codex/usage`。
|
||||
- 其他 base URL:优先 `/api/codex/usage`,再 `/codex/usage`。
|
||||
|
||||
调用方差异:
|
||||
|
||||
- 业务请求前验证、管理端主动刷新、管理端健康检查和后台定时刷新都会按账号为 `/codex/usage` 读取 Cloudflare cookie,并传给 usage 请求头。
|
||||
- 账号导入补全 usage 发生在新账号入库前,通常没有可复用的本地 cookie;当前仍按无 cookie usage 请求处理。
|
||||
- `usage_request_headers` 当前支持可选 `cookie`,但是否携带取决于调用方传入。
|
||||
|
||||
### 真实数据验证
|
||||
|
||||
2026-06-25 使用本地 `.runtime` active free 账号调用 `GET /api/admin/accounts/quota`,经项目真实解密、请求官方 `/wham/usage`、归一化并落库。
|
||||
@@ -341,14 +388,22 @@ WebSocket 事件:
|
||||
- `monthly_limit.source = "rate_limit"`
|
||||
- 管理端账号列表返回 `quota.windows[0].group = "monthly"`、`labelDisplay = "月限额"`、`windowSeconds = 2592000`
|
||||
|
||||
2026-06-26 使用临时 18080 本地实例复用 `.runtime` 数据,登录管理端后调用 `GET /api/admin/accounts/quota?id=acct_fbed55cb574b4aee81760d0e8ac84ba6`。请求真实官方 `/wham/usage` 返回 200,`quota_fetched_at` 从 `2026-06-25T16:32:12.697845388+00:00` 更新为 `2026-06-25T16:53:08.898420541+00:00`。
|
||||
|
||||
本地 `account_cookies` 当前没有该 active 账号的 `cf_clearance`,所以这次真实链路只验证 usage 主刷新路径未被 cookie 改动破坏;cookie 头携带行为由 `tests/admin/accounts/quota.rs` 和 `tests/upstream/accounts/quota_refresh.rs` 的 route/mock 请求头测试覆盖。
|
||||
|
||||
结论:本地真实 free 账号确认为 30d/月限额单独展示。当前本地没有 active 的 5h+7d 账号可实时验证组合返回;该组合路径由官方 bundle 证据和项目测试覆盖。
|
||||
|
||||
## 当前差距
|
||||
|
||||
- `x-oai-attestation` 只确认官方二进制里存在,未确认生成算法、触发条件、签名材料和失败行为;当前项目不实现该头。
|
||||
- `requestAttestation` 说明官方 app-server 协议支持 attestation 生成,但当前没有动态运行态证据说明每个 upstream 请求都必须生成该头。
|
||||
- `traceparent` / `tracestate`、`x-openai-internal-codex-responses-lite`、`x-codex-ws-stream-request-start-ms` 只确认官方存在,当前项目未实现。
|
||||
- 原生二进制不能在当前 Linux 环境执行,无法做动态断点或真实 app-server 调用链追踪。
|
||||
- 官方包里没有 source map,JS bundle 为压缩产物;函数名和局部变量名不能视为稳定 API。
|
||||
- WebSocket 降级策略目前按现有项目实现理解为连接/transport 错误可降级,显式上游状态错误不降级;官方二进制字符串能证明有降级路径,但不能单独证明完整条件矩阵。
|
||||
- 官方有 WebSocket 60 分钟连接限制文案;当前项目 WebSocket pool 只按本地 pool TTL 控制,未显式实现官方 60 分钟上限。
|
||||
- 官方有 `cyber_policy` 风控语义;当前项目还没有把对应文案/错误码拆成独立错误类型。
|
||||
|
||||
## 实现建议
|
||||
|
||||
@@ -362,3 +417,4 @@ WebSocket 事件:
|
||||
- `WebSocketPreferred` 可降级 HTTP SSE,`WebSocketRequired` 不降级。
|
||||
- `codex.rate_limits` 内部事件不下发给客户端,但用于刷新账号限额状态。
|
||||
- Cloudflare cookie 捕获、重放、challenge cooldown、path-block 清 cookie/禁用账号。
|
||||
- usage 请求头保持轻量:不要把 Responses 专用 beta/residency/turn header 全部塞到 `/api/codex/usage`。
|
||||
|
||||
@@ -1,6 +1,14 @@
|
||||
use super::{types::*, AdminAccountService};
|
||||
|
||||
impl AdminAccountService {
|
||||
pub(crate) async fn usage_cookie_header(&self, account_id: &str) -> Option<String> {
|
||||
self.cookies
|
||||
.cookie_header_for_request(account_id, "chatgpt.com", "/codex/usage")
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
}
|
||||
|
||||
pub async fn cookies(&self, account_id: &str) -> Result<Option<String>, AdminAccountError> {
|
||||
self.ensure_cookie_account_exists(account_id).await?;
|
||||
self.cookies
|
||||
|
||||
@@ -19,6 +19,7 @@ impl AdminAccountService {
|
||||
|
||||
let request_id = uuid::Uuid::new_v4().to_string();
|
||||
let token = stored.access_token.expose_secret().to_string();
|
||||
let cookie_header = self.usage_cookie_header(account_id).await;
|
||||
let context = crate::upstream::transport::CodexRequestContext {
|
||||
access_token: &token,
|
||||
account_id: stored.account_id.as_deref(),
|
||||
@@ -30,7 +31,7 @@ impl AdminAccountService {
|
||||
version: None,
|
||||
codex_window_id: None,
|
||||
parent_thread_id: None,
|
||||
cookie_header: None,
|
||||
cookie_header: cookie_header.as_deref(),
|
||||
installation_id: self.installation_id.as_deref(),
|
||||
session_id: None,
|
||||
};
|
||||
@@ -143,6 +144,7 @@ impl AdminAccountService {
|
||||
let request_id = uuid::Uuid::new_v4().to_string();
|
||||
for account in &accounts {
|
||||
let token = account.access_token.expose_secret().to_string();
|
||||
let cookie_header = self.usage_cookie_header(&account.id).await;
|
||||
let start = std::time::Instant::now();
|
||||
let context = crate::upstream::transport::CodexRequestContext {
|
||||
access_token: &token,
|
||||
@@ -155,7 +157,7 @@ impl AdminAccountService {
|
||||
version: None,
|
||||
codex_window_id: None,
|
||||
parent_thread_id: None,
|
||||
cookie_header: None,
|
||||
cookie_header: cookie_header.as_deref(),
|
||||
installation_id: self.installation_id.as_deref(),
|
||||
session_id: None,
|
||||
};
|
||||
|
||||
@@ -101,6 +101,7 @@ impl TaskCoordinator {
|
||||
30 * 60,
|
||||
)
|
||||
.with_installation_id(services.installation_id.clone())
|
||||
.with_cookie_store(stores.cookies.clone())
|
||||
.start(),
|
||||
);
|
||||
coordinator.push(
|
||||
|
||||
@@ -7,6 +7,7 @@ use tracing::{debug, info, warn};
|
||||
|
||||
use crate::{
|
||||
upstream::accounts::{
|
||||
cookies::SqliteCookieStore,
|
||||
quota::{QuotaRefreshSummary, RuntimeQuotaRefreshService},
|
||||
store::SqliteAccountStore,
|
||||
},
|
||||
@@ -55,6 +56,12 @@ impl QuotaRefreshTask {
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 usage 请求可复用的账号 Cookie 存储。
|
||||
pub fn with_cookie_store(mut self, cookie_store: SqliteCookieStore) -> Self {
|
||||
self.service = self.service.with_cookie_store(cookie_store);
|
||||
self
|
||||
}
|
||||
|
||||
/// 启动后台刷新任务。
|
||||
pub fn start(self) -> SchedulerHandle {
|
||||
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
|
||||
|
||||
@@ -14,7 +14,10 @@ use crate::upstream::accounts::{
|
||||
},
|
||||
store::{AccountStore, AccountStoreError, SqliteAccountStore, SqliteAccountStoreError},
|
||||
};
|
||||
use crate::upstream::transport::{CodexBackendClient, CodexClientError, CodexRequestContext};
|
||||
use crate::upstream::{
|
||||
accounts::cookies::SqliteCookieStore,
|
||||
transport::{CodexBackendClient, CodexClientError, CodexRequestContext},
|
||||
};
|
||||
|
||||
const MIN_REFRESH_INTERVAL_SECS: u64 = 30 * 60;
|
||||
const DEFAULT_REQUEST_SPACING_SECS: u64 = 3;
|
||||
@@ -26,6 +29,7 @@ pub struct RuntimeQuotaRefreshService {
|
||||
min_refresh_interval_secs: u64,
|
||||
request_spacing: Duration,
|
||||
installation_id: Option<String>,
|
||||
cookie_store: Option<SqliteCookieStore>,
|
||||
}
|
||||
|
||||
impl RuntimeQuotaRefreshService {
|
||||
@@ -37,6 +41,7 @@ impl RuntimeQuotaRefreshService {
|
||||
min_refresh_interval_secs: MIN_REFRESH_INTERVAL_SECS,
|
||||
request_spacing: Self::default_request_spacing(),
|
||||
installation_id: None,
|
||||
cookie_store: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -57,6 +62,7 @@ impl RuntimeQuotaRefreshService {
|
||||
min_refresh_interval_secs,
|
||||
request_spacing: Self::default_request_spacing(),
|
||||
installation_id: None,
|
||||
cookie_store: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -72,6 +78,12 @@ impl RuntimeQuotaRefreshService {
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 usage 请求可复用的账号 Cookie 存储。
|
||||
pub fn with_cookie_store(mut self, cookie_store: SqliteCookieStore) -> Self {
|
||||
self.cookie_store = Some(cookie_store);
|
||||
self
|
||||
}
|
||||
|
||||
/// 执行一次配额锁定或待验证账号刷新。
|
||||
pub async fn refresh_locked_accounts_once(
|
||||
&self,
|
||||
@@ -117,6 +129,7 @@ impl RuntimeQuotaRefreshService {
|
||||
last_refreshed.insert(account.id.clone(), now);
|
||||
|
||||
let request_id = uuid::Uuid::new_v4().to_string();
|
||||
let cookie_header = self.usage_cookie_header(&account.id).await;
|
||||
match self
|
||||
.codex
|
||||
.fetch_usage(CodexRequestContext {
|
||||
@@ -130,7 +143,7 @@ impl RuntimeQuotaRefreshService {
|
||||
version: None,
|
||||
codex_window_id: None,
|
||||
parent_thread_id: None,
|
||||
cookie_header: None,
|
||||
cookie_header: cookie_header.as_deref(),
|
||||
installation_id: self.installation_id.as_deref(),
|
||||
session_id: None,
|
||||
})
|
||||
@@ -181,6 +194,15 @@ impl RuntimeQuotaRefreshService {
|
||||
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
async fn usage_cookie_header(&self, account_id: &str) -> Option<String> {
|
||||
self.cookie_store
|
||||
.as_ref()?
|
||||
.cookie_header_for_request(account_id, "chatgpt.com", "/codex/usage")
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
}
|
||||
}
|
||||
|
||||
/// 单次配额刷新摘要。
|
||||
|
||||
@@ -806,6 +806,7 @@ impl CodexBackendClient {
|
||||
HeaderValue::from_str(&self.fingerprint.originator)?,
|
||||
);
|
||||
insert_optional_header(&mut headers, "chatgpt-account-id", context.account_id)?;
|
||||
insert_optional_header(&mut headers, "cookie", context.cookie_header)?;
|
||||
headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
@@ -1,4 +1,8 @@
|
||||
use super::*;
|
||||
use wiremock::{
|
||||
matchers::{method, path},
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_usage_stats_should_return_page_and_summary() {
|
||||
@@ -206,6 +210,155 @@ async fn admin_usage_stats_should_cursor_page_account_usage() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_account_quota_should_send_usage_cookie() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/backend-api/wham/usage"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"plan_type": "plus",
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary_window": {
|
||||
"used_percent": 12,
|
||||
"reset_at": 1_800_000_000,
|
||||
"limit_window_seconds": 18_000
|
||||
}
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let (app, _state, pool, _dir, secret_box) = admin_accounts_test_app_with_api_base_url(
|
||||
"admin-account-quota-cookie.sqlite",
|
||||
88,
|
||||
format!("{}/backend-api", server.uri()),
|
||||
)
|
||||
.await;
|
||||
seed_encrypted_account(
|
||||
&pool,
|
||||
secret_box.clone(),
|
||||
NewAccount {
|
||||
id: "acct_quota_cookie".to_string(),
|
||||
email: Some("quota-cookie@example.com".to_string()),
|
||||
account_id: Some("chatgpt-quota-cookie".to_string()),
|
||||
user_id: None,
|
||||
label: None,
|
||||
plan_type: Some("plus".to_string()),
|
||||
access_token: SecretString::new("access-quota-cookie".to_string().into()),
|
||||
refresh_token: None,
|
||||
access_token_expires_at: None,
|
||||
status: AccountStatus::Active,
|
||||
added_at: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
SqliteCookieStore::new(pool.clone(), secret_box)
|
||||
.set_cookie_header("acct_quota_cookie", "cf_clearance=admin-quota")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/accounts/quota?id=acct_quota_cookie")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_admin_quota_cookie")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
let cookie_header = requests
|
||||
.iter()
|
||||
.find(|request| request.url.path() == "/backend-api/wham/usage")
|
||||
.and_then(|request| request.headers.get("cookie"))
|
||||
.and_then(|value| value.to_str().ok());
|
||||
|
||||
assert_eq!(
|
||||
(response.status(), cookie_header),
|
||||
(StatusCode::OK, Some("cf_clearance=admin-quota"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_account_health_check_should_send_usage_cookie() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/backend-api/wham/usage"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"plan_type": "plus",
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary_window": {
|
||||
"used_percent": 15,
|
||||
"reset_at": 1_800_000_000,
|
||||
"limit_window_seconds": 18_000
|
||||
}
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let (app, _state, pool, _dir, secret_box) = admin_accounts_test_app_with_api_base_url(
|
||||
"admin-account-health-cookie.sqlite",
|
||||
89,
|
||||
format!("{}/backend-api", server.uri()),
|
||||
)
|
||||
.await;
|
||||
seed_encrypted_account(
|
||||
&pool,
|
||||
secret_box.clone(),
|
||||
NewAccount {
|
||||
id: "acct_health_cookie".to_string(),
|
||||
email: Some("health-cookie@example.com".to_string()),
|
||||
account_id: Some("chatgpt-health-cookie".to_string()),
|
||||
user_id: None,
|
||||
label: None,
|
||||
plan_type: Some("plus".to_string()),
|
||||
access_token: SecretString::new("access-health-cookie".to_string().into()),
|
||||
refresh_token: None,
|
||||
access_token_expires_at: None,
|
||||
status: AccountStatus::Active,
|
||||
added_at: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
SqliteCookieStore::new(pool.clone(), secret_box)
|
||||
.set_cookie_header("acct_health_cookie", "cf_clearance=admin-health")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
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_admin_health_cookie")
|
||||
.body(Body::from(
|
||||
json!({ "ids": ["acct_health_cookie"] }).to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
let cookie_header = requests
|
||||
.iter()
|
||||
.find(|request| request.url.path() == "/backend-api/wham/usage")
|
||||
.and_then(|request| request.headers.get("cookie"))
|
||||
.and_then(|value| value.to_str().ok());
|
||||
|
||||
assert_eq!(
|
||||
(response.status(), cookie_header),
|
||||
(StatusCode::OK, Some("cf_clearance=admin-health"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_account_quota_warnings_should_return_threshold_matches_from_cached_quota() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
@@ -3,6 +3,7 @@ use std::{sync::Arc, time::Duration as StdDuration};
|
||||
use chrono::{Duration, Utc};
|
||||
use codex_proxy_rs::infra::crypto::SecretBox;
|
||||
use codex_proxy_rs::infra::database::connect_sqlite;
|
||||
use codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore;
|
||||
use codex_proxy_rs::upstream::accounts::model::AccountStatus;
|
||||
use codex_proxy_rs::upstream::accounts::quota::{quota_from_usage, RuntimeQuotaRefreshService};
|
||||
use codex_proxy_rs::upstream::accounts::store::{NewAccount, SqliteAccountStore};
|
||||
@@ -137,6 +138,67 @@ fn quota_from_usage_should_preserve_additional_limit_name_and_metered_feature()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_refresh_service_should_send_usage_cookie_when_cookie_store_is_configured() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/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": 12.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-cookie.sqlite");
|
||||
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
|
||||
.await
|
||||
.expect("sqlite pool");
|
||||
let secret_box = SecretBox::new([14u8; 32]);
|
||||
let store = SqliteAccountStore::new(pool.clone(), secret_box.clone());
|
||||
let cookies = SqliteCookieStore::new(pool.clone(), secret_box);
|
||||
insert_quota_locked_account(&store, &pool, "acct-quota-cookie", "access-token-cookie").await;
|
||||
cookies
|
||||
.set_cookie_header("acct-quota-cookie", "cf_clearance=quota-refresh")
|
||||
.await
|
||||
.expect("cookie should be stored");
|
||||
let codex = CodexBackendClient::new(
|
||||
reqwest::Client::new(),
|
||||
server.uri(),
|
||||
crate::support::fingerprint::test_fingerprint(),
|
||||
);
|
||||
let service =
|
||||
RuntimeQuotaRefreshService::new(store, Arc::new(codex)).with_cookie_store(cookies);
|
||||
|
||||
let summary = service
|
||||
.refresh_locked_accounts_once()
|
||||
.await
|
||||
.expect("quota refresh should succeed");
|
||||
let requests = server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("received requests should load");
|
||||
let cookie_header = requests
|
||||
.iter()
|
||||
.find(|request| request.url.path() == "/api/codex/usage")
|
||||
.and_then(|request| request.headers.get("cookie"))
|
||||
.and_then(|value| value.to_str().ok());
|
||||
|
||||
assert_eq!(
|
||||
(summary.refreshed, cookie_header),
|
||||
(1, Some("cf_clearance=quota-refresh"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_refresh_service_should_fetch_usage_for_quota_locked_accounts_and_store_quota() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
@@ -522,7 +522,10 @@ async fn codex_backend_client_usage_should_use_wham_usage_headers() {
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("Codex Desktop")
|
||||
);
|
||||
assert!(headers.get("cookie").is_none());
|
||||
assert_eq!(
|
||||
headers.get("cookie").and_then(|value| value.to_str().ok()),
|
||||
Some("cf_clearance=old")
|
||||
);
|
||||
assert!(headers.get("content-type").is_none());
|
||||
assert!(headers.get("sec-ch-ua").is_none());
|
||||
assert!(headers.get("openai-beta").is_none());
|
||||
|
||||
Reference in New Issue
Block a user