Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions codex-rs/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions codex-rs/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ codex-app-server-protocol = { path = "app-server-protocol" }
codex-apply-patch = { path = "apply-patch" }
codex-arg0 = { path = "arg0" }
codex-async-utils = { path = "async-utils" }
codex-backend-client = { path = "backend-client" }
codex-chatgpt = { path = "chatgpt" }
codex-common = { path = "common" }
codex-core = { path = "core" }
Expand Down
29 changes: 29 additions & 0 deletions codex-rs/app-server-protocol/src/protocol.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ use codex_protocol::parse_command::ParsedCommand;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::FileChange;
use codex_protocol::protocol::RateLimitSnapshot;
use codex_protocol::protocol::ReviewDecision;
use codex_protocol::protocol::SandboxPolicy;
use codex_protocol::protocol::TurnAbortReason;
Expand Down Expand Up @@ -183,6 +184,12 @@ client_request_definitions! {
params: ExecOneOffCommandParams,
response: ExecOneOffCommandResponse,
},
#[serde(rename = "account/rateLimits/read")]
#[ts(rename = "account/rateLimits/read")]
GetAccountRateLimits {
params: #[ts(type = "undefined")] #[serde(skip_serializing_if = "Option::is_none")] Option<()>,
response: GetAccountRateLimitsResponse,
},
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema, TS)]
Expand Down Expand Up @@ -420,6 +427,12 @@ pub struct ExecOneOffCommandResponse {
pub stderr: String,
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
pub struct GetAccountRateLimitsResponse {
pub rate_limits: RateLimitSnapshot,
}

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
pub struct GetAuthStatusResponse {
Expand Down Expand Up @@ -970,4 +983,20 @@ mod tests {
assert_eq!(payload.request_with_id(RequestId::Integer(7)), request);
Ok(())
}

#[test]
fn serialize_get_account_rate_limits() -> Result<()> {
let request = ClientRequest::GetAccountRateLimits {
request_id: RequestId::Integer(1),
params: None,
};
assert_eq!(
json!({
"method": "account/rateLimits/read",
"id": 1,
}),
serde_json::to_value(&request)?,
);
Ok(())
}
}
2 changes: 2 additions & 0 deletions codex-rs/app-server/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,13 @@ anyhow = { workspace = true }
codex-arg0 = { workspace = true }
codex-common = { workspace = true, features = ["cli"] }
codex-core = { workspace = true }
codex-backend-client = { workspace = true }
codex-file-search = { workspace = true }
codex-login = { workspace = true }
codex-protocol = { workspace = true }
codex-app-server-protocol = { workspace = true }
codex-utils-json-to-toml = { workspace = true }
chrono = { workspace = true }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
tokio = { workspace = true, features = [
Expand Down
57 changes: 57 additions & 0 deletions codex-rs/app-server/src/codex_message_processor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ use codex_app_server_protocol::ApplyPatchApprovalParams;
use codex_app_server_protocol::ApplyPatchApprovalResponse;
use codex_app_server_protocol::ArchiveConversationParams;
use codex_app_server_protocol::ArchiveConversationResponse;
use codex_app_server_protocol::AuthMode;
use codex_app_server_protocol::AuthStatusChangeNotification;
use codex_app_server_protocol::ClientRequest;
use codex_app_server_protocol::ConversationSummary;
Expand All @@ -18,6 +19,7 @@ use codex_app_server_protocol::ExecOneOffCommandParams;
use codex_app_server_protocol::ExecOneOffCommandResponse;
use codex_app_server_protocol::FuzzyFileSearchParams;
use codex_app_server_protocol::FuzzyFileSearchResponse;
use codex_app_server_protocol::GetAccountRateLimitsResponse;
use codex_app_server_protocol::GetUserAgentResponse;
use codex_app_server_protocol::GetUserSavedConfigResponse;
use codex_app_server_protocol::GitDiffToRemoteResponse;
Expand Down Expand Up @@ -49,6 +51,7 @@ use codex_app_server_protocol::SetDefaultModelParams;
use codex_app_server_protocol::SetDefaultModelResponse;
use codex_app_server_protocol::UserInfoResponse;
use codex_app_server_protocol::UserSavedConfig;
use codex_backend_client::Client as BackendClient;
use codex_core::AuthManager;
use codex_core::CodexConversation;
use codex_core::ConversationManager;
Expand Down Expand Up @@ -88,6 +91,7 @@ use codex_protocol::config_types::ForcedLoginMethod;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::InputMessageKind;
use codex_protocol::protocol::RateLimitSnapshot;
use codex_protocol::protocol::USER_MESSAGE_BEGIN;
use codex_utils_json_to_toml::json_to_toml;
use std::collections::HashMap;
Expand Down Expand Up @@ -240,6 +244,12 @@ impl CodexMessageProcessor {
ClientRequest::ExecOneOffCommand { request_id, params } => {
self.exec_one_off_command(request_id, params).await;
}
ClientRequest::GetAccountRateLimits {
request_id,
params: _,
} => {
self.get_account_rate_limits(request_id).await;
}
}
}

Expand Down Expand Up @@ -527,6 +537,53 @@ impl CodexMessageProcessor {
self.outgoing.send_response(request_id, response).await;
}

async fn get_account_rate_limits(&self, request_id: RequestId) {
match self.fetch_account_rate_limits().await {
Ok(rate_limits) => {
let response = GetAccountRateLimitsResponse { rate_limits };
self.outgoing.send_response(request_id, response).await;
}
Err(error) => {
self.outgoing.send_error(request_id, error).await;
}
}
}

async fn fetch_account_rate_limits(&self) -> Result<RateLimitSnapshot, JSONRPCErrorError> {
let Some(auth) = self.auth_manager.auth() else {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we have some helper for internal errors? We use them quite a lot and they are verbose.

Copy link
Contributor Author

@owenlin0 owenlin0 Oct 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

return Err(JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: "codex account authentication required to read rate limits".to_string(),
data: None,
});
};

if auth.mode != AuthMode::ChatGPT {
return Err(JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: "chatgpt authentication required to read rate limits".to_string(),
data: None,
});
}

let client = BackendClient::from_auth(self.config.chatgpt_base_url.clone(), &auth)
.await
.map_err(|err| JSONRPCErrorError {
code: INTERNAL_ERROR_CODE,
message: format!("failed to construct backend client: {err}"),
data: None,
})?;

client
.get_rate_limits()
.await
.map_err(|err| JSONRPCErrorError {
code: INTERNAL_ERROR_CODE,
message: format!("failed to fetch codex rate limits: {err}"),
data: None,
})
}

