| use super::*; |
| use crate::session::tests::make_session_configuration_for_tests; |
| use crate::state::AutoCompactWindowSnapshot; |
| use codex_protocol::SessionId; |
| use codex_protocol::ThreadId; |
| use codex_protocol::protocol::CreditsSnapshot; |
| use codex_protocol::protocol::RateLimitWindow; |
| use codex_protocol::protocol::SpendControlLimitSnapshot; |
| use codex_protocol::protocol::TokenUsage; |
| use codex_protocol::protocol::TokenUsageRecord; |
| use pretty_assertions::assert_eq; |
|
|
| #[tokio::test] |
| async fn record_token_usage_continues_restored_totals() { |
| let thread_id = ThreadId::new(); |
| let session_id = SessionId::from(ThreadId::new()); |
| let usage = |total_tokens| TokenUsage { |
| total_tokens, |
| ..TokenUsage::default() |
| }; |
| let mut restored = SessionState::new(make_session_configuration_for_tests().await); |
| restored.latest_token_usage_record = Some(TokenUsageRecord { |
| thread_id, |
| turn_id: "turn-b".to_string(), |
| session_id, |
| root_turn_id: "root-turn".to_string(), |
| response_id: "response-c".to_string(), |
| usage: usage(30), |
| turn_token_usage: usage(30), |
| thread_token_usage: usage(230), |
| }); |
| let after_resume = restored.record_token_usage( |
| thread_id, |
| "turn-b", |
| session_id, |
| "root-turn".to_string(), |
| "response-d".to_string(), |
| &usage(20), |
| ); |
| assert_eq!( |
| after_resume, |
| TokenUsageRecord { |
| thread_id, |
| turn_id: "turn-b".to_string(), |
| session_id, |
| root_turn_id: "root-turn".to_string(), |
| response_id: "response-d".to_string(), |
| usage: usage(20), |
| turn_token_usage: usage(50), |
| thread_token_usage: usage(250), |
| } |
| ); |
| } |
|
|
| #[tokio::test] |
| |
| async fn merge_connector_selection_deduplicates_entries() { |
| let session_configuration = make_session_configuration_for_tests().await; |
| let mut state = SessionState::new(session_configuration); |
| let merged = state.merge_connector_selection([ |
| "calendar".to_string(), |
| "calendar".to_string(), |
| "drive".to_string(), |
| ]); |
|
|
| assert_eq!( |
| merged, |
| HashSet::from(["calendar".to_string(), "drive".to_string()]) |
| ); |
| } |
|
|
| #[tokio::test] |
| |
| async fn clear_connector_selection_removes_entries() { |
| let session_configuration = make_session_configuration_for_tests().await; |
| let mut state = SessionState::new(session_configuration); |
| state.merge_connector_selection(["calendar".to_string()]); |
|
|
| state.clear_connector_selection(); |
|
|
| assert_eq!(state.get_connector_selection(), HashSet::new()); |
| } |
|
|
| #[tokio::test] |
| async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { |
| let session_configuration = make_session_configuration_for_tests().await; |
| let mut state = SessionState::new(session_configuration); |
|
|
| state.set_rate_limits(RateLimitSnapshot { |
| limit_id: None, |
| limit_name: None, |
| normal_model_slug: None, |
| primary: Some(RateLimitWindow { |
| used_percent: 12.0, |
| window_minutes: Some(60), |
| resets_at: Some(100), |
| }), |
| secondary: None, |
| credits: None, |
| individual_limit: None, |
| spend_control_reached: None, |
| plan_type: None, |
| rate_limit_reached_type: None, |
| }); |
|
|
| assert_eq!( |
| state |
| .latest_rate_limits |
| .as_ref() |
| .and_then(|v| v.limit_id.clone()), |
| Some("codex".to_string()) |
| ); |
| } |
|
|
| #[tokio::test] |
| async fn replace_history_clears_auto_compact_window_prefill() { |
| let session_configuration = make_session_configuration_for_tests().await; |
| let mut state = SessionState::new(session_configuration); |
|
|
| state.set_auto_compact_window_estimated_prefill( 100); |
| state.replace_history(Vec::new(), None); |
|
|
| assert_eq!( |
| state.auto_compact_window_snapshot(), |
| AutoCompactWindowSnapshot { |
| prefill_input_tokens: None, |
| } |
| ); |
| } |
|
|
| #[tokio::test] |
| async fn set_rate_limits_defaults_to_codex_when_limit_id_missing_after_other_bucket() { |
| let session_configuration = make_session_configuration_for_tests().await; |
| let mut state = SessionState::new(session_configuration); |
|
|
| state.set_rate_limits(RateLimitSnapshot { |
| limit_id: Some("codex_other".to_string()), |
| limit_name: Some("codex_other".to_string()), |
| normal_model_slug: None, |
| primary: Some(RateLimitWindow { |
| used_percent: 20.0, |
| window_minutes: Some(60), |
| resets_at: Some(200), |
| }), |
| secondary: None, |
| credits: None, |
| individual_limit: None, |
| spend_control_reached: None, |
| plan_type: None, |
| rate_limit_reached_type: None, |
| }); |
| state.set_rate_limits(RateLimitSnapshot { |
| limit_id: None, |
| limit_name: None, |
| normal_model_slug: None, |
| primary: Some(RateLimitWindow { |
| used_percent: 30.0, |
| window_minutes: Some(60), |
| resets_at: Some(300), |
| }), |
| secondary: None, |
| credits: None, |
| individual_limit: None, |
| spend_control_reached: None, |
| plan_type: None, |
| rate_limit_reached_type: None, |
| }); |
|
|
| assert_eq!( |
| state |
| .latest_rate_limits |
| .as_ref() |
| .and_then(|v| v.limit_id.clone()), |
| Some("codex".to_string()) |
| ); |
| } |
|
|
| #[tokio::test] |
| async fn set_rate_limits_carries_account_metadata_from_codex_to_codex_other() { |
| let session_configuration = make_session_configuration_for_tests().await; |
| let mut state = SessionState::new(session_configuration); |
|
|
| state.set_rate_limits(RateLimitSnapshot { |
| limit_id: Some("codex".to_string()), |
| limit_name: Some("codex".to_string()), |
| normal_model_slug: None, |
| primary: Some(RateLimitWindow { |
| used_percent: 10.0, |
| window_minutes: Some(60), |
| resets_at: Some(100), |
| }), |
| secondary: None, |
| credits: Some(CreditsSnapshot { |
| has_credits: true, |
| unlimited: false, |
| balance: Some("50".to_string()), |
| }), |
| individual_limit: Some(SpendControlLimitSnapshot { |
| limit: "25000".to_string(), |
| used: "8000".to_string(), |
| remaining_percent: 68, |
| resets_at: 300, |
| }), |
| spend_control_reached: Some(true), |
| plan_type: Some(codex_protocol::account::PlanType::Plus), |
| rate_limit_reached_type: None, |
| }); |
|
|
| state.set_rate_limits(RateLimitSnapshot { |
| limit_id: Some("codex_other".to_string()), |
| limit_name: None, |
| normal_model_slug: None, |
| primary: Some(RateLimitWindow { |
| used_percent: 30.0, |
| window_minutes: Some(120), |
| resets_at: Some(200), |
| }), |
| secondary: None, |
| credits: None, |
| individual_limit: None, |
| spend_control_reached: None, |
| plan_type: None, |
| rate_limit_reached_type: None, |
| }); |
|
|
| assert_eq!( |
| state.latest_rate_limits, |
| Some(RateLimitSnapshot { |
| limit_id: Some("codex_other".to_string()), |
| limit_name: None, |
| normal_model_slug: None, |
| primary: Some(RateLimitWindow { |
| used_percent: 30.0, |
| window_minutes: Some(120), |
| resets_at: Some(200), |
| }), |
| secondary: None, |
| credits: Some(CreditsSnapshot { |
| has_credits: true, |
| unlimited: false, |
| balance: Some("50".to_string()), |
| }), |
| individual_limit: Some(SpendControlLimitSnapshot { |
| limit: "25000".to_string(), |
| used: "8000".to_string(), |
| remaining_percent: 68, |
| resets_at: 300, |
| }), |
| spend_control_reached: Some(true), |
| plan_type: Some(codex_protocol::account::PlanType::Plus), |
| rate_limit_reached_type: None, |
| }) |
| ); |
|
|
| state.set_rate_limits(RateLimitSnapshot { |
| limit_id: Some("codex_other".to_string()), |
| limit_name: None, |
| normal_model_slug: None, |
| primary: None, |
| secondary: None, |
| credits: None, |
| individual_limit: None, |
| spend_control_reached: Some(false), |
| plan_type: None, |
| rate_limit_reached_type: None, |
| }); |
|
|
| assert_eq!( |
| state |
| .latest_rate_limits |
| .as_ref() |
| .and_then(|snapshot| snapshot.spend_control_reached), |
| Some(false) |
| ); |
| } |
|
|
