Upstream merge: openai/codex@main into main

codex-rs/core/src/chat_completions.rs

@@ -869,6 +869,9 @@ where
                     // Not an assistant message – forward immediately.
                     return Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone { item, sequence_number: None, output_index: None })));
                 }
+                Poll::Ready(Some(Ok(ResponseEvent::RateLimits(snapshot)))) => {
+                    return Poll::Ready(Some(Ok(ResponseEvent::RateLimits(snapshot))));
+                }
                 Poll::Ready(Some(Ok(ResponseEvent::Completed {
                     response_id,
                     token_usage,

codex-rs/core/src/client.rs

@@ -11,6 +11,7 @@ use eventsource_stream::Eventsource;
 use futures::prelude::*;
 // use regex_lite::Regex;
 use reqwest::StatusCode;
+use reqwest::header::HeaderMap;
 use serde::Deserialize;
 use serde::Serialize;
 use serde_json::Value;
@@ -44,6 +45,7 @@ use crate::model_family::ModelFamily;
 use crate::model_provider_info::ModelProviderInfo;
 use crate::model_provider_info::WireApi;
 use crate::openai_tools::create_tools_json_for_responses_api;
+use codex_protocol::protocol::RateLimitSnapshotEvent;
 use crate::protocol::TokenUsage;
 use crate::util::backoff;
 use std::sync::Arc;
@@ -359,6 +361,15 @@ impl ModelClient {
                     }
                     let (tx_event, rx_event) = mpsc::channel::<Result<ResponseEvent>>(1600);
 
+                    if let Some(snapshot) = parse_rate_limit_snapshot(resp.headers())
+                        && tx_event
+                            .send(Ok(ResponseEvent::RateLimits(snapshot)))
+                            .await
+                            .is_err()
+                    {
+                        debug!("receiver dropped rate limit snapshot event");
+                    }
+
                     // spawn task to process SSE
                     let stream = resp.bytes_stream().map_err(CodexErr::Reqwest);
                     let debug_logger = Arc::clone(&self.debug_logger);
@@ -633,6 +644,38 @@ fn attach_item_ids(payload_json: &mut Value, original_items: &[ResponseItem]) {
     }
 }
 
+fn parse_rate_limit_snapshot(headers: &HeaderMap) -> Option<RateLimitSnapshotEvent> {
+    let primary_used_percent = parse_header_f64(headers, "x-codex-primary-used-percent")?;
+    let weekly_used_percent = parse_header_f64(headers, "x-codex-protection-used-percent")?;
+    let primary_to_weekly_ratio_percent =
+        parse_header_f64(headers, "x-codex-primary-over-protection-limit-percent")?;
+    let primary_window_minutes = parse_header_u64(headers, "x-codex-primary-window-minutes")?;
+    let weekly_window_minutes = parse_header_u64(headers, "x-codex-protection-window-minutes")?;
+
+    Some(RateLimitSnapshotEvent {
+        primary_used_percent,
+        weekly_used_percent,
+        primary_to_weekly_ratio_percent,
+        primary_window_minutes,
+        weekly_window_minutes,
+    })
+}
+
+fn parse_header_f64(headers: &HeaderMap, name: &str) -> Option<f64> {
+    parse_header_str(headers, name)?
+        .parse::<f64>()
+        .ok()
+        .filter(|v| v.is_finite())
+}
+
+fn parse_header_u64(headers: &HeaderMap, name: &str) -> Option<u64> {
+    parse_header_str(headers, name)?.parse::<u64>().ok()
+}
+
+fn parse_header_str<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> {
+    headers.get(name)?.to_str().ok()
+}
+
 async fn process_sse<S>(
     stream: S,
     tx_event: mpsc::Sender<Result<ResponseEvent>>,

codex-rs/core/src/client_common.rs

@@ -5,6 +5,7 @@ use crate::environment_context::EnvironmentContext;
 use crate::error::Result;
 use crate::model_family::ModelFamily;
 use crate::openai_tools::OpenAiTool;
+use codex_protocol::protocol::RateLimitSnapshotEvent;
 use crate::protocol::TokenUsage;
 use codex_apply_patch::APPLY_PATCH_TOOL_INSTRUCTIONS;
 use codex_protocol::models::ContentItem;
@@ -251,6 +252,7 @@ pub enum ResponseEvent {
         call_id: String,
         query: Option<String>,
     },
+    RateLimits(RateLimitSnapshotEvent),
 }
 
 #[derive(Debug, Serialize)]

codex-rs/core/src/codex.rs

@@ -3268,6 +3268,9 @@ async fn try_run_turn(
                     sess.tx_event.send(stamped).await.ok();
                 }
             }
+            ResponseEvent::RateLimits(_rl) => {
+                // Upstream emits periodic rate limit snapshots. No-op for now.
+            }
             // Note: ReasoningSummaryPartAdded handled above without scratchpad mutation.
         }
     }

codex-rs/core/src/exec.rs

@@ -53,7 +53,7 @@ const AGGREGATE_BUFFER_INITIAL_CAPACITY: usize = 8 * 1024; // 8 KiB
 /// Aggregation still collects full output; only the live event stream is capped.
 pub(crate) const MAX_EXEC_OUTPUT_DELTAS_PER_CALL: usize = 10_000;
 
-#[derive(Debug, Clone)]
+#[derive(Clone, Debug)]
 pub struct ExecParams {
     pub command: Vec<String>,
     pub cwd: PathBuf,

codex-rs/core/tests/common/Cargo.toml

@@ -11,3 +11,4 @@ codex-core = { path = "../.." }
 serde_json = "1"
 tempfile = "3"
 tokio = { version = "1", features = ["time"] }
+wiremock = "0.6"

codex-rs/core/tests/common/lib.rs

@@ -7,6 +7,8 @@ use codex_core::config::Config;
 use codex_core::config::ConfigOverrides;
 use codex_core::config::ConfigToml;
 
+pub mod responses;
+
 /// Returns a default `Config` whose on-disk state is confined to the provided
 /// temporary directory. Using a per-test directory keeps tests hermetic and
 /// avoids clobbering a developer’s real `~/.code` directory (legacy `~/.codex`

codex-rs/core/tests/common/responses.rs

@@ -0,0 +1,133 @@
+use serde_json::Value;
+use wiremock::BodyPrintLimit;
+use wiremock::Mock;
+use wiremock::MockServer;
+use wiremock::ResponseTemplate;
+use wiremock::matchers::method;
+use wiremock::matchers::path;
+
+/// Build an SSE stream body from a list of JSON events.
+pub fn sse(events: Vec<Value>) -> String {
+    use std::fmt::Write as _;
+    let mut out = String::new();
+    for ev in events {
+        let kind = ev.get("type").and_then(|v| v.as_str()).unwrap();
+        writeln!(&mut out, "event: {kind}").unwrap();
+        if !ev.as_object().map(|o| o.len() == 1).unwrap_or(false) {
+            write!(&mut out, "data: {ev}\n\n").unwrap();
+        } else {
+            out.push('\n');
+        }
+    }
+    out
+}
+
+/// Convenience: SSE event for a completed response with a specific id.
+pub fn ev_completed(id: &str) -> Value {
+    serde_json::json!({
+        "type": "response.completed",
+        "response": {
+            "id": id,
+            "usage": {"input_tokens":0,"input_tokens_details":null,"output_tokens":0,"output_tokens_details":null,"total_tokens":0}
+        }
+    })
+}
+
+pub fn ev_completed_with_tokens(id: &str, total_tokens: u64) -> Value {
+    serde_json::json!({
+        "type": "response.completed",
+        "response": {
+            "id": id,
+            "usage": {
+                "input_tokens": total_tokens,
+                "input_tokens_details": null,
+                "output_tokens": 0,
+                "output_tokens_details": null,
+                "total_tokens": total_tokens
+            }
+        }
+    })
+}
+
+/// Convenience: SSE event for a single assistant message output item.
+pub fn ev_assistant_message(id: &str, text: &str) -> Value {
+    serde_json::json!({
+        "type": "response.output_item.done",
+        "item": {
+            "type": "message",
+            "role": "assistant",
+            "id": id,
+            "content": [{"type": "output_text", "text": text}]
+        }
+    })
+}
+
+pub fn ev_function_call(call_id: &str, name: &str, arguments: &str) -> Value {
+    serde_json::json!({
+        "type": "response.output_item.done",
+        "item": {
+            "type": "function_call",
+            "call_id": call_id,
+            "name": name,
+            "arguments": arguments
+        }
+    })
+}
+
+/// Convenience: SSE event for an `apply_patch` custom tool call with raw patch
+/// text. This mirrors the payload produced by the Responses API when the model
+/// invokes `apply_patch` directly (before we convert it to a function call).
+pub fn ev_apply_patch_custom_tool_call(call_id: &str, patch: &str) -> Value {
+    serde_json::json!({
+        "type": "response.output_item.done",
+        "item": {
+            "type": "custom_tool_call",
+            "name": "apply_patch",
+            "input": patch,
+            "call_id": call_id
+        }
+    })
+}
+
+/// Convenience: SSE event for an `apply_patch` function call. The Responses API
+/// wraps the patch content in a JSON string under the `input` key; we recreate
+/// the same structure so downstream code exercises the full parsing path.
+pub fn ev_apply_patch_function_call(call_id: &str, patch: &str) -> Value {
+    let arguments = serde_json::json!({ "input": patch });
+    let arguments = serde_json::to_string(&arguments).expect("serialize apply_patch arguments");
+
+    serde_json::json!({
+        "type": "response.output_item.done",
+        "item": {
+            "type": "function_call",
+            "name": "apply_patch",
+            "arguments": arguments,
+            "call_id": call_id
+        }
+    })
+}
+
+pub fn sse_response(body: String) -> ResponseTemplate {
+    ResponseTemplate::new(200)
+        .insert_header("content-type", "text/event-stream")
+        .set_body_raw(body, "text/event-stream")
+}
+
+pub async fn mount_sse_once<M>(server: &MockServer, matcher: M, body: String)
+where
+    M: wiremock::Match + Send + Sync + 'static,
+{
+    Mock::given(method("POST"))
+        .and(path("/v1/responses"))
+        .and(matcher)
+        .respond_with(sse_response(body))
+        .mount(server)
+        .await;
+}
+
+pub async fn start_mock_server() -> MockServer {
+    MockServer::builder()
+        .body_print_limit(BodyPrintLimit::Limited(80_000))
+        .start()
+        .await
+}

codex-rs/core/tests/suite/client.rs

@@ -22,6 +22,7 @@ use codex_protocol::models::ReasoningItemReasoningSummary;
 use codex_protocol::models::WebSearchAction;
 use core_test_support::load_default_config_for_test;
 use core_test_support::load_sse_fixture_with_id;
+use core_test_support::responses;
 use core_test_support::wait_for_event;
 use futures::StreamExt;
 use serde_json::json;
@@ -875,6 +876,100 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() {
     assert_eq!(body["input"][5]["id"].as_str(), Some("custom-tool-id"));
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn token_count_includes_rate_limits_snapshot() {
+    let server = MockServer::start().await;
+
+    let sse_body = responses::sse(vec![responses::ev_completed_with_tokens("resp_rate", 123)]);
+
+    let response = ResponseTemplate::new(200)
+        .insert_header("content-type", "text/event-stream")
+        .insert_header("x-codex-primary-used-percent", "12.5")
+        .insert_header("x-codex-protection-used-percent", "40.0")
+        .insert_header("x-codex-primary-over-protection-limit-percent", "75.0")
+        .insert_header("x-codex-primary-window-minutes", "10")
+        .insert_header("x-codex-protection-window-minutes", "60")
+        .set_body_raw(sse_body, "text/event-stream");
+
+    Mock::given(method("POST"))
+        .and(path("/v1/responses"))
+        .respond_with(response)
+        .expect(1)
+        .mount(&server)
+        .await;
+
+    let mut provider = built_in_model_providers()["openai"].clone();
+    provider.base_url = Some(format!("{}/v1", server.uri()));
+
+    let home = TempDir::new().unwrap();
+    let mut config = load_default_config_for_test(&home);
+    config.model_provider = provider;
+
+    let conversation_manager = ConversationManager::with_auth(CodexAuth::from_api_key("test"));
+    let codex = conversation_manager
+        .new_conversation(config)
+        .await
+        .expect("create conversation")
+        .conversation;
+
+    codex
+        .submit(Op::UserInput {
+            items: vec![InputItem::Text {
+                text: "hello".into(),
+            }],
+        })
+        .await
+        .unwrap();
+
+    let token_event = wait_for_event(&codex, |msg| matches!(msg, EventMsg::TokenCount(_))).await;
+    let final_payload = match token_event {
+        EventMsg::TokenCount(ev) => ev,
+        _ => unreachable!(),
+    };
+    // Assert full JSON for the final token count event (usage + rate limits)
+    let final_json = serde_json::to_value(&final_payload).unwrap();
+    pretty_assertions::assert_eq!(
+        final_json,
+        json!({
+            "info": {
+                "total_token_usage": {
+                    "input_tokens": 123,
+                    "cached_input_tokens": 0,
+                    "output_tokens": 0,
+                    "reasoning_output_tokens": 0,
+                    "total_tokens": 123
+                },
+                "last_token_usage": {
+                    "input_tokens": 123,
+                    "cached_input_tokens": 0,
+                    "output_tokens": 0,
+                    "reasoning_output_tokens": 0,
+                    "total_tokens": 123
+                },
+                // Default model is gpt-5 in tests → 272000 context window
+                "model_context_window": 272000
+            },
+            "rate_limits": {
+                "primary_used_percent": 12.5,
+                "weekly_used_percent": 40.0,
+                "primary_to_weekly_ratio_percent": 75.0,
+                "primary_window_minutes": 10,
+                "weekly_window_minutes": 60
+            }
+        })
+    );
+    let usage = final_payload
+        .info
+        .expect("token usage info should be recorded after completion");
+    assert_eq!(usage.total_token_usage.total_tokens, 123);
+    let final_snapshot = final_payload
+        .rate_limits
+        .expect("latest rate limit snapshot should be retained");
+    assert_eq!(final_snapshot.primary_used_percent, 12.5);
+
+    wait_for_event(&codex, |msg| matches!(msg, EventMsg::TaskComplete(_))).await;
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn azure_overrides_assign_properties_used_for_responses_url() {
     let existing_env_var_with_random_value = if cfg!(windows) { "USERNAME" } else { "USER" };