From 583a61da7d352c08cbe16434aedc52add7587d54 Mon Sep 17 00:00:00 2001 From: zyycn Date: Sat, 4 Jul 2026 19:55:51 +0800 Subject: [PATCH] fix: harden account scheduling and client context Co-authored-by: Codex --- backend/src/admin/auth/session.rs | 22 +-- backend/src/http/middleware/request_id.rs | 65 ++++---- .../src/proxy/dispatch/responses/service.rs | 134 ++++++++++++++++ backend/src/proxy/openai/responses.rs | 53 ------- backend/src/upstream/accounts/pool.rs | 143 +++++++++++++++--- backend/src/upstream/accounts/store/mod.rs | 77 ++++++++++ .../src/upstream/accounts/store/queries.rs | 57 +++++++ backend/tests/http/trace_middleware/mod.rs | 139 ++++++++++++----- .../chat_upstream/responses_websocket.rs | 86 +++++++++++ backend/tests/runtime/tasks/mod.rs | 8 + backend/tests/runtime/tasks/model_refresh.rs | 8 + .../accounts/account_pool/selection.rs | 124 +++++++++++++++ .../upstream/accounts/account_repository.rs | 93 ++++++++++++ 13 files changed, 855 insertions(+), 154 deletions(-) diff --git a/backend/src/admin/auth/session.rs b/backend/src/admin/auth/session.rs index 244bedae4..75e899f1d 100644 --- a/backend/src/admin/auth/session.rs +++ b/backend/src/admin/auth/session.rs @@ -4,7 +4,7 @@ use axum::{ extract::{FromRequestParts, State}, http::{header::SET_COOKIE, request::Parts, HeaderMap, HeaderValue, StatusCode}, response::{IntoResponse, Response}, - Json, + Extension, Json, }; use serde::{Deserialize, Serialize}; @@ -13,6 +13,7 @@ use crate::{ auth::service::AdminSessionError, response::{AdminEnvelope, AdminError, AdminResponse}, }, + http::middleware::request_id::ClientIp, infra::time::china_rfc3339, runtime::state::AppState, }; @@ -57,10 +58,10 @@ struct AdminSessionStatusData { /// `POST /api/admin/login` pub(crate) async fn login( State(state): State, - headers: HeaderMap, + client_ip: Option>, Json(payload): Json, ) -> Result { - let source = admin_login_source(&headers); + let source = admin_login_source(optional_client_ip(client_ip)); let login_result = state .services .admin_sessions @@ -191,12 +192,11 @@ fn admin_api_key_header(headers: &HeaderMap) -> Option { (!value.is_empty()).then(|| value.to_string()) } -fn admin_login_source(headers: &HeaderMap) -> String { - headers - .get("x-real-ip") - .and_then(|value| value.to_str().ok()) - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or("unknown") - .to_string() +fn optional_client_ip(client_ip: Option>) -> Option { + let Extension(client_ip) = client_ip?; + Some(client_ip.as_str().to_string()) +} + +fn admin_login_source(fallback_client_ip: Option) -> String { + fallback_client_ip.unwrap_or_else(|| "unknown".to_string()) } diff --git a/backend/src/http/middleware/request_id.rs b/backend/src/http/middleware/request_id.rs index 6edfee8fa..578427906 100644 --- a/backend/src/http/middleware/request_id.rs +++ b/backend/src/http/middleware/request_id.rs @@ -14,14 +14,8 @@ use uuid::Uuid; /// 请求 ID 头。 pub const REQUEST_ID_HEADER: &str = "x-request-id"; const MAX_REQUEST_ID_LEN: usize = 128; -const CF_CONNECTING_IP_HEADER: &str = "cf-connecting-ip"; const REAL_IP_HEADER: &str = "x-real-ip"; const X_FORWARDED_FOR_HEADER: &str = "x-forwarded-for"; -const UNTRUSTED_CLIENT_IP_HEADERS: &[&str] = &[ - CF_CONNECTING_IP_HEADER, - REAL_IP_HEADER, - X_FORWARDED_FOR_HEADER, -]; /// 可信反向代理配置。 #[derive(Debug, Clone, Default, PartialEq, Eq)] @@ -77,6 +71,17 @@ impl RequestId { } } +/// 由连接来源和可信代理链解析出的客户端 IP。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ClientIp(String); + +impl ClientIp { + /// 返回 IP 字符串。 + pub fn as_str(&self) -> &str { + &self.0 + } +} + fn valid_request_id(value: &str) -> bool { !value.is_empty() && value.len() <= MAX_REQUEST_ID_LEN @@ -117,10 +122,10 @@ async fn run_with_request_id(mut request: Request, next: Next) -> Response { } fn attach_request_context(trusted_proxies: &TrustedProxyConfig, request: &mut Request) { - attach_real_ip_from_connection(trusted_proxies, request); + attach_client_ip_from_connection(trusted_proxies, request); } -fn attach_real_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request: &mut Request) { +fn attach_client_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request: &mut Request) { let peer_ip = request .extensions() .get::>() @@ -128,47 +133,45 @@ fn attach_real_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request: let client_ip = peer_ip.map(|peer_ip| client_ip_from_peer(peer_ip, request.headers(), trusted_proxies)); - for header in UNTRUSTED_CLIENT_IP_HEADERS { - request.headers_mut().remove(*header); - } - let Some(client_ip) = client_ip else { return; }; - if let Ok(value) = HeaderValue::from_str(&client_ip.to_string()) { - request.headers_mut().insert(REAL_IP_HEADER, value); - } + request.extensions_mut().insert(ClientIp(client_ip)); } fn client_ip_from_peer( peer_ip: IpAddr, headers: &HeaderMap, trusted_proxies: &TrustedProxyConfig, -) -> IpAddr { +) -> String { if trusted_proxies.is_trusted(peer_ip) { - trusted_forwarded_client_ip(headers).unwrap_or(peer_ip) + trusted_forwarded_client_ip(headers).unwrap_or_else(|| peer_ip.to_string()) } else { - peer_ip + peer_ip.to_string() } } -fn trusted_forwarded_client_ip(headers: &HeaderMap) -> Option { - header_ip(headers, CF_CONNECTING_IP_HEADER) - .or_else(|| first_forwarded_for_ip(headers)) - .or_else(|| header_ip(headers, REAL_IP_HEADER)) +fn trusted_forwarded_client_ip(headers: &HeaderMap) -> Option { + first_forwarded_for_value(headers).or_else(|| header_string(headers, REAL_IP_HEADER)) } -fn header_ip(headers: &HeaderMap, name: &str) -> Option { - parse_ip(headers.get(name)?.to_str().ok()?) +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 first_forwarded_for_ip(headers: &HeaderMap) -> Option { - let value = headers.get(X_FORWARDED_FOR_HEADER)?.to_str().ok()?; - parse_ip(value.split(',').next()?) -} - -fn parse_ip(value: &str) -> Option { - value.trim().parse().ok() +fn first_forwarded_for_value(headers: &HeaderMap) -> Option { + let value = header_string(headers, X_FORWARDED_FOR_HEADER)?; + value + .split(',') + .next() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) } fn parse_trusted_proxy_entry(entry: &str) -> Result { diff --git a/backend/src/proxy/dispatch/responses/service.rs b/backend/src/proxy/dispatch/responses/service.rs index c1a3b4586..6fb4d2164 100644 --- a/backend/src/proxy/dispatch/responses/service.rs +++ b/backend/src/proxy/dispatch/responses/service.rs @@ -1231,6 +1231,140 @@ impl ResponseDispatchService { history_recovery_used = true; continue; } + Err(ResponseDispatchError::Upstream(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); + continue; + } + Err(ResponseDispatchError::Upstream(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); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_auth_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + let upstream_error = upstream_error_body(&error); + let account_status = auth_failure_account_status(&error); + match account_status { + AccountStatus::Disabled => { + disabled_count += 1; + last_disabled_auth_error = Some(upstream_error); + last_exhausted_account_class = + Some(ExhaustedAccountClass::Disabled); + } + AccountStatus::Banned => { + banned_count += 1; + last_banned_status_code = + Some(upstream_error_http_status(&error)); + last_banned_auth_error = Some(upstream_error); + last_exhausted_account_class = + Some(ExhaustedAccountClass::Banned); + } + _ => { + expired_count += 1; + last_auth_error = Some(upstream_error); + last_exhausted_account_class = + Some(ExhaustedAccountClass::Expired); + } + } + self.account_pool + .set_status(&release_account_id, account_status) + .await; + trigger_refresh_after_auth_failure( + &self.token_refresh, + &release_account_id, + account_status, + ); + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(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); + continue; + } + Err(ResponseDispatchError::Upstream(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); + continue; + } + Err(ResponseDispatchError::Upstream(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_stream_dispatch_error!( + ResponseDispatchError::ModelUnsupported { + count: model_unsupported_count + 1, + upstream_error, + }, + account_id: Some(&release_account_id), + transport: Some(backend_transport_name(transport)) + ); + } + 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; + } + Err(ResponseDispatchError::Upstream(error)) + if is_banned_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + banned_count += 1; + last_banned_status_code = Some(upstream_error_http_status(&error)); + last_banned_auth_error = Some(upstream_error_body(&error)); + last_exhausted_account_class = Some(ExhaustedAccountClass::Banned); + self.account_pool + .set_status(&release_account_id, AccountStatus::Banned) + .await; + excluded_account_ids.push(release_account_id); + continue; + } Err(error) => { self.account_pool.release(&release_account_id).await; if let ResponseDispatchError::Upstream(upstream_error) = &error { diff --git a/backend/src/proxy/openai/responses.rs b/backend/src/proxy/openai/responses.rs index 5ddda300f..95445e1d7 100644 --- a/backend/src/proxy/openai/responses.rs +++ b/backend/src/proxy/openai/responses.rs @@ -9,7 +9,6 @@ use axum::{ }; use serde::{Deserialize, Serialize}; use serde_json::{json, Map, Value}; -use std::net::IpAddr; use crate::{ http::middleware::request_id::RequestId, @@ -612,7 +611,6 @@ async fn handle_compact_responses( } }; let mut compact_request = translate_response_to_compact(openai_request); - compact_request.client_ip = client_ip_from_headers(&headers); compact_request.client_user_agent = user_agent_from_headers(&headers); match state @@ -694,7 +692,6 @@ fn fill_string_from_header(field: &mut Option, headers: &HeaderMap, name } pub(crate) fn attach_client_context(request: &mut CodexResponsesRequest, headers: &HeaderMap) { - request.client_ip = client_ip_from_headers(headers); request.client_user_agent = user_agent_from_headers(headers); } @@ -702,56 +699,6 @@ pub(crate) fn user_agent_from_headers(headers: &HeaderMap) -> Option { header_string(headers, "user-agent") } -pub(crate) fn client_ip_from_headers(headers: &HeaderMap) -> Option { - header_string(headers, "cf-connecting-ip") - .or_else(|| header_string(headers, "x-real-ip")) - .or_else(|| forwarded_client_ip(headers)) - .and_then(|value| normalize_client_ip(&value)) -} - -fn forwarded_client_ip(headers: &HeaderMap) -> Option { - let value = header_string(headers, "x-forwarded-for")?; - let parts = value - .split(',') - .map(str::trim) - .filter(|part| !part.is_empty()) - .collect::>(); - parts - .iter() - .find(|part| !is_private_ip(part)) - .or_else(|| parts.first()) - .map(|part| (*part).to_string()) -} - -fn normalize_client_ip(value: &str) -> Option { - let trimmed = value.trim(); - if trimmed.is_empty() { - return None; - } - - trimmed - .parse::() - .map(|addr| addr.ip().to_string()) - .ok() - .or_else(|| Some(trimmed.to_string())) -} - -fn is_private_ip(value: &str) -> bool { - let Some(ip) = normalize_client_ip(value).and_then(|value| value.parse::().ok()) else { - return false; - }; - - match ip { - IpAddr::V4(ip) => { - ip.is_private() - || ip.is_loopback() - || ip.is_link_local() - || ip.octets()[0] == 100 && (64..=127).contains(&ip.octets()[1]) - } - IpAddr::V6(ip) => ip.is_loopback() || ip.is_unique_local() || ip.is_unicast_link_local(), - } -} - fn openai_subagent_from_headers(headers: &HeaderMap) -> Option { header_string(headers, OPENAI_SUBAGENT_HEADER).and_then(|value| { if matches!( diff --git a/backend/src/upstream/accounts/pool.rs b/backend/src/upstream/accounts/pool.rs index 25f5ace60..c0706ab3e 100644 --- a/backend/src/upstream/accounts/pool.rs +++ b/backend/src/upstream/accounts/pool.rs @@ -154,13 +154,19 @@ pub struct DistinctPlanAccount { #[derive(Debug)] struct AccountAcquireWithStatusRefresh { acquired: Option, - expired_account_ids: Vec, + refreshed_accounts: Vec, } #[derive(Debug)] struct DistinctPlanAccountsWithStatusRefresh { accounts: Vec, - expired_account_ids: Vec, + refreshed_accounts: Vec, +} + +#[derive(Debug)] +struct AccountCapacitySummaryWithStatusRefresh { + summary: AccountCapacitySummary, + refreshed_accounts: Vec, } /// 账号池容量摘要。 @@ -176,6 +182,71 @@ pub struct AccountCapacitySummary { pub available_slots: usize, } +#[derive(Debug, Clone, PartialEq, Eq)] +struct RuntimeAccountStateSnapshot { + status: AccountStatus, + quota_limit_reached: bool, + quota_verify_required: bool, + quota_cooldown_until: Option>, + cloudflare_cooldown_until: Option>, + window_request_count: u64, + window_input_tokens: u64, + window_output_tokens: u64, + window_cached_tokens: u64, + window_image_input_tokens: u64, + window_image_output_tokens: u64, + window_image_request_count: u64, + window_image_request_failed_count: u64, + window_started_at: Option>, + window_reset_at: Option>, + limit_window_seconds: Option, +} + +impl From<&Account> for RuntimeAccountStateSnapshot { + fn from(account: &Account) -> Self { + Self { + status: account.status, + quota_limit_reached: account.quota_limit_reached, + quota_verify_required: account.quota_verify_required, + quota_cooldown_until: account.quota_cooldown_until, + cloudflare_cooldown_until: account.cloudflare_cooldown_until, + window_request_count: account.window_request_count, + window_input_tokens: account.window_input_tokens, + window_output_tokens: account.window_output_tokens, + window_cached_tokens: account.window_cached_tokens, + window_image_input_tokens: account.window_image_input_tokens, + window_image_output_tokens: account.window_image_output_tokens, + window_image_request_count: account.window_image_request_count, + window_image_request_failed_count: account.window_image_request_failed_count, + window_started_at: account.window_started_at, + window_reset_at: account.window_reset_at, + limit_window_seconds: account.limit_window_seconds, + } + } +} + +impl RuntimeAccountStateSnapshot { + fn usage_window_changed(&self, other: &Self) -> bool { + self.window_request_count != other.window_request_count + || self.window_input_tokens != other.window_input_tokens + || self.window_output_tokens != other.window_output_tokens + || self.window_cached_tokens != other.window_cached_tokens + || self.window_image_input_tokens != other.window_image_input_tokens + || self.window_image_output_tokens != other.window_image_output_tokens + || self.window_image_request_count != other.window_image_request_count + || self.window_image_request_failed_count != other.window_image_request_failed_count + || self.window_started_at != other.window_started_at + || self.window_reset_at != other.window_reset_at + || self.limit_window_seconds != other.limit_window_seconds + } +} + +#[derive(Debug, Clone)] +struct RefreshedAccountState { + account: Account, + sync_usage_window: bool, +} + /// 运行时窗口用量增量。 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct AccountWindowUsageDelta { @@ -419,7 +490,7 @@ impl AccountPool { request: &AccountAcquireRequest, ) -> AccountAcquireWithStatusRefresh { self.cleanup_stale_slots(request.now); - let expired_account_ids = self.refresh_account_statuses(request.now); + let refreshed_accounts = self.refresh_account_statuses(request.now); let candidates = self.candidates(request); let selected = if let Some(preferred_account_id) = &request.preferred_account_id { let preferred = candidates @@ -458,7 +529,7 @@ impl AccountPool { let Some((selected, previous_slot_at, selection_source)) = selected else { return AccountAcquireWithStatusRefresh { acquired: None, - expired_account_ids, + refreshed_accounts, }; }; tracing::info!( @@ -485,13 +556,15 @@ impl AccountPool { account: selected, previous_slot_at, }), - expired_account_ids, + refreshed_accounts, } } - fn refresh_account_statuses(&mut self, now: DateTime) -> Vec { + fn refresh_account_statuses(&mut self, now: DateTime) -> Vec { + let mut refreshed_accounts = Vec::new(); let mut expired_account_ids = Vec::new(); for (account_id, account) in &mut self.accounts { + let before = RuntimeAccountStateSnapshot::from(&*account); if account.status == AccountStatus::Active && access_token_expired(account, now) { account.status = AccountStatus::Expired; expired_account_ids.push(account_id.clone()); @@ -499,11 +572,18 @@ impl AccountPool { refresh_quota_window(account, now); refresh_cloudflare_cooldown(account, now); } + let after = RuntimeAccountStateSnapshot::from(&*account); + if before != after { + refreshed_accounts.push(RefreshedAccountState { + account: account.clone(), + sync_usage_window: before.usage_window_changed(&after), + }); + } } for account_id in &expired_account_ids { self.slots.remove(account_id); } - expired_account_ids + refreshed_accounts } /// 释放指定账号的一个在途槽位。 @@ -527,8 +607,15 @@ impl AccountPool { /// 计算账号池容量摘要。 pub fn capacity_summary(&mut self, now: DateTime) -> AccountCapacitySummary { + self.capacity_summary_with_status_refresh(now).summary + } + + fn capacity_summary_with_status_refresh( + &mut self, + now: DateTime, + ) -> AccountCapacitySummaryWithStatusRefresh { self.cleanup_stale_slots(now); - self.refresh_account_statuses(now); + let refreshed_accounts = self.refresh_account_statuses(now); let is_capacity_account = |account: &Account| { account.status == AccountStatus::Active && AccountService::quota_available_at(account, now, self.options.skip_quota_limited) @@ -550,11 +637,14 @@ impl AccountPool { .map(|(_, slots)| slots.len().min(self.options.max_concurrent_per_account)) .sum(); - AccountCapacitySummary { - max_concurrent_per_account: self.options.max_concurrent_per_account, - total_slots, - used_slots, - available_slots: total_slots.saturating_sub(used_slots), + AccountCapacitySummaryWithStatusRefresh { + summary: AccountCapacitySummary { + max_concurrent_per_account: self.options.max_concurrent_per_account, + total_slots, + used_slots, + available_slots: total_slots.saturating_sub(used_slots), + }, + refreshed_accounts, } } @@ -569,7 +659,7 @@ impl AccountPool { now: DateTime, ) -> DistinctPlanAccountsWithStatusRefresh { self.cleanup_stale_slots(now); - let expired_account_ids = self.refresh_account_statuses(now); + let refreshed_accounts = self.refresh_account_statuses(now); let mut by_plan = IndexMap::>::new(); for account in self.accounts.values() { @@ -615,7 +705,7 @@ impl AccountPool { DistinctPlanAccountsWithStatusRefresh { accounts, - expired_account_ids, + refreshed_accounts, } } @@ -874,7 +964,14 @@ impl RuntimeAccountPoolService { /// 读取账号池容量摘要。 pub async fn capacity_summary(&self, now: DateTime) -> AccountCapacitySummary { - self.pool.lock().await.capacity_summary(now) + let refresh = self + .pool + .lock() + .await + .capacity_summary_with_status_refresh(now); + self.persist_runtime_account_states(refresh.refreshed_accounts) + .await; + refresh.summary } /// 使用当前时间读取账号池容量摘要。 @@ -964,7 +1061,7 @@ impl RuntimeAccountPoolService { /// 按请求上下文获取可用账号。 pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option { let refresh = self.pool.lock().await.acquire_with_status_refresh(request); - self.persist_expired_statuses(refresh.expired_account_ids) + self.persist_runtime_account_states(refresh.refreshed_accounts) .await; refresh.acquired } @@ -976,22 +1073,22 @@ impl RuntimeAccountPoolService { .lock() .await .distinct_plan_accounts_with_status_refresh(now); - self.persist_expired_statuses(refresh.expired_account_ids) + self.persist_runtime_account_states(refresh.refreshed_accounts) .await; refresh.accounts } - async fn persist_expired_statuses(&self, account_ids: Vec) { - for account_id in account_ids { + async fn persist_runtime_account_states(&self, refreshed_states: Vec) { + for refreshed in refreshed_states { if let Err(error) = self .store - .set_status(&account_id, AccountStatus::Expired) + .sync_runtime_account_state(&refreshed.account, refreshed.sync_usage_window) .await { tracing::warn!( - account_id, + account_id = %refreshed.account.id, error = %error, - "failed to persist runtime-expired account status" + "failed to persist refreshed runtime account state" ); } } diff --git a/backend/src/upstream/accounts/store/mod.rs b/backend/src/upstream/accounts/store/mod.rs index 4fd2d27fc..1ee19e786 100644 --- a/backend/src/upstream/accounts/store/mod.rs +++ b/backend/src/upstream/accounts/store/mod.rs @@ -338,6 +338,13 @@ pub trait AccountStore: Send + Sync + 'static { cooldown_until: Option>, ) -> AccountStoreResult; + /// 同步运行时自然刷新出来的账号状态。 + async fn sync_runtime_account_state( + &self, + account: &Account, + sync_usage_window: bool, + ) -> AccountStoreResult; + /// 同步账号当前 rate-limit 统计窗口。 async fn sync_rate_limit_window( &self, @@ -711,6 +718,66 @@ impl SqliteAccountStore { Ok(result.rows_affected() > 0) } + /// 同步运行时自然刷新出来的账号状态。 + pub async fn sync_runtime_account_state( + &self, + account: &Account, + sync_usage_window: bool, + ) -> SqliteAccountStoreResult { + let now = Utc::now().to_rfc3339(); + let quota_limit_reached = if account.quota_limit_reached { 1 } else { 0 }; + let quota_verify_required = if account.quota_verify_required { 1 } else { 0 }; + let quota_cooldown_until = account.quota_cooldown_until.map(|dt| dt.to_rfc3339()); + let cloudflare_cooldown_until = account.cloudflare_cooldown_until.map(|dt| dt.to_rfc3339()); + let mut tx = self.pool.begin().await?; + let result = sqlx::query(SYNC_RUNTIME_ACCOUNT_STATE_SQL) + .bind(status_to_db(account.status)) + .bind(quota_limit_reached) + .bind(&now) + .bind(quota_limit_reached) + .bind(quota_verify_required) + .bind(&now) + .bind(quota_verify_required) + .bind("a_cooldown_until) + .bind(&now) + .bind("a_cooldown_until) + .bind(&cloudflare_cooldown_until) + .bind(&now) + .bind(&cloudflare_cooldown_until) + .bind(&now) + .bind(&account.id) + .execute(&mut *tx) + .await?; + if result.rows_affected() == 0 { + tx.commit().await?; + return Ok(false); + } + + if sync_usage_window { + sqlx::query(SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL) + .bind(&account.id) + .bind(u64_to_i64_saturating(account.window_request_count)) + .bind(u64_to_i64_saturating(account.window_input_tokens)) + .bind(u64_to_i64_saturating(account.window_output_tokens)) + .bind(u64_to_i64_saturating(account.window_cached_tokens)) + .bind(u64_to_i64_saturating(account.window_image_input_tokens)) + .bind(u64_to_i64_saturating(account.window_image_output_tokens)) + .bind(u64_to_i64_saturating(account.window_image_request_count)) + .bind(u64_to_i64_saturating( + account.window_image_request_failed_count, + )) + .bind(account.window_started_at.map(|dt| dt.to_rfc3339())) + .bind(account.window_reset_at.map(|dt| dt.to_rfc3339())) + .bind(account.limit_window_seconds.map(u64_to_i64_saturating)) + .bind(&now) + .execute(&mut *tx) + .await?; + } + + tx.commit().await?; + Ok(true) + } + /// 标记账号 Cloudflare 冷却期。 pub async fn set_cloudflare_cooldown_until( &self, @@ -1082,6 +1149,16 @@ impl AccountStore for SqliteAccountStore { .map_err(|error| map_account_store_error(&error)) } + async fn sync_runtime_account_state( + &self, + account: &Account, + sync_usage_window: bool, + ) -> AccountStoreResult { + SqliteAccountStore::sync_runtime_account_state(self, account, sync_usage_window) + .await + .map_err(|error| map_account_store_error(&error)) + } + async fn sync_rate_limit_window( &self, account_id: &str, diff --git a/backend/src/upstream/accounts/store/queries.rs b/backend/src/upstream/accounts/store/queries.rs index 47982fd61..400cc5dac 100644 --- a/backend/src/upstream/accounts/store/queries.rs +++ b/backend/src/upstream/accounts/store/queries.rs @@ -432,6 +432,63 @@ set updated_at = ? where id = ?"; +pub(super) const SYNC_RUNTIME_ACCOUNT_STATE_SQL: &str = r" +update accounts +set + status = case + when status in ('disabled', 'banned', 'refreshing') then status + else ? + end, + quota_limit_reached = case + when ? = 0 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_limit_reached + else ? + end, + quota_verify_required = case + when ? = 1 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_verify_required + else ? + end, + quota_cooldown_until = case + when ? is null and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_cooldown_until + else ? + end, + cloudflare_cooldown_until = case + when ? is null and cloudflare_cooldown_until is not null and cloudflare_cooldown_until > ? then cloudflare_cooldown_until + else ? + end, + updated_at = ? +where id = ?"; + +pub(super) const SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +on conflict(account_id) do update set + window_request_count = excluded.window_request_count, + window_input_tokens = excluded.window_input_tokens, + window_output_tokens = excluded.window_output_tokens, + window_cached_tokens = excluded.window_cached_tokens, + window_image_input_tokens = excluded.window_image_input_tokens, + window_image_output_tokens = excluded.window_image_output_tokens, + window_image_request_count = excluded.window_image_request_count, + window_image_request_failed_count = excluded.window_image_request_failed_count, + window_started_at = excluded.window_started_at, + window_reset_at = excluded.window_reset_at, + limit_window_seconds = excluded.limit_window_seconds +where account_usage.window_reset_at is null + or account_usage.window_reset_at <= excluded.window_reset_at + or account_usage.window_reset_at <= ?"; + pub(super) const SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL: &str = r" update accounts set diff --git a/backend/tests/http/trace_middleware/mod.rs b/backend/tests/http/trace_middleware/mod.rs index c3a382408..47186d64a 100644 --- a/backend/tests/http/trace_middleware/mod.rs +++ b/backend/tests/http/trace_middleware/mod.rs @@ -11,10 +11,12 @@ use axum::{ http::{Request, StatusCode}, middleware::{from_fn, from_fn_with_state}, routing::get, - Router, + Extension, Router, }; use codex_proxy_rs::http::middleware::{ - request_id::{attach_request_id, attach_request_id_with_proxy_config, TrustedProxyConfig}, + request_id::{ + attach_request_id, attach_request_id_with_proxy_config, ClientIp, TrustedProxyConfig, + }, trace::http_trace_layer, }; use tokio::time::{sleep, Duration}; @@ -219,16 +221,12 @@ async fn request_context_should_generate_new_request_id_for_invalid_header() { } #[tokio::test] -async fn request_context_should_attach_real_ip_from_connection() { +async fn request_context_should_attach_client_ip_from_connection() { let app = Router::new() .route( "/ip", - get(|headers: HeaderMap| async move { - headers - .get("x-real-ip") - .and_then(|value| value.to_str().ok()) - .unwrap_or_default() - .to_string() + get(|Extension(client_ip): Extension| async move { + client_ip.as_str().to_string() }), ) .layer(from_fn(attach_request_id)); @@ -247,17 +245,30 @@ async fn request_context_should_attach_real_ip_from_connection() { } #[tokio::test] -async fn request_context_should_replace_untrusted_forwarded_ip_headers() { +async fn request_context_should_preserve_forwarded_headers_from_untrusted_peer() { let app = Router::new() .route( "/ip", - get(|headers: HeaderMap| async move { - headers - .get("x-real-ip") - .and_then(|value| value.to_str().ok()) - .unwrap_or("missing") - .to_string() - }), + get( + |Extension(client_ip): Extension, headers: HeaderMap| async move { + let cf_ip = headers + .get("cf-connecting-ip") + .and_then(|value| value.to_str().ok()) + .unwrap_or("missing"); + let forwarded_for = headers + .get("x-forwarded-for") + .and_then(|value| value.to_str().ok()) + .unwrap_or("missing"); + let real_ip = headers + .get("x-real-ip") + .and_then(|value| value.to_str().ok()) + .unwrap_or("missing"); + format!( + "{}|cf={cf_ip}|xff={forwarded_for}|real={real_ip}", + client_ip.as_str() + ) + }, + ), ) .layer(from_fn(attach_request_id)); let mut request = Request::builder() @@ -274,7 +285,10 @@ async fn request_context_should_replace_untrusted_forwarded_ip_headers() { let response = app.oneshot(request).await.expect("response"); let bytes = to_bytes(response.into_body(), 1024).await.unwrap(); - assert_eq!(std::str::from_utf8(&bytes).unwrap(), "127.0.0.1"); + assert_eq!( + std::str::from_utf8(&bytes).unwrap(), + "127.0.0.1|cf=203.0.113.43|xff=203.0.113.42, 10.0.0.2|real=203.0.113.44" + ); } #[tokio::test] @@ -284,14 +298,12 @@ async fn request_context_should_accept_forwarded_ip_from_trusted_proxy() { let app = Router::new() .route( "/ip", - get(|headers: HeaderMap| async move { - let real_ip = headers - .get("x-real-ip") - .and_then(|value| value.to_str().ok()) - .unwrap_or("missing"); - let forwarded_present = headers.contains_key("x-forwarded-for"); - format!("{real_ip}|forwarded={forwarded_present}") - }), + get( + |Extension(client_ip): Extension, headers: HeaderMap| async move { + let forwarded_present = headers.contains_key("x-forwarded-for"); + format!("{}|forwarded={forwarded_present}", client_ip.as_str()) + }, + ), ) .layer(from_fn_with_state( trusted_proxies, @@ -311,23 +323,19 @@ async fn request_context_should_accept_forwarded_ip_from_trusted_proxy() { assert_eq!( std::str::from_utf8(&bytes).unwrap(), - "203.0.113.42|forwarded=false" + "203.0.113.42|forwarded=true" ); } #[tokio::test] -async fn request_context_should_prefer_cloudflare_ip_from_trusted_proxy() { +async fn request_context_should_accept_real_ip_from_trusted_proxy_when_forwarded_for_missing() { let trusted_proxies = TrustedProxyConfig::from_entries(&["127.0.0.1".to_string()]).expect("trusted proxy"); let app = Router::new() .route( "/ip", - get(|headers: HeaderMap| async move { - headers - .get("x-real-ip") - .and_then(|value| value.to_str().ok()) - .unwrap_or("missing") - .to_string() + get(|Extension(client_ip): Extension| async move { + client_ip.as_str().to_string() }), ) .layer(from_fn_with_state( @@ -336,8 +344,7 @@ async fn request_context_should_prefer_cloudflare_ip_from_trusted_proxy() { )); let mut request = Request::builder() .uri("/ip") - .header("cf-connecting-ip", "198.51.100.23") - .header("x-forwarded-for", "203.0.113.42, 10.0.0.2") + .header("x-real-ip", "198.51.100.23") .body(Body::empty()) .expect("request should build"); request @@ -350,6 +357,66 @@ async fn request_context_should_prefer_cloudflare_ip_from_trusted_proxy() { assert_eq!(std::str::from_utf8(&bytes).unwrap(), "198.51.100.23"); } +#[tokio::test] +async fn request_context_should_ignore_forwarded_ip_from_untrusted_peer() { + let trusted_proxies = + TrustedProxyConfig::from_entries(&["10.0.0.0/8".to_string()]).expect("trusted proxy"); + let app = Router::new() + .route( + "/ip", + get(|Extension(client_ip): Extension| async move { + client_ip.as_str().to_string() + }), + ) + .layer(from_fn_with_state( + trusted_proxies, + attach_request_id_with_proxy_config, + )); + let mut request = Request::builder() + .uri("/ip") + .header("x-forwarded-for", "203.0.113.42") + .body(Body::empty()) + .expect("request should build"); + request + .extensions_mut() + .insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 61234)))); + + let response = app.oneshot(request).await.expect("response"); + let bytes = to_bytes(response.into_body(), 1024).await.unwrap(); + + assert_eq!(std::str::from_utf8(&bytes).unwrap(), "127.0.0.1"); +} + +#[tokio::test] +async fn request_context_should_not_inject_real_ip_header() { + let app = Router::new() + .route( + "/ip", + get(|headers: HeaderMap| async move { + headers + .get("x-real-ip") + .and_then(|value| value.to_str().ok()) + .unwrap_or("missing") + .to_string() + }), + ) + .layer(from_fn(attach_request_id)); + let mut request = Request::builder() + .uri("/ip") + .header("x-forwarded-for", "203.0.113.42, 10.0.0.2") + .header("cf-connecting-ip", "203.0.113.43") + .body(Body::empty()) + .expect("request should build"); + request + .extensions_mut() + .insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 61234)))); + + let response = app.oneshot(request).await.expect("response"); + let bytes = to_bytes(response.into_body(), 1024).await.unwrap(); + + assert_eq!(std::str::from_utf8(&bytes).unwrap(), "missing"); +} + #[test] fn trusted_proxy_config_should_reject_invalid_entries() { let error = TrustedProxyConfig::from_entries(&["not-an-ip".to_string()]).unwrap_err(); diff --git a/backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs b/backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs index 957a2db86..acfed5a0b 100644 --- a/backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs +++ b/backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs @@ -133,6 +133,92 @@ async fn responses_websocket_stream_should_synthesize_response_failed_when_close assert!(body.contains("stream_disconnected")); } +#[tokio::test] +async fn responses_websocket_stream_first_error_429_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": "error", + "error": { + "type": "usage_limit_reached", + "message": "The usage limit has been reached", + "retry_after_seconds": 77 + }, + "status_code": 429 + }) + .to_string(), + ) + .await; + first_websocket.close(None).await.unwrap(); + + let second_payload = accept_successful_websocket_response_with_authorization( + &listener, + "Bearer access-secondary", + "resp_after_ws_stream_rate_limit", + ) + .await; + (first_payload, second_payload) + }); + 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 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 (first_payload, second_payload) = upstream.await.unwrap(); + 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,) = + 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_eq!(status, StatusCode::OK); + assert!(content_type.starts_with("text/event-stream")); + assert!(body.contains("resp_after_ws_stream_rate_limit")); + assert_eq!(first_payload["type"], "response.create"); + assert_eq!(second_payload["type"], "response.create"); + assert_eq!(primary_quota_state.0, 1); + assert!(primary_quota_state.1.is_some()); + assert_eq!(primary_usage.0, 1); + assert_eq!(secondary_usage.0, 1); +} + #[tokio::test] async fn responses_websocket_should_reuse_connection_for_recorded_conversation() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/backend/tests/runtime/tasks/mod.rs b/backend/tests/runtime/tasks/mod.rs index 7a4eb88ca..ed781131e 100644 --- a/backend/tests/runtime/tasks/mod.rs +++ b/backend/tests/runtime/tasks/mod.rs @@ -97,6 +97,14 @@ impl AccountStore for FakeAccountStore { Ok(false) } + async fn sync_runtime_account_state( + &self, + _account: &Account, + _sync_usage_window: bool, + ) -> AccountStoreResult { + Ok(false) + } + async fn sync_rate_limit_window( &self, _account_id: &str, diff --git a/backend/tests/runtime/tasks/model_refresh.rs b/backend/tests/runtime/tasks/model_refresh.rs index dd8c6eac8..89451d84c 100644 --- a/backend/tests/runtime/tasks/model_refresh.rs +++ b/backend/tests/runtime/tasks/model_refresh.rs @@ -210,6 +210,14 @@ impl AccountStore for PlanAccountStore { Ok(false) } + async fn sync_runtime_account_state( + &self, + _account: &Account, + _sync_usage_window: bool, + ) -> AccountStoreResult { + Ok(false) + } + async fn sync_rate_limit_window( &self, _account_id: &str, diff --git a/backend/tests/upstream/accounts/account_pool/selection.rs b/backend/tests/upstream/accounts/account_pool/selection.rs index cabf7cf1a..cfaf5d206 100644 --- a/backend/tests/upstream/accounts/account_pool/selection.rs +++ b/backend/tests/upstream/accounts/account_pool/selection.rs @@ -210,6 +210,130 @@ async fn runtime_account_pool_should_persist_expired_status_when_jwt_expiry_is_d assert_eq!(status.0, "expired"); } +#[tokio::test] +async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_expires() { + let now = fixed_time(); + let (pool, _dir) = + crate::support::sqlite::init_test_db("runtime-pool-quota-window-reset.sqlite").await; + let store = codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()); + store + .insert(codex_proxy_rs::upstream::accounts::store::NewAccount { + id: "acct_quota_reset".to_string(), + email: None, + account_id: Some("chatgpt-quota-reset".to_string()), + user_id: None, + label: None, + plan_type: Some("plus".to_string()), + access_token: secrecy::SecretString::new("access-token".to_string().into()), + refresh_token: None, + access_token_expires_at: Some(now + Duration::hours(1)), + status: AccountStatus::QuotaExhausted, + added_at: None, + }) + .await + .unwrap(); + sqlx::query( + r#" + update accounts + set + quota_limit_reached = 1, + quota_verify_required = 0, + quota_cooldown_until = ?, + cloudflare_cooldown_until = ? + where id = ? + "#, + ) + .bind((now - Duration::seconds(1)).to_rfc3339()) + .bind((now - Duration::seconds(1)).to_rfc3339()) + .bind("acct_quota_reset") + .execute(&pool) + .await + .unwrap(); + sqlx::query( + r#" + insert into account_usage ( + account_id, + request_count, + window_request_count, + window_input_tokens, + window_output_tokens, + window_cached_tokens, + window_started_at, + window_reset_at, + limit_window_seconds + ) values (?, 9, 7, 11, 13, 17, ?, ?, 60) + "#, + ) + .bind("acct_quota_reset") + .bind((now - Duration::seconds(120)).to_rfc3339()) + .bind((now - Duration::seconds(1)).to_rfc3339()) + .execute(&pool) + .await + .unwrap(); + let runtime_pool = codex_proxy_rs::upstream::accounts::pool::RuntimeAccountPoolService::new( + std::sync::Arc::new(store), + AccountPoolOptions::default(), + 0, + ); + runtime_pool.restore_from_repository().await.unwrap(); + + let acquired = runtime_pool + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) + .await + .unwrap(); + let stored: ( + String, + i64, + i64, + Option, + Option, + i64, + i64, + i64, + i64, + Option, + Option, + ) = sqlx::query_as( + r#" + select + a.status, + a.quota_limit_reached, + a.quota_verify_required, + a.quota_cooldown_until, + a.cloudflare_cooldown_until, + au.window_request_count, + au.window_input_tokens, + au.window_output_tokens, + au.window_cached_tokens, + au.window_reset_at, + au.limit_window_seconds + from accounts a + left join account_usage au on au.account_id = a.id + where a.id = ? + "#, + ) + .bind("acct_quota_reset") + .fetch_one(&pool) + .await + .unwrap(); + let persisted_reset_at = stored + .9 + .as_deref() + .and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok()) + .map(|value| value.with_timezone(&Utc)); + + assert_eq!(acquired.account.id, "acct_quota_reset"); + assert!(acquired.account.quota_verify_required); + assert_eq!(stored.0, "active"); + assert_eq!(stored.1, 0); + assert_eq!(stored.2, 1); + assert!(stored.3.is_none()); + assert!(stored.4.is_none()); + assert_eq!((stored.5, stored.6, stored.7, stored.8), (0, 0, 0, 0)); + assert!(persisted_reset_at.is_some_and(|reset_at| reset_at > now)); + assert_eq!(stored.10, Some(60)); +} + #[test] fn account_pool_should_prefer_configured_tier_priority() { let mut pool = AccountPool::with_options(AccountPoolOptions { diff --git a/backend/tests/upstream/accounts/account_repository.rs b/backend/tests/upstream/accounts/account_repository.rs index 00d5d8f5b..89d72e44a 100644 --- a/backend/tests/upstream/accounts/account_repository.rs +++ b/backend/tests/upstream/accounts/account_repository.rs @@ -518,6 +518,99 @@ async fn account_repository_should_map_quota_snapshot_to_account_status() { assert!(restored_status.2.is_none()); } +#[tokio::test] +async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quota_cooldown() { + let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 42).await; + let repo = SqliteAccountStore::new(pool.clone()); + seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; + let future_cooldown = Utc::now() + Duration::hours(1); + sqlx::query( + "update accounts set quota_limit_reached = 1, quota_verify_required = 0, quota_cooldown_until = ? where id = ?", + ) + .bind(future_cooldown.to_rfc3339()) + .bind("acct_a") + .execute(&pool) + .await + .unwrap(); + let mut stale_account = AccountStore::list_pool_accounts(&repo) + .await + .unwrap() + .into_iter() + .find(|account| account.id == "acct_a") + .unwrap(); + stale_account.quota_limit_reached = false; + stale_account.quota_verify_required = true; + stale_account.quota_cooldown_until = None; + + repo.sync_runtime_account_state(&stale_account, false) + .await + .unwrap(); + let stored: (i64, i64, Option) = sqlx::query_as( + "select quota_limit_reached, quota_verify_required, quota_cooldown_until from accounts where id = ?", + ) + .bind("acct_a") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(stored.0, 1); + assert_eq!(stored.1, 0); + assert_eq!(stored.2, Some(future_cooldown.to_rfc3339())); +} + +#[tokio::test] +async fn account_repository_runtime_state_sync_should_not_regress_newer_usage_window() { + let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 43).await; + let repo = SqliteAccountStore::new(pool.clone()); + seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; + let now = Utc::now(); + let future_reset = now + Duration::hours(1); + let stale_reset = now + Duration::minutes(1); + sqlx::query( + r#" + insert into account_usage ( + account_id, + window_request_count, + window_input_tokens, + window_reset_at, + limit_window_seconds + ) values (?, 5, 17, ?, 3600) + "#, + ) + .bind("acct_a") + .bind(future_reset.to_rfc3339()) + .execute(&pool) + .await + .unwrap(); + let mut stale_account = AccountStore::list_pool_accounts(&repo) + .await + .unwrap() + .into_iter() + .find(|account| account.id == "acct_a") + .unwrap(); + stale_account.window_request_count = 0; + stale_account.window_input_tokens = 0; + stale_account.window_reset_at = Some(stale_reset); + stale_account.window_started_at = Some(now); + stale_account.limit_window_seconds = Some(60); + + repo.sync_runtime_account_state(&stale_account, true) + .await + .unwrap(); + let stored: (i64, i64, Option, Option) = sqlx::query_as( + "select window_request_count, window_input_tokens, window_reset_at, limit_window_seconds from account_usage where account_id = ?", + ) + .bind("acct_a") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(stored.0, 5); + assert_eq!(stored.1, 17); + assert_eq!(stored.2, Some(future_reset.to_rfc3339())); + assert_eq!(stored.3, Some(3600)); +} + #[tokio::test] async fn account_repository_should_not_override_disabled_status_from_quota_snapshot() { let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 38).await;