diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 4cf9e520eb..365bda8da7 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -3271,7 +3271,6 @@ impl Session { 1, &[("mode", "server_side")], ); - self.recompute_token_usage(turn_context).await; } pub(crate) async fn record_model_warning(&self, message: impl Into, ctx: &TurnContext) { @@ -7370,6 +7369,7 @@ async fn try_run_sampling_request( if let Some(PendingServerSideCompactionCheckpoint { history_at_checkpoint, item, + turn_item, }) = pending_server_side_compaction_checkpoint.take() { sess.apply_server_side_compaction_checkpoint( @@ -7381,6 +7381,17 @@ async fn try_run_sampling_request( history_at_checkpoint.as_slice(), ) .await; + sess.emit_turn_item_started(&turn_context, &turn_item).await; + sess.emit_turn_item_completed(&turn_context, turn_item) + .await; + if let Some(token_usage) = token_usage.as_ref() { + let mut state = sess.state.lock().await; + state.update_token_info_from_usage( + token_usage, + turn_context.model_context_window(), + ); + } + sess.recompute_token_usage(&turn_context).await; applied_server_side_compaction_checkpoint = true; } if !applied_server_side_compaction_checkpoint { diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 1536ef8bac..5484de5756 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -4078,6 +4078,62 @@ async fn task_finish_emits_turn_item_lifecycle_for_leftover_pending_user_input() )); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn inline_compaction_output_item_buffers_checkpoint_without_committing_turn_item() { + let (sess, tc, rx) = make_session_and_context_with_rx().await; + while rx.try_recv().is_ok() {} + + let router = Arc::new(ToolRouter::from_config( + &tc.tools_config, + None, + None, + tc.dynamic_tools.as_slice(), + )); + let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); + let mut ctx = crate::stream_events_utils::HandleOutputCtx { + sess: Arc::clone(&sess), + turn_context: Arc::clone(&tc), + tool_runtime: crate::tools::parallel::ToolCallRuntime::new( + router, + Arc::clone(&sess), + Arc::clone(&tc), + tracker, + ), + cancellation_token: tokio_util::sync::CancellationToken::new(), + }; + let item = ResponseItem::Compaction { + encrypted_content: "INLINE_SERVER_SUMMARY".to_string(), + }; + + let output = crate::stream_events_utils::handle_output_item_done( + &mut ctx, + item.clone(), + Some(TurnItem::ContextCompaction( + codex_protocol::items::ContextCompactionItem::new(), + )), + ) + .await + .expect("handle output item"); + + assert!(output.pending_server_side_compaction.is_some()); + assert!(output.tool_future.is_none()); + assert!(!output.needs_follow_up); + assert!(output.last_agent_message.is_none()); + + let first = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv()) + .await + .expect("expected raw response item event") + .expect("channel open"); + assert!(matches!( + first.msg, + EventMsg::RawResponseItem(raw) if raw.item == item + )); + assert!( + rx.try_recv().is_err(), + "expected no committed compaction item lifecycle before response.completed" + ); +} + #[tokio::test] async fn steer_input_requires_active_turn() { let (sess, _tc, _rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 71d798f614..a172ab1f3d 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -154,6 +154,7 @@ pub(crate) struct OutputItemResult { pub(crate) struct PendingServerSideCompactionCheckpoint { pub history_at_checkpoint: Vec, pub item: ResponseItem, + pub turn_item: TurnItem, } pub(crate) struct HandleOutputCtx { @@ -179,14 +180,8 @@ pub(crate) async fn handle_output_item_done( }); debug!( turn_id = %ctx.turn_context.sub_id, - "emitting streamed server-side compaction item and buffering history rewrite until response.completed" + "emitting streamed server-side raw compaction item and buffering committed checkpoint until response.completed" ); - ctx.sess - .emit_turn_item_started(&ctx.turn_context, &turn_item) - .await; - ctx.sess - .emit_turn_item_completed(&ctx.turn_context, turn_item) - .await; ctx.sess .send_event( &ctx.turn_context, @@ -196,6 +191,7 @@ pub(crate) async fn handle_output_item_done( output.pending_server_side_compaction = Some(PendingServerSideCompactionCheckpoint { history_at_checkpoint: ctx.sess.clone_history().await.raw_items().to_vec(), item, + turn_item, }); return Ok(output); } diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 8e05408a72..7dd0fa7e94 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -547,7 +547,7 @@ async fn auto_server_side_compaction_uses_inline_context_management() -> Result< } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn auto_server_side_compaction_emits_events_before_later_streamed_items() -> Result<()> { +async fn auto_server_side_compaction_commits_events_after_later_streamed_items() -> Result<()> { skip_if_no_network!(Ok(())); let compact_threshold = 120; @@ -628,9 +628,9 @@ async fn auto_server_side_compaction_emits_events_before_later_streamed_items() } assert!( - context_compacted_index.expect("context compacted event") - < assistant_completed_index.expect("assistant completed event"), - "expected the inline compaction event to arrive before later streamed assistant items" + assistant_completed_index.expect("assistant completed event") + < context_compacted_index.expect("context compacted event"), + "expected the committed inline compaction event to arrive only after later streamed assistant items" ); Ok(()) @@ -901,12 +901,11 @@ async fn auto_server_side_compaction_preserves_recomputed_token_estimate() -> Re let event = codex.next_event().await.expect("event"); match event.msg { EventMsg::TokenCount(token_count) => { - if let Some(last_token_usage) = token_count - .info - .as_ref() - .map(|info| info.last_token_usage.total_tokens) - { - token_usage_events.push(last_token_usage); + if let Some(info) = token_count.info.as_ref() { + token_usage_events.push(( + info.last_token_usage.total_tokens, + info.total_token_usage.total_tokens, + )); } } EventMsg::TurnComplete(_) => { @@ -926,8 +925,8 @@ async fn auto_server_side_compaction_preserves_recomputed_token_estimate() -> Re token_usage_events .iter() .copied() - .any(|tokens| tokens > 500), - "expected a post-compaction token event to keep the recomputed local estimate instead of reverting to the provider-reported total: {token_usage_events:?}" + .any(|(last_tokens, total_tokens)| last_tokens > 500 && total_tokens >= 1_000), + "expected a post-compaction token event to keep the recomputed local estimate and preserve cumulative provider accounting: {token_usage_events:?}" ); Ok(())