async fn get_user_saved_config(&self, request_id: RequestId) {
let toml_value = match load_config_as_toml(&self.config.codex_home).await {
Ok(val) => val,
Expand Down
3 changes: 3 additions & 0 deletions codex-rs/app-server/tests/common/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,10 @@ path = "lib.rs"
[dependencies]
anyhow = { workspace = true }
assert_cmd = { workspace = true }
base64 = { workspace = true }
chrono = { workspace = true }
codex-app-server-protocol = { workspace = true }
codex-core = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true, features = [
Expand Down
131 changes: 131 additions & 0 deletions codex-rs/app-server/tests/common/auth_fixtures.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,131 @@
use std::path::Path;

use anyhow::Context;
use anyhow::Result;
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use chrono::DateTime;
use chrono::Utc;
use codex_core::auth::AuthDotJson;
use codex_core::auth::get_auth_file;
use codex_core::auth::write_auth_json;
use codex_core::token_data::TokenData;
use codex_core::token_data::parse_id_token;
use serde_json::json;

/// Builder for writing a fake ChatGPT auth.json in tests.
#[derive(Debug, Clone)]
pub struct ChatGptAuthFixture {
access_token: String,
refresh_token: String,
account_id: Option<String>,
claims: ChatGptIdTokenClaims,
last_refresh: Option<Option<DateTime<Utc>>>,
}

impl ChatGptAuthFixture {
pub fn new(access_token: impl Into<String>) -> Self {
Self {
access_token: access_token.into(),
refresh_token: "refresh-token".to_string(),
account_id: None,
claims: ChatGptIdTokenClaims::default(),
last_refresh: None,
}
}

pub fn refresh_token(mut self, refresh_token: impl Into<String>) -> Self {
self.refresh_token = refresh_token.into();
self
}

pub fn account_id(mut self, account_id: impl Into<String>) -> Self {
self.account_id = Some(account_id.into());
self
}

pub fn plan_type(mut self, plan_type: impl Into<String>) -> Self {
self.claims.plan_type = Some(plan_type.into());
self
}

pub fn email(mut self, email: impl Into<String>) -> Self {
self.claims.email = Some(email.into());
self
}

pub fn last_refresh(mut self, last_refresh: Option<DateTime<Utc>>) -> Self {
self.last_refresh = Some(last_refresh);
self
}

pub fn claims(mut self, claims: ChatGptIdTokenClaims) -> Self {
self.claims = claims;
self
}
}

#[derive(Debug, Clone, Default)]
pub struct ChatGptIdTokenClaims {
pub email: Option<String>,
pub plan_type: Option<String>,
}

impl ChatGptIdTokenClaims {
pub fn new() -> Self {
Self::default()
}

pub fn email(mut self, email: impl Into<String>) -> Self {
self.email = Some(email.into());
self
}

pub fn plan_type(mut self, plan_type: impl Into<String>) -> Self {
self.plan_type = Some(plan_type.into());
self
}
}

pub fn encode_id_token(claims: &ChatGptIdTokenClaims) -> Result<String> {
let header = json!({ "alg": "none", "typ": "JWT" });
let mut payload = serde_json::Map::new();
if let Some(email) = &claims.email {
payload.insert("email".to_string(), json!(email));
}
if let Some(plan_type) = &claims.plan_type {
payload.insert(
"https://api.openai.com/auth".to_string(),
json!({ "chatgpt_plan_type": plan_type }),
);
}
let payload = serde_json::Value::Object(payload);

let header_b64 =
URL_SAFE_NO_PAD.encode(serde_json::to_vec(&header).context("serialize jwt header")?);
let payload_b64 =
URL_SAFE_NO_PAD.encode(serde_json::to_vec(&payload).context("serialize jwt payload")?);
let signature_b64 = URL_SAFE_NO_PAD.encode(b"signature");
Ok(format!("{header_b64}.{payload_b64}.{signature_b64}"))
}

pub fn write_chatgpt_auth(codex_home: &Path, fixture: ChatGptAuthFixture) -> Result<()> {
let id_token_raw = encode_id_token(&fixture.claims)?;
let id_token = parse_id_token(&id_token_raw).context("parse id token")?;
let tokens = TokenData {
id_token,
access_token: fixture.access_token,
refresh_token: fixture.refresh_token,
account_id: fixture.account_id,
};

let last_refresh = fixture.last_refresh.unwrap_or_else(|| Some(Utc::now()));

let auth = AuthDotJson {
openai_api_key: None,
tokens: Some(tokens),
last_refresh,
};

write_auth_json(&get_auth_file(codex_home), &auth).context("write auth.json")
}
5 changes: 5 additions & 0 deletions codex-rs/app-server/tests/common/lib.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
mod auth_fixtures;
mod mcp_process;
mod mock_model_server;
mod responses;

pub use auth_fixtures::ChatGptAuthFixture;
pub use auth_fixtures::ChatGptIdTokenClaims;
pub use auth_fixtures::encode_id_token;
pub use auth_fixtures::write_chatgpt_auth;
use codex_app_server_protocol::JSONRPCResponse;
pub use mcp_process::McpProcess;
pub use mock_model_server::create_mock_chat_completions_server;
Expand Down
5 changes: 5 additions & 0 deletions codex-rs/app-server/tests/common/mcp_process.rs
Original file line number Diff line number Diff line change
Expand Up @@ -236,6 +236,11 @@ impl McpProcess {
self.send_request("getUserAgent", None).await
}

/// Send an `account/rateLimits/read` JSON-RPC request.
pub async fn send_get_account_rate_limits_request(&mut self) -> anyhow::Result<i64> {
self.send_request("account/rateLimits/read", None).await
}

/// Send a `userInfo` JSON-RPC request.
pub async fn send_user_info_request(&mut self) -> anyhow::Result<i64> {
self.send_request("userInfo", None).await
Expand Down
1 change: 1 addition & 0 deletions codex-rs/app-server/tests/suite/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ mod fuzzy_file_search;
mod interrupt;
mod list_resume;
mod login;
mod rate_limits;
mod send_message;
mod set_default_model;
mod user_agent;
Expand Down
Loading
Loading