mirror of
https://github.com/openai/codex.git
synced 2026-10-20 17:38:45 +00:00
dhruvgupta-oai
5037919777
## Why When ordinary included usage is exhausted, eligible users need a way to keep working with the backend-provided Luna Reserve quota and return to their previous model once ordinary usage recovers. ## What changed - Automatically switch eligible TUI tasks to Luna Reserve after an authoritative account usage read, preserving the prior model and reasoning effort per task. - Restrict model selection while on Reserve to its supported reasoning levels, keep queued turns on the accepted model, and restore the saved model after a fresh read confirms recovery. - Surface Reserve usage in the composer and `/status`, and expose `normalModelSlug` in app-server rate-limit snapshots so clients can use the associated model's display metadata without changing the request model. - Poll account limits more frequently near exhaustion and retain compatibility with app servers that only accept the legacy null request parameters. ## Testing - Added coverage for Reserve entry, queued-turn handling, task reconstruction, recovery authorization, account changes, model selection, usage rendering, and banner dismissal. GitOrigin-RevId: 46a435e9531051f8cec3f6a46c4b9f7cfa479747
276 lines
8.5 KiB
Rust
276 lines
8.5 KiB
Rust
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]
|
|
// Verifies connector merging deduplicates repeated IDs.
|
|
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]
|
|
// Verifies clearing connector selection removes all saved IDs.
|
|
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(/*tokens*/ 100);
|
|
state.replace_history(Vec::new(), /*reference_context_item*/ 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)
|
|
);
|
|
}
|