diff --git a/.github/blob-size-allowlist.txt b/.github/blob-size-allowlist.txt index 9375b49c2739..a1ef72e0fed3 100644 --- a/.github/blob-size-allowlist.txt +++ b/.github/blob-size-allowlist.txt @@ -8,3 +8,5 @@ codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.js codex-rs/tui/tests/fixtures/oss-story.jsonl codex-rs/tui_app_server/tests/fixtures/oss-story.jsonl codex-rs/tui/src/app.rs +code-rs/core/src/codex/streaming.rs +code-rs/tui/src/chatwidget.rs diff --git a/code-rs/app-server-protocol/schema/json/EventMsg.json b/code-rs/app-server-protocol/schema/json/EventMsg.json index d248208ee824..d7bf723c4904 100644 --- a/code-rs/app-server-protocol/schema/json/EventMsg.json +++ b/code-rs/app-server-protocol/schema/json/EventMsg.json @@ -3687,6 +3687,16 @@ ], "type": "string" }, + "RateLimitReachedType": { + "enum": [ + "rate_limit_reached", + "workspace_owner_credits_depleted", + "workspace_member_credits_depleted", + "workspace_owner_usage_limit_reached", + "workspace_member_usage_limit_reached" + ], + "type": "string" + }, "RateLimitSnapshot": { "properties": { "credits": { @@ -3733,6 +3743,17 @@ } ] }, + "rate_limit_reached_type": { + "anyOf": [ + { + "$ref": "#/definitions/RateLimitReachedType" + }, + { + "type": "null" + } + ], + "default": null + }, "secondary": { "anyOf": [ { diff --git a/code-rs/app-server-protocol/schema/json/ServerNotification.json b/code-rs/app-server-protocol/schema/json/ServerNotification.json index 87bba8110433..eeab7a3d99a3 100644 --- a/code-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/code-rs/app-server-protocol/schema/json/ServerNotification.json @@ -4649,6 +4649,26 @@ ], "type": "string" }, + "RateLimitReachedType": { + "enum": [ + "rate_limit_reached", + "workspace_owner_credits_depleted", + "workspace_member_credits_depleted", + "workspace_owner_usage_limit_reached", + "workspace_member_usage_limit_reached" + ], + "type": "string" + }, + "RateLimitReachedType2": { + "enum": [ + "rate_limit_reached", + "workspace_owner_credits_depleted", + "workspace_member_credits_depleted", + "workspace_owner_usage_limit_reached", + "workspace_member_usage_limit_reached" + ], + "type": "string" + }, "RateLimitSnapshot": { "properties": { "credits": { @@ -4693,6 +4713,17 @@ } ] }, + "rateLimitReachedType": { + "anyOf": [ + { + "$ref": "#/definitions/RateLimitReachedType" + }, + { + "type": "null" + } + ], + "default": null + }, "secondary": { "anyOf": [ { @@ -4752,6 +4783,17 @@ } ] }, + "rate_limit_reached_type": { + "anyOf": [ + { + "$ref": "#/definitions/RateLimitReachedType2" + }, + { + "type": "null" + } + ], + "default": null + }, "secondary": { "anyOf": [ { diff --git a/code-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/code-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 226db3d72d9e..0edb86fb4f37 100644 --- a/code-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/code-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -6955,6 +6955,16 @@ }, "type": "object" }, + "RateLimitReachedType": { + "enum": [ + "rate_limit_reached", + "workspace_owner_credits_depleted", + "workspace_member_credits_depleted", + "workspace_owner_usage_limit_reached", + "workspace_member_usage_limit_reached" + ], + "type": "string" + }, "RateLimitSnapshot": { "properties": { "credits": { @@ -7001,6 +7011,17 @@ } ] }, + "rate_limit_reached_type": { + "anyOf": [ + { + "$ref": "#/definitions/RateLimitReachedType" + }, + { + "type": "null" + } + ], + "default": null + }, "secondary": { "anyOf": [ { @@ -13758,6 +13779,16 @@ }, "type": "object" }, + "RateLimitReachedType": { + "enum": [ + "rate_limit_reached", + "workspace_owner_credits_depleted", + "workspace_member_credits_depleted", + "workspace_owner_usage_limit_reached", + "workspace_member_usage_limit_reached" + ], + "type": "string" + }, "RateLimitSnapshot": { "properties": { "credits": { @@ -13802,6 +13833,17 @@ } ] }, + "rateLimitReachedType": { + "anyOf": [ + { + "$ref": "#/definitions/v2/RateLimitReachedType" + }, + { + "type": "null" + } + ], + "default": null + }, "secondary": { "anyOf": [ { diff --git a/code-rs/app-server-protocol/schema/json/v1/ForkConversationResponse.json b/code-rs/app-server-protocol/schema/json/v1/ForkConversationResponse.json index 010e4fee0473..6bc91666fb4e 100644 --- a/code-rs/app-server-protocol/schema/json/v1/ForkConversationResponse.json +++ b/code-rs/app-server-protocol/schema/json/v1/ForkConversationResponse.json @@ -3687,6 +3687,16 @@ ], "type": "string" }, + "RateLimitReachedType": { + "enum": [ + "rate_limit_reached", + "workspace_owner_credits_depleted", + "workspace_member_credits_depleted", + "workspace_owner_usage_limit_reached", + "workspace_member_usage_limit_reached" + ], + "type": "string" + }, "RateLimitSnapshot": { "properties": { "credits": { @@ -3733,6 +3743,17 @@ } ] }, + "rate_limit_reached_type": { + "anyOf": [ + { + "$ref": "#/definitions/RateLimitReachedType" + }, + { + "type": "null" + } + ], + "default": null + }, "secondary": { "anyOf": [ { diff --git a/code-rs/app-server-protocol/schema/json/v1/ResumeConversationResponse.json b/code-rs/app-server-protocol/schema/json/v1/ResumeConversationResponse.json index a703f18f76cc..2309ca0b4b06 100644 --- a/code-rs/app-server-protocol/schema/json/v1/ResumeConversationResponse.json +++ b/code-rs/app-server-protocol/schema/json/v1/ResumeConversationResponse.json @@ -3687,6 +3687,16 @@ ], "type": "string" }, + "RateLimitReachedType": { + "enum": [ + "rate_limit_reached", + "workspace_owner_credits_depleted", + "workspace_member_credits_depleted", + "workspace_owner_usage_limit_reached", + "workspace_member_usage_limit_reached" + ], + "type": "string" + }, "RateLimitSnapshot": { "properties": { "credits": { @@ -3733,6 +3743,17 @@ } ] }, + "rate_limit_reached_type": { + "anyOf": [ + { + "$ref": "#/definitions/RateLimitReachedType" + }, + { + "type": "null" + } + ], + "default": null + }, "secondary": { "anyOf": [ { diff --git a/code-rs/app-server-protocol/schema/json/v1/SessionConfiguredNotification.json b/code-rs/app-server-protocol/schema/json/v1/SessionConfiguredNotification.json index 2f01a04c18bd..efb3610926da 100644 --- a/code-rs/app-server-protocol/schema/json/v1/SessionConfiguredNotification.json +++ b/code-rs/app-server-protocol/schema/json/v1/SessionConfiguredNotification.json @@ -3687,6 +3687,16 @@ ], "type": "string" }, + "RateLimitReachedType": { + "enum": [ + "rate_limit_reached", + "workspace_owner_credits_depleted", + "workspace_member_credits_depleted", + "workspace_owner_usage_limit_reached", + "workspace_member_usage_limit_reached" + ], + "type": "string" + }, "RateLimitSnapshot": { "properties": { "credits": { @@ -3733,6 +3743,17 @@ } ] }, + "rate_limit_reached_type": { + "anyOf": [ + { + "$ref": "#/definitions/RateLimitReachedType" + }, + { + "type": "null" + } + ], + "default": null + }, "secondary": { "anyOf": [ { diff --git a/code-rs/app-server-protocol/schema/json/v2/AccountRateLimitsUpdatedNotification.json b/code-rs/app-server-protocol/schema/json/v2/AccountRateLimitsUpdatedNotification.json index 2569b422d5e3..c4e92d2f380a 100644 --- a/code-rs/app-server-protocol/schema/json/v2/AccountRateLimitsUpdatedNotification.json +++ b/code-rs/app-server-protocol/schema/json/v2/AccountRateLimitsUpdatedNotification.json @@ -36,6 +36,16 @@ ], "type": "string" }, + "RateLimitReachedType": { + "enum": [ + "rate_limit_reached", + "workspace_owner_credits_depleted", + "workspace_member_credits_depleted", + "workspace_owner_usage_limit_reached", + "workspace_member_usage_limit_reached" + ], + "type": "string" + }, "RateLimitSnapshot": { "properties": { "credits": { @@ -80,6 +90,17 @@ } ] }, + "rateLimitReachedType": { + "anyOf": [ + { + "$ref": "#/definitions/RateLimitReachedType" + }, + { + "type": "null" + } + ], + "default": null + }, "secondary": { "anyOf": [ { diff --git a/code-rs/app-server-protocol/schema/json/v2/GetAccountRateLimitsResponse.json b/code-rs/app-server-protocol/schema/json/v2/GetAccountRateLimitsResponse.json index e9495d90567a..713b2002b836 100644 --- a/code-rs/app-server-protocol/schema/json/v2/GetAccountRateLimitsResponse.json +++ b/code-rs/app-server-protocol/schema/json/v2/GetAccountRateLimitsResponse.json @@ -36,6 +36,16 @@ ], "type": "string" }, + "RateLimitReachedType": { + "enum": [ + "rate_limit_reached", + "workspace_owner_credits_depleted", + "workspace_member_credits_depleted", + "workspace_owner_usage_limit_reached", + "workspace_member_usage_limit_reached" + ], + "type": "string" + }, "RateLimitSnapshot": { "properties": { "credits": { @@ -80,6 +90,17 @@ } ] }, + "rateLimitReachedType": { + "anyOf": [ + { + "$ref": "#/definitions/RateLimitReachedType" + }, + { + "type": "null" + } + ], + "default": null + }, "secondary": { "anyOf": [ { diff --git a/code-rs/app-server-protocol/schema/typescript/RateLimitReachedType.ts b/code-rs/app-server-protocol/schema/typescript/RateLimitReachedType.ts new file mode 100644 index 000000000000..78f106c905d1 --- /dev/null +++ b/code-rs/app-server-protocol/schema/typescript/RateLimitReachedType.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type RateLimitReachedType = "rate_limit_reached" | "workspace_owner_credits_depleted" | "workspace_member_credits_depleted" | "workspace_owner_usage_limit_reached" | "workspace_member_usage_limit_reached"; diff --git a/code-rs/app-server-protocol/schema/typescript/RateLimitSnapshot.ts b/code-rs/app-server-protocol/schema/typescript/RateLimitSnapshot.ts index 8604128b4e47..5da995e0273f 100644 --- a/code-rs/app-server-protocol/schema/typescript/RateLimitSnapshot.ts +++ b/code-rs/app-server-protocol/schema/typescript/RateLimitSnapshot.ts @@ -3,6 +3,7 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { CreditsSnapshot } from "./CreditsSnapshot"; import type { PlanType } from "./PlanType"; +import type { RateLimitReachedType } from "./RateLimitReachedType"; import type { RateLimitWindow } from "./RateLimitWindow"; -export type RateLimitSnapshot = { limit_id: string | null, limit_name: string | null, primary: RateLimitWindow | null, secondary: RateLimitWindow | null, credits: CreditsSnapshot | null, plan_type: PlanType | null, }; +export type RateLimitSnapshot = { limit_id: string | null, limit_name: string | null, primary: RateLimitWindow | null, secondary: RateLimitWindow | null, credits: CreditsSnapshot | null, plan_type: PlanType | null, rate_limit_reached_type: RateLimitReachedType | null, }; diff --git a/code-rs/app-server-protocol/schema/typescript/index.ts b/code-rs/app-server-protocol/schema/typescript/index.ts index a236bb9a2b45..94110512dc55 100644 --- a/code-rs/app-server-protocol/schema/typescript/index.ts +++ b/code-rs/app-server-protocol/schema/typescript/index.ts @@ -146,6 +146,7 @@ export type { PlanItem } from "./PlanItem"; export type { PlanItemArg } from "./PlanItemArg"; export type { PlanType } from "./PlanType"; export type { Profile } from "./Profile"; +export type { RateLimitReachedType } from "./RateLimitReachedType"; export type { RateLimitSnapshot } from "./RateLimitSnapshot"; export type { RateLimitWindow } from "./RateLimitWindow"; export type { RawResponseItemEvent } from "./RawResponseItemEvent"; diff --git a/code-rs/app-server-protocol/schema/typescript/v2/RateLimitReachedType.ts b/code-rs/app-server-protocol/schema/typescript/v2/RateLimitReachedType.ts new file mode 100644 index 000000000000..78f106c905d1 --- /dev/null +++ b/code-rs/app-server-protocol/schema/typescript/v2/RateLimitReachedType.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type RateLimitReachedType = "rate_limit_reached" | "workspace_owner_credits_depleted" | "workspace_member_credits_depleted" | "workspace_owner_usage_limit_reached" | "workspace_member_usage_limit_reached"; diff --git a/code-rs/app-server-protocol/schema/typescript/v2/RateLimitSnapshot.ts b/code-rs/app-server-protocol/schema/typescript/v2/RateLimitSnapshot.ts index 0c2ebe1893f7..dc8417a30402 100644 --- a/code-rs/app-server-protocol/schema/typescript/v2/RateLimitSnapshot.ts +++ b/code-rs/app-server-protocol/schema/typescript/v2/RateLimitSnapshot.ts @@ -3,6 +3,7 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { PlanType } from "../PlanType"; import type { CreditsSnapshot } from "./CreditsSnapshot"; +import type { RateLimitReachedType } from "./RateLimitReachedType"; import type { RateLimitWindow } from "./RateLimitWindow"; -export type RateLimitSnapshot = { limitId: string | null, limitName: string | null, primary: RateLimitWindow | null, secondary: RateLimitWindow | null, credits: CreditsSnapshot | null, planType: PlanType | null, }; +export type RateLimitSnapshot = { limitId: string | null, limitName: string | null, primary: RateLimitWindow | null, secondary: RateLimitWindow | null, credits: CreditsSnapshot | null, planType: PlanType | null, rateLimitReachedType: RateLimitReachedType | null, }; diff --git a/code-rs/app-server-protocol/schema/typescript/v2/index.ts b/code-rs/app-server-protocol/schema/typescript/v2/index.ts index dc7f88ce2c56..08da9a846212 100644 --- a/code-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/code-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -111,6 +111,7 @@ export type { PatchApplyStatus } from "./PatchApplyStatus"; export type { PatchChangeKind } from "./PatchChangeKind"; export type { PlanDeltaNotification } from "./PlanDeltaNotification"; export type { ProfileV2 } from "./ProfileV2"; +export type { RateLimitReachedType } from "./RateLimitReachedType"; export type { RateLimitSnapshot } from "./RateLimitSnapshot"; export type { RateLimitWindow } from "./RateLimitWindow"; export type { RawResponseItemCompletedNotification } from "./RawResponseItemCompletedNotification"; diff --git a/code-rs/app-server-protocol/src/protocol/v2.rs b/code-rs/app-server-protocol/src/protocol/v2.rs index c5317ba1ec00..838b92c940ac 100644 --- a/code-rs/app-server-protocol/src/protocol/v2.rs +++ b/code-rs/app-server-protocol/src/protocol/v2.rs @@ -39,6 +39,7 @@ use code_protocol::protocol::RejectConfig as CoreRejectConfig; use code_protocol::protocol::CodexErrorInfo as CoreCodexErrorInfo; use code_protocol::protocol::CreditsSnapshot as CoreCreditsSnapshot; use code_protocol::protocol::NetworkAccess as CoreNetworkAccess; +use code_protocol::protocol::RateLimitReachedType as CoreRateLimitReachedType; use code_protocol::protocol::RateLimitSnapshot as CoreRateLimitSnapshot; use code_protocol::protocol::RateLimitWindow as CoreRateLimitWindow; use code_protocol::protocol::SessionSource as CoreSessionSource; @@ -3418,6 +3419,8 @@ pub struct RateLimitSnapshot { pub secondary: Option, pub credits: Option, pub plan_type: Option, + #[serde(default)] + pub rate_limit_reached_type: Option, } impl From for RateLimitSnapshot { @@ -3429,6 +3432,41 @@ impl From for RateLimitSnapshot { secondary: value.secondary.map(RateLimitWindow::from), credits: value.credits.map(CreditsSnapshot::from), plan_type: value.plan_type, + rate_limit_reached_type: value + .rate_limit_reached_type + .map(RateLimitReachedType::from), + } + } +} + +#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "snake_case")] +#[ts(rename_all = "snake_case")] +#[ts(export_to = "v2/")] +pub enum RateLimitReachedType { + RateLimitReached, + WorkspaceOwnerCreditsDepleted, + WorkspaceMemberCreditsDepleted, + WorkspaceOwnerUsageLimitReached, + WorkspaceMemberUsageLimitReached, +} + +impl From for RateLimitReachedType { + fn from(value: CoreRateLimitReachedType) -> Self { + match value { + CoreRateLimitReachedType::RateLimitReached => Self::RateLimitReached, + CoreRateLimitReachedType::WorkspaceOwnerCreditsDepleted => { + Self::WorkspaceOwnerCreditsDepleted + } + CoreRateLimitReachedType::WorkspaceMemberCreditsDepleted => { + Self::WorkspaceMemberCreditsDepleted + } + CoreRateLimitReachedType::WorkspaceOwnerUsageLimitReached => { + Self::WorkspaceOwnerUsageLimitReached + } + CoreRateLimitReachedType::WorkspaceMemberUsageLimitReached => { + Self::WorkspaceMemberUsageLimitReached + } } } } diff --git a/code-rs/app-server/src/code_message_processor.rs b/code-rs/app-server/src/code_message_processor.rs index 9a0b8a1e626b..b1ac75a75fe1 100644 --- a/code-rs/app-server/src/code_message_processor.rs +++ b/code-rs/app-server/src/code_message_processor.rs @@ -118,6 +118,7 @@ use code_protocol::mcp_protocol::CancelLoginChatGptResponse; use code_protocol::mcp_protocol::LogoutChatGptResponse; use code_protocol::account::PlanType; use code_protocol::protocol::RateLimitSnapshot as CoreRateLimitSnapshot; +use code_protocol::protocol::RateLimitReachedType as CoreRateLimitReachedType; use code_protocol::protocol::RateLimitWindow as CoreRateLimitWindow; // Removed deprecated ChatGPT login support scaffolding @@ -2143,6 +2144,7 @@ mod tests { secondary_window_minutes: 1440, primary_reset_after_seconds: Some(12), secondary_reset_after_seconds: Some(34), + rate_limit_reached_type: None, }; let snapshot = rate_limit_snapshot_from_event(&event, Some(PlanType::Pro)); @@ -2338,6 +2340,31 @@ fn rate_limit_snapshot_from_event( secondary: Some(secondary), credits: None, plan_type, + rate_limit_reached_type: snapshot + .rate_limit_reached_type + .map(rate_limit_reached_type_to_protocol), + } +} + +fn rate_limit_reached_type_to_protocol( + reached: code_core::protocol::RateLimitReachedType, +) -> CoreRateLimitReachedType { + match reached { + code_core::protocol::RateLimitReachedType::RateLimitReached => { + CoreRateLimitReachedType::RateLimitReached + } + code_core::protocol::RateLimitReachedType::WorkspaceOwnerCreditsDepleted => { + CoreRateLimitReachedType::WorkspaceOwnerCreditsDepleted + } + code_core::protocol::RateLimitReachedType::WorkspaceMemberCreditsDepleted => { + CoreRateLimitReachedType::WorkspaceMemberCreditsDepleted + } + code_core::protocol::RateLimitReachedType::WorkspaceOwnerUsageLimitReached => { + CoreRateLimitReachedType::WorkspaceOwnerUsageLimitReached + } + code_core::protocol::RateLimitReachedType::WorkspaceMemberUsageLimitReached => { + CoreRateLimitReachedType::WorkspaceMemberUsageLimitReached + } } } diff --git a/code-rs/code-auto-drive-core/src/auto_coordinator.rs b/code-rs/code-auto-drive-core/src/auto_coordinator.rs index 9ca3be54a9a6..e81755d084d5 100644 --- a/code-rs/code-auto-drive-core/src/auto_coordinator.rs +++ b/code-rs/code-auto-drive-core/src/auto_coordinator.rs @@ -1571,6 +1571,7 @@ mod tests { let err = anyhow!(CodexErr::UsageLimitReached(UsageLimitReachedError { plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, })); match classify_model_error(&err) { RetryDecision::Fatal(e) => { diff --git a/code-rs/code-auto-drive-core/src/faults.rs b/code-rs/code-auto-drive-core/src/faults.rs index 08a82aaae5d2..d502e6d188ce 100644 --- a/code-rs/code-auto-drive-core/src/faults.rs +++ b/code-rs/code-auto-drive-core/src/faults.rs @@ -130,6 +130,7 @@ pub fn fault_to_error(fault: InjectedFault) -> anyhow::Error { Some(FaultReset::Seconds(secs)) => anyhow!(CodexErr::UsageLimitReached(UsageLimitReachedError { plan_type: None, resets_in_seconds: Some(secs), + rate_limit_reached_type: None, })), Some(FaultReset::Timestamp(instant)) => { let reset_at = chrono::Utc::now() diff --git a/code-rs/core/src/account_switching.rs b/code-rs/core/src/account_switching.rs index b07c4d903e20..e14abed038cb 100644 --- a/code-rs/core/src/account_switching.rs +++ b/code-rs/core/src/account_switching.rs @@ -8,6 +8,7 @@ use code_app_server_protocol::AuthMode; use crate::auth; use crate::account_usage; use crate::auth_accounts; +use crate::protocol::RateLimitReachedType; #[derive(Debug, Default)] pub struct RateLimitSwitchState { @@ -66,6 +67,29 @@ fn account_has_credentials(account: &auth_accounts::StoredAccount) -> bool { fn usage_reset_blocked_until( snapshot: &account_usage::StoredRateLimitSnapshot, ) -> Option> { + if snapshot + .snapshot + .as_ref() + .and_then(|snapshot| snapshot.rate_limit_reached_type) + .is_some_and(|reached| { + matches!( + reached, + RateLimitReachedType::RateLimitReached + | RateLimitReachedType::WorkspaceOwnerCreditsDepleted + | RateLimitReachedType::WorkspaceMemberCreditsDepleted + | RateLimitReachedType::WorkspaceOwnerUsageLimitReached + | RateLimitReachedType::WorkspaceMemberUsageLimitReached + ) + }) + { + return snapshot + .primary_next_reset_at + .into_iter() + .chain(snapshot.secondary_next_reset_at) + .max() + .or(snapshot.last_usage_limit_hit_at); + } + snapshot .primary_next_reset_at .into_iter() @@ -278,6 +302,19 @@ mod tests { Utc.with_ymd_and_hms(2025, 12, 22, 12, 0, 0).unwrap() } + fn sample_snapshot(used_percent: f64) -> crate::protocol::RateLimitSnapshotEvent { + crate::protocol::RateLimitSnapshotEvent { + primary_used_percent: used_percent, + secondary_used_percent: used_percent, + primary_to_secondary_ratio_percent: 25.0, + primary_window_minutes: 300, + secondary_window_minutes: 10_080, + primary_reset_after_seconds: Some(600), + secondary_reset_after_seconds: Some(3_600), + rate_limit_reached_type: None, + } + } + #[test] fn selects_another_chatgpt_account_when_available() { let home = tempdir().expect("tmp"); @@ -344,6 +381,72 @@ mod tests { assert!(next.is_none()); } + #[test] + fn typed_usage_limit_snapshot_blocks_only_that_account() { + let home = tempdir().expect("tmp"); + let a = auth_accounts::upsert_chatgpt_account( + home.path(), + chatgpt_tokens("acct-a", "a@example.com"), + Utc::now(), + None, + true, + ) + .expect("insert a"); + let b = auth_accounts::upsert_chatgpt_account( + home.path(), + chatgpt_tokens("acct-b", "b@example.com"), + Utc::now(), + None, + false, + ) + .expect("insert b"); + let c = auth_accounts::upsert_chatgpt_account( + home.path(), + chatgpt_tokens("acct-c", "c@example.com"), + Utc::now(), + None, + false, + ) + .expect("insert c"); + + let now = fixed_now(); + let mut limited_snapshot = sample_snapshot(95.0); + limited_snapshot.rate_limit_reached_type = Some( + RateLimitReachedType::WorkspaceMemberUsageLimitReached, + ); + account_usage::record_rate_limit_snapshot( + home.path(), + &b.id, + Some("Pro"), + &limited_snapshot, + now, + ) + .expect("limited snapshot"); + let mut available_snapshot = sample_snapshot(20.0); + available_snapshot.primary_reset_after_seconds = None; + available_snapshot.secondary_reset_after_seconds = None; + account_usage::record_rate_limit_snapshot( + home.path(), + &c.id, + Some("Pro"), + &available_snapshot, + now, + ) + .expect("available snapshot"); + + let mut state = RateLimitSwitchState::default(); + state.mark_limited(&a.id, AuthMode::ChatGPT, None); + let next = select_next_account_id( + home.path(), + &state, + false, + now, + Some(a.id.as_str()), + ) + .expect("select"); + assert_eq!(next.as_deref(), Some(c.id.as_str())); + } + #[test] fn api_key_fallback_requires_all_chatgpt_limited() { let home = tempdir().expect("tmp"); diff --git a/code-rs/core/src/account_usage.rs b/code-rs/core/src/account_usage.rs index cacadbd2a750..fe1ed2e815fd 100644 --- a/code-rs/core/src/account_usage.rs +++ b/code-rs/core/src/account_usage.rs @@ -4,7 +4,7 @@ use std::io::{ErrorKind, Read, Seek, SeekFrom, Write}; use std::path::{Path, PathBuf}; use chrono::{DateTime, Datelike, Duration, NaiveDate, TimeZone, Timelike, Utc}; -use crate::protocol::RateLimitSnapshotEvent; +use crate::protocol::{RateLimitReachedType, RateLimitSnapshotEvent}; use fs2::FileExt; use serde::{Deserialize, Serialize}; use serde_json; @@ -516,11 +516,33 @@ pub fn record_usage_limit_hint( resets_in_seconds: Option, observed_at: DateTime, ) -> std::io::Result<()> { + record_usage_limit_hint_with_type( + code_home, + account_id, + plan, + resets_in_seconds, + observed_at, + Some(RateLimitReachedType::RateLimitReached), + ) +} + +pub fn record_usage_limit_hint_with_type( + code_home: &Path, + account_id: &str, + plan: Option<&str>, + resets_in_seconds: Option, + observed_at: DateTime, + reached_type: Option, +) -> std::io::Result<()> { + let reached_type = reached_type.or(Some(RateLimitReachedType::RateLimitReached)); if resets_in_seconds.is_none() { return with_usage_file(code_home, account_id, plan, |data| { data.last_updated = observed_at; let mut info = data.rate_limit.take().unwrap_or_default(); info.last_usage_limit_hit_at = Some(observed_at); + if let Some(snapshot) = info.snapshot.as_mut() { + snapshot.rate_limit_reached_type = reached_type; + } data.rate_limit = Some(info); }); } @@ -529,6 +551,9 @@ pub fn record_usage_limit_hint( data.last_updated = observed_at; let mut info = data.rate_limit.take().unwrap_or_default(); info.last_usage_limit_hit_at = Some(observed_at); + if let Some(snapshot) = info.snapshot.as_mut() { + snapshot.rate_limit_reached_type = reached_type; + } if let Some(seconds) = resets_in_seconds { let reset_at = observed_at + Duration::seconds(seconds as i64); info.primary_next_reset_at = Some(reset_at); @@ -826,6 +851,7 @@ mod tests { //! The helper tests below construct scenarios targeting each rule so the state //! machine in `record_threshold_log` can be refactored confidently. use super::*; + use crate::protocol::RateLimitReachedType; use std::fs::File; use crate::protocol::TokenUsage; use tempfile::TempDir; @@ -840,6 +866,19 @@ mod tests { } } + fn sample_snapshot() -> RateLimitSnapshotEvent { + RateLimitSnapshotEvent { + primary_used_percent: 50.0, + secondary_used_percent: 60.0, + primary_to_secondary_ratio_percent: 25.0, + primary_window_minutes: 300, + secondary_window_minutes: 10_080, + primary_reset_after_seconds: Some(600), + secondary_reset_after_seconds: Some(3_600), + rate_limit_reached_type: None, + } + } + #[test] fn usage_limit_hint_updates_last_hit_and_resets() { let home = TempDir::new().expect("tempdir"); @@ -854,11 +893,74 @@ mod tests { assert_eq!(snapshot.account_id, "acct-1"); assert_eq!(snapshot.plan.as_deref(), Some("Team")); assert_eq!(snapshot.last_usage_limit_hit_at, Some(now)); + assert_eq!( + snapshot + .snapshot + .as_ref() + .and_then(|snapshot| snapshot.rate_limit_reached_type), + None, + "a hint without a prior snapshot should not fabricate usage bars" + ); let expected_reset = now + Duration::seconds(300); assert_eq!(snapshot.primary_next_reset_at, Some(expected_reset)); assert_eq!(snapshot.secondary_next_reset_at, Some(expected_reset)); } + #[test] + fn usage_limit_hint_marks_existing_snapshot_with_reached_type() { + let home = TempDir::new().expect("tempdir"); + let now = Utc::now(); + let mut snapshot = sample_snapshot(); + + record_rate_limit_snapshot(home.path(), "acct-1", Some("Team"), &snapshot, now) + .expect("snapshot recorded"); + record_usage_limit_hint_with_type( + home.path(), + "acct-1", + Some("Team"), + Some(300), + now, + Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached), + ) + .expect("hint recorded"); + + let snapshots = list_rate_limit_snapshots(home.path()).expect("snapshot listing"); + let stored = snapshots + .iter() + .find(|snapshot| snapshot.account_id == "acct-1") + .expect("account snapshot"); + assert_eq!( + stored + .snapshot + .as_ref() + .and_then(|snapshot| snapshot.rate_limit_reached_type), + Some(RateLimitReachedType::WorkspaceMemberUsageLimitReached) + ); + + snapshot.primary_used_percent = 1.0; + record_rate_limit_snapshot( + home.path(), + "acct-1", + Some("Team"), + &snapshot, + now + Duration::minutes(1), + ) + .expect("fresh snapshot recorded"); + let snapshots = list_rate_limit_snapshots(home.path()).expect("snapshot listing"); + let stored = snapshots + .iter() + .find(|snapshot| snapshot.account_id == "acct-1") + .expect("account snapshot"); + assert_eq!( + stored + .snapshot + .as_ref() + .and_then(|snapshot| snapshot.rate_limit_reached_type), + None, + "fresh successful snapshots clear stale reached classification" + ); + } + #[test] fn token_usage_compacts_old_hourly_entries_into_buckets() { let home = TempDir::new().expect("tempdir"); diff --git a/code-rs/core/src/client.rs b/code-rs/core/src/client.rs index b9609840d2ed..7bf86d739fb9 100644 --- a/code-rs/core/src/client.rs +++ b/code-rs/core/src/client.rs @@ -78,6 +78,7 @@ use crate::openai_tools::create_tools_json_for_responses_api; use crate::openai_tools::ConfigShellToolType; use crate::openai_tools::ToolsConfig; use crate::protocol::RateLimitSnapshotEvent; +use crate::protocol::RateLimitReachedType; use crate::protocol::SandboxPolicy; use crate::protocol::TokenUsage; use crate::reasoning::clamp_reasoning_effort_for_model; @@ -257,6 +258,7 @@ struct Error { // Optional fields available on "usage_limit_reached" and "usage_not_included" errors plan_type: Option, resets_in_seconds: Option, + rate_limit_reached_type: Option, } #[derive(Serialize)] @@ -1881,16 +1883,20 @@ impl ModelClient { .map(|s| s.to_string()); let resets_in_seconds = body.as_ref().and_then(|err| err.error.resets_in_seconds); + let reached_type = body + .as_ref() + .and_then(|err| err.error.rate_limit_reached_type); let code_home = self.code_home().to_path_buf(); let account_id = current_account_id.clone(); tokio::task::spawn_blocking(move || { let observed_at = Utc::now(); - if let Err(err) = account_usage::record_usage_limit_hint( + if let Err(err) = account_usage::record_usage_limit_hint_with_type( &code_home, &account_id, plan_type.as_deref(), resets_in_seconds, observed_at, + reached_type, ) { tracing::warn!("Failed to persist usage limit hint: {err}"); } @@ -2023,6 +2029,7 @@ impl ModelClient { return Err(CodexErr::UsageLimitReached(UsageLimitReachedError { plan_type, resets_in_seconds, + rate_limit_reached_type: error.rate_limit_reached_type, })); } else if error.r#type.as_deref() == Some("usage_not_included") { return Err(CodexErr::UsageNotIncluded); @@ -2399,6 +2406,7 @@ fn map_wrapped_websocket_error_event(event: WrappedWebsocketErrorEvent) -> Optio return Some(CodexErr::UsageLimitReached(UsageLimitReachedError { plan_type: error.plan_type, resets_in_seconds: error.resets_in_seconds, + rate_limit_reached_type: error.rate_limit_reached_type, })); } @@ -2646,6 +2654,7 @@ fn parse_rate_limit_snapshot(headers: &HeaderMap) -> Option( UsageLimitReachedError { plan_type: error.plan_type, resets_in_seconds: error.resets_in_seconds, + rate_limit_reached_type: error.rate_limit_reached_type, }, )); } else if error.r#type.as_deref() == Some("usage_not_included") { @@ -4043,6 +4053,7 @@ mod tests { param: None, plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, }; let retry_after = try_parse_retry_after(&err, now).expect("retry"); @@ -4060,6 +4071,7 @@ mod tests { param: None, plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, }; let retry_after = try_parse_retry_after(&err, now).expect("retry"); assert_eq!(retry_after.delay, Duration::from_secs_f64(1.898)); @@ -4075,6 +4087,7 @@ mod tests { param: None, plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, }; let retry_after = try_parse_retry_after(&err, now).expect("retry"); assert_eq!(retry_after.delay, Duration::from_secs(35)); @@ -4090,6 +4103,7 @@ mod tests { param: None, plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, }; assert!(try_parse_retry_after(&err, now).is_none()); @@ -4143,6 +4157,7 @@ mod tests { param: None, plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, }; chosen = try_parse_retry_after(&err, now); } @@ -4169,6 +4184,7 @@ mod tests { param: None, plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, }; for status in [ @@ -4194,6 +4210,7 @@ mod tests { param: None, plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, }; assert!(!is_quota_exceeded_http_error(StatusCode::BAD_REQUEST, &error)); @@ -4208,6 +4225,7 @@ mod tests { param: Some("reasoning.summary".to_string()), plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, }; assert!(is_reasoning_summary_rejected(&error_with_param)); @@ -4219,6 +4237,7 @@ mod tests { param: None, plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, }; assert!(is_reasoning_summary_rejected(&error_by_message)); @@ -4232,6 +4251,7 @@ mod tests { param: Some("reasoning.summary".to_string()), plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, }; assert!(!is_reasoning_summary_rejected(&rate_limit_error)); diff --git a/code-rs/core/src/codex/streaming.rs b/code-rs/core/src/codex/streaming.rs index d5dd4ae86f79..01e3cc75dbe2 100644 --- a/code-rs/core/src/codex/streaming.rs +++ b/code-rs/core/src/codex/streaming.rs @@ -3107,13 +3107,15 @@ async fn run_turn( let usage_account = ctx.account_id.clone(); let usage_plan = ctx.plan.clone(); let resets = limit_err.resets_in_seconds; + let reached_type = limit_err.rate_limit_reached_type; spawn_usage_task(move || { - if let Err(err) = account_usage::record_usage_limit_hint( + if let Err(err) = account_usage::record_usage_limit_hint_with_type( &usage_home, &usage_account, usage_plan.as_deref(), resets, Utc::now(), + reached_type, ) { warn!("Failed to persist usage limit hint: {err}"); } diff --git a/code-rs/core/src/error.rs b/code-rs/core/src/error.rs index 3b2b9088408d..9dfcbfd9f41d 100644 --- a/code-rs/core/src/error.rs +++ b/code-rs/core/src/error.rs @@ -1,4 +1,5 @@ use crate::exec::ExecToolCallOutput; +use crate::protocol::RateLimitReachedType; use chrono::{DateTime, Duration as ChronoDuration, Utc}; use reqwest::StatusCode; use serde_json; @@ -239,6 +240,7 @@ impl std::fmt::Display for RetryLimitReachedError { pub struct UsageLimitReachedError { pub plan_type: Option, pub resets_in_seconds: Option, + pub rate_limit_reached_type: Option, } impl UsageLimitReachedError { @@ -418,6 +420,7 @@ mod tests { let err = UsageLimitReachedError { plan_type: Some("plus".to_string()), resets_in_seconds: None, + rate_limit_reached_type: None, }; assert_eq!( err.to_string(), @@ -430,6 +433,7 @@ mod tests { let err = UsageLimitReachedError { plan_type: None, resets_in_seconds: None, + rate_limit_reached_type: None, }; assert_eq!( err.to_string(), @@ -442,6 +446,7 @@ mod tests { let err = UsageLimitReachedError { plan_type: Some("pro".to_string()), resets_in_seconds: None, + rate_limit_reached_type: None, }; assert_eq!( err.to_string(), @@ -454,6 +459,7 @@ mod tests { let err = UsageLimitReachedError { plan_type: None, resets_in_seconds: Some(5 * 60), + rate_limit_reached_type: None, }; assert_eq!( err.to_string(), @@ -466,6 +472,7 @@ mod tests { let err = UsageLimitReachedError { plan_type: Some("plus".to_string()), resets_in_seconds: Some(3 * 3600 + 32 * 60), + rate_limit_reached_type: None, }; assert_eq!( err.to_string(), @@ -478,6 +485,7 @@ mod tests { let err = UsageLimitReachedError { plan_type: None, resets_in_seconds: Some(2 * 86_400 + 3 * 3600 + 5 * 60), + rate_limit_reached_type: None, }; assert_eq!( err.to_string(), @@ -490,6 +498,7 @@ mod tests { let err = UsageLimitReachedError { plan_type: None, resets_in_seconds: Some(30), + rate_limit_reached_type: None, }; assert_eq!( err.to_string(), diff --git a/code-rs/core/src/history/state.rs b/code-rs/core/src/history/state.rs index ce69d22cfaa5..1b6272e2e072 100644 --- a/code-rs/core/src/history/state.rs +++ b/code-rs/core/src/history/state.rs @@ -3130,6 +3130,7 @@ mod tests { secondary_window_minutes: 5, primary_reset_after_seconds: Some(30), secondary_reset_after_seconds: Some(60), + rate_limit_reached_type: None, }, legend: vec![RateLimitLegendEntry { label: "primary".into(), diff --git a/code-rs/core/src/protocol.rs b/code-rs/core/src/protocol.rs index 63325d1dc155..b66f7cf63327 100644 --- a/code-rs/core/src/protocol.rs +++ b/code-rs/core/src/protocol.rs @@ -747,6 +747,53 @@ fn rate_limit_snapshot_to_protocol( secondary: Some(secondary), credits: None, plan_type: None, + rate_limit_reached_type: snapshot + .rate_limit_reached_type + .map(rate_limit_reached_type_to_protocol), + } +} + +fn rate_limit_reached_type_to_protocol( + reached: RateLimitReachedType, +) -> code_protocol::protocol::RateLimitReachedType { + match reached { + RateLimitReachedType::RateLimitReached => { + code_protocol::protocol::RateLimitReachedType::RateLimitReached + } + RateLimitReachedType::WorkspaceOwnerCreditsDepleted => { + code_protocol::protocol::RateLimitReachedType::WorkspaceOwnerCreditsDepleted + } + RateLimitReachedType::WorkspaceMemberCreditsDepleted => { + code_protocol::protocol::RateLimitReachedType::WorkspaceMemberCreditsDepleted + } + RateLimitReachedType::WorkspaceOwnerUsageLimitReached => { + code_protocol::protocol::RateLimitReachedType::WorkspaceOwnerUsageLimitReached + } + RateLimitReachedType::WorkspaceMemberUsageLimitReached => { + code_protocol::protocol::RateLimitReachedType::WorkspaceMemberUsageLimitReached + } + } +} + +fn rate_limit_reached_type_from_protocol( + reached: code_protocol::protocol::RateLimitReachedType, +) -> RateLimitReachedType { + match reached { + code_protocol::protocol::RateLimitReachedType::RateLimitReached => { + RateLimitReachedType::RateLimitReached + } + code_protocol::protocol::RateLimitReachedType::WorkspaceOwnerCreditsDepleted => { + RateLimitReachedType::WorkspaceOwnerCreditsDepleted + } + code_protocol::protocol::RateLimitReachedType::WorkspaceMemberCreditsDepleted => { + RateLimitReachedType::WorkspaceMemberCreditsDepleted + } + code_protocol::protocol::RateLimitReachedType::WorkspaceOwnerUsageLimitReached => { + RateLimitReachedType::WorkspaceOwnerUsageLimitReached + } + code_protocol::protocol::RateLimitReachedType::WorkspaceMemberUsageLimitReached => { + RateLimitReachedType::WorkspaceMemberUsageLimitReached + } } } @@ -800,6 +847,9 @@ fn rate_limit_snapshot_from_protocol( secondary_window_minutes, primary_reset_after_seconds, secondary_reset_after_seconds, + rate_limit_reached_type: snapshot + .rate_limit_reached_type + .map(rate_limit_reached_type_from_protocol), } } @@ -1159,6 +1209,19 @@ pub struct RateLimitSnapshotEvent { /// Seconds until the secondary window resets, if reported by the API. #[serde(skip_serializing_if = "Option::is_none")] pub secondary_reset_after_seconds: Option, + /// Backend classification for a reached limit, when provided. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub rate_limit_reached_type: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum RateLimitReachedType { + RateLimitReached, + WorkspaceOwnerCreditsDepleted, + WorkspaceMemberCreditsDepleted, + WorkspaceOwnerUsageLimitReached, + WorkspaceMemberUsageLimitReached, } #[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] diff --git a/code-rs/protocol/src/protocol.rs b/code-rs/protocol/src/protocol.rs index 0a896854271b..56179fc26f91 100644 --- a/code-rs/protocol/src/protocol.rs +++ b/code-rs/protocol/src/protocol.rs @@ -1356,6 +1356,19 @@ pub struct RateLimitSnapshot { pub secondary: Option, pub credits: Option, pub plan_type: Option, + #[serde(default)] + pub rate_limit_reached_type: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize, JsonSchema, TS)] +#[serde(rename_all = "snake_case")] +#[ts(rename_all = "snake_case")] +pub enum RateLimitReachedType { + RateLimitReached, + WorkspaceOwnerCreditsDepleted, + WorkspaceMemberCreditsDepleted, + WorkspaceOwnerUsageLimitReached, + WorkspaceMemberUsageLimitReached, } #[derive(Debug, Clone, PartialEq, Deserialize, Serialize, JsonSchema, TS)] diff --git a/code-rs/tui/src/chatwidget.rs b/code-rs/tui/src/chatwidget.rs index 3bde033b5751..4825051cfe0d 100644 --- a/code-rs/tui/src/chatwidget.rs +++ b/code-rs/tui/src/chatwidget.rs @@ -134,7 +134,7 @@ use code_auto_drive_core::{ TurnConfig, TurnDescriptor, }; -use self::limits_overlay::{LimitsOverlayContent, LimitsTab}; +use self::limits_overlay::{LimitsOverlayContent, LimitsTab, LimitsTabBody}; use crate::chrome_launch::ChromeLaunchOption; use crate::insert_history::word_wrap_lines; use self::rate_limit_refresh::{ @@ -1051,7 +1051,7 @@ use code_cloud_tasks_client::{ApplyOutcome, CloudTaskError, CreatedTask, TaskSum use code_protocol::models::ContentItem; use code_protocol::models::ResponseItem; use code_core::config_types::{validation_tool_category, ValidationCategory}; -use code_core::protocol::RateLimitSnapshotEvent; +use code_core::protocol::{RateLimitReachedType, RateLimitSnapshotEvent}; use code_core::protocol::ValidationGroup; use crate::rate_limits_view::{ build_limits_view, RateLimitDisplayConfig, RateLimitResetInfo, DEFAULT_DISPLAY_CONFIG, @@ -9109,13 +9109,26 @@ impl ChatWidget<'_> { } } - fn set_limits_overlay_tabs(&mut self, tabs: Vec) { - let content = if tabs.is_empty() { + fn build_limits_overlay_content( + &self, + snapshot: Option, + ) -> LimitsOverlayContent { + let reset_info = self.rate_limit_reset_info(); + let tabs = self.build_limits_tabs(snapshot, reset_info); + if tabs.is_empty() { LimitsOverlayContent::Placeholder } else { LimitsOverlayContent::Tabs(tabs) - }; - self.set_limits_overlay_content(content); + } + } + + fn limits_content_has_snapshot(content: &LimitsOverlayContent) -> bool { + match content { + LimitsOverlayContent::Tabs(tabs) => tabs + .iter() + .any(|tab| matches!(tab.body, LimitsTabBody::View(_))), + _ => false, + } } fn build_limits_tabs( @@ -9410,6 +9423,17 @@ impl ChatWidget<'_> { RtSpan::raw(status_field_prefix("Plan")), RtSpan::styled(plan.to_string(), value_style), ])); + if let Some(snapshot) = record.and_then(|r| r.snapshot.as_ref()) { + if let Some(reached_type) = snapshot.rate_limit_reached_type { + lines.push(RtLine::from(vec![ + RtSpan::raw(status_field_prefix("Status")), + RtSpan::styled( + Self::rate_limit_reached_label(reached_type), + Style::default().fg(crate::colors::warning()), + ), + ])); + } + } let tokens_prefix = status_field_prefix("Tokens"); let tokens_summary = format!("{formatted_total} total {cost_suffix}"); lines.push(RtLine::from(vec![ @@ -9440,6 +9464,24 @@ impl ChatWidget<'_> { lines } + fn rate_limit_reached_label(reached_type: RateLimitReachedType) -> String { + match reached_type { + RateLimitReachedType::RateLimitReached => "Usage limit reached".to_string(), + RateLimitReachedType::WorkspaceOwnerCreditsDepleted => { + "Workspace owner credits depleted".to_string() + } + RateLimitReachedType::WorkspaceMemberCreditsDepleted => { + "Workspace member credits depleted".to_string() + } + RateLimitReachedType::WorkspaceOwnerUsageLimitReached => { + "Workspace owner usage limit reached".to_string() + } + RateLimitReachedType::WorkspaceMemberUsageLimitReached => { + "Workspace member usage limit reached".to_string() + } + } + } + fn hourly_usage_lines( summary: Option<&StoredUsageSummary>, is_api_key_account: bool, @@ -16878,12 +16920,12 @@ impl ChatWidget<'_> { let snapshot = self.rate_limit_snapshot.clone(); let needs_refresh = self.should_refresh_limits(); - if self.rate_limit_fetch_inflight || needs_refresh { + let content = self.build_limits_overlay_content(snapshot.clone()); + let has_snapshot = Self::limits_content_has_snapshot(&content); + if (self.rate_limit_fetch_inflight || needs_refresh) && !has_snapshot { self.set_limits_overlay_content(LimitsOverlayContent::Loading); } else { - let reset_info = self.rate_limit_reset_info(); - let tabs = self.build_limits_tabs(snapshot.clone(), reset_info); - self.set_limits_overlay_tabs(tabs); + self.set_limits_overlay_content(content); } self.request_redraw(); @@ -16955,7 +16997,12 @@ impl ChatWidget<'_> { } if show_loading { - self.set_limits_overlay_content(LimitsOverlayContent::Loading); + let content = self.build_limits_overlay_content(self.rate_limit_snapshot.clone()); + if Self::limits_content_has_snapshot(&content) { + self.set_limits_overlay_content(content); + } else { + self.set_limits_overlay_content(LimitsOverlayContent::Loading); + } self.request_redraw(); } @@ -16989,7 +17036,11 @@ impl ChatWidget<'_> { pub(crate) fn on_rate_limit_refresh_failed(&mut self, message: String) { self.rate_limit_fetch_inflight = false; - let content = if self.rate_limit_snapshot.is_some() { + let has_live_snapshot = self.rate_limit_snapshot.is_some(); + let content = self.build_limits_overlay_content(self.rate_limit_snapshot.clone()); + let content = if Self::limits_content_has_snapshot(&content) { + content + } else if has_live_snapshot { LimitsOverlayContent::Error(message.clone()) } else { LimitsOverlayContent::Placeholder @@ -16997,7 +17048,7 @@ impl ChatWidget<'_> { self.set_limits_overlay_content(content); self.request_redraw(); - if self.rate_limit_snapshot.is_some() { + if has_live_snapshot { self.history_push_plain_state(history_cell::new_warning_event(message)); } } @@ -25414,16 +25465,12 @@ Have we met every part of this goal and is there no further work to do?"# let snapshot = self.rate_limit_snapshot.clone(); let needs_refresh = self.should_refresh_limits(); - let content = if self.rate_limit_fetch_inflight || needs_refresh { + let content = self.build_limits_overlay_content(snapshot.clone()); + let has_snapshot = Self::limits_content_has_snapshot(&content); + let content = if (self.rate_limit_fetch_inflight || needs_refresh) && !has_snapshot { LimitsOverlayContent::Loading } else { - let reset_info = self.rate_limit_reset_info(); - let tabs = self.build_limits_tabs(snapshot.clone(), reset_info); - if tabs.is_empty() { - LimitsOverlayContent::Placeholder - } else { - LimitsOverlayContent::Tabs(tabs) - } + content }; if needs_refresh { @@ -31517,6 +31564,7 @@ impl Drop for AutoReviewStubGuard { use crate::chatwidget::message::UserMessage; use crate::chatwidget::smoke_helpers::{enter_test_runtime_guard, ChatWidgetHarness}; use crate::history_cell::{self, ExploreAggregationCell, HistoryCellType}; + use base64::Engine; use code_common::model_presets::ReasoningEffortPreset; use code_auto_drive_core::{ AutoContinueMode, @@ -31592,6 +31640,68 @@ use code_core::protocol::OrderMeta; } } + fn test_rate_limit_snapshot() -> RateLimitSnapshotEvent { + RateLimitSnapshotEvent { + primary_used_percent: 12.0, + secondary_used_percent: 34.0, + primary_to_secondary_ratio_percent: 25.0, + primary_window_minutes: 300, + secondary_window_minutes: 10_080, + primary_reset_after_seconds: Some(600), + secondary_reset_after_seconds: Some(3_600), + rate_limit_reached_type: None, + } + } + + fn write_test_chatgpt_account( + code_home: &std::path::Path, + account_id: &str, + email: &str, + ) { + let payload = serde_json::json!({ + "email": email, + "https://api.openai.com/auth": { + "chatgpt_plan_type": "pro", + "chatgpt_account_id": account_id, + "chatgpt_user_id": "user-12345", + "user_id": "user-12345" + } + }); + let header = serde_json::json!({ "alg": "none", "typ": "JWT" }); + let encode = |value: serde_json::Value| { + base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(serde_json::to_vec(&value).expect("json bytes")) + }; + let id_token = format!( + "{}.{}.{}", + encode(header), + encode(payload), + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b"sig") + ); + let accounts = serde_json::json!({ + "version": 1, + "active_account_id": account_id, + "accounts": [ + { + "id": account_id, + "mode": "chatgpt", + "label": email, + "tokens": { + "id_token": id_token, + "access_token": "access", + "refresh_token": "refresh", + "account_id": account_id + } + } + ] + }); + std::fs::write( + code_home.join("auth_accounts.json"), + serde_json::to_string_pretty(&accounts).expect("accounts json"), + ) + .expect("write auth accounts"); + } + fn expect_post_turn_developer_input(code_op_rx: &mut UnboundedReceiver) -> String { loop { match code_op_rx.try_recv().expect("post-turn developer op") { @@ -32592,6 +32702,52 @@ use code_core::protocol::OrderMeta; chat.maybe_schedule_rate_limit_refresh(); } + #[test] + fn limits_overlay_uses_cached_snapshot_while_refresh_due() { + let _runtime_guard = enter_test_runtime_guard(); + let code_home = tempdir().expect("temp code home"); + let mut harness = ChatWidgetHarness::new(); + harness.with_chat(|chat| { + chat.config.code_home = code_home.path().to_path_buf(); + write_test_chatgpt_account( + &chat.config.code_home, + "limits-account", + "limits@example.com", + ); + account_usage::record_rate_limit_snapshot( + &chat.config.code_home, + "limits-account", + Some("Pro"), + &test_rate_limit_snapshot(), + Utc::now(), + ) + .expect("snapshot stored"); + chat.rate_limit_last_fetch_at = None; + chat.rate_limit_fetch_inflight = false; + + let tabs = chat.build_limits_tabs(None, RateLimitResetInfo::default()); + assert!( + tabs.iter() + .any(|tab| matches!(tab.body, LimitsTabBody::View(_))), + "test fixture should produce a cached snapshot tab" + ); + + chat.show_limits_settings_ui(); + + let content = chat + .settings + .overlay + .as_ref() + .and_then(|overlay| overlay.limits_content()) + .expect("limits content"); + assert!( + content.has_snapshot_view(), + "cached snapshot should remain visible while refresh runs" + ); + assert!(chat.rate_limit_fetch_inflight); + }); + } + #[test] fn apply_context_mode_selection_persists_disabled_override() { let _runtime_guard = enter_test_runtime_guard(); diff --git a/code-rs/tui/src/chatwidget/limits_overlay.rs b/code-rs/tui/src/chatwidget/limits_overlay.rs index fb955e18d3fc..1bda9f41638b 100644 --- a/code-rs/tui/src/chatwidget/limits_overlay.rs +++ b/code-rs/tui/src/chatwidget/limits_overlay.rs @@ -132,6 +132,17 @@ impl LimitsOverlay { } } + #[cfg(any(test, feature = "test-helpers"))] + #[allow(dead_code)] + pub(crate) fn has_snapshot_view(&self) -> bool { + match &self.content { + LimitsOverlayContent::Tabs(tabs) => tabs + .iter() + .any(|tab| matches!(tab.body, LimitsTabBody::View(_))), + _ => false, + } + } + pub(crate) fn lines_for_width(&self, width: u16) -> Vec> { let mut lines = match &self.content { LimitsOverlayContent::Loading => loading_lines(), diff --git a/code-rs/tui/src/chatwidget/settings_overlay.rs b/code-rs/tui/src/chatwidget/settings_overlay.rs index df9977c8150f..a14fe3d40fc7 100644 --- a/code-rs/tui/src/chatwidget/settings_overlay.rs +++ b/code-rs/tui/src/chatwidget/settings_overlay.rs @@ -992,6 +992,12 @@ impl LimitsSettingsContent { self.overlay.set_content(content); } + #[cfg(any(test, feature = "test-helpers"))] + #[allow(dead_code)] + pub(crate) fn has_snapshot_view(&self) -> bool { + self.overlay.has_snapshot_view() + } + fn render_tabs(&self, area: Rect, buf: &mut Buffer) { use ratatui::widgets::Paragraph; @@ -1525,6 +1531,12 @@ impl SettingsOverlayView { self.limits_content.as_mut() } + #[cfg(any(test, feature = "test-helpers"))] + #[allow(dead_code)] + pub(crate) fn limits_content(&self) -> Option<&LimitsSettingsContent> { + self.limits_content.as_ref() + } + pub(crate) fn set_section(&mut self, section: SettingsSection) -> bool { if self.active_section() == section { return false; diff --git a/code-rs/tui/src/rate_limits_view.rs b/code-rs/tui/src/rate_limits_view.rs index 4ea7215ef200..7357251c10cb 100644 --- a/code-rs/tui/src/rate_limits_view.rs +++ b/code-rs/tui/src/rate_limits_view.rs @@ -966,6 +966,7 @@ mod tests { secondary_window_minutes: 60, primary_reset_after_seconds: Some(900), secondary_reset_after_seconds: Some(3600), + rate_limit_reached_type: None, } }