diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 5fa3ff1f67..a04c409e03 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -4,7 +4,9 @@ use std::fmt::Debug; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use crate::AuthManager; use crate::CodexAuth; @@ -668,6 +670,7 @@ pub(crate) struct Session { pub(crate) active_turn: Mutex>, pub(crate) services: SessionServices, js_repl: Arc, + inline_server_side_compaction_incompatible: AtomicBool, next_internal_sub_id: AtomicU64, } @@ -1667,6 +1670,7 @@ impl Session { active_turn: Mutex::new(None), services, js_repl, + inline_server_side_compaction_incompatible: AtomicBool::new(false), next_internal_sub_id: AtomicU64::new(0), }); if let Some(network_policy_decider_session) = network_policy_decider_session { @@ -3239,20 +3243,14 @@ impl Session { &self, turn_context: &TurnContext, item: ResponseItem, + history_before_turn: &[ResponseItem], ) { - let ghost_snapshots = self - .clone_history() - .await - .raw_items() - .iter() - .filter(|history_item| matches!(history_item, ResponseItem::GhostSnapshot { .. })) - .cloned() - .collect::>(); - - let mut replacement_history = vec![item.clone()]; - if !ghost_snapshots.is_empty() { - replacement_history.extend(ghost_snapshots); - } + let current_history = self.clone_history().await; + let replacement_history = build_server_side_compaction_replacement_history( + item.clone(), + history_before_turn, + current_history.raw_items(), + ); let reference_context_item = Some(turn_context.to_turn_context_item()); let compacted_item = CompactedItem { @@ -3376,6 +3374,18 @@ impl Session { self.features.clone() } + fn inline_server_side_compaction_supported(&self) -> bool { + !self + .inline_server_side_compaction_incompatible + .load(Ordering::Relaxed) + } + + fn disable_inline_server_side_compaction(&self) -> bool { + !self + .inline_server_side_compaction_incompatible + .swap(true, Ordering::Relaxed) + } + pub(crate) async fn collaboration_mode(&self) -> CollaborationMode { let state = self.state.lock().await; state.session_configuration.collaboration_mode.clone() @@ -5710,6 +5720,7 @@ pub(crate) async fn run_turn( &mut client_session, turn_metadata_header.as_deref(), sampling_request_input, + &preturn_inline_compaction_state.history_before_turn, inline_compaction_for_request.map(|pending| pending.threshold), &turn_enabled_connectors, skills_outcome, @@ -6004,6 +6015,29 @@ fn collect_new_ghost_snapshots_since( .collect() } +fn build_server_side_compaction_replacement_history( + compaction_item: ResponseItem, + history_before_turn: &[ResponseItem], + current_history: &[ResponseItem], +) -> Vec { + let mut replacement_history = vec![compaction_item]; + replacement_history.extend( + current_history + .strip_prefix(history_before_turn) + .unwrap_or(current_history) + .iter() + .filter(|item| !matches!(item, ResponseItem::GhostSnapshot { .. })) + .cloned(), + ); + replacement_history.extend( + current_history + .iter() + .filter(|item| matches!(item, ResponseItem::GhostSnapshot { .. })) + .cloned(), + ); + replacement_history +} + fn record_compaction_metric( sess: &Session, mode: &'static str, @@ -6052,6 +6086,9 @@ fn inline_server_side_compaction_threshold( if !sess.enabled(Feature::ServerSideCompaction) { return None; } + if !sess.inline_server_side_compaction_supported() { + return None; + } if !should_use_remote_compact_task(&turn_context.provider) { return None; } @@ -6068,6 +6105,8 @@ fn record_inline_compaction_skip( ) { let reason = if !sess.enabled(Feature::ServerSideCompaction) { "flag_off" + } else if !sess.inline_server_side_compaction_supported() { + "backend_incompatible" } else if !should_use_remote_compact_task(&turn_context.provider) { "non_openai" } else if has_custom_compact_prompt(turn_context) { @@ -6116,6 +6155,14 @@ async fn downgrade_known_inline_compaction_error( return Ok(false); } + if sess.disable_inline_server_side_compaction() { + tracing::warn!( + turn_id = %turn_context.sub_id, + trigger = pending_compaction.trigger.as_str(), + "disabling inline server-side compaction for this session after compatibility failure" + ); + } + tracing::warn!( turn_id = %turn_context.sub_id, trigger = pending_compaction.trigger.as_str(), @@ -6522,6 +6569,7 @@ async fn run_sampling_request( client_session: &mut ModelClientSession, turn_metadata_header: Option<&str>, input: Vec, + history_before_turn: &[ResponseItem], inline_compaction_threshold: Option, explicitly_enabled_connectors: &HashSet, skills_outcome: Option<&SkillLoadOutcome>, @@ -6558,6 +6606,7 @@ async fn run_sampling_request( Arc::clone(&turn_diff_tracker), server_model_warning_emitted_for_turn, &prompt, + history_before_turn, cancellation_token.child_token(), ) .await @@ -7281,6 +7330,7 @@ async fn try_run_sampling_request( turn_diff_tracker: SharedTurnDiffTracker, server_model_warning_emitted_for_turn: &mut bool, prompt: &Prompt, + history_before_turn: &[ResponseItem], cancellation_token: CancellationToken, ) -> CodexResult { feedback_tags!( @@ -7321,6 +7371,7 @@ async fn try_run_sampling_request( let plan_mode = turn_context.collaboration_mode.mode == ModeKind::Plan; let mut assistant_message_stream_parsers = AssistantMessageStreamParsers::new(plan_mode); let mut plan_mode_state = plan_mode.then(|| PlanModeStreamState::new(&turn_context.sub_id)); + let history_before_turn = Arc::new(history_before_turn.to_vec()); let receiving_span = trace_span!("receiving_stream"); let outcome: CodexResult = loop { let handle_responses = trace_span!( @@ -7393,6 +7444,7 @@ async fn try_run_sampling_request( turn_context: turn_context.clone(), tool_runtime: tool_runtime.clone(), cancellation_token: cancellation_token.child_token(), + history_before_turn: history_before_turn.clone(), }; let output_result = handle_output_item_done(&mut ctx, item, previously_active_item) diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index c6be5cbf46..4a0523104d 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -278,6 +278,45 @@ fn collect_new_ghost_snapshots_since_returns_only_snapshots_added_after_turn_sta assert_eq!(new_snapshots, vec![same_turn_snapshot]); } +#[test] +fn build_server_side_compaction_replacement_history_keeps_current_turn_inputs() { + let prior_snapshot = ghost_snapshot("ghost-before"); + let same_turn_snapshot = ghost_snapshot("ghost-during"); + let history_before_turn = vec![user_message("earlier"), prior_snapshot.clone()]; + let current_turn_user = user_message("current turn"); + let current_turn_tool_output = ResponseItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_text("tool result".to_string()), + }; + let current_history = vec![ + user_message("earlier"), + prior_snapshot.clone(), + current_turn_user.clone(), + current_turn_tool_output.clone(), + same_turn_snapshot.clone(), + ]; + let compaction_item = ResponseItem::Compaction { + encrypted_content: "INLINE_SUMMARY".to_string(), + }; + + let replacement_history = build_server_side_compaction_replacement_history( + compaction_item.clone(), + &history_before_turn, + ¤t_history, + ); + + assert_eq!( + replacement_history, + vec![ + compaction_item, + current_turn_user, + current_turn_tool_output, + prior_snapshot, + same_turn_snapshot, + ] + ); +} + fn make_mcp_tool( server_name: &str, tool_name: &str, @@ -2293,6 +2332,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { active_turn: Mutex::new(None), services, js_repl, + inline_server_side_compaction_incompatible: std::sync::atomic::AtomicBool::new(false), next_internal_sub_id: AtomicU64::new(0), }; @@ -2853,6 +2893,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( active_turn: Mutex::new(None), services, js_repl, + inline_server_side_compaction_incompatible: std::sync::atomic::AtomicBool::new(false), next_internal_sub_id: AtomicU64::new(0), }); diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 9894883f5e..fd964305c9 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -153,6 +153,7 @@ pub(crate) struct HandleOutputCtx { pub turn_context: Arc, pub tool_runtime: ToolCallRuntime, pub cancellation_token: CancellationToken, + pub history_before_turn: Arc>, } #[instrument(level = "trace", skip_all)] @@ -186,7 +187,11 @@ pub(crate) async fn handle_output_item_done( "processing streamed server-side compaction item" ); ctx.sess - .apply_server_side_compaction_checkpoint(ctx.turn_context.as_ref(), item) + .apply_server_side_compaction_checkpoint( + ctx.turn_context.as_ref(), + item, + ctx.history_before_turn.as_slice(), + ) .await; ctx.sess .emit_turn_item_completed( diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index e28aabaccb..ae44c940ae 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -516,6 +516,82 @@ async fn auto_server_side_compaction_uses_inline_context_management() -> Result< Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn auto_server_side_compaction_keeps_current_turn_inputs_for_follow_ups() -> Result<()> { + skip_if_no_network!(Ok(())); + + let compact_threshold = 120; + let first_turn_text = "inline compact turn one"; + let second_turn_text = "inline compact turn two"; + let inline_summary = summary_with_prefix("INLINE_SERVER_SUMMARY"); + + let harness = TestCodexHarness::with_builder( + test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + config + .features + .enable(Feature::ServerSideCompaction) + .expect("enable server-side compaction"); + config.model_auto_compact_token_limit = Some(compact_threshold); + }), + ) + .await?; + let codex = harness.test().codex.clone(); + + let responses_mock = responses::mount_sse_sequence( + harness.server(), + vec![ + responses::sse(vec![ + responses::ev_assistant_message("m1", "FIRST_REMOTE_REPLY"), + responses::ev_completed_with_tokens("resp-1", 500), + ]), + responses::sse(vec![ + responses::ev_compaction(&inline_summary), + responses::ev_function_call("call-inline-mid-turn", DUMMY_FUNCTION_NAME, "{}"), + responses::ev_completed_with_tokens("resp-2", 80), + ]), + responses::sse(vec![ + responses::ev_assistant_message("m3", "AFTER_INLINE_TOOL_REPLY"), + responses::ev_completed("resp-3"), + ]), + ], + ) + .await; + + submit_text_turn_and_wait(&codex, first_turn_text).await?; + submit_text_turn_and_wait(&codex, second_turn_text).await?; + + let requests = responses_mock.requests(); + assert_eq!( + requests.len(), + 3, + "expected initial request, inline-compacted tool call, and same-turn follow-up" + ); + + let follow_up_request = &requests[2]; + assert!( + follow_up_request.body_contains_text(&inline_summary), + "expected same-turn follow-up to include the inline compaction item" + ); + assert!( + follow_up_request.body_contains_text(second_turn_text), + "expected same-turn follow-up to retain the current turn user input" + ); + assert!( + !follow_up_request.body_contains_text(first_turn_text), + "expected same-turn follow-up to drop pre-compaction history" + ); + assert!( + follow_up_request + .function_call_output_text("call-inline-mid-turn") + .is_some(), + "expected same-turn follow-up to include the tool output" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn auto_server_side_compaction_uses_legacy_remote_path_with_custom_prompt() -> Result<()> { skip_if_no_network!(Ok(())); @@ -610,25 +686,33 @@ async fn auto_server_side_compaction_downgrades_known_compat_errors_once() -> Re })), responses::sse_response(sse(vec![ responses::ev_assistant_message("m2", "AFTER_DOWNGRADE_REPLY"), - responses::ev_completed("resp-2"), + responses::ev_completed_with_tokens("resp-2", 500), + ])), + responses::sse_response(sse(vec![ + responses::ev_assistant_message("m3", "AFTER_COMPAT_SKIP_REPLY"), + responses::ev_completed("resp-3"), ])), ], ) .await; - let compact_mock = responses::mount_compact_user_history_with_summary_once( + let compact_mock = responses::mount_compact_user_history_with_summary_sequence( harness.server(), - "DOWNGRADE_REMOTE_SUMMARY", + vec![ + "DOWNGRADE_REMOTE_SUMMARY".to_string(), + "POST_DOWNGRADE_REMOTE_SUMMARY".to_string(), + ], ) .await; submit_text_turn_and_wait(&codex, "downgrade turn one").await?; submit_text_turn_and_wait(&codex, "downgrade turn two").await?; + submit_text_turn_and_wait(&codex, "downgrade turn three").await?; let requests = responses_mock.requests(); assert_eq!( requests.len(), - 3, - "expected initial turn, failed inline attempt, and downgraded follow-up request" + 4, + "expected the initial turn, one failed inline attempt, one downgraded retry, and a later direct legacy request" ); let inline_attempt = requests[1].body_json(); @@ -649,10 +733,18 @@ async fn auto_server_side_compaction_downgrades_known_compat_errors_once() -> Re requests[2].body_contains_text("DOWNGRADE_REMOTE_SUMMARY"), "downgraded retry should reuse the client-side compaction output" ); + assert!( + requests[3].body_json().get("context_management").is_none(), + "future auto-compaction requests should skip inline compaction after a known compat error" + ); + assert!( + requests[3].body_contains_text("POST_DOWNGRADE_REMOTE_SUMMARY"), + "future auto-compaction requests should go straight to the legacy compaction output" + ); assert_eq!( compact_mock.requests().len(), - 1, - "expected one legacy remote compact request after the compat error" + 2, + "expected later auto-compactions to use the legacy path directly after the compat error" ); Ok(())