diff --git a/codex-rs/app-server-protocol/src/protocol.rs b/codex-rs/app-server-protocol/src/protocol.rs index b4cd358b22..bca57c7233 100644 --- a/codex-rs/app-server-protocol/src/protocol.rs +++ b/codex-rs/app-server-protocol/src/protocol.rs @@ -875,6 +875,11 @@ pub struct AuthStatusChangeNotification { #[serde(tag = "method", content = "params", rename_all = "camelCase")] #[strum(serialize_all = "camelCase")] pub enum ServerNotification { + #[serde(rename = "account/rateLimits/updated")] + #[ts(rename = "account/rateLimits/updated")] + #[strum(serialize = "account/rateLimits/updated")] + AccountRateLimitsUpdated(RateLimitSnapshot), + /// Authentication status changed AuthStatusChange(AuthStatusChangeNotification), @@ -888,6 +893,7 @@ pub enum ServerNotification { impl ServerNotification { pub fn to_params(self) -> Result { match self { + ServerNotification::AccountRateLimitsUpdated(params) => serde_json::to_value(params), ServerNotification::AuthStatusChange(params) => serde_json::to_value(params), ServerNotification::LoginChatGptComplete(params) => serde_json::to_value(params), ServerNotification::SessionConfigured(params) => serde_json::to_value(params), diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 39992e8391..f256b78d1a 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -1436,6 +1436,15 @@ async fn apply_bespoke_event_handling( on_exec_approval_response(event_id, rx, conversation).await; }); } + EventMsg::TokenCount(token_count_event) => { + if let Some(rate_limits) = token_count_event.rate_limits { + outgoing + .send_server_notification(ServerNotification::AccountRateLimitsUpdated( + rate_limits, + )) + .await; + } + } // If this is a TurnAborted, reply to any pending interrupt requests. EventMsg::TurnAborted(turn_aborted_event) => { let pending = { diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index 96a2c5a96c..feaa19e37a 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -142,6 +142,8 @@ pub(crate) struct OutgoingError { #[cfg(test)] mod tests { use codex_app_server_protocol::LoginChatGptCompleteNotification; + use codex_protocol::protocol::RateLimitSnapshot; + use codex_protocol::protocol::RateLimitWindow; use pretty_assertions::assert_eq; use serde_json::json; use uuid::Uuid; @@ -171,4 +173,34 @@ mod tests { "ensure the strum macros serialize the method field correctly" ); } + + #[test] + fn verify_account_rate_limits_notification_serialization() { + let notification = ServerNotification::AccountRateLimitsUpdated(RateLimitSnapshot { + primary: Some(RateLimitWindow { + used_percent: 25.0, + window_minutes: Some(15), + resets_at: Some(123), + }), + secondary: None, + }); + + let jsonrpc_notification = OutgoingMessage::AppServerNotification(notification); + assert_eq!( + json!({ + "method": "account/rateLimits/updated", + "params": { + "primary": { + "used_percent": 25.0, + "window_minutes": 15, + "resets_at": 123, + }, + "secondary": null, + }, + }), + serde_json::to_value(jsonrpc_notification) + .expect("ensure the notification serializes correctly"), + "ensure the notification serializes correctly" + ); + } }