[app-server] send account/rateLimits/updated notifications (#5477)
Codex will now send an `account/rateLimits/updated` notification whenever the user's rate limits are updated. This is implemented by just transforming the existing TokenCount event.
This commit is contained in:
@@ -875,6 +875,11 @@ pub struct AuthStatusChangeNotification {
|
|||||||
#[serde(tag = "method", content = "params", rename_all = "camelCase")]
|
#[serde(tag = "method", content = "params", rename_all = "camelCase")]
|
||||||
#[strum(serialize_all = "camelCase")]
|
#[strum(serialize_all = "camelCase")]
|
||||||
pub enum ServerNotification {
|
pub enum ServerNotification {
|
||||||
|
#[serde(rename = "account/rateLimits/updated")]
|
||||||
|
#[ts(rename = "account/rateLimits/updated")]
|
||||||
|
#[strum(serialize = "account/rateLimits/updated")]
|
||||||
|
AccountRateLimitsUpdated(RateLimitSnapshot),
|
||||||
|
|
||||||
/// Authentication status changed
|
/// Authentication status changed
|
||||||
AuthStatusChange(AuthStatusChangeNotification),
|
AuthStatusChange(AuthStatusChangeNotification),
|
||||||
|
|
||||||
@@ -888,6 +893,7 @@ pub enum ServerNotification {
|
|||||||
impl ServerNotification {
|
impl ServerNotification {
|
||||||
pub fn to_params(self) -> Result<serde_json::Value, serde_json::Error> {
|
pub fn to_params(self) -> Result<serde_json::Value, serde_json::Error> {
|
||||||
match self {
|
match self {
|
||||||
|
ServerNotification::AccountRateLimitsUpdated(params) => serde_json::to_value(params),
|
||||||
ServerNotification::AuthStatusChange(params) => serde_json::to_value(params),
|
ServerNotification::AuthStatusChange(params) => serde_json::to_value(params),
|
||||||
ServerNotification::LoginChatGptComplete(params) => serde_json::to_value(params),
|
ServerNotification::LoginChatGptComplete(params) => serde_json::to_value(params),
|
||||||
ServerNotification::SessionConfigured(params) => serde_json::to_value(params),
|
ServerNotification::SessionConfigured(params) => serde_json::to_value(params),
|
||||||
|
|||||||
@@ -1436,6 +1436,15 @@ async fn apply_bespoke_event_handling(
|
|||||||
on_exec_approval_response(event_id, rx, conversation).await;
|
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.
|
// If this is a TurnAborted, reply to any pending interrupt requests.
|
||||||
EventMsg::TurnAborted(turn_aborted_event) => {
|
EventMsg::TurnAborted(turn_aborted_event) => {
|
||||||
let pending = {
|
let pending = {
|
||||||
|
|||||||
@@ -142,6 +142,8 @@ pub(crate) struct OutgoingError {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use codex_app_server_protocol::LoginChatGptCompleteNotification;
|
use codex_app_server_protocol::LoginChatGptCompleteNotification;
|
||||||
|
use codex_protocol::protocol::RateLimitSnapshot;
|
||||||
|
use codex_protocol::protocol::RateLimitWindow;
|
||||||
use pretty_assertions::assert_eq;
|
use pretty_assertions::assert_eq;
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
@@ -171,4 +173,34 @@ mod tests {
|
|||||||
"ensure the strum macros serialize the method field correctly"
|
"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"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user