mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-13 16:07:29 +00:00
fix: align account quota usage tracking
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -7,7 +7,7 @@ use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::infra::{
|
||||
format::{format_compact_percent, format_tokens},
|
||||
format::{format_compact_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64},
|
||||
time::{china_datetime, china_relative_time},
|
||||
};
|
||||
|
||||
@@ -31,13 +31,29 @@ struct AdminAccountQuotaWindowData {
|
||||
label_display: String,
|
||||
used_percent: Option<f64>,
|
||||
used_percent_display: String,
|
||||
token_usage_display: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
local_usage: Option<AdminAccountQuotaWindowLocalUsageData>,
|
||||
reset_at_display: String,
|
||||
window_used_display: String,
|
||||
#[serde(skip)]
|
||||
reset_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AdminAccountQuotaWindowLocalUsageData {
|
||||
request_count: i64,
|
||||
request_count_display: String,
|
||||
input_tokens: i64,
|
||||
input_tokens_display: String,
|
||||
output_tokens: i64,
|
||||
output_tokens_display: String,
|
||||
cached_tokens: i64,
|
||||
cached_tokens_display: String,
|
||||
total_tokens: i64,
|
||||
total_tokens_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct AdminAccountQuotaUsageWindow {
|
||||
pub(crate) key: String,
|
||||
@@ -56,17 +72,37 @@ impl AdminAccountQuotaUsageWindow {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub(crate) struct AdminAccountQuotaWindowTokenUsage {
|
||||
total_tokens: u64,
|
||||
pub(crate) struct AdminAccountQuotaWindowLocalUsage {
|
||||
pub(crate) request_count: i64,
|
||||
pub(crate) input_tokens: i64,
|
||||
pub(crate) output_tokens: i64,
|
||||
pub(crate) cached_tokens: i64,
|
||||
}
|
||||
|
||||
impl AdminAccountQuotaWindowTokenUsage {
|
||||
pub(crate) fn add_tokens(&mut self, tokens: u64) {
|
||||
self.total_tokens = self.total_tokens.saturating_add(tokens);
|
||||
impl AdminAccountQuotaWindowLocalUsage {
|
||||
pub(crate) fn add(&mut self, other: Self) {
|
||||
self.request_count = self.request_count.saturating_add(other.request_count);
|
||||
self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
|
||||
self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
|
||||
self.cached_tokens = self.cached_tokens.saturating_add(other.cached_tokens);
|
||||
}
|
||||
}
|
||||
|
||||
fn display(self) -> String {
|
||||
format_tokens(self.total_tokens)
|
||||
impl From<AdminAccountQuotaWindowLocalUsage> for AdminAccountQuotaWindowLocalUsageData {
|
||||
fn from(usage: AdminAccountQuotaWindowLocalUsage) -> Self {
|
||||
let total_tokens = usage.input_tokens.saturating_add(usage.output_tokens);
|
||||
Self {
|
||||
request_count: usage.request_count,
|
||||
request_count_display: format_plain_number(nonnegative_i64_to_u64(usage.request_count)),
|
||||
input_tokens: usage.input_tokens,
|
||||
input_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.input_tokens)),
|
||||
output_tokens: usage.output_tokens,
|
||||
output_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.output_tokens)),
|
||||
cached_tokens: usage.cached_tokens,
|
||||
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.cached_tokens)),
|
||||
total_tokens,
|
||||
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -93,17 +129,18 @@ impl AdminAccountQuotaData {
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn apply_token_usage(
|
||||
pub(crate) fn apply_local_usage(
|
||||
&mut self,
|
||||
usage_by_window: &HashMap<String, AdminAccountQuotaWindowTokenUsage>,
|
||||
usage_by_window: &HashMap<String, AdminAccountQuotaWindowLocalUsage>,
|
||||
) {
|
||||
for window in &mut self.windows {
|
||||
window.token_usage_display =
|
||||
match (window.usage_window(), usage_by_window.get(&window.key)) {
|
||||
(Some(_), Some(usage)) => usage.display(),
|
||||
(Some(_), None) => AdminAccountQuotaWindowTokenUsage::default().display(),
|
||||
(None, _) => "-".to_string(),
|
||||
};
|
||||
window.local_usage = window.usage_window().map(|_| {
|
||||
usage_by_window
|
||||
.get(&window.key)
|
||||
.copied()
|
||||
.unwrap_or_default()
|
||||
.into()
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -185,7 +222,7 @@ fn push_monthly_quota_window(
|
||||
label_display: "月限额".to_string(),
|
||||
used_percent,
|
||||
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
|
||||
token_usage_display: "-".to_string(),
|
||||
local_usage: None,
|
||||
reset_at_display: reset_at
|
||||
.as_ref()
|
||||
.map_or_else(|| "-".to_string(), china_datetime),
|
||||
@@ -274,7 +311,7 @@ fn push_quota_window(
|
||||
label_display,
|
||||
used_percent,
|
||||
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
|
||||
token_usage_display: "-".to_string(),
|
||||
local_usage: None,
|
||||
reset_at_display: reset_at
|
||||
.as_ref()
|
||||
.map_or_else(|| "-".to_string(), china_datetime),
|
||||
|
||||
@@ -27,7 +27,7 @@ use crate::{
|
||||
admin::{
|
||||
accounts::quota_view::{
|
||||
quota_data, AdminAccountQuotaData, AdminAccountQuotaUsageWindow,
|
||||
AdminAccountQuotaWindowTokenUsage,
|
||||
AdminAccountQuotaWindowLocalUsage,
|
||||
},
|
||||
accounts::service::{
|
||||
AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata,
|
||||
@@ -303,14 +303,6 @@ struct AccountModelUsageRecord {
|
||||
last_used_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
struct AccountUsageWindowStats {
|
||||
request_count: i64,
|
||||
input_tokens: i64,
|
||||
output_tokens: i64,
|
||||
cached_tokens: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct AccountQuotaUsageWindow {
|
||||
account_id: String,
|
||||
@@ -691,7 +683,7 @@ pub(crate) async fn account_quota(
|
||||
let raw = data.get("raw").cloned().unwrap_or(Value::Null);
|
||||
let quota_json = quota.to_string();
|
||||
let mut quota_data = quota_data("a_json, Some(Utc::now()));
|
||||
apply_account_quota_window_token_usage(&state, &account_id, &mut quota_data).await;
|
||||
apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await;
|
||||
let account =
|
||||
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
|
||||
Ok(AdminResponse::new(
|
||||
@@ -922,10 +914,11 @@ async fn account_list_stats(
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let mut quota_by_account = quota_by_account.clone();
|
||||
apply_quota_window_token_usage(state, accounts, &mut quota_by_account).await;
|
||||
let selected_quota_windows = selected_quota_windows_by_account(accounts, "a_by_account);
|
||||
let all_quota_windows = quota_usage_windows_by_account(accounts, "a_by_account);
|
||||
let quota_window_usage_stats =
|
||||
account_usage_stats_by_quota_window(state, &selected_quota_windows).await;
|
||||
quota_window_local_usage_by_account(state, &all_quota_windows).await;
|
||||
apply_quota_window_local_usage(accounts, &mut quota_by_account, "a_window_usage_stats);
|
||||
let selected_quota_windows = selected_quota_windows_by_account(accounts, "a_by_account);
|
||||
let usage_records = usage_records
|
||||
.into_iter()
|
||||
.map(|usage| {
|
||||
@@ -972,105 +965,11 @@ fn selected_quota_window(quota: &AdminAccountQuotaData) -> Option<AdminAccountQu
|
||||
})
|
||||
}
|
||||
|
||||
async fn account_usage_stats_by_quota_window(
|
||||
state: &AppState,
|
||||
windows_by_account: &HashMap<String, AdminAccountQuotaUsageWindow>,
|
||||
) -> HashMap<String, AccountUsageWindowStats> {
|
||||
if windows_by_account.is_empty() {
|
||||
return HashMap::new();
|
||||
}
|
||||
|
||||
let mut builder = QueryBuilder::<Sqlite>::new(
|
||||
"select
|
||||
account_id,
|
||||
coalesce(sum(request_count), 0) as request_count,
|
||||
coalesce(sum(input_tokens), 0) as input_tokens,
|
||||
coalesce(sum(output_tokens), 0) as output_tokens,
|
||||
coalesce(sum(cached_tokens), 0) as cached_tokens
|
||||
from usage_time_buckets
|
||||
where ",
|
||||
);
|
||||
for (index, (account_id, window)) in windows_by_account.iter().enumerate() {
|
||||
if index > 0 {
|
||||
builder.push(" or ");
|
||||
}
|
||||
let bucket_start = china_quarter_hour_start(window.start);
|
||||
builder.push("(account_id = ");
|
||||
builder.push_bind(account_id);
|
||||
builder.push(" and bucket_start >= ");
|
||||
builder.push_bind(bucket_start.to_rfc3339());
|
||||
builder.push(" and bucket_start <= ");
|
||||
builder.push_bind(window.end.to_rfc3339());
|
||||
builder.push(")");
|
||||
}
|
||||
builder.push(" group by account_id");
|
||||
|
||||
let rows = builder
|
||||
.build()
|
||||
.fetch_all(state.services.background_tasks.accounts.pool())
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
rows.into_iter()
|
||||
.map(|row| {
|
||||
(
|
||||
row.get::<String, _>("account_id"),
|
||||
AccountUsageWindowStats {
|
||||
request_count: row.get("request_count"),
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn apply_selected_quota_window_usage(
|
||||
mut usage: AdminUsageRecord,
|
||||
selected_quota_windows: &HashMap<String, AdminAccountQuotaUsageWindow>,
|
||||
quota_window_usage_stats: &HashMap<String, AccountUsageWindowStats>,
|
||||
) -> AdminUsageRecord {
|
||||
let Some(window) = selected_quota_windows.get(&usage.account_id) else {
|
||||
return usage;
|
||||
};
|
||||
let stats = quota_window_usage_stats
|
||||
.get(&usage.account_id)
|
||||
.copied()
|
||||
.unwrap_or_default();
|
||||
usage.window_request_count = stats.request_count;
|
||||
usage.window_input_tokens = stats.input_tokens;
|
||||
usage.window_output_tokens = stats.output_tokens;
|
||||
usage.window_cached_tokens = stats.cached_tokens;
|
||||
usage.window_started_at = Some(window.start);
|
||||
usage.window_reset_at = Some(window.end);
|
||||
usage.limit_window_seconds = Some(window.window_seconds);
|
||||
usage
|
||||
}
|
||||
|
||||
async fn apply_account_quota_window_token_usage(
|
||||
state: &AppState,
|
||||
account_id: &str,
|
||||
quota: &mut AdminAccountQuotaData,
|
||||
) {
|
||||
let windows = quota
|
||||
.usage_windows()
|
||||
.into_iter()
|
||||
.map(|window| AccountQuotaUsageWindow {
|
||||
account_id: account_id.to_string(),
|
||||
window,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let usage_by_account = quota_window_token_usage_by_account(state, &windows).await;
|
||||
let empty_usage = HashMap::new();
|
||||
quota.apply_token_usage(usage_by_account.get(account_id).unwrap_or(&empty_usage));
|
||||
}
|
||||
|
||||
async fn apply_quota_window_token_usage(
|
||||
state: &AppState,
|
||||
fn quota_usage_windows_by_account(
|
||||
accounts: &[AdminAccountMetadata],
|
||||
quota_by_account: &mut HashMap<String, AdminAccountQuotaData>,
|
||||
) {
|
||||
let windows = accounts
|
||||
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
|
||||
) -> Vec<AccountQuotaUsageWindow> {
|
||||
accounts
|
||||
.iter()
|
||||
.filter_map(|account| {
|
||||
quota_by_account
|
||||
@@ -1086,21 +985,44 @@ async fn apply_quota_window_token_usage(
|
||||
window,
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let usage_by_account = quota_window_token_usage_by_account(state, &windows).await;
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn apply_account_quota_window_local_usage(
|
||||
state: &AppState,
|
||||
account_id: &str,
|
||||
quota: &mut AdminAccountQuotaData,
|
||||
) {
|
||||
let windows = quota
|
||||
.usage_windows()
|
||||
.into_iter()
|
||||
.map(|window| AccountQuotaUsageWindow {
|
||||
account_id: account_id.to_string(),
|
||||
window,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let usage_by_account = quota_window_local_usage_by_account(state, &windows).await;
|
||||
let empty_usage = HashMap::new();
|
||||
quota.apply_local_usage(usage_by_account.get(account_id).unwrap_or(&empty_usage));
|
||||
}
|
||||
|
||||
fn apply_quota_window_local_usage(
|
||||
accounts: &[AdminAccountMetadata],
|
||||
quota_by_account: &mut HashMap<String, AdminAccountQuotaData>,
|
||||
usage_by_account: &HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>,
|
||||
) {
|
||||
for account in accounts {
|
||||
if let Some(quota) = quota_by_account.get_mut(&account.id) {
|
||||
let empty_usage = HashMap::new();
|
||||
quota.apply_token_usage(usage_by_account.get(&account.id).unwrap_or(&empty_usage));
|
||||
quota.apply_local_usage(usage_by_account.get(&account.id).unwrap_or(&empty_usage));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn quota_window_token_usage_by_account(
|
||||
async fn quota_window_local_usage_by_account(
|
||||
state: &AppState,
|
||||
windows: &[AccountQuotaUsageWindow],
|
||||
) -> HashMap<String, HashMap<String, AdminAccountQuotaWindowTokenUsage>> {
|
||||
) -> HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>> {
|
||||
let Some(min_start) = windows.iter().map(|window| window.window.start).min() else {
|
||||
return HashMap::new();
|
||||
};
|
||||
@@ -1123,8 +1045,10 @@ async fn quota_window_token_usage_by_account(
|
||||
"select
|
||||
account_id,
|
||||
bucket_start,
|
||||
coalesce(sum(request_count), 0) as request_count,
|
||||
coalesce(sum(input_tokens), 0) as input_tokens,
|
||||
coalesce(sum(output_tokens), 0) as output_tokens
|
||||
coalesce(sum(output_tokens), 0) as output_tokens,
|
||||
coalesce(sum(cached_tokens), 0) as cached_tokens
|
||||
from usage_time_buckets
|
||||
where account_id in (",
|
||||
);
|
||||
@@ -1134,7 +1058,7 @@ async fn quota_window_token_usage_by_account(
|
||||
}
|
||||
separated.push_unseparated(")");
|
||||
builder.push(" and bucket_start >= ");
|
||||
builder.push_bind(min_start.to_rfc3339());
|
||||
builder.push_bind(china_quarter_hour_start(min_start).to_rfc3339());
|
||||
builder.push(" and bucket_start <= ");
|
||||
builder.push_bind(max_end.to_rfc3339());
|
||||
builder.push(" group by account_id, bucket_start");
|
||||
@@ -1145,35 +1069,61 @@ async fn quota_window_token_usage_by_account(
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let mut usage_by_account =
|
||||
HashMap::<String, HashMap<String, AdminAccountQuotaWindowTokenUsage>>::new();
|
||||
HashMap::<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>::new();
|
||||
|
||||
for row in rows {
|
||||
let account_id: String = row.get("account_id");
|
||||
let Ok(bucket_start) = parse_rfc3339_utc(row.get::<&str, _>("bucket_start")) else {
|
||||
continue;
|
||||
};
|
||||
let total_tokens = nonnegative_i64_to_u64(row.get::<i64, _>("input_tokens"))
|
||||
.saturating_add(nonnegative_i64_to_u64(row.get::<i64, _>("output_tokens")));
|
||||
let bucket_usage = AdminAccountQuotaWindowLocalUsage {
|
||||
request_count: row.get("request_count"),
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
};
|
||||
let Some(account_windows) = windows_by_account.get(account_id.as_str()) else {
|
||||
continue;
|
||||
};
|
||||
|
||||
for window in account_windows
|
||||
.iter()
|
||||
.filter(|window| bucket_start >= window.start && bucket_start <= window.end)
|
||||
{
|
||||
for window in account_windows.iter().filter(|window| {
|
||||
bucket_start >= china_quarter_hour_start(window.start) && bucket_start <= window.end
|
||||
}) {
|
||||
usage_by_account
|
||||
.entry(account_id.clone())
|
||||
.or_default()
|
||||
.entry(window.key.clone())
|
||||
.or_default()
|
||||
.add_tokens(total_tokens);
|
||||
.add(bucket_usage);
|
||||
}
|
||||
}
|
||||
|
||||
usage_by_account
|
||||
}
|
||||
|
||||
fn apply_selected_quota_window_usage(
|
||||
mut usage: AdminUsageRecord,
|
||||
selected_quota_windows: &HashMap<String, AdminAccountQuotaUsageWindow>,
|
||||
quota_window_usage_stats: &HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>,
|
||||
) -> AdminUsageRecord {
|
||||
let Some(window) = selected_quota_windows.get(&usage.account_id) else {
|
||||
return usage;
|
||||
};
|
||||
let stats = quota_window_usage_stats
|
||||
.get(&usage.account_id)
|
||||
.and_then(|usage_by_window| usage_by_window.get(&window.key))
|
||||
.copied()
|
||||
.unwrap_or_default();
|
||||
usage.window_request_count = stats.request_count;
|
||||
usage.window_input_tokens = stats.input_tokens;
|
||||
usage.window_output_tokens = stats.output_tokens;
|
||||
usage.window_cached_tokens = stats.cached_tokens;
|
||||
usage.window_started_at = Some(window.start);
|
||||
usage.window_reset_at = Some(window.end);
|
||||
usage.limit_window_seconds = Some(window.window_seconds);
|
||||
usage
|
||||
}
|
||||
|
||||
async fn account_summary_data(
|
||||
state: &AppState,
|
||||
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
|
||||
|
||||
@@ -755,7 +755,7 @@ async fn execute_fresh_pooled_response_create_request(
|
||||
key,
|
||||
PooledWebSocketConnection {
|
||||
websocket,
|
||||
metadata,
|
||||
metadata: reusable_websocket_metadata(metadata),
|
||||
created_at: now,
|
||||
},
|
||||
)
|
||||
@@ -827,7 +827,7 @@ async fn execute_pooled_response_create_request(
|
||||
let created_at = connection.created_at;
|
||||
match collect_websocket_response(
|
||||
connection.websocket,
|
||||
connection.metadata,
|
||||
reusable_websocket_metadata(connection.metadata),
|
||||
true,
|
||||
started_at,
|
||||
first_token_started_at,
|
||||
@@ -840,7 +840,7 @@ async fn execute_pooled_response_create_request(
|
||||
key,
|
||||
PooledWebSocketConnection {
|
||||
websocket,
|
||||
metadata,
|
||||
metadata: reusable_websocket_metadata(metadata),
|
||||
created_at,
|
||||
},
|
||||
)
|
||||
@@ -861,7 +861,7 @@ async fn execute_pooled_response_create_request_stream(
|
||||
connection: PooledWebSocketConnection,
|
||||
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
|
||||
let mut websocket = connection.websocket;
|
||||
let mut metadata = connection.metadata;
|
||||
let mut metadata = reusable_websocket_metadata(connection.metadata);
|
||||
let created_at = connection.created_at;
|
||||
let first_token_started_at = Instant::now();
|
||||
if let Err(error) = websocket
|
||||
@@ -1144,6 +1144,13 @@ fn websocket_connection_metadata(
|
||||
}
|
||||
}
|
||||
|
||||
fn reusable_websocket_metadata(
|
||||
mut metadata: CodexWebSocketConnectionMetadata,
|
||||
) -> CodexWebSocketConnectionMetadata {
|
||||
metadata.rate_limit_headers.clear();
|
||||
metadata
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Streaming
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -1406,7 +1413,7 @@ async fn finish_stream_websocket(
|
||||
pool_return.key,
|
||||
PooledWebSocketConnection {
|
||||
websocket,
|
||||
metadata,
|
||||
metadata: reusable_websocket_metadata(metadata),
|
||||
created_at: pool_return.created_at,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -4,6 +4,29 @@ use codex_proxy_rs::upstream::{
|
||||
transport::websocket::write_websocket_audit_artifact_for_dir,
|
||||
};
|
||||
|
||||
fn rate_limit_event(used_percent: u64) -> String {
|
||||
json!({
|
||||
"type": "codex.rate_limits",
|
||||
"rate_limits": {
|
||||
"primary": {
|
||||
"used_percent": used_percent,
|
||||
"window_minutes": 43200,
|
||||
"reset_at": 1893456000 + used_percent as i64,
|
||||
}
|
||||
}
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn primary_used_percent_values(headers: &[(String, String)]) -> Vec<&str> {
|
||||
headers
|
||||
.iter()
|
||||
.filter_map(|(name, value)| {
|
||||
(name == "x-codex-primary-used-percent").then_some(value.as_str())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_audit_artifact_should_require_explicit_directory() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
@@ -663,6 +686,74 @@ async fn websocket_execute_response_create_request_should_capture_internal_metad
|
||||
.any(|(name, value)| name == "x-codex-primary-reset-at" && value == "1893456300"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn codex_backend_client_should_not_reuse_websocket_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_codex_test_websocket(stream).await;
|
||||
|
||||
let _first_message = websocket.next().await.unwrap().unwrap();
|
||||
websocket
|
||||
.send(Message::Text(rate_limit_event(10).into()))
|
||||
.await
|
||||
.unwrap();
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_rate_limit_first", 3, 1).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let _second_message = websocket.next().await.unwrap().unwrap();
|
||||
websocket
|
||||
.send(Message::Text(rate_limit_event(20).into()))
|
||||
.await
|
||||
.unwrap();
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
completed_websocket_response("resp_rate_limit_second", 5, 2).into(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
websocket.close(None).await.unwrap();
|
||||
});
|
||||
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
|
||||
let backend = CodexBackendClient::new(
|
||||
reqwest::Client::builder().no_proxy().build().unwrap(),
|
||||
format!("http://{addr}"),
|
||||
crate::support::fingerprint::runtime_test_fingerprint(),
|
||||
)
|
||||
.with_websocket_pool(pool);
|
||||
let request = pooled_websocket_request("conversation-rate-limit-reuse");
|
||||
|
||||
let first = backend
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_rate_limit_first", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("first pooled websocket response should succeed");
|
||||
let second = backend
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_rate_limit_second", Some("chatgpt-account")),
|
||||
)
|
||||
.await
|
||||
.expect("second pooled websocket response should reuse connection");
|
||||
server.await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
primary_used_percent_values(&first.rate_limit_headers),
|
||||
vec!["10"]
|
||||
);
|
||||
assert_eq!(
|
||||
primary_used_percent_values(&second.rate_limit_headers),
|
||||
vec!["20"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_execute_response_create_request_should_surface_incomplete_response() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
|
||||
@@ -42,6 +42,29 @@ function quotaWindowBarClass(window?: any) {
|
||||
}
|
||||
return 'bg-(--cp-success)'
|
||||
}
|
||||
|
||||
function quotaWindowPercentTextClass(window?: any) {
|
||||
if (window?.usedPercent === null || window?.usedPercent === undefined) {
|
||||
return 'text-(--cp-text-muted)'
|
||||
}
|
||||
if (window.usedPercent >= 95) {
|
||||
return 'text-(--cp-danger-text)'
|
||||
}
|
||||
if (window.usedPercent >= 80) {
|
||||
return 'text-(--cp-warning-text)'
|
||||
}
|
||||
return 'text-(--cp-success-text)'
|
||||
}
|
||||
|
||||
function quotaWindowLocalUsageDisplay(window?: any) {
|
||||
const display = window?.localUsage?.totalTokensDisplay
|
||||
return typeof display === 'string' && display.trim() ? display : '-'
|
||||
}
|
||||
|
||||
function shouldShowQuotaWindowLocalUsage(window?: any) {
|
||||
const totalTokens = window?.localUsage?.totalTokens
|
||||
return typeof totalTokens === 'number' && totalTokens > 0
|
||||
}
|
||||
</script>
|
||||
|
||||
<template>
|
||||
@@ -80,7 +103,17 @@ function quotaWindowBarClass(window?: any) {
|
||||
>
|
||||
<div class="flex items-center justify-between gap-3 text-[12px] font-[720]">
|
||||
<span class="text-(--cp-text-secondary)">{{ window.labelDisplay }}</span>
|
||||
<span class="text-(--cp-text-primary)">{{ window.usedPercentDisplay }}</span>
|
||||
<span class="flex shrink-0 items-baseline justify-end gap-1.5 font-mono tabular-nums">
|
||||
<span
|
||||
v-if="shouldShowQuotaWindowLocalUsage(window)"
|
||||
class="text-[12px] font-[680] text-(--cp-text-muted)"
|
||||
>
|
||||
{{ quotaWindowLocalUsageDisplay(window) }}
|
||||
</span>
|
||||
<span class="text-[12px] font-[780]" :class="quotaWindowPercentTextClass(window)">
|
||||
{{ window.usedPercentDisplay }}
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<div class="mt-2 h-2 overflow-hidden rounded-full bg-(--cp-default-border)">
|
||||
<div
|
||||
@@ -105,7 +138,17 @@ function quotaWindowBarClass(window?: any) {
|
||||
>
|
||||
<div class="flex items-center justify-between gap-3 text-[12px] font-[720]">
|
||||
<span class="text-(--cp-text-secondary)">{{ window.labelDisplay }}</span>
|
||||
<span class="text-(--cp-text-primary)">{{ window.usedPercentDisplay }}</span>
|
||||
<span class="flex shrink-0 items-baseline justify-end gap-1.5 font-mono tabular-nums">
|
||||
<span
|
||||
v-if="shouldShowQuotaWindowLocalUsage(window)"
|
||||
class="text-[12px] font-[680] text-(--cp-text-muted)"
|
||||
>
|
||||
{{ quotaWindowLocalUsageDisplay(window) }}
|
||||
</span>
|
||||
<span class="text-[12px] font-[780]" :class="quotaWindowPercentTextClass(window)">
|
||||
{{ window.usedPercentDisplay }}
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<div class="mt-2 h-2 overflow-hidden rounded-full bg-(--cp-default-border)">
|
||||
<div
|
||||
@@ -130,7 +173,17 @@ function quotaWindowBarClass(window?: any) {
|
||||
>
|
||||
<div class="flex items-center justify-between gap-3 text-[12px] font-[720]">
|
||||
<span class="text-(--cp-text-secondary)">{{ window.labelDisplay }}</span>
|
||||
<span class="text-(--cp-text-primary)">{{ window.usedPercentDisplay }}</span>
|
||||
<span class="flex shrink-0 items-baseline justify-end gap-1.5 font-mono tabular-nums">
|
||||
<span
|
||||
v-if="shouldShowQuotaWindowLocalUsage(window)"
|
||||
class="text-[12px] font-[680] text-(--cp-text-muted)"
|
||||
>
|
||||
{{ quotaWindowLocalUsageDisplay(window) }}
|
||||
</span>
|
||||
<span class="text-[12px] font-[780]" :class="quotaWindowPercentTextClass(window)">
|
||||
{{ window.usedPercentDisplay }}
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<div class="mt-2 h-2 overflow-hidden rounded-full bg-(--cp-default-border)">
|
||||
<div
|
||||
|
||||
@@ -50,16 +50,21 @@ function quotaWindowBarClass(window?: any) {
|
||||
return 'bg-(--cp-success)'
|
||||
}
|
||||
|
||||
function windowUsageDisplay(window?: any) {
|
||||
const display = window?.tokenUsageDisplay
|
||||
return typeof display === 'string' && display.trim() ? display : '-'
|
||||
}
|
||||
|
||||
function windowPercentDisplay(window?: any) {
|
||||
const display = window?.usedPercentDisplay
|
||||
return typeof display === 'string' && display.trim() ? display : '-'
|
||||
}
|
||||
|
||||
function windowLocalUsageDisplay(window?: any) {
|
||||
const display = window?.localUsage?.totalTokensDisplay
|
||||
return typeof display === 'string' && display.trim() ? display : '-'
|
||||
}
|
||||
|
||||
function shouldShowWindowLocalUsage(window?: any) {
|
||||
const totalTokens = window?.localUsage?.totalTokens
|
||||
return typeof totalTokens === 'number' && totalTokens > 0
|
||||
}
|
||||
|
||||
function quotaWindowPercentTextClass(window?: any) {
|
||||
if (window?.usedPercent === null || window?.usedPercent === undefined) {
|
||||
return 'text-(--cp-text-muted)'
|
||||
@@ -99,15 +104,18 @@ function quotaWindowMatches(actual: unknown, expected: number) {
|
||||
{{ window.labelDisplay }}
|
||||
</span>
|
||||
<span class="flex shrink-0 items-baseline justify-end gap-1.5 font-mono tabular-nums">
|
||||
<span
|
||||
v-if="shouldShowWindowLocalUsage(window)"
|
||||
class="text-[10px] leading-none font-[680] text-(--cp-text-muted)"
|
||||
>
|
||||
{{ windowLocalUsageDisplay(window) }}
|
||||
</span>
|
||||
<span
|
||||
class="text-[10px] leading-none font-[780]"
|
||||
:class="quotaWindowPercentTextClass(window)"
|
||||
>
|
||||
{{ windowPercentDisplay(window) }}
|
||||
</span>
|
||||
<span class="text-[10px] leading-none font-[680] text-(--cp-text-muted)">
|
||||
{{ windowUsageDisplay(window) }}
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<div class="h-1 w-full overflow-hidden rounded-full bg-(--cp-default-border)">
|
||||
|
||||
Reference in New Issue
Block a user