From 3fa70ca0a87acf33deab2bf4ecc9ca76eb2ee524 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 25 Jul 2026 23:30:53 +0300 Subject: [PATCH] feat: A tool for the agent to see all accounts' usage/rate-limits (like /profile) to pick/switch profile; low-usage system-pro --- codex-rs/core/src/auth_profile_usage.rs | 3 +- codex-rs/core/src/session/session.rs | 6 + codex-rs/core/src/session/turn.rs | 91 ++++++- codex-rs/core/src/session/turn_tests.rs | 70 +++++ .../tools/handlers/auth_profile_control.rs | 116 ++++++-- .../handlers/auth_profile_control_spec.rs | 48 +++- .../handlers/auth_profile_usage_control.rs | 247 +++++++++++++++++- .../auth_profile_usage_control_spec.rs | 32 ++- 8 files changed, 569 insertions(+), 44 deletions(-) diff --git a/codex-rs/core/src/auth_profile_usage.rs b/codex-rs/core/src/auth_profile_usage.rs index ff4542880e..57ae207ead 100644 --- a/codex-rs/core/src/auth_profile_usage.rs +++ b/codex-rs/core/src/auth_profile_usage.rs @@ -42,7 +42,8 @@ pub struct AuthProfileUsageRecommendation { } /// Stable recommendation reason codes. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] pub enum AuthProfileUsageRecommendationReason { CurrentProfileHealthy, CurrentProfileUnknown, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index cb2c78549b..6e3c9d9714 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -332,6 +332,11 @@ impl SessionConfiguration { config.selected_auth_profile = auth_profile; next_configuration.original_config_do_not_use = Arc::new(config); } + if let Some(enabled) = updates.auth_profile_auto_switch_enabled { + let mut config = (*next_configuration.original_config_do_not_use).clone(); + config.auth_profile_auto_switch.enabled = enabled; + next_configuration.original_config_do_not_use = Arc::new(config); + } if let Some(personality) = updates.personality { next_configuration.personality = Some(personality); } @@ -561,6 +566,7 @@ pub(crate) struct SessionSettingsUpdate { pub(crate) permission_profile: Option, pub(crate) active_permission_profile: Option, pub(crate) auth_profile: Option>, + pub(crate) auth_profile_auto_switch_enabled: Option, pub(crate) windows_sandbox_level: Option, pub(crate) model_provider_id: Option, pub(crate) collaboration_mode: Option, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 3b9da0497c..fdc11cc925 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -5,6 +5,8 @@ use std::sync::Arc; use std::sync::atomic::Ordering; use crate::SkillInjections; +use crate::auth_profile_usage::AuthProfileUsageHealth; +use crate::auth_profile_usage::usage_health_for_snapshots; use crate::build_skill_injections; use crate::client::ModelClientSession; use crate::client_common::Prompt; @@ -15,6 +17,7 @@ use crate::compact::run_inline_auto_compact_task; use crate::compact::should_use_remote_compact_task; use crate::compact_remote::run_inline_remote_auto_compact_task; use crate::compact_remote_v2::run_inline_remote_auto_compact_task as run_inline_remote_auto_compact_task_v2; +use crate::config::AuthProfileAutoSwitchConfig; use crate::connectors; use crate::context::ContextualUserFragment; use crate::context_manager::estimate_response_item_model_visible_bytes; @@ -56,6 +59,8 @@ use crate::stream_events_utils::record_completed_response_item_with_finalized_fa use crate::tasks::emit_compact_metric; use crate::tools::ToolRouter; use crate::tools::context::SharedTurnDiffTracker; +use crate::tools::handlers::auth_profile_control_spec::MANAGE_AUTH_PROFILES_TOOL_NAME; +use crate::tools::handlers::auth_profile_usage_control_spec::GET_USAGE_TOOL_NAME; use crate::tools::parallel::ToolCallRuntime; use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::router::ToolRouterParams; @@ -97,6 +102,7 @@ use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::PlanDeltaEvent; +use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::RawResponseCompletedEvent; use codex_protocol::protocol::ReasoningContentDeltaEvent; use codex_protocol::protocol::ReasoningRawContentDeltaEvent; @@ -104,6 +110,7 @@ use codex_protocol::protocol::TurnDiffEvent; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; use codex_tools::ToolName; +use codex_tools::ToolSpec; use codex_tools::filter_request_plugin_install_discoverable_tools_for_client; use codex_utils_output_truncation::approx_token_count; use codex_utils_output_truncation::approx_tokens_from_byte_count_i64; @@ -125,6 +132,12 @@ use tracing::trace; use tracing::trace_span; use tracing::warn; +const AUTH_PROFILE_USAGE_PROMPT_NUDGE: &str = "When account usage or rate limits may affect progress, use `get_usage` with `scope: \"all_accounts\"` to compare auth profiles before switching; use `manage_auth_profiles` to switch profiles or set `auto_switch_enabled` with `action: \"set_auto_switch\"`."; + +/// Remaining-capacity threshold (percent) below which the auth-profile usage nudge is added to the +/// developer instructions. +const AUTH_PROFILE_USAGE_LOW_REMAINING_PERCENT: f64 = 20.0; + /// Takes initial turn input and runs a loop where, at each sampling request, /// the model replies with either: /// @@ -1193,9 +1206,10 @@ pub(crate) fn build_prompt( turn_context: &TurnContext, base_instructions: BaseInstructions, ) -> Prompt { + let tools = router.model_visible_specs(); Prompt { input, - tools: router.model_visible_specs(), + tools, parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls, base_instructions, personality: turn_context.personality, @@ -1204,6 +1218,75 @@ pub(crate) fn build_prompt( } } +/// Appends the auth-profile usage nudge only when the current profile is actually running low +/// on Codex capacity. The nudge is deliberately usage-gated so healthy sessions keep byte-identical +/// developer instructions (and therefore keep their prompt cache). +async fn with_auth_profile_usage_prompt_nudge( + sess: &Session, + base_instructions: BaseInstructions, + tools: &[ToolSpec], +) -> BaseInstructions { + if !auth_profile_usage_tools_visible(tools) { + return base_instructions; + } + let (rate_limits, auto_switch_config) = { + let state = sess.state.lock().await; + let config = &state.session_configuration.original_config_do_not_use; + ( + state.latest_rate_limits.clone(), + config.auth_profile_auto_switch.clone(), + ) + }; + let Some(rate_limits) = rate_limits else { + return base_instructions; + }; + if !auth_profile_usage_is_low(&rate_limits, &auto_switch_config) { + return base_instructions; + } + append_auth_profile_usage_prompt_nudge(base_instructions) +} + +fn auth_profile_usage_is_low( + rate_limits: &RateLimitSnapshot, + auto_switch_config: &AuthProfileAutoSwitchConfig, +) -> bool { + match usage_health_for_snapshots(std::slice::from_ref(rate_limits), auto_switch_config) { + AuthProfileUsageHealth::Exhausted { .. } => true, + AuthProfileUsageHealth::Healthy { + remaining_percent, .. + } => remaining_percent <= AUTH_PROFILE_USAGE_LOW_REMAINING_PERCENT, + AuthProfileUsageHealth::Unknown => false, + } +} + +fn append_auth_profile_usage_prompt_nudge( + mut base_instructions: BaseInstructions, +) -> BaseInstructions { + if base_instructions + .text + .contains(AUTH_PROFILE_USAGE_PROMPT_NUDGE) + { + return base_instructions; + } + + if !base_instructions.text.ends_with('\n') { + base_instructions.text.push('\n'); + } + base_instructions.text.push('\n'); + base_instructions + .text + .push_str(AUTH_PROFILE_USAGE_PROMPT_NUDGE); + base_instructions +} + +fn auth_profile_usage_tools_visible(tools: &[ToolSpec]) -> bool { + let has_usage = tools.iter().any(|tool| tool.name() == GET_USAGE_TOOL_NAME); + let has_profile_control = tools + .iter() + .any(|tool| tool.name() == MANAGE_AUTH_PROFILES_TOOL_NAME); + has_usage && has_profile_control +} + #[allow(clippy::too_many_arguments)] #[allow(deprecated)] #[instrument(level = "trace", @@ -1227,6 +1310,12 @@ async fn run_sampling_request( let router = built_tools(sess.as_ref(), turn_context.as_ref(), &cancellation_token).await?; let base_instructions = sess.get_base_instructions().await; + let base_instructions = with_auth_profile_usage_prompt_nudge( + sess.as_ref(), + base_instructions, + &router.model_visible_specs(), + ) + .await; let tool_runtime = ToolCallRuntime::new( Arc::clone(&router), diff --git a/codex-rs/core/src/session/turn_tests.rs b/codex-rs/core/src/session/turn_tests.rs index b18ec97426..f3e6c32e92 100644 --- a/codex-rs/core/src/session/turn_tests.rs +++ b/codex-rs/core/src/session/turn_tests.rs @@ -81,6 +81,76 @@ fn user_input_text(text: &str) -> ResponseItem { } } +fn auth_profile_usage_snapshot(used_percent: f64) -> RateLimitSnapshot { + RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: None, + primary: Some(codex_protocol::protocol::RateLimitWindow { + used_percent, + window_minutes: Some(5 * 60), + resets_at: Some(123), + }), + secondary: None, + credits: None, + individual_limit: None, + plan_type: None, + rate_limit_reached_type: None, + } +} + +#[test] +fn auth_profile_usage_prompt_nudge_requires_both_visible_tools() { + let usage_tool = + crate::tools::handlers::auth_profile_usage_control_spec::create_get_usage_tool(); + let profile_tool = + crate::tools::handlers::auth_profile_control_spec::create_manage_auth_profiles_tool(); + + assert!(!auth_profile_usage_tools_visible(&[])); + assert!(!auth_profile_usage_tools_visible(std::slice::from_ref( + &usage_tool + ))); + assert!(!auth_profile_usage_tools_visible(std::slice::from_ref( + &profile_tool + ))); + assert!(auth_profile_usage_tools_visible(&[ + usage_tool, + profile_tool + ])); +} + +#[test] +fn auth_profile_usage_prompt_nudge_only_fires_when_remaining_capacity_is_low() { + let config = AuthProfileAutoSwitchConfig::default(); + + assert!(!auth_profile_usage_is_low( + &auth_profile_usage_snapshot(/*used_percent*/ 10.0), + &config + )); + assert!(auth_profile_usage_is_low( + &auth_profile_usage_snapshot(/*used_percent*/ 90.0), + &config + )); + assert!(auth_profile_usage_is_low( + &auth_profile_usage_snapshot(/*used_percent*/ 100.0), + &config + )); +} + +#[test] +fn auth_profile_usage_prompt_nudge_is_appended_once() { + let base = BaseInstructions { + text: "base instructions".to_string(), + }; + + let nudged = append_auth_profile_usage_prompt_nudge(base); + assert!(nudged.text.starts_with("base instructions\n\n")); + assert!(nudged.text.contains("scope: \"all_accounts\"")); + assert!(nudged.text.contains("action: \"set_auto_switch\"")); + + let duplicate = append_auth_profile_usage_prompt_nudge(nudged.clone()); + assert_eq!(duplicate.text, nudged.text); +} + #[tokio::test] async fn plan_mode_uses_contributed_turn_item_for_last_agent_message() { let (mut session, turn_context) = crate::session::tests::make_session_and_context().await; diff --git a/codex-rs/core/src/tools/handlers/auth_profile_control.rs b/codex-rs/core/src/tools/handlers/auth_profile_control.rs index 4bbe15063b..3dcdc7e9ac 100644 --- a/codex-rs/core/src/tools/handlers/auth_profile_control.rs +++ b/codex-rs/core/src/tools/handlers/auth_profile_control.rs @@ -1,3 +1,5 @@ +use crate::config::AuthProfileAutoSwitchConfig; +use crate::config::AuthProfileAutoSwitchStrategy; use crate::function_tool::FunctionCallError; use crate::session::session::SessionSettingsUpdate; use crate::tools::context::FunctionToolOutput; @@ -25,6 +27,7 @@ enum ManageAuthProfilesAction { List, Current, Switch, + SetAutoSwitch, } #[derive(Debug, Deserialize)] @@ -33,6 +36,8 @@ struct ManageAuthProfilesArgs { action: ManageAuthProfilesAction, #[serde(default)] profile: Option, + #[serde(default)] + auto_switch_enabled: Option, } #[derive(Debug, Serialize)] @@ -42,6 +47,7 @@ struct ManageAuthProfilesResponse { current_profile: Option, switched_to: Option, profiles: Vec, + auto_switch: AuthProfileAutoSwitchState, } #[derive(Debug, Serialize)] @@ -57,6 +63,16 @@ struct AuthProfileSummary { current: bool, } +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct AuthProfileAutoSwitchState { + enabled: bool, + on_5h_limit: bool, + on_weekly_limit: bool, + strategy: &'static str, + profiles: Vec, +} + #[async_trait::async_trait] impl ToolExecutor for ManageAuthProfilesHandler { fn tool_name(&self) -> ToolName { @@ -91,28 +107,47 @@ impl ToolExecutor for ManageAuthProfilesHandler { let mut current_profile = session.selected_auth_profile().await; let mut switched_to = None; - if matches!(args.action, ManageAuthProfilesAction::Switch) { - let profile = normalize_requested_profile(args.profile)?; - session - .update_settings(SessionSettingsUpdate { - auth_profile: Some(profile.clone()), - ..Default::default() - }) - .await - .map_err(|err| FunctionCallError::RespondToModel(err.to_string()))?; - let turn_context = session - .new_default_turn_with_sub_id(turn.sub_id.clone()) - .await; - session - .refresh_token_context_window_for_profile_switch(&turn_context) - .await; - current_profile = profile.clone(); - switched_to = profile; + match args.action { + ManageAuthProfilesAction::Switch => { + let profile = normalize_requested_profile(args.profile)?; + session + .update_settings(SessionSettingsUpdate { + auth_profile: Some(profile.clone()), + ..Default::default() + }) + .await + .map_err(|err| FunctionCallError::RespondToModel(err.to_string()))?; + let turn_context = session + .new_default_turn_with_sub_id(turn.sub_id.clone()) + .await; + session + .refresh_token_context_window_for_profile_switch(&turn_context) + .await; + current_profile = profile.clone(); + switched_to = profile; + } + ManageAuthProfilesAction::SetAutoSwitch => { + let enabled = args.auto_switch_enabled.ok_or_else(|| { + FunctionCallError::RespondToModel( + "auto_switch_enabled is required when action is set_auto_switch" + .to_string(), + ) + })?; + session + .update_settings(SessionSettingsUpdate { + auth_profile_auto_switch_enabled: Some(enabled), + ..Default::default() + }) + .await + .map_err(|err| FunctionCallError::RespondToModel(err.to_string()))?; + } + ManageAuthProfilesAction::List | ManageAuthProfilesAction::Current => {} } + let config = session.get_config().await; let profiles = codex_login::list_auth_profiles( - &turn.config.codex_home, - turn.config.cli_auth_credentials_store_mode, + &config.codex_home, + config.cli_auth_credentials_store_mode, ) .map_err(|err| FunctionCallError::RespondToModel(err.to_string()))?; let response = ManageAuthProfilesResponse { @@ -120,6 +155,7 @@ impl ToolExecutor for ManageAuthProfilesHandler { current_profile: current_profile.clone(), switched_to, profiles: summarize_profiles(profiles, current_profile.as_deref()), + auto_switch: summarize_auto_switch(&config.auth_profile_auto_switch), }; let response = serde_json::to_string_pretty(&response) .map_err(|err| FunctionCallError::Fatal(err.to_string()))?; @@ -176,6 +212,23 @@ fn summarize_profiles( summaries } +fn summarize_auto_switch(config: &AuthProfileAutoSwitchConfig) -> AuthProfileAutoSwitchState { + AuthProfileAutoSwitchState { + enabled: config.enabled, + on_5h_limit: config.on_5h_limit, + on_weekly_limit: config.on_weekly_limit, + strategy: auth_profile_auto_switch_strategy_name(config.strategy), + profiles: config.profiles.clone(), + } +} + +fn auth_profile_auto_switch_strategy_name(strategy: AuthProfileAutoSwitchStrategy) -> &'static str { + match strategy { + AuthProfileAutoSwitchStrategy::HighestAvailable => "highest_available", + AuthProfileAutoSwitchStrategy::Ordered => "ordered", + } +} + #[cfg(test)] mod tests { use super::*; @@ -238,4 +291,29 @@ mod tests { ]) ); } + + #[test] + fn summarize_auto_switch_serializes_session_state() { + let state = summarize_auto_switch(&AuthProfileAutoSwitchConfig { + enabled: true, + profiles: vec!["work".to_string(), "spare".to_string()], + on_5h_limit: true, + on_weekly_limit: false, + strategy: AuthProfileAutoSwitchStrategy::Ordered, + heartbeat_interval_secs: 60, + heartbeat_freshness_secs: 120, + }); + + let response = serde_json::to_value(state).expect("serialize auto switch state"); + assert_eq!( + response, + serde_json::json!({ + "enabled": true, + "on5hLimit": true, + "onWeeklyLimit": false, + "strategy": "ordered", + "profiles": ["work", "spare"] + }) + ); + } } diff --git a/codex-rs/core/src/tools/handlers/auth_profile_control_spec.rs b/codex-rs/core/src/tools/handlers/auth_profile_control_spec.rs index ecc9a1f749..15830d13b6 100644 --- a/codex-rs/core/src/tools/handlers/auth_profile_control_spec.rs +++ b/codex-rs/core/src/tools/handlers/auth_profile_control_spec.rs @@ -13,25 +13,49 @@ pub fn create_manage_auth_profiles_tool() -> ToolSpec { ( "action".to_string(), JsonSchema::string_enum( - vec![json!("list"), json!("current"), json!("switch")], + vec![ + json!("list"), + json!("current"), + json!("switch"), + json!("set_auto_switch"), + ], Some( - "Required. Use list to inspect available auth profiles, current to inspect the active auth profile, and switch to request a profile switch through the session settings path." + "Required. Use list to inspect available auth profiles, current to inspect the active auth profile, switch to request a profile switch through the session settings path, and set_auto_switch to toggle session-local auth-profile auto-switching." .to_string(), ), ), ), ( "profile".to_string(), - JsonSchema::string(Some( - "Profile name for switch. Omit or pass null to switch back to the default root auth profile." - .to_string(), - )), + JsonSchema::any_of( + vec![ + JsonSchema::string(/*description*/ None), + JsonSchema::null(/*description*/ None), + ], + Some( + "Profile name for switch. Omit or pass null to switch back to the default root auth profile." + .to_string(), + ), + ), + ), + ( + "auto_switch_enabled".to_string(), + JsonSchema::any_of( + vec![ + JsonSchema::boolean(/*description*/ None), + JsonSchema::null(/*description*/ None), + ], + Some( + "Required when action is set_auto_switch. Enables or disables auth-profile auto-switching for this session." + .to_string(), + ), + ), ), ]); ToolSpec::Function(ResponsesApiTool { name: MANAGE_AUTH_PROFILES_TOOL_NAME.to_string(), - description: "List Codewith auth profiles, inspect the current auth profile, or request a safe auth profile switch for the current session. Profile switches use the same session settings path as the /profile picker." + description: "List Codewith auth profiles, inspect the current auth profile, request a safe auth profile switch, or toggle auth-profile auto-switching for the current session. Profile switches use the same session settings path as the /profile picker." .to_string(), strict: false, defer_loading: None, @@ -60,7 +84,8 @@ mod tests { tool.parameters .properties .as_ref() - .is_some_and(|properties| properties.contains_key("profile")) + .is_some_and(|properties| properties.contains_key("profile") + && properties.contains_key("auto_switch_enabled")) ); let action = tool .parameters @@ -70,7 +95,12 @@ mod tests { .expect("action property"); assert_eq!( action.enum_values.as_ref(), - Some(&vec![json!("list"), json!("current"), json!("switch")]) + Some(&vec![ + json!("list"), + json!("current"), + json!("switch"), + json!("set_auto_switch") + ]) ); } } diff --git a/codex-rs/core/src/tools/handlers/auth_profile_usage_control.rs b/codex-rs/core/src/tools/handlers/auth_profile_usage_control.rs index 7fccd28fa4..ce9d7cc888 100644 --- a/codex-rs/core/src/tools/handlers/auth_profile_usage_control.rs +++ b/codex-rs/core/src/tools/handlers/auth_profile_usage_control.rs @@ -1,7 +1,10 @@ use crate::auth_profile_usage::AuthProfileUsageHealth; +use crate::auth_profile_usage::AuthProfileUsageRecommendationReason; use crate::auth_profile_usage::TokenUsageProfileResponse; +use crate::auth_profile_usage::recommend_auth_profile; use crate::auth_profile_usage::usage_capture_is_stale; use crate::auth_profile_usage::usage_health_for_snapshots; +use crate::config::AuthProfileAutoSwitchConfig; use crate::function_tool::FunctionCallError; use crate::session::session::Session; use crate::session::turn_context::TurnContext; @@ -25,6 +28,7 @@ use codex_tools::ToolSpec; use serde::Deserialize; use serde::Serialize; use std::collections::BTreeMap; +use std::collections::HashSet; use std::sync::LazyLock; use std::time::Duration; use tokio::sync::Mutex; @@ -42,6 +46,7 @@ pub struct GetUsageHandler; enum GetUsageScope { Session, Account, + AllAccounts, Both, } @@ -63,6 +68,10 @@ struct GetUsageResponse { session: Option, #[serde(skip_serializing_if = "Option::is_none")] account: Option, + #[serde(skip_serializing_if = "Option::is_none")] + accounts: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + recommendation: Option, } #[derive(Debug, Serialize)] @@ -76,6 +85,7 @@ struct SessionUsageResponse { #[serde(rename_all = "camelCase")] struct AccountUsageResponse { target: AccountUsageTarget, + current: bool, include_token_profile: bool, spend_status: UsageSpendSummary, #[serde(skip_serializing_if = "Option::is_none")] @@ -88,6 +98,15 @@ struct AccountUsageResponse { error: Option, } +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountUsageRecommendationResponse { + profile_name: Option, + display_name: String, + current: bool, + reason: AuthProfileUsageRecommendationReason, +} + #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] struct AccountUsageTarget { @@ -215,11 +234,11 @@ impl ToolExecutor for GetUsageHandler { }; let args: GetUsageArgs = parse_arguments(&arguments)?; - let target_profile = normalize_requested_profile( - args.auth_profile.clone(), - session.selected_auth_profile().await, - )?; - let response = get_usage_response(&session, &turn, args, target_profile).await?; + let current_profile = session.selected_auth_profile().await; + let target_profile = + normalize_requested_profile(args.auth_profile.clone(), current_profile.clone())?; + let response = + get_usage_response(&session, &turn, args, current_profile, target_profile).await?; let response = serde_json::to_string_pretty(&response) .map_err(|err| FunctionCallError::Fatal(err.to_string()))?; Ok(boxed_tool_output(FunctionToolOutput::from_text( @@ -235,10 +254,12 @@ async fn get_usage_response( session: &Session, turn: &TurnContext, args: GetUsageArgs, + current_profile: Option, target_profile: Option, ) -> Result { let include_session = matches!(args.scope, GetUsageScope::Session | GetUsageScope::Both); let include_account = matches!(args.scope, GetUsageScope::Account | GetUsageScope::Both); + let include_accounts = matches!(args.scope, GetUsageScope::AllAccounts); let session_usage = if include_session { Some(SessionUsageResponse { token_usage: session.token_usage_info().await, @@ -248,14 +269,42 @@ async fn get_usage_response( None }; let account_usage = if include_account { - Some(fetch_account_usage(session, turn, target_profile, args.include_token_profile).await?) + Some( + fetch_account_usage( + session, + turn, + target_profile, + current_profile.as_deref(), + args.include_token_profile, + ) + .await?, + ) } else { None }; + let (accounts_usage, recommendation) = if include_accounts { + let accounts = fetch_all_account_usages( + session, + turn, + current_profile.as_deref(), + args.include_token_profile, + ) + .await?; + let recommendation = Some(account_usage_recommendation( + &accounts, + current_profile.as_deref(), + &turn.config.auth_profile_auto_switch, + )); + (Some(accounts), recommendation) + } else { + (None, None) + }; Ok(GetUsageResponse { scope: args.scope, session: session_usage, account: account_usage, + accounts: accounts_usage, + recommendation, }) } @@ -279,9 +328,28 @@ async fn fetch_account_usage( session: &Session, turn: &TurnContext, target_profile: Option, + current_profile: Option<&str>, include_token_profile: bool, ) -> Result { validate_target_profile(turn, target_profile.as_deref())?; + fetch_account_usage_unvalidated( + session, + turn, + target_profile, + current_profile, + include_token_profile, + ) + .await +} + +async fn fetch_account_usage_unvalidated( + session: &Session, + turn: &TurnContext, + target_profile: Option, + current_profile: Option<&str>, + include_token_profile: bool, +) -> Result { + let current = target_profile.as_deref() == current_profile; let scoped_auth_manager = session .services .auth_manager @@ -290,6 +358,7 @@ async fn fetch_account_usage( let Some(auth) = scoped_auth_manager.auth().await else { return Ok(account_unavailable_response( target_profile, + current, include_token_profile, AuthProfileUsageStatusReason::NoAuth, )); @@ -298,6 +367,7 @@ async fn fetch_account_usage( if !auth.uses_codex_backend() { return Ok(AccountUsageResponse { target, + current, include_token_profile, spend_status: UsageSpendSummary::account_without_backend_credits(), rate_limits: None, @@ -313,6 +383,7 @@ async fn fetch_account_usage( Err(_) => { return Ok(AccountUsageResponse { target, + current, include_token_profile, spend_status: UsageSpendSummary::account_without_backend_credits(), rate_limits: None, @@ -372,6 +443,7 @@ async fn fetch_account_usage( Ok(AccountUsageResponse { target, + current, include_token_profile, spend_status, rate_limits, @@ -407,6 +479,7 @@ fn validate_target_profile( fn account_unavailable_response( target_profile: Option, + current: bool, include_token_profile: bool, reason: AuthProfileUsageStatusReason, ) -> AccountUsageResponse { @@ -417,6 +490,7 @@ fn account_unavailable_response( plan: None, redacted_account_id: None, }, + current, include_token_profile, spend_status: UsageSpendSummary::account_without_backend_credits(), rate_limits: None, @@ -426,6 +500,97 @@ fn account_unavailable_response( } } +async fn fetch_all_account_usages( + session: &Session, + turn: &TurnContext, + current_profile: Option<&str>, + include_token_profile: bool, +) -> Result, FunctionCallError> { + let mut targets = vec![None]; + let mut seen = HashSet::from([None]); + for profile in codex_login::list_auth_profiles( + &turn.config.codex_home, + turn.config.cli_auth_credentials_store_mode, + ) + .map_err(|err| FunctionCallError::RespondToModel(err.to_string()))? + { + let target = Some(profile.name); + if seen.insert(target.clone()) { + targets.push(target); + } + } + + let usages = futures::future::join_all(targets.into_iter().map(|target_profile| { + fetch_account_usage_unvalidated( + session, + turn, + target_profile, + current_profile, + include_token_profile, + ) + })) + .await; + usages.into_iter().collect() +} + +fn account_usage_recommendation( + accounts: &[AccountUsageResponse], + current_profile: Option<&str>, + config: &AuthProfileAutoSwitchConfig, +) -> AccountUsageRecommendationResponse { + let ordered_profiles = accounts + .iter() + .map(|account| account.target.profile_name.clone()) + .collect::>(); + let health_by_profile = accounts + .iter() + .filter_map(|account| { + let limits = account.rate_limits.as_ref()?; + Some(( + account.target.profile_name.clone(), + auth_profile_usage_health_from_summary(&limits.health), + )) + }) + .collect::>(); + let recommendation = recommend_auth_profile( + current_profile, + config.strategy, + &ordered_profiles, + &health_by_profile, + ); + let profile_name = recommendation.profile; + AccountUsageRecommendationResponse { + display_name: display_name_for_profile(profile_name.as_deref()), + current: profile_name.as_deref() == current_profile, + profile_name, + reason: recommendation.reason, + } +} + +fn auth_profile_usage_health_from_summary( + summary: &AuthProfileUsageSummary, +) -> AuthProfileUsageHealth { + if summary.stale { + return AuthProfileUsageHealth::Unknown; + } + match summary.status { + AuthProfileUsageStatus::Healthy => AuthProfileUsageHealth::Healthy { + remaining_percent: summary.remaining_percent.unwrap_or(0.0), + resets_at: summary.resets_at, + }, + AuthProfileUsageStatus::Exhausted => AuthProfileUsageHealth::Exhausted { + retry_at: summary.resets_at, + }, + AuthProfileUsageStatus::Unknown => AuthProfileUsageHealth::Unknown, + #[cfg(test)] + AuthProfileUsageStatus::Unavailable => AuthProfileUsageHealth::Unknown, + } +} + +fn display_name_for_profile(profile_name: Option<&str>) -> String { + profile_name.unwrap_or("Default").to_string() +} + async fn fetch_rate_limit_snapshots( turn: &TurnContext, target_profile: Option, @@ -756,6 +921,76 @@ mod tests { ); } + fn account_usage( + profile_name: Option<&str>, + current: bool, + health: AuthProfileUsageSummary, + ) -> AccountUsageResponse { + AccountUsageResponse { + target: AccountUsageTarget { + profile_name: profile_name.map(str::to_string), + auth_mode: None, + plan: None, + redacted_account_id: None, + }, + current, + include_token_profile: false, + spend_status: UsageSpendSummary::account_without_backend_credits(), + rate_limits: Some(AccountRateLimitUsage { + captured_at: 123, + stale_after_secs: 120, + health, + snapshots: Vec::new(), + }), + token_profile: None, + token_profile_error: None, + error: None, + } + } + + #[test] + fn account_usage_recommendation_serializes_shared_reason_codes() { + let accounts = vec![ + account_usage( + Some("work"), + true, + AuthProfileUsageSummary { + status: AuthProfileUsageStatus::Exhausted, + remaining_percent: Some(0.0), + resets_at: Some(100), + captured_at: Some(123), + stale: false, + reason: None, + }, + ), + account_usage( + Some("spare"), + false, + AuthProfileUsageSummary { + status: AuthProfileUsageStatus::Healthy, + remaining_percent: Some(80.0), + resets_at: Some(200), + captured_at: Some(123), + stale: false, + reason: None, + }, + ), + ]; + + let recommendation = account_usage_recommendation(&accounts, Some("work"), &config()); + let response = serde_json::to_value(recommendation).expect("serialize recommendation"); + + assert_eq!( + response, + serde_json::json!({ + "profileName": "spare", + "displayName": "spare", + "current": false, + "reason": "selected_highest_remaining" + }) + ); + } + #[tokio::test] async fn cached_rate_limit_snapshots_reuses_fresh_entries_and_expires_stale_entries() { let key = AuthProfileUsageCacheKey { diff --git a/codex-rs/core/src/tools/handlers/auth_profile_usage_control_spec.rs b/codex-rs/core/src/tools/handlers/auth_profile_usage_control_spec.rs index 60bf21f4ab..d43245c3a7 100644 --- a/codex-rs/core/src/tools/handlers/auth_profile_usage_control_spec.rs +++ b/codex-rs/core/src/tools/handlers/auth_profile_usage_control_spec.rs @@ -13,19 +13,30 @@ pub fn create_get_usage_tool() -> ToolSpec { ( "scope".to_string(), JsonSchema::string_enum( - vec![json!("session"), json!("account"), json!("both")], + vec![ + json!("session"), + json!("account"), + json!("all_accounts"), + json!("both"), + ], Some( - "Required. Use session for current conversation token usage, account for the selected auth profile's Codex account usage, and both for both surfaces." + "Required. Use session for current conversation token usage, account for the selected auth profile's Codex account usage, all_accounts for every saved auth profile plus the default root account, and both for session plus selected account." .to_string(), ), ), ), ( "auth_profile".to_string(), - JsonSchema::string(Some( - "Optional auth profile name for account usage. Omit to inspect the current session profile; pass an empty string to inspect the default root auth without switching." - .to_string(), - )), + JsonSchema::any_of( + vec![ + JsonSchema::string(/*description*/ None), + JsonSchema::null(/*description*/ None), + ], + Some( + "Optional auth profile name for account usage. Omit to inspect the current session profile; pass an empty string to inspect the default root auth without switching. Ignored when scope is all_accounts." + .to_string(), + ), + ), ), ( "include_token_profile".to_string(), @@ -38,7 +49,7 @@ pub fn create_get_usage_tool() -> ToolSpec { ToolSpec::Function(ResponsesApiTool { name: GET_USAGE_TOOL_NAME.to_string(), - description: "Read current session usage and scoped Codewith auth-profile account usage. This tool is read-only, does not enumerate every saved profile, does not switch profiles, and never returns auth tokens or raw auth storage details." + description: "Read current session usage and scoped Codewith auth-profile account usage. This tool is read-only, does not switch profiles, and never returns auth tokens or raw auth storage details." .to_string(), strict: false, defer_loading: None, @@ -70,7 +81,12 @@ mod tests { properties .get("scope") .and_then(|scope| scope.enum_values.as_ref()), - Some(&vec![json!("session"), json!("account"), json!("both")]) + Some(&vec![ + json!("session"), + json!("account"), + json!("all_accounts"), + json!("both") + ]) ); } }