mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
fix: taking plan type from usage endpoint instead of thru auth token (#7610)
pull plan type from the usage endpoint, persist it in session state / tui state, and propagate through rate limit snapshots
This commit is contained in:
committed by
GitHub
Unverified
parent
b1c918d8f7
commit
b8eab7ce90
@@ -1524,6 +1524,7 @@ pub struct RateLimitSnapshot {
|
||||
pub primary: Option<RateLimitWindow>,
|
||||
pub secondary: Option<RateLimitWindow>,
|
||||
pub credits: Option<CreditsSnapshot>,
|
||||
pub plan_type: Option<PlanType>,
|
||||
}
|
||||
|
||||
impl From<CoreRateLimitSnapshot> for RateLimitSnapshot {
|
||||
@@ -1532,6 +1533,7 @@ impl From<CoreRateLimitSnapshot> for RateLimitSnapshot {
|
||||
primary: value.primary.map(RateLimitWindow::from),
|
||||
secondary: value.secondary.map(RateLimitWindow::from),
|
||||
credits: value.credits.map(CreditsSnapshot::from),
|
||||
plan_type: value.plan_type,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1499,6 +1499,7 @@ mod tests {
|
||||
unlimited: false,
|
||||
balance: Some("5".to_string()),
|
||||
}),
|
||||
plan_type: None,
|
||||
};
|
||||
|
||||
handle_token_count_event(
|
||||
|
||||
@@ -16,6 +16,9 @@ use tracing::warn;
|
||||
|
||||
use crate::error_code::INTERNAL_ERROR_CODE;
|
||||
|
||||
#[cfg(test)]
|
||||
use codex_protocol::account::PlanType;
|
||||
|
||||
/// Sends messages to the client and manages request callbacks.
|
||||
pub(crate) struct OutgoingMessageSender {
|
||||
next_request_id: AtomicI64,
|
||||
@@ -230,6 +233,7 @@ mod tests {
|
||||
}),
|
||||
secondary: None,
|
||||
credits: None,
|
||||
plan_type: Some(PlanType::Plus),
|
||||
},
|
||||
});
|
||||
|
||||
@@ -245,7 +249,8 @@ mod tests {
|
||||
"resetsAt": 123
|
||||
},
|
||||
"secondary": null,
|
||||
"credits": null
|
||||
"credits": null,
|
||||
"planType": "plus"
|
||||
}
|
||||
},
|
||||
}),
|
||||
|
||||
@@ -11,6 +11,7 @@ use codex_app_server_protocol::RateLimitSnapshot;
|
||||
use codex_app_server_protocol::RateLimitWindow;
|
||||
use codex_app_server_protocol::RequestId;
|
||||
use codex_core::auth::AuthCredentialsStoreMode;
|
||||
use codex_protocol::account::PlanType as AccountPlanType;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use std::path::Path;
|
||||
@@ -153,6 +154,7 @@ async fn get_account_rate_limits_returns_snapshot() -> Result<()> {
|
||||
resets_at: Some(secondary_reset_timestamp),
|
||||
}),
|
||||
credits: None,
|
||||
plan_type: Some(AccountPlanType::Pro),
|
||||
},
|
||||
};
|
||||
assert_eq!(received, expected);
|
||||
|
||||
@@ -7,6 +7,7 @@ use crate::types::TurnAttemptsSiblingTurnsResponse;
|
||||
use anyhow::Result;
|
||||
use codex_core::auth::CodexAuth;
|
||||
use codex_core::default_client::get_codex_user_agent;
|
||||
use codex_protocol::account::PlanType as AccountPlanType;
|
||||
use codex_protocol::protocol::CreditsSnapshot;
|
||||
use codex_protocol::protocol::RateLimitSnapshot;
|
||||
use codex_protocol::protocol::RateLimitWindow;
|
||||
@@ -291,6 +292,7 @@ impl Client {
|
||||
primary,
|
||||
secondary,
|
||||
credits: Self::map_credits(payload.credits),
|
||||
plan_type: Some(Self::map_plan_type(payload.plan_type)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -325,6 +327,23 @@ impl Client {
|
||||
})
|
||||
}
|
||||
|
||||
fn map_plan_type(plan_type: crate::types::PlanType) -> AccountPlanType {
|
||||
match plan_type {
|
||||
crate::types::PlanType::Free => AccountPlanType::Free,
|
||||
crate::types::PlanType::Plus => AccountPlanType::Plus,
|
||||
crate::types::PlanType::Pro => AccountPlanType::Pro,
|
||||
crate::types::PlanType::Team => AccountPlanType::Team,
|
||||
crate::types::PlanType::Business => AccountPlanType::Business,
|
||||
crate::types::PlanType::Enterprise => AccountPlanType::Enterprise,
|
||||
crate::types::PlanType::Edu | crate::types::PlanType::Education => AccountPlanType::Edu,
|
||||
crate::types::PlanType::Guest
|
||||
| crate::types::PlanType::Go
|
||||
| crate::types::PlanType::FreeWorkspace
|
||||
| crate::types::PlanType::Quorum
|
||||
| crate::types::PlanType::K12 => AccountPlanType::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
fn window_minutes_from_seconds(seconds: i32) -> Option<i64> {
|
||||
if seconds <= 0 {
|
||||
return None;
|
||||
|
||||
@@ -37,6 +37,7 @@ pub fn parse_rate_limit(headers: &HeaderMap) -> Option<RateLimitSnapshot> {
|
||||
primary,
|
||||
secondary,
|
||||
credits,
|
||||
plan_type: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -227,23 +227,6 @@ impl CodexAuth {
|
||||
})
|
||||
}
|
||||
|
||||
/// Raw plan string from the ID token (including unknown/new plan types).
|
||||
pub fn raw_plan_type(&self) -> Option<String> {
|
||||
self.get_plan_type().map(|plan| match plan {
|
||||
InternalPlanType::Known(k) => format!("{k:?}"),
|
||||
InternalPlanType::Unknown(raw) => raw,
|
||||
})
|
||||
}
|
||||
|
||||
/// Raw internal plan value from the ID token.
|
||||
/// Exposes the underlying `token_data::PlanType` without mapping it to the
|
||||
/// public `AccountPlanType`. Use this when downstream code needs to inspect
|
||||
/// internal/unknown plan strings exactly as issued in the token.
|
||||
pub(crate) fn get_plan_type(&self) -> Option<InternalPlanType> {
|
||||
self.get_current_token_data()
|
||||
.and_then(|t| t.id_token.chatgpt_plan_type)
|
||||
}
|
||||
|
||||
fn get_current_auth_json(&self) -> Option<AuthDotJson> {
|
||||
#[expect(clippy::unwrap_used)]
|
||||
self.auth_dot_json.lock().unwrap().clone()
|
||||
@@ -1041,10 +1024,6 @@ mod tests {
|
||||
.expect("auth available");
|
||||
|
||||
pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Pro));
|
||||
pretty_assertions::assert_eq!(
|
||||
auth.get_plan_type(),
|
||||
Some(InternalPlanType::Known(InternalKnownPlan::Pro))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1065,10 +1044,6 @@ mod tests {
|
||||
.expect("auth available");
|
||||
|
||||
pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Unknown));
|
||||
pretty_assertions::assert_eq!(
|
||||
auth.get_plan_type(),
|
||||
Some(InternalPlanType::Unknown("mystery-tier".to_string()))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2598,6 +2598,7 @@ mod tests {
|
||||
unlimited: false,
|
||||
balance: Some("10.00".to_string()),
|
||||
}),
|
||||
plan_type: Some(codex_protocol::account::PlanType::Plus),
|
||||
};
|
||||
state.set_rate_limits(initial.clone());
|
||||
|
||||
@@ -2613,6 +2614,7 @@ mod tests {
|
||||
resets_at: Some(1_900),
|
||||
}),
|
||||
credits: None,
|
||||
plan_type: None,
|
||||
};
|
||||
state.set_rate_limits(update.clone());
|
||||
|
||||
@@ -2622,6 +2624,78 @@ mod tests {
|
||||
primary: update.primary.clone(),
|
||||
secondary: update.secondary,
|
||||
credits: initial.credits,
|
||||
plan_type: initial.plan_type,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_rate_limits_updates_plan_type_when_present() {
|
||||
let codex_home = tempfile::tempdir().expect("create temp dir");
|
||||
let config = Config::load_from_base_config_with_overrides(
|
||||
ConfigToml::default(),
|
||||
ConfigOverrides::default(),
|
||||
codex_home.path().to_path_buf(),
|
||||
)
|
||||
.expect("load default test config");
|
||||
let config = Arc::new(config);
|
||||
let session_configuration = SessionConfiguration {
|
||||
provider: config.model_provider.clone(),
|
||||
model: config.model.clone(),
|
||||
model_reasoning_effort: config.model_reasoning_effort,
|
||||
model_reasoning_summary: config.model_reasoning_summary,
|
||||
developer_instructions: config.developer_instructions.clone(),
|
||||
user_instructions: config.user_instructions.clone(),
|
||||
base_instructions: config.base_instructions.clone(),
|
||||
compact_prompt: config.compact_prompt.clone(),
|
||||
approval_policy: config.approval_policy,
|
||||
sandbox_policy: config.sandbox_policy.clone(),
|
||||
cwd: config.cwd.clone(),
|
||||
original_config_do_not_use: Arc::clone(&config),
|
||||
exec_policy: Arc::new(RwLock::new(ExecPolicy::empty())),
|
||||
session_source: SessionSource::Exec,
|
||||
};
|
||||
|
||||
let mut state = SessionState::new(session_configuration);
|
||||
let initial = RateLimitSnapshot {
|
||||
primary: Some(RateLimitWindow {
|
||||
used_percent: 15.0,
|
||||
window_minutes: Some(20),
|
||||
resets_at: Some(1_600),
|
||||
}),
|
||||
secondary: Some(RateLimitWindow {
|
||||
used_percent: 5.0,
|
||||
window_minutes: Some(45),
|
||||
resets_at: Some(1_650),
|
||||
}),
|
||||
credits: Some(CreditsSnapshot {
|
||||
has_credits: true,
|
||||
unlimited: false,
|
||||
balance: Some("15.00".to_string()),
|
||||
}),
|
||||
plan_type: Some(codex_protocol::account::PlanType::Plus),
|
||||
};
|
||||
state.set_rate_limits(initial.clone());
|
||||
|
||||
let update = RateLimitSnapshot {
|
||||
primary: Some(RateLimitWindow {
|
||||
used_percent: 35.0,
|
||||
window_minutes: Some(25),
|
||||
resets_at: Some(1_700),
|
||||
}),
|
||||
secondary: None,
|
||||
credits: None,
|
||||
plan_type: Some(codex_protocol::account::PlanType::Pro),
|
||||
};
|
||||
state.set_rate_limits(update.clone());
|
||||
|
||||
assert_eq!(
|
||||
state.latest_rate_limits,
|
||||
Some(RateLimitSnapshot {
|
||||
primary: update.primary,
|
||||
secondary: update.secondary,
|
||||
credits: initial.credits,
|
||||
plan_type: update.plan_type,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
@@ -560,6 +560,7 @@ mod tests {
|
||||
resets_at: Some(secondary_reset_at),
|
||||
}),
|
||||
credits: None,
|
||||
plan_type: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -62,7 +62,7 @@ impl SessionState {
|
||||
}
|
||||
|
||||
pub(crate) fn set_rate_limits(&mut self, snapshot: RateLimitSnapshot) {
|
||||
self.latest_rate_limits = Some(merge_rate_limit_credits(
|
||||
self.latest_rate_limits = Some(merge_rate_limit_fields(
|
||||
self.latest_rate_limits.as_ref(),
|
||||
snapshot,
|
||||
));
|
||||
@@ -83,13 +83,16 @@ impl SessionState {
|
||||
}
|
||||
}
|
||||
|
||||
// Sometimes new snapshots don't include credits
|
||||
fn merge_rate_limit_credits(
|
||||
// Sometimes new snapshots don't include credits or plan information.
|
||||
fn merge_rate_limit_fields(
|
||||
previous: Option<&RateLimitSnapshot>,
|
||||
mut snapshot: RateLimitSnapshot,
|
||||
) -> RateLimitSnapshot {
|
||||
if snapshot.credits.is_none() {
|
||||
snapshot.credits = previous.and_then(|prior| prior.credits.clone());
|
||||
}
|
||||
if snapshot.plan_type.is_none() {
|
||||
snapshot.plan_type = previous.and_then(|prior| prior.plan_type);
|
||||
}
|
||||
snapshot
|
||||
}
|
||||
|
||||
@@ -1195,7 +1195,8 @@ async fn token_count_includes_rate_limits_snapshot() {
|
||||
"window_minutes": 60,
|
||||
"resets_at": 1704074400
|
||||
},
|
||||
"credits": null
|
||||
"credits": null,
|
||||
"plan_type": null
|
||||
}
|
||||
})
|
||||
);
|
||||
@@ -1243,7 +1244,8 @@ async fn token_count_includes_rate_limits_snapshot() {
|
||||
"window_minutes": 60,
|
||||
"resets_at": 1704074400
|
||||
},
|
||||
"credits": null
|
||||
"credits": null,
|
||||
"plan_type": null
|
||||
}
|
||||
})
|
||||
);
|
||||
@@ -1314,7 +1316,8 @@ async fn usage_limit_error_emits_rate_limit_event() -> anyhow::Result<()> {
|
||||
"window_minutes": 60,
|
||||
"resets_at": null
|
||||
},
|
||||
"credits": null
|
||||
"credits": null,
|
||||
"plan_type": null
|
||||
});
|
||||
|
||||
let submission_id = codex
|
||||
|
||||
@@ -847,6 +847,7 @@ pub struct RateLimitSnapshot {
|
||||
pub primary: Option<RateLimitWindow>,
|
||||
pub secondary: Option<RateLimitWindow>,
|
||||
pub credits: Option<CreditsSnapshot>,
|
||||
pub plan_type: Option<crate::account::PlanType>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, JsonSchema, TS)]
|
||||
|
||||
@@ -58,6 +58,7 @@ use codex_core::protocol::WebSearchBeginEvent;
|
||||
use codex_core::protocol::WebSearchEndEvent;
|
||||
use codex_core::skills::model::SkillMetadata;
|
||||
use codex_protocol::ConversationId;
|
||||
use codex_protocol::account::PlanType;
|
||||
use codex_protocol::approvals::ElicitationRequestEvent;
|
||||
use codex_protocol::parse_command::ParsedCommand;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
@@ -282,6 +283,7 @@ pub(crate) struct ChatWidget {
|
||||
initial_user_message: Option<UserMessage>,
|
||||
token_info: Option<TokenUsageInfo>,
|
||||
rate_limit_snapshot: Option<RateLimitSnapshotDisplay>,
|
||||
plan_type: Option<PlanType>,
|
||||
rate_limit_warnings: RateLimitWarningState,
|
||||
rate_limit_switch_prompt: RateLimitSwitchPromptState,
|
||||
rate_limit_poller: Option<JoinHandle<()>>,
|
||||
@@ -580,6 +582,8 @@ impl ChatWidget {
|
||||
});
|
||||
}
|
||||
|
||||
self.plan_type = snapshot.plan_type.or(self.plan_type);
|
||||
|
||||
let warnings = self.rate_limit_warnings.take_warnings(
|
||||
snapshot
|
||||
.secondary
|
||||
@@ -1275,6 +1279,7 @@ impl ChatWidget {
|
||||
),
|
||||
token_info: None,
|
||||
rate_limit_snapshot: None,
|
||||
plan_type: None,
|
||||
rate_limit_warnings: RateLimitWarningState::default(),
|
||||
rate_limit_switch_prompt: RateLimitSwitchPromptState::default(),
|
||||
rate_limit_poller: None,
|
||||
@@ -1357,6 +1362,7 @@ impl ChatWidget {
|
||||
),
|
||||
token_info: None,
|
||||
rate_limit_snapshot: None,
|
||||
plan_type: None,
|
||||
rate_limit_warnings: RateLimitWarningState::default(),
|
||||
rate_limit_switch_prompt: RateLimitSwitchPromptState::default(),
|
||||
rate_limit_poller: None,
|
||||
@@ -2001,6 +2007,7 @@ impl ChatWidget {
|
||||
context_usage,
|
||||
&self.conversation_id,
|
||||
self.rate_limit_snapshot.as_ref(),
|
||||
self.plan_type,
|
||||
Local::now(),
|
||||
));
|
||||
}
|
||||
|
||||
@@ -47,6 +47,7 @@ use codex_core::protocol::UndoStartedEvent;
|
||||
use codex_core::protocol::ViewImageToolCallEvent;
|
||||
use codex_core::protocol::WarningEvent;
|
||||
use codex_protocol::ConversationId;
|
||||
use codex_protocol::account::PlanType;
|
||||
use codex_protocol::openai_models::ModelPreset;
|
||||
use codex_protocol::openai_models::ReasoningEffortPreset;
|
||||
use codex_protocol::parse_command::ParsedCommand;
|
||||
@@ -90,6 +91,7 @@ fn snapshot(percent: f64) -> RateLimitSnapshot {
|
||||
}),
|
||||
secondary: None,
|
||||
credits: None,
|
||||
plan_type: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -398,6 +400,7 @@ fn make_chatwidget_manual() -> (
|
||||
initial_user_message: None,
|
||||
token_info: None,
|
||||
rate_limit_snapshot: None,
|
||||
plan_type: None,
|
||||
rate_limit_warnings: RateLimitWarningState::default(),
|
||||
rate_limit_switch_prompt: RateLimitSwitchPromptState::default(),
|
||||
rate_limit_poller: None,
|
||||
@@ -546,6 +549,7 @@ fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() {
|
||||
unlimited: false,
|
||||
balance: Some("17.5".to_string()),
|
||||
}),
|
||||
plan_type: None,
|
||||
}));
|
||||
let initial_balance = chat
|
||||
.rate_limit_snapshot
|
||||
@@ -562,6 +566,7 @@ fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() {
|
||||
}),
|
||||
secondary: None,
|
||||
credits: None,
|
||||
plan_type: None,
|
||||
}));
|
||||
|
||||
let display = chat
|
||||
@@ -581,6 +586,59 @@ fn rate_limit_snapshot_keeps_prior_credits_when_missing_from_headers() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rate_limit_snapshot_updates_and_retains_plan_type() {
|
||||
let (mut chat, _rx, _op_rx) = make_chatwidget_manual();
|
||||
|
||||
chat.on_rate_limit_snapshot(Some(RateLimitSnapshot {
|
||||
primary: Some(RateLimitWindow {
|
||||
used_percent: 10.0,
|
||||
window_minutes: Some(60),
|
||||
resets_at: None,
|
||||
}),
|
||||
secondary: Some(RateLimitWindow {
|
||||
used_percent: 5.0,
|
||||
window_minutes: Some(300),
|
||||
resets_at: None,
|
||||
}),
|
||||
credits: None,
|
||||
plan_type: Some(PlanType::Plus),
|
||||
}));
|
||||
assert_eq!(chat.plan_type, Some(PlanType::Plus));
|
||||
|
||||
chat.on_rate_limit_snapshot(Some(RateLimitSnapshot {
|
||||
primary: Some(RateLimitWindow {
|
||||
used_percent: 25.0,
|
||||
window_minutes: Some(30),
|
||||
resets_at: Some(123),
|
||||
}),
|
||||
secondary: Some(RateLimitWindow {
|
||||
used_percent: 15.0,
|
||||
window_minutes: Some(300),
|
||||
resets_at: Some(234),
|
||||
}),
|
||||
credits: None,
|
||||
plan_type: Some(PlanType::Pro),
|
||||
}));
|
||||
assert_eq!(chat.plan_type, Some(PlanType::Pro));
|
||||
|
||||
chat.on_rate_limit_snapshot(Some(RateLimitSnapshot {
|
||||
primary: Some(RateLimitWindow {
|
||||
used_percent: 30.0,
|
||||
window_minutes: Some(60),
|
||||
resets_at: Some(456),
|
||||
}),
|
||||
secondary: Some(RateLimitWindow {
|
||||
used_percent: 18.0,
|
||||
window_minutes: Some(300),
|
||||
resets_at: Some(567),
|
||||
}),
|
||||
credits: None,
|
||||
plan_type: None,
|
||||
}));
|
||||
assert_eq!(chat.plan_type, Some(PlanType::Pro));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rate_limit_switch_prompt_skips_when_on_lower_cost_model() {
|
||||
let (mut chat, _, _) = make_chatwidget_manual();
|
||||
|
||||
@@ -10,6 +10,7 @@ use codex_core::config::Config;
|
||||
use codex_core::protocol::SandboxPolicy;
|
||||
use codex_core::protocol::TokenUsage;
|
||||
use codex_protocol::ConversationId;
|
||||
use codex_protocol::account::PlanType;
|
||||
use ratatui::prelude::*;
|
||||
use ratatui::style::Stylize;
|
||||
use std::collections::BTreeSet;
|
||||
@@ -65,6 +66,7 @@ struct StatusHistoryCell {
|
||||
rate_limits: StatusRateLimitData,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn new_status_output(
|
||||
config: &Config,
|
||||
auth_manager: &AuthManager,
|
||||
@@ -72,6 +74,7 @@ pub(crate) fn new_status_output(
|
||||
context_usage: Option<&TokenUsage>,
|
||||
session_id: &Option<ConversationId>,
|
||||
rate_limits: Option<&RateLimitSnapshotDisplay>,
|
||||
plan_type: Option<PlanType>,
|
||||
now: DateTime<Local>,
|
||||
) -> CompositeHistoryCell {
|
||||
let command = PlainHistoryCell::new(vec!["/status".magenta().into()]);
|
||||
@@ -82,6 +85,7 @@ pub(crate) fn new_status_output(
|
||||
context_usage,
|
||||
session_id,
|
||||
rate_limits,
|
||||
plan_type,
|
||||
now,
|
||||
);
|
||||
|
||||
@@ -89,6 +93,7 @@ pub(crate) fn new_status_output(
|
||||
}
|
||||
|
||||
impl StatusHistoryCell {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn new(
|
||||
config: &Config,
|
||||
auth_manager: &AuthManager,
|
||||
@@ -96,6 +101,7 @@ impl StatusHistoryCell {
|
||||
context_usage: Option<&TokenUsage>,
|
||||
session_id: &Option<ConversationId>,
|
||||
rate_limits: Option<&RateLimitSnapshotDisplay>,
|
||||
plan_type: Option<PlanType>,
|
||||
now: DateTime<Local>,
|
||||
) -> Self {
|
||||
let config_entries = create_config_summary_entries(config);
|
||||
@@ -111,7 +117,7 @@ impl StatusHistoryCell {
|
||||
SandboxPolicy::WorkspaceWrite { .. } => "workspace-write".to_string(),
|
||||
};
|
||||
let agents_summary = compose_agents_summary(config);
|
||||
let account = compose_account_display(auth_manager);
|
||||
let account = compose_account_display(auth_manager, plan_type);
|
||||
let session_id = session_id.as_ref().map(std::string::ToString::to_string);
|
||||
let context_window = config.model_context_window.and_then(|window| {
|
||||
context_usage.map(|usage| StatusContextWindowData {
|
||||
|
||||
@@ -6,6 +6,7 @@ use codex_app_server_protocol::AuthMode;
|
||||
use codex_core::AuthManager;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::project_doc::discover_project_doc_paths;
|
||||
use codex_protocol::account::PlanType;
|
||||
use std::path::Path;
|
||||
use unicode_width::UnicodeWidthStr;
|
||||
|
||||
@@ -83,13 +84,18 @@ pub(crate) fn compose_agents_summary(config: &Config) -> String {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn compose_account_display(auth_manager: &AuthManager) -> Option<StatusAccountDisplay> {
|
||||
pub(crate) fn compose_account_display(
|
||||
auth_manager: &AuthManager,
|
||||
plan: Option<PlanType>,
|
||||
) -> Option<StatusAccountDisplay> {
|
||||
let auth = auth_manager.auth()?;
|
||||
|
||||
match auth.mode {
|
||||
AuthMode::ChatGPT => {
|
||||
let email = auth.get_account_email();
|
||||
let plan = auth.raw_plan_type().map(|plan| title_case(plan.as_str()));
|
||||
let plan = plan
|
||||
.map(|plan_type| title_case(format!("{plan_type:?}").as_str()))
|
||||
.or_else(|| Some("Unknown".to_string()));
|
||||
Some(StatusAccountDisplay::ChatGpt { email, plan })
|
||||
}
|
||||
AuthMode::ApiKey => Some(StatusAccountDisplay::ApiKey),
|
||||
|
||||
@@ -120,6 +120,7 @@ fn status_snapshot_includes_reasoning_details() {
|
||||
resets_at: Some(reset_at_from(&captured_at, 1_200)),
|
||||
}),
|
||||
credits: None,
|
||||
plan_type: None,
|
||||
};
|
||||
let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
|
||||
|
||||
@@ -130,6 +131,7 @@ fn status_snapshot_includes_reasoning_details() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
Some(&rate_display),
|
||||
None,
|
||||
captured_at,
|
||||
);
|
||||
let mut rendered_lines = render_lines(&composite.display_lines(80));
|
||||
@@ -171,6 +173,7 @@ fn status_snapshot_includes_monthly_limit() {
|
||||
}),
|
||||
secondary: None,
|
||||
credits: None,
|
||||
plan_type: None,
|
||||
};
|
||||
let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
|
||||
|
||||
@@ -181,6 +184,7 @@ fn status_snapshot_includes_monthly_limit() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
Some(&rate_display),
|
||||
None,
|
||||
captured_at,
|
||||
);
|
||||
let mut rendered_lines = render_lines(&composite.display_lines(80));
|
||||
@@ -211,6 +215,7 @@ fn status_snapshot_shows_unlimited_credits() {
|
||||
unlimited: true,
|
||||
balance: None,
|
||||
}),
|
||||
plan_type: None,
|
||||
};
|
||||
let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
|
||||
let composite = new_status_output(
|
||||
@@ -220,6 +225,7 @@ fn status_snapshot_shows_unlimited_credits() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
Some(&rate_display),
|
||||
None,
|
||||
captured_at,
|
||||
);
|
||||
let rendered = render_lines(&composite.display_lines(120));
|
||||
@@ -249,6 +255,7 @@ fn status_snapshot_shows_positive_credits() {
|
||||
unlimited: false,
|
||||
balance: Some("12.5".to_string()),
|
||||
}),
|
||||
plan_type: None,
|
||||
};
|
||||
let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
|
||||
let composite = new_status_output(
|
||||
@@ -258,6 +265,7 @@ fn status_snapshot_shows_positive_credits() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
Some(&rate_display),
|
||||
None,
|
||||
captured_at,
|
||||
);
|
||||
let rendered = render_lines(&composite.display_lines(120));
|
||||
@@ -287,6 +295,7 @@ fn status_snapshot_hides_zero_credits() {
|
||||
unlimited: false,
|
||||
balance: Some("0".to_string()),
|
||||
}),
|
||||
plan_type: None,
|
||||
};
|
||||
let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
|
||||
let composite = new_status_output(
|
||||
@@ -296,6 +305,7 @@ fn status_snapshot_hides_zero_credits() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
Some(&rate_display),
|
||||
None,
|
||||
captured_at,
|
||||
);
|
||||
let rendered = render_lines(&composite.display_lines(120));
|
||||
@@ -323,6 +333,7 @@ fn status_snapshot_hides_when_has_no_credits_flag() {
|
||||
unlimited: true,
|
||||
balance: None,
|
||||
}),
|
||||
plan_type: None,
|
||||
};
|
||||
let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
|
||||
let composite = new_status_output(
|
||||
@@ -332,6 +343,7 @@ fn status_snapshot_hides_when_has_no_credits_flag() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
Some(&rate_display),
|
||||
None,
|
||||
captured_at,
|
||||
);
|
||||
let rendered = render_lines(&composite.display_lines(120));
|
||||
@@ -369,6 +381,7 @@ fn status_card_token_usage_excludes_cached_tokens() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
None,
|
||||
None,
|
||||
now,
|
||||
);
|
||||
let rendered = render_lines(&composite.display_lines(120));
|
||||
@@ -410,6 +423,7 @@ fn status_snapshot_truncates_in_narrow_terminal() {
|
||||
}),
|
||||
secondary: None,
|
||||
credits: None,
|
||||
plan_type: None,
|
||||
};
|
||||
let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
|
||||
|
||||
@@ -420,6 +434,7 @@ fn status_snapshot_truncates_in_narrow_terminal() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
Some(&rate_display),
|
||||
None,
|
||||
captured_at,
|
||||
);
|
||||
let mut rendered_lines = render_lines(&composite.display_lines(70));
|
||||
@@ -461,6 +476,7 @@ fn status_snapshot_shows_missing_limits_message() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
None,
|
||||
None,
|
||||
now,
|
||||
);
|
||||
let mut rendered_lines = render_lines(&composite.display_lines(80));
|
||||
@@ -509,6 +525,7 @@ fn status_snapshot_includes_credits_and_limits() {
|
||||
unlimited: false,
|
||||
balance: Some("37.5".to_string()),
|
||||
}),
|
||||
plan_type: None,
|
||||
};
|
||||
let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
|
||||
|
||||
@@ -519,6 +536,7 @@ fn status_snapshot_includes_credits_and_limits() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
Some(&rate_display),
|
||||
None,
|
||||
captured_at,
|
||||
);
|
||||
let mut rendered_lines = render_lines(&composite.display_lines(80));
|
||||
@@ -551,6 +569,7 @@ fn status_snapshot_shows_empty_limits_message() {
|
||||
primary: None,
|
||||
secondary: None,
|
||||
credits: None,
|
||||
plan_type: None,
|
||||
};
|
||||
let captured_at = chrono::Local
|
||||
.with_ymd_and_hms(2024, 6, 7, 8, 9, 10)
|
||||
@@ -565,6 +584,7 @@ fn status_snapshot_shows_empty_limits_message() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
Some(&rate_display),
|
||||
None,
|
||||
captured_at,
|
||||
);
|
||||
let mut rendered_lines = render_lines(&composite.display_lines(80));
|
||||
@@ -609,6 +629,7 @@ fn status_snapshot_shows_stale_limits_message() {
|
||||
resets_at: Some(reset_at_from(&captured_at, 1_800)),
|
||||
}),
|
||||
credits: None,
|
||||
plan_type: None,
|
||||
};
|
||||
let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
|
||||
let now = captured_at + ChronoDuration::minutes(20);
|
||||
@@ -620,6 +641,7 @@ fn status_snapshot_shows_stale_limits_message() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
Some(&rate_display),
|
||||
None,
|
||||
now,
|
||||
);
|
||||
let mut rendered_lines = render_lines(&composite.display_lines(80));
|
||||
@@ -668,6 +690,7 @@ fn status_snapshot_cached_limits_hide_credits_without_flag() {
|
||||
unlimited: false,
|
||||
balance: Some("80".to_string()),
|
||||
}),
|
||||
plan_type: None,
|
||||
};
|
||||
let rate_display = rate_limit_snapshot_display(&snapshot, captured_at);
|
||||
let now = captured_at + ChronoDuration::minutes(20);
|
||||
@@ -679,6 +702,7 @@ fn status_snapshot_cached_limits_hide_credits_without_flag() {
|
||||
Some(&usage),
|
||||
&None,
|
||||
Some(&rate_display),
|
||||
None,
|
||||
now,
|
||||
);
|
||||
let mut rendered_lines = render_lines(&composite.display_lines(80));
|
||||
@@ -725,6 +749,7 @@ fn status_context_window_uses_last_usage() {
|
||||
Some(&last_usage),
|
||||
&None,
|
||||
None,
|
||||
None,
|
||||
now,
|
||||
);
|
||||
let rendered_lines = render_lines(&composite.display_lines(80));
|
||||
|
||||
Reference in New Issue
Block a user