diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 9fbbcd312d..2db71e3655 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -114,6 +114,7 @@ use crate::config::resolve_web_search_mode_for_turn; use crate::config::types::McpServerConfig; use crate::config::types::ShellEnvironmentPolicy; use crate::context_manager::ContextManager; +use crate::context_manager::TotalTokenUsageBreakdown; use crate::environment_context::EnvironmentContext; use crate::error::CodexErr; use crate::error::Result as CodexResult; @@ -1214,12 +1215,20 @@ impl Session { format!("auto-compact-{id}") } - async fn get_total_token_usage(&self) -> i64 { + pub(crate) async fn get_total_token_usage(&self) -> i64 { let state = self.state.lock().await; state.get_total_token_usage(state.server_reasoning_included()) } - async fn get_estimated_token_count(&self, turn_context: &TurnContext) -> Option { + pub(crate) async fn get_total_token_usage_breakdown(&self) -> TotalTokenUsageBreakdown { + let state = self.state.lock().await; + state.get_total_token_usage_breakdown() + } + + pub(crate) async fn get_estimated_token_count( + &self, + turn_context: &TurnContext, + ) -> Option { let state = self.state.lock().await; state.history.estimate_token_count(turn_context) } diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index eac3188cd5..ebcb01712e 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -4,16 +4,20 @@ use crate::Prompt; use crate::codex::Session; use crate::codex::TurnContext; use crate::context_manager::ContextManager; +use crate::context_manager::TotalTokenUsageBreakdown; use crate::context_manager::is_codex_generated_item; +use crate::error::CodexErr; use crate::error::Result as CodexResult; use crate::protocol::CompactedItem; use crate::protocol::EventMsg; use crate::protocol::RolloutItem; use crate::protocol::TurnStartedEvent; +use codex_api::CompactionInput as ApiCompactionInput; use codex_protocol::items::ContextCompactionItem; use codex_protocol::items::TurnItem; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ResponseItem; +use tracing::error; use tracing::info; pub(crate) async fn run_inline_remote_auto_compact_task( @@ -90,7 +94,7 @@ async fn run_remote_compact_task_inner_impl( output_schema: None, }; - let mut new_history = sess + let compact_result = sess .services .model_client .compact_conversation_history( @@ -98,7 +102,25 @@ async fn run_remote_compact_task_inner_impl( &turn_context.model_info, &turn_context.otel_manager, ) - .await?; + .await; + let mut new_history = match compact_result { + Ok(history) => history, + Err(err) => { + let total_usage_breakdown = sess.get_total_token_usage_breakdown().await; + let compact_request_metrics = build_compact_request_metrics( + &turn_context.model_info.slug, + &prompt.input, + &prompt.base_instructions.text, + ); + log_remote_compact_failure( + turn_context, + &compact_request_metrics, + total_usage_breakdown, + &err, + ); + return Err(err); + } + }; new_history = sess .process_compacted_history(turn_context, new_history) .await; @@ -121,6 +143,66 @@ async fn run_remote_compact_task_inner_impl( Ok(()) } +#[derive(Debug)] +struct CompactRequestMetrics { + failing_compaction_request_body_json: String, + failing_compaction_request_body_bytes: usize, +} + +fn build_compact_request_metrics( + model: &str, + input: &[ResponseItem], + instructions: &str, +) -> CompactRequestMetrics { + let payload = ApiCompactionInput { + model, + input, + instructions, + }; + let failing_compaction_request_body_json = serde_json::to_string(&payload) + .unwrap_or_else(|err| format!("{{\"compact_request_serialization_error\":\"{err}\"}}")); + let failing_compaction_request_body_bytes = failing_compaction_request_body_json.len(); + + CompactRequestMetrics { + failing_compaction_request_body_json, + failing_compaction_request_body_bytes, + } +} + +fn compact_error_status_code(err: &CodexErr) -> Option { + match err { + CodexErr::InvalidRequest(_) => Some(400), + CodexErr::UnexpectedStatus(status) => Some(status.status.as_u16()), + CodexErr::ContextWindowExceeded => Some(400), + _ => None, + } +} + +fn log_remote_compact_failure( + turn_context: &TurnContext, + metrics: &CompactRequestMetrics, + total_usage_breakdown: TotalTokenUsageBreakdown, + err: &CodexErr, +) { + error!( + turn_id = %turn_context.sub_id, + failing_compaction_request_body_json = %metrics.failing_compaction_request_body_json, + "remote compaction request payload before failure" + ); + error!( + turn_id = %turn_context.sub_id, + compact_error_status = ?compact_error_status_code(err), + last_api_response_total_tokens = total_usage_breakdown.last_api_response_total_tokens, + last_api_response_total_bytes_estimate = total_usage_breakdown.last_api_response_total_bytes_estimate, + estimated_tokens_of_items_added_since_last_successful_api_response = total_usage_breakdown.estimated_tokens_of_items_added_since_last_successful_api_response, + estimated_bytes_of_items_added_since_last_successful_api_response = total_usage_breakdown.estimated_bytes_of_items_added_since_last_successful_api_response, + model_context_window_tokens = ?turn_context.model_context_window(), + failing_compaction_request_body_bytes = metrics.failing_compaction_request_body_bytes, + compact_error = %err, + "remote compaction failed" + ); +} + fn trim_function_call_history_to_fit_context_window( history: &mut ContextManager, turn_context: &TurnContext, diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index a67ea90d3d..4c430c460a 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -4,6 +4,7 @@ use crate::instructions::SkillInstructions; use crate::instructions::UserInstructions; use crate::session_prefix::is_session_prefix; use crate::truncate::TruncationPolicy; +use crate::truncate::approx_bytes_for_tokens; use crate::truncate::approx_token_count; use crate::truncate::approx_tokens_from_byte_count; use crate::truncate::truncate_function_output_items_with_policy; @@ -27,6 +28,14 @@ pub(crate) struct ContextManager { token_info: Option, } +#[derive(Debug, Clone, Copy, Default)] +pub(crate) struct TotalTokenUsageBreakdown { + pub last_api_response_total_tokens: i64, + pub last_api_response_total_bytes_estimate: usize, + pub estimated_tokens_of_items_added_since_last_successful_api_response: i64, + pub estimated_bytes_of_items_added_since_last_successful_api_response: usize, +} + impl ContextManager { pub(crate) fn new() -> Self { Self { @@ -247,6 +256,21 @@ impl ContextManager { total } + fn get_trailing_codex_generated_items_bytes(&self) -> usize { + let mut total = 0usize; + for item in self.items.iter().rev() { + if !is_codex_generated_item(item) { + break; + } + total = total.saturating_add( + serde_json::to_vec(item) + .map(|bytes| bytes.len()) + .unwrap_or_default(), + ); + } + total + } + /// When true, the server already accounted for past reasoning tokens and /// the client should not re-estimate them. pub(crate) fn get_total_token_usage(&self, server_reasoning_included: bool) -> i64 { @@ -265,6 +289,29 @@ impl ContextManager { } } + pub(crate) fn get_total_token_usage_breakdown(&self) -> TotalTokenUsageBreakdown { + let last_usage = self + .token_info + .as_ref() + .map(|info| info.last_token_usage.clone()) + .unwrap_or_default(); + + TotalTokenUsageBreakdown { + last_api_response_total_tokens: last_usage.total_tokens, + last_api_response_total_bytes_estimate: if last_usage.total_tokens <= 0 { + 0 + } else { + usize::try_from(last_usage.total_tokens) + .map(approx_bytes_for_tokens) + .unwrap_or(usize::MAX) + }, + estimated_tokens_of_items_added_since_last_successful_api_response: self + .get_trailing_codex_generated_items_tokens(), + estimated_bytes_of_items_added_since_last_successful_api_response: self + .get_trailing_codex_generated_items_bytes(), + } + } + /// This function enforces a couple of invariants on the in-memory history: /// 1. every call (function/custom) has a corresponding output entry /// 2. every output has a corresponding call entry diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index 22e9682fe3..4ea399e167 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -2,5 +2,6 @@ mod history; mod normalize; pub(crate) use history::ContextManager; +pub(crate) use history::TotalTokenUsageBreakdown; pub(crate) use history::is_codex_generated_item; pub(crate) use history::is_user_turn_boundary; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 29bc14d1cc..621cbd5010 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -6,6 +6,7 @@ use std::collections::HashSet; use crate::codex::SessionConfiguration; use crate::context_manager::ContextManager; +use crate::context_manager::TotalTokenUsageBreakdown; use crate::protocol::RateLimitSnapshot; use crate::protocol::TokenUsage; use crate::protocol::TokenUsageInfo; @@ -100,6 +101,10 @@ impl SessionState { .get_total_token_usage(server_reasoning_included) } + pub(crate) fn get_total_token_usage_breakdown(&self) -> TotalTokenUsageBreakdown { + self.history.get_total_token_usage_breakdown() + } + pub(crate) fn set_server_reasoning_included(&mut self, included: bool) { self.server_reasoning_included = included; }