From 9ce2a2e99284a289fa64baff85477a83c4110622 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Tue, 21 Oct 2025 11:46:28 -0700 Subject: [PATCH] [app-server] account rate limits updated event --- codex-rs/app-server-protocol/src/protocol.rs | 5 + .../app-server/src/codex_message_processor.rs | 143 ++++++++++++++++++ codex-rs/app-server/tests/suite/user_agent.rs | 9 +- codex-rs/docs/codex_mcp_interface.md | 3 +- 4 files changed, 156 insertions(+), 4 deletions(-) diff --git a/codex-rs/app-server-protocol/src/protocol.rs b/codex-rs/app-server-protocol/src/protocol.rs index b4cd358b22..c6d4fe8842 100644 --- a/codex-rs/app-server-protocol/src/protocol.rs +++ b/codex-rs/app-server-protocol/src/protocol.rs @@ -883,6 +883,10 @@ pub enum ServerNotification { /// The special session configured event for a new or resumed conversation. SessionConfigured(SessionConfiguredNotification), + #[serde(rename = "account/rateLimits/updated")] + #[ts(rename = "account/rateLimits/updated")] + #[strum(serialize = "account/rateLimits/updated")] + AccountRateLimitsUpdated(RateLimitSnapshot), } impl ServerNotification { @@ -891,6 +895,7 @@ impl ServerNotification { ServerNotification::AuthStatusChange(params) => serde_json::to_value(params), ServerNotification::LoginChatGptComplete(params) => serde_json::to_value(params), ServerNotification::SessionConfigured(params) => serde_json::to_value(params), + ServerNotification::AccountRateLimitsUpdated(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 a5dd0d92fc..87f1598ae0 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -137,6 +137,7 @@ pub(crate) struct CodexMessageProcessor { // Queue of pending interrupt requests per conversation. We reply when TurnAborted arrives. pending_interrupts: Arc>>>, pending_fuzzy_searches: Arc>>>, + latest_rate_limits: Arc>>, } impl CodexMessageProcessor { @@ -157,6 +158,7 @@ impl CodexMessageProcessor { active_login: Arc::new(Mutex::new(None)), pending_interrupts: Arc::new(Mutex::new(HashMap::new())), pending_fuzzy_searches: Arc::new(Mutex::new(HashMap::new())), + latest_rate_limits: Arc::new(Mutex::new(None)), } } @@ -1257,6 +1259,7 @@ impl CodexMessageProcessor { .insert(subscription_id, cancel_tx); let outgoing_for_task = self.outgoing.clone(); let pending_interrupts = self.pending_interrupts.clone(); + let latest_rate_limits = self.latest_rate_limits.clone(); tokio::spawn(async move { loop { tokio::select! { @@ -1298,6 +1301,7 @@ impl CodexMessageProcessor { .await; apply_bespoke_event_handling(event.clone(), conversation_id, conversation.clone(), outgoing_for_task.clone(), pending_interrupts.clone()).await; + apply_rate_limit_notification(event, outgoing_for_task.clone(), latest_rate_limits.clone()).await; } } } @@ -1466,6 +1470,39 @@ async fn apply_bespoke_event_handling( } } +async fn apply_rate_limit_notification( + event: Event, + outgoing: Arc, + latest_rate_limits: Arc>>, +) -> bool { + let Event { msg, .. } = event; + let EventMsg::TokenCount(token_count) = msg else { + return false; + }; + let Some(snapshot) = token_count.rate_limits else { + return false; + }; + + // Only notify when the snapshot changes to avoid redundant traffic. + let should_send = { + let mut guard = latest_rate_limits.lock().await; + if guard.as_ref() == Some(&snapshot) { + false + } else { + *guard = Some(snapshot.clone()); + true + } + }; + + if should_send { + outgoing + .send_server_notification(ServerNotification::AccountRateLimitsUpdated(snapshot)) + .await; + } + + should_send +} + async fn derive_config_from_params( params: NewConversationParams, codex_linux_sandbox_exe: Option, @@ -1633,9 +1670,16 @@ fn extract_conversation_summary( #[cfg(test)] mod tests { use super::*; + use anyhow::Context as _; use anyhow::Result; use pretty_assertions::assert_eq; use serde_json::json; + use std::sync::Arc; + use tokio::sync::mpsc; + + use crate::outgoing_message::OutgoingMessage; + use codex_protocol::protocol::RateLimitWindow; + use codex_protocol::protocol::TokenCountEvent; #[test] fn extract_conversation_summary_prefers_plain_user_messages() -> Result<()> { @@ -1681,4 +1725,103 @@ mod tests { assert_eq!(summary.preview, "Count to 5"); Ok(()) } + + #[tokio::test] + async fn apply_rate_limit_notification_emits_on_change_only() -> Result<()> { + let (tx, mut rx) = mpsc::unbounded_channel(); + let outgoing = Arc::new(OutgoingMessageSender::new(tx)); + let latest_rate_limits = Arc::new(Mutex::new(None)); + + let snapshot_one = RateLimitSnapshot { + primary: Some(RateLimitWindow { + used_percent: 10.0, + window_minutes: Some(60), + resets_at: Some(1_111), + }), + secondary: None, + }; + let event_one = Event { + id: "evt-1".to_string(), + msg: EventMsg::TokenCount(TokenCountEvent { + info: None, + rate_limits: Some(snapshot_one.clone()), + }), + }; + + assert!( + apply_rate_limit_notification(event_one, outgoing.clone(), latest_rate_limits.clone()) + .await + ); + + let message = rx.recv().await.context("receive first notification")?; + match message { + OutgoingMessage::AppServerNotification( + ServerNotification::AccountRateLimitsUpdated(received), + ) => assert_eq!(received, snapshot_one), + other => panic!("unexpected message: {other:?}"), + } + + let event_dup = Event { + id: "evt-2".to_string(), + msg: EventMsg::TokenCount(TokenCountEvent { + info: None, + rate_limits: Some(snapshot_one.clone()), + }), + }; + + assert!( + !apply_rate_limit_notification(event_dup, outgoing.clone(), latest_rate_limits.clone()) + .await + ); + + if let Ok(unexpected) = rx.try_recv() { + panic!("unexpected extra message after second snapshot: {unexpected:?}"); + } + + let snapshot_two = RateLimitSnapshot { + primary: Some(RateLimitWindow { + used_percent: 20.0, + window_minutes: Some(120), + resets_at: Some(2_222), + }), + secondary: None, + }; + let event_two = Event { + id: "evt-3".to_string(), + msg: EventMsg::TokenCount(TokenCountEvent { + info: None, + rate_limits: Some(snapshot_two.clone()), + }), + }; + + assert!( + apply_rate_limit_notification( + event_two.clone(), + outgoing.clone(), + latest_rate_limits.clone(), + ) + .await + ); + + let message = rx.recv().await.context("receive second notification")?; + match message { + OutgoingMessage::AppServerNotification( + ServerNotification::AccountRateLimitsUpdated(received), + ) => assert_eq!(received, snapshot_two), + other => panic!("unexpected message: {other:?}"), + } + + if let Ok(unexpected) = rx.try_recv() { + panic!("unexpected extra message after duplicate snapshot: {unexpected:?}"); + } + + // Re-emitting the same snapshot, even if a different conversation produced it, + // should be suppressed by the global tracker. + assert!(!apply_rate_limit_notification(event_two, outgoing, latest_rate_limits).await); + if let Ok(unexpected) = rx.try_recv() { + panic!("unexpected extra message after duplicate snapshot: {unexpected:?}"); + } + + Ok(()) + } } diff --git a/codex-rs/app-server/tests/suite/user_agent.rs b/codex-rs/app-server/tests/suite/user_agent.rs index 95a0b1a3e0..4d1156527c 100644 --- a/codex-rs/app-server/tests/suite/user_agent.rs +++ b/codex-rs/app-server/tests/suite/user_agent.rs @@ -13,9 +13,12 @@ const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs async fn get_user_agent_returns_current_codex_user_agent() { let codex_home = TempDir::new().unwrap_or_else(|err| panic!("create tempdir: {err}")); - let mut mcp = McpProcess::new(codex_home.path()) - .await - .expect("spawn mcp process"); + let mut mcp = McpProcess::new_with_env( + codex_home.path(), + &[("CODEX_INTERNAL_ORIGINATOR_OVERRIDE", Some("codex_cli_rs"))], + ) + .await + .expect("spawn mcp process"); timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()) .await .expect("initialize timeout") diff --git a/codex-rs/docs/codex_mcp_interface.md b/codex-rs/docs/codex_mcp_interface.md index aaaa0f4cad..32423735e5 100644 --- a/codex-rs/docs/codex_mcp_interface.md +++ b/codex-rs/docs/codex_mcp_interface.md @@ -27,7 +27,7 @@ At a glance: - Approvals (server → client requests) - `applyPatchApproval`, `execCommandApproval` - Notifications (server → client) - - `loginChatGptComplete`, `authStatusChange` + - `loginChatGptComplete`, `authStatusChange`, `account/rateLimits/updated` - `codex/event` stream with agent events See code for full type definitions and exact shapes: `protocol/src/mcp_protocol.rs`. @@ -97,6 +97,7 @@ Each response yields: While a conversation runs, the server sends notifications: - `codex/event` with the serialized Codex event payload. The shape matches `core/src/protocol.rs`’s `Event` and `EventMsg` types. Some notifications include a `_meta.requestId` to correlate with the originating request. +- `account/rateLimits/updated` whenever Codex observes a new rate-limit snapshot for the active account. The payload matches `RateLimitSnapshot`. - Auth notifications via method names `loginChatGptComplete` and `authStatusChange`. Clients should render events and, when present, surface approval requests (see next section).