diff --git a/backend/src/admin/accounts/quota_view.rs b/backend/src/admin/accounts/quota_view.rs index 42206fa4c..8a008d855 100644 --- a/backend/src/admin/accounts/quota_view.rs +++ b/backend/src/admin/accounts/quota_view.rs @@ -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, used_percent_display: String, - token_usage_display: String, + #[serde(skip_serializing_if = "Option::is_none")] + local_usage: Option, reset_at_display: String, window_used_display: String, #[serde(skip)] reset_at: Option>, } +#[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 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, + usage_by_window: &HashMap, ) { 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), diff --git a/backend/src/admin/accounts/routes.rs b/backend/src/admin/accounts/routes.rs index 9f66ede2d..35a158f18 100644 --- a/backend/src/admin/accounts/routes.rs +++ b/backend/src/admin/accounts/routes.rs @@ -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>, } -#[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, -) -> HashMap { - if windows_by_account.is_empty() { - return HashMap::new(); - } - - let mut builder = QueryBuilder::::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::("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, - quota_window_usage_stats: &HashMap, -) -> 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::>(); - 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, -) { - let windows = accounts + quota_by_account: &HashMap, +) -> Vec { + accounts .iter() .filter_map(|account| { quota_by_account @@ -1086,21 +985,44 @@ async fn apply_quota_window_token_usage( window, }) }) - .collect::>(); - 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::>(); + 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, + usage_by_account: &HashMap>, +) { 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> { +) -> HashMap> { 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::>::new(); + HashMap::>::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::("input_tokens")) - .saturating_add(nonnegative_i64_to_u64(row.get::("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, + quota_window_usage_stats: &HashMap>, +) -> 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, diff --git a/backend/src/upstream/transport/websocket.rs b/backend/src/upstream/transport/websocket.rs index 837cc26f8..f5e2e528e 100644 --- a/backend/src/upstream/transport/websocket.rs +++ b/backend/src/upstream/transport/websocket.rs @@ -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 { 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, }, ) diff --git a/backend/tests/upstream/transport/websocket.rs b/backend/tests/upstream/transport/websocket.rs index 5a98eb7e1..94965e658 100644 --- a/backend/tests/upstream/transport/websocket.rs +++ b/backend/tests/upstream/transport/websocket.rs @@ -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(); diff --git a/frontend/src/views/accounts/components/AccountQuotaPanel.vue b/frontend/src/views/accounts/components/AccountQuotaPanel.vue index 21d565fdd..c2449b45b 100644 --- a/frontend/src/views/accounts/components/AccountQuotaPanel.vue +++ b/frontend/src/views/accounts/components/AccountQuotaPanel.vue @@ -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 +}