diff --git a/codex-rs/codex-api/src/api_bridge.rs b/codex-rs/codex-api/src/api_bridge.rs index 7f429562b..505fcea66 100644 --- a/codex-rs/codex-api/src/api_bridge.rs +++ b/codex-rs/codex-api/src/api_bridge.rs @@ -13,6 +13,7 @@ use codex_protocol::error::CodexErr; use codex_protocol::error::RetryLimitReachedError; use codex_protocol::error::UnexpectedResponseError; use codex_protocol::error::UsageLimitReachedError; +use codex_protocol::protocol::RateLimitReachedType; use http::HeaderMap; use serde::Deserialize; use serde_json::Value; @@ -96,29 +97,39 @@ pub fn map_api_error(err: ApiError) -> CodexErr { } else if status == http::StatusCode::INTERNAL_SERVER_ERROR { CodexErr::InternalServerError } else if status == http::StatusCode::TOO_MANY_REQUESTS { - if let Ok(err) = serde_json::from_str::(&body_text) { - if err.error.error_type.as_deref() == Some("usage_limit_reached") { - let limit_id = extract_header(headers.as_ref(), ACTIVE_LIMIT_HEADER); - let rate_limits = headers.as_ref().and_then(|map| { - parse_rate_limit_for_limit(map, limit_id.as_deref()) - }); - let promo_message = headers.as_ref().and_then(parse_promo_message); - let rate_limit_reached_type = - headers.as_ref().and_then(parse_rate_limit_reached_type); - let resets_at = err - .error + let usage_error = serde_json::from_str::(&body_text).ok(); + let rate_limit_reached_type = + headers.as_ref().and_then(parse_rate_limit_reached_type); + + if usage_error.as_ref().is_some_and(|err| { + err.error.error_type.as_deref() == Some("usage_limit_reached") + }) || is_account_scoped_rate_limit_reached(rate_limit_reached_type) + { + let limit_id = extract_header(headers.as_ref(), ACTIVE_LIMIT_HEADER); + let rate_limits = headers + .as_ref() + .and_then(|map| parse_rate_limit_for_limit(map, limit_id.as_deref())); + let promo_message = headers.as_ref().and_then(parse_promo_message); + let resets_at = usage_error.as_ref().and_then(|err| { + err.error .resets_at - .and_then(|seconds| DateTime::::from_timestamp(seconds, 0)); - return CodexErr::UsageLimitReached(UsageLimitReachedError { - plan_type: err.error.plan_type, - resets_at, - rate_limits: rate_limits.map(Box::new), - promo_message, - rate_limit_reached_type, - }); - } else if err.error.error_type.as_deref() == Some("usage_not_included") { - return CodexErr::UsageNotIncluded; - } + .and_then(|seconds| DateTime::::from_timestamp(seconds, 0)) + }); + return CodexErr::UsageLimitReached(UsageLimitReachedError { + plan_type: usage_error + .as_ref() + .and_then(|err| err.error.plan_type.clone()), + resets_at, + rate_limits: rate_limits.map(Box::new), + promo_message, + rate_limit_reached_type, + }); + } + + if usage_error.as_ref().is_some_and(|err| { + err.error.error_type.as_deref() == Some("usage_not_included") + }) { + return CodexErr::UsageNotIncluded; } CodexErr::RetryLimit(RetryLimitReachedError { @@ -218,3 +229,17 @@ struct UsageErrorBody { plan_type: Option, resets_at: Option, } + +fn is_account_scoped_rate_limit_reached( + rate_limit_reached_type: Option, +) -> bool { + matches!( + rate_limit_reached_type, + Some( + RateLimitReachedType::WorkspaceOwnerCreditsDepleted + | RateLimitReachedType::WorkspaceMemberCreditsDepleted + | RateLimitReachedType::WorkspaceOwnerUsageLimitReached + | RateLimitReachedType::WorkspaceMemberUsageLimitReached + ) + ) +} diff --git a/codex-rs/codex-api/src/api_bridge_tests.rs b/codex-rs/codex-api/src/api_bridge_tests.rs index 2d5da6198..ef2dfd8a7 100644 --- a/codex-rs/codex-api/src/api_bridge_tests.rs +++ b/codex-rs/codex-api/src/api_bridge_tests.rs @@ -265,6 +265,78 @@ fn map_api_error_ignores_unparseable_rate_limit_reached_type_headers() { } } +#[test] +fn map_api_error_maps_account_scoped_rate_limit_header_without_usage_error_type() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-codex-primary-used-percent", + http::HeaderValue::from_static("100.0"), + ); + headers.insert( + "x-codex-primary-window-minutes", + http::HeaderValue::from_static("300"), + ); + headers.insert( + "x-codex-rate-limit-reached-type", + http::HeaderValue::from_static("workspace_member_usage_limit_reached"), + ); + let body = serde_json::json!({ + "error": { + "type": "rate_limit_reached", + "message": "workspace spend cap reached", + } + }) + .to_string(); + + let err = map_api_error(ApiError::Transport(TransportError::Http { + status: http::StatusCode::TOO_MANY_REQUESTS, + url: Some("http://example.com/v1/responses".to_string()), + headers: Some(headers), + body: Some(body), + })); + + let CodexErr::UsageLimitReached(usage_limit) = err else { + panic!("expected CodexErr::UsageLimitReached, got {err:?}"); + }; + assert_eq!( + usage_limit.rate_limit_reached_type, + Some(codex_protocol::protocol::RateLimitReachedType::WorkspaceMemberUsageLimitReached) + ); + assert_eq!( + usage_limit + .rate_limits + .as_ref() + .and_then(|snapshot| snapshot.primary.as_ref()) + .map(|window| window.used_percent), + Some(100.0) + ); +} + +#[test] +fn map_api_error_keeps_generic_rate_limit_header_retryable() { + let mut headers = HeaderMap::new(); + headers.insert( + "x-codex-rate-limit-reached-type", + http::HeaderValue::from_static("rate_limit_reached"), + ); + let body = serde_json::json!({ + "error": { + "type": "rate_limit_reached", + "message": "retry later", + } + }) + .to_string(); + + let err = map_api_error(ApiError::Transport(TransportError::Http { + status: http::StatusCode::TOO_MANY_REQUESTS, + url: Some("http://example.com/v1/responses".to_string()), + headers: Some(headers), + body: Some(body), + })); + + assert!(matches!(err, CodexErr::RetryLimit(_))); +} + #[test] fn map_api_error_extracts_identity_auth_details_from_headers() { let mut headers = HeaderMap::new(); diff --git a/codex-rs/core/tests/suite/auth_profile_auto_switch.rs b/codex-rs/core/tests/suite/auth_profile_auto_switch.rs index ad000b441..1af38da31 100644 --- a/codex-rs/core/tests/suite/auth_profile_auto_switch.rs +++ b/codex-rs/core/tests/suite/auth_profile_auto_switch.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use codex_app_server_protocol::AuthMode; use codex_config::types::AuthCredentialsStoreMode; +use codex_core::compact::SUMMARY_PREFIX; use codex_core::config::AuthProfileAutoSwitchStrategy; use codex_login::AuthDotJson; use codex_login::CodexAuth; @@ -9,7 +10,9 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::user_input::UserInput; use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_completed_with_tokens; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_compact_response_sequence; use core_test_support::responses::mount_response_sequence; use core_test_support::responses::sse; use core_test_support::responses::sse_response; @@ -48,6 +51,35 @@ fn usage_limit_response() -> ResponseTemplate { })) } +fn account_scoped_usage_limit_response() -> ResponseTemplate { + ResponseTemplate::new(429) + .insert_header( + "x-codex-rate-limit-reached-type", + "workspace_member_usage_limit_reached", + ) + .insert_header("x-codex-primary-used-percent", "100.0") + .insert_header("x-codex-primary-window-minutes", "300") + .set_body_json(json!({ + "error": { + "type": "rate_limit_reached", + "message": "workspace usage limit reached", + "resets_at": 1704067242, + "plan_type": "pro" + } + })) +} + +fn compact_success_response() -> ResponseTemplate { + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .set_body_json(json!({ + "output": [{ + "type": "compaction", + "encrypted_content": format!("{SUMMARY_PREFIX}\ncompacted summary") + }] + })) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn auth_profile_auto_switch_retries_until_profiles_are_exhausted_or_one_succeeds() -> anyhow::Result<()> { @@ -67,7 +99,7 @@ async fn auth_profile_auto_switch_retries_until_profiles_are_exhausted_or_one_su let request_log = mount_response_sequence( &server, vec![ - usage_limit_response(), + account_scoped_usage_limit_response(), usage_limit_response(), sse_response(sse(vec![ ev_response_created("resp-success"), @@ -189,3 +221,89 @@ async fn auth_profile_auto_switch_stops_after_all_profiles_are_exhausted() -> an Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn auth_profile_auto_switch_retries_pre_turn_compaction_after_account_exhaustion() +-> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let codex_home = Arc::new(TempDir::new()?); + + for profile in ["account001", "account002"] { + codex_login::save_auth_profile( + codex_home.path(), + AuthCredentialsStoreMode::File, + profile, + &api_key_auth(&format!("test-key-{profile}")), + )?; + } + + let responses_mock = mount_response_sequence( + &server, + vec![ + sse_response(sse(vec![ + ev_response_created("response-one"), + ev_completed_with_tokens("response-one", /*total_tokens*/ 60), + ])), + sse_response(sse(vec![ + ev_response_created("response-two"), + ev_completed_with_tokens("response-two", /*total_tokens*/ 500), + ])), + sse_response(sse(vec![ + ev_response_created("response-three"), + ev_completed_with_tokens("response-three", /*total_tokens*/ 80), + ])), + ], + ) + .await; + let compact_mock = mount_compact_response_sequence( + &server, + vec![ + account_scoped_usage_limit_response(), + compact_success_response(), + ], + ) + .await; + + let mut builder = test_codex() + .with_home(Arc::clone(&codex_home)) + .with_auth(CodexAuth::from_api_key("test-key-root")) + .with_config(|config| { + config.model_auto_compact_token_limit = Some(200); + config.selected_auth_profile = Some("account001".to_string()); + config.auth_profile_auto_switch.enabled = true; + config.auth_profile_auto_switch.strategy = AuthProfileAutoSwitchStrategy::Ordered; + config.auth_profile_auto_switch.profiles = + vec!["account001".to_string(), "account002".to_string()]; + }); + let test = builder.build(&server).await?; + + for text in ["USER_ONE", "USER_TWO", "USER_THREE"] { + test.codex + .submit(Op::UserInput { + environments: None, + items: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + } + + assert_eq!(3, responses_mock.requests().len()); + assert_eq!(2, compact_mock.requests().len()); + assert_eq!( + Some("account002".to_string()), + test.codex.config_snapshot().await.selected_auth_profile + ); + + Ok(()) +}