diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 30221d944..0157101cf 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -3341,6 +3341,8 @@ mod tests { message: String::new(), replacement_history: None, window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, }), RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 66b5b5c4c..8dbcab338 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1160,6 +1160,8 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { message: String::new(), replacement_history: Some(replacement_history), window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, }), RolloutItem::TurnContext(turn_context.to_turn_context_item()), diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 7a2eaf935..8a5efd594 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -302,7 +302,7 @@ async fn run_compact_task_inner_impl( let user_messages = collect_user_messages(history_items); let mut new_history = build_compacted_history(Vec::new(), &user_messages, &summary_text); - let (window_number, window_id) = sess.advance_auto_compact_window().await; + let (window_number, window_ids) = sess.advance_auto_compact_window().await; if matches!( initial_context_injection, @@ -320,7 +320,9 @@ async fn run_compact_task_inner_impl( message: summary_text.clone(), replacement_history: Some(new_history.clone()), window_number: Some(window_number), - window_id: Some(window_id), + first_window_id: Some(window_ids.first_window_id.to_string()), + previous_window_id: window_ids.previous_window_id.map(|id| id.to_string()), + window_id: Some(window_ids.window_id.to_string()), }; sess.replace_compacted_history( turn_context.as_ref(), diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 26469c201..7a47881e4 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -258,7 +258,7 @@ async fn run_remote_compact_task_inner_impl( &responses_metadata, ) .await?; - let (new_window_number, new_window_id) = sess.advance_auto_compact_window().await; + let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await; new_history = process_compacted_history( sess.as_ref(), turn_context.as_ref(), @@ -275,7 +275,9 @@ async fn run_remote_compact_task_inner_impl( message: String::new(), replacement_history: Some(new_history.clone()), window_number: Some(new_window_number), - window_id: Some(new_window_id), + first_window_id: Some(new_window_ids.first_window_id.to_string()), + previous_window_id: new_window_ids.previous_window_id.map(|id| id.to_string()), + window_id: Some(new_window_ids.window_id.to_string()), }; // Install is the semantic boundary where the compact endpoint's output becomes live // thread history. Keep it distinct from the later inference request so the reducer can diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 6f10a474a..0b382700a 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -292,7 +292,7 @@ async fn run_remote_compact_task_inner_impl( let (compacted_history, retained_images) = build_v2_compacted_history(&prompt_input, compaction_output); analytics_details.retained_image_count = Some(retained_images); - let (new_window_number, new_window_id) = sess.advance_auto_compact_window().await; + let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await; let new_history = process_compacted_history( sess.as_ref(), turn_context.as_ref(), @@ -309,7 +309,9 @@ async fn run_remote_compact_task_inner_impl( message: String::new(), replacement_history: Some(new_history.clone()), window_number: Some(new_window_number), - window_id: Some(new_window_id), + first_window_id: Some(new_window_ids.first_window_id.to_string()), + previous_window_id: new_window_ids.previous_window_id.map(|id| id.to_string()), + window_id: Some(new_window_ids.window_id.to_string()), }; compaction_trace.record_installed(&CompactionCheckpointTracePayload { input_history: &trace_input_history, diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index b21b61735..bac1caafd 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -5,14 +5,24 @@ use uuid::Uuid; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct TokenBudgetContext { thread_id: ThreadId, + first_window_id: Uuid, + previous_window_id: Option, window_id: Uuid, tokens_left: i64, } impl TokenBudgetContext { - pub(crate) fn new(thread_id: ThreadId, window_id: Uuid, tokens_left: i64) -> Self { + pub(crate) fn new( + thread_id: ThreadId, + first_window_id: Uuid, + previous_window_id: Option, + window_id: Uuid, + tokens_left: i64, + ) -> Self { Self { thread_id, + first_window_id, + previous_window_id, window_id, tokens_left, } @@ -34,10 +44,14 @@ impl ContextualUserFragment for TokenBudgetContext { fn body(&self) -> String { let thread_id = self.thread_id; + let first_window_id = self.first_window_id; + let previous_window_id = self + .previous_window_id + .map_or_else(|| "none".to_string(), |window_id| window_id.to_string()); let window_id = self.window_id; let tokens_left = self.tokens_left; format!( - "Thread id {thread_id}.\nCurrent context window id {window_id}.\nYou have {tokens_left} tokens left in this context window." + "Thread id {thread_id}.\nFirst context window id {first_window_id}.\nPrevious context window id {previous_window_id}.\nCurrent context window id {window_id}.\nYou have {tokens_left} tokens left in this context window." ) } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 0792334b3..3cc746912 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -311,6 +311,7 @@ use crate::session_startup_prewarm::SessionStartupPrewarmHandle; use crate::shell; #[cfg(test)] use crate::skills::SkillLoadOutcome; +use crate::state::AutoCompactWindowIds; use crate::state::AutoCompactWindowSnapshot; use crate::state::PendingRequestPermissions; use crate::state::SessionServices; @@ -1364,6 +1365,8 @@ impl Session { previous_turn_settings, reference_context_item, window_number, + first_window_id, + previous_window_id, window_id, } = self .reconstruct_history_from_rollout(turn_context, rollout_items) @@ -1382,8 +1385,16 @@ impl Session { { let mut state = self.state.lock().await; state.replace_history(history, reference_context_item); - let window_id = window_id.unwrap_or_else(|| state.auto_compact_window_id()); - state.restore_auto_compact_window(window_number, window_id); + let fallback_ids = state.auto_compact_window_ids(); + let window_id = window_id.unwrap_or(fallback_ids.window_id); + state.restore_auto_compact_window( + window_number, + AutoCompactWindowIds { + first_window_id: first_window_id.unwrap_or(window_id), + previous_window_id, + window_id, + }, + ); state.set_previous_turn_settings(previous_turn_settings.clone()); } let prefix_tokens = if matches!( @@ -3011,7 +3022,7 @@ impl Session { collaboration_mode, base_instructions, session_source, - auto_compact_window_id, + auto_compact_window_ids, ) = { let state = self.state.lock().await; ( @@ -3020,7 +3031,7 @@ impl Session { state.session_configuration.collaboration_mode.clone(), state.session_configuration.base_instructions.clone(), state.session_configuration.session_source.clone(), - state.auto_compact_window_id(), + state.auto_compact_window_ids(), ) }; if let Some(model_switch_message) = @@ -3210,7 +3221,9 @@ impl Session { developer_sections.push( crate::context::TokenBudgetContext::new( self.thread_id(), - auto_compact_window_id, + auto_compact_window_ids.first_window_id, + auto_compact_window_ids.previous_window_id, + auto_compact_window_ids.window_id, model_context_window, ) .render(), @@ -3309,10 +3322,9 @@ impl Session { format!("{thread_id}:{window_number}") } - pub(crate) async fn advance_auto_compact_window(&self) -> (u64, String) { + pub(crate) async fn advance_auto_compact_window(&self) -> (u64, AutoCompactWindowIds) { let mut state = self.state.lock().await; - let (window_number, window_id) = state.advance_auto_compact_window(); - (window_number, window_id.to_string()) + state.advance_auto_compact_window() } pub(crate) async fn request_new_context_window(&self) { @@ -3328,7 +3340,7 @@ impl Session { let mut state = self.state.lock().await; state.start_new_context_window_if_requested() }; - let (window_number, window_id) = window?; + let (window_number, window_ids) = window?; let context_items = self.build_initial_context(turn_context).await; let turn_context_item = turn_context.to_turn_context_item(); let replacement_history = context_items; @@ -3341,7 +3353,9 @@ impl Session { message: String::new(), replacement_history: Some(replacement_history), window_number: Some(window_number), - window_id: Some(window_id.to_string()), + first_window_id: Some(window_ids.first_window_id.to_string()), + previous_window_id: window_ids.previous_window_id.map(|id| id.to_string()), + window_id: Some(window_ids.window_id.to_string()), }), RolloutItem::TurnContext(turn_context_item), ]) diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index 22d0990e3..bf534c1e8 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -10,12 +10,16 @@ pub(super) struct RolloutReconstruction { pub(super) previous_turn_settings: Option, pub(super) reference_context_item: Option, pub(super) window_number: u64, + pub(super) first_window_id: Option, + pub(super) previous_window_id: Option, pub(super) window_id: Option, } #[derive(Debug, Clone, Copy)] struct ReconstructedWindow { number: u64, + first_id: Option, + previous_id: Option, id: Option, } @@ -133,6 +137,11 @@ impl Session { { active_segment.window = Some(ReconstructedWindow { number: window_number, + first_id: compacted.first_window_id.as_deref().and_then(parse_uuid_v7), + previous_id: compacted + .previous_window_id + .as_deref() + .and_then(parse_uuid_v7), id: compacted.window_id.as_deref().and_then(parse_uuid_v7), }); } @@ -341,6 +350,8 @@ impl Session { let window = window.unwrap_or(ReconstructedWindow { number: fallback_window_number, + first_id: None, + previous_id: None, id: None, }); RolloutReconstruction { @@ -348,6 +359,8 @@ impl Session { previous_turn_settings, reference_context_item, window_number: window.number, + first_window_id: window.first_id, + previous_window_id: window.previous_id, window_id: window.id, } } diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index 43563d233..1823b28fd 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -829,6 +829,8 @@ async fn record_initial_history_resumed_rollback_drops_incomplete_user_turn_comp message: String::new(), replacement_history: Some(Vec::new()), window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, }), RolloutItem::EventMsg(EventMsg::ThreadRolledBack( @@ -887,6 +889,8 @@ async fn record_initial_history_resumed_does_not_seed_reference_context_item_aft message: String::new(), replacement_history: Some(Vec::new()), window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, }), ]; @@ -914,6 +918,8 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_does_ message: "legacy summary".to_string(), replacement_history: None, window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, }), ]; @@ -947,6 +953,8 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_clear message: "legacy summary".to_string(), replacement_history: None, window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, }), RolloutItem::EventMsg(EventMsg::TurnStarted( @@ -1043,6 +1051,8 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis message: String::new(), replacement_history: Some(Vec::new()), window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, }), RolloutItem::TurnContext(previous_context_item), @@ -1196,6 +1206,8 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu message: String::new(), replacement_history: Some(Vec::new()), window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, }), ]; @@ -1433,6 +1445,8 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea message: String::new(), replacement_history: Some(Vec::new()), window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, }), ]; @@ -1599,6 +1613,8 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear message: String::new(), replacement_history: Some(Vec::new()), window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, }), // A newer TurnStarted replaces the incomplete compacted turn without a matching diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index bca6be6a7..6f84da3c7 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1653,11 +1653,15 @@ async fn reconstruct_history_uses_replacement_history_verbatim() { metadata: None, }, ]; + let first_window_id = Uuid::now_v7(); + let previous_window_id = Uuid::now_v7(); let window_id = Uuid::now_v7(); let rollout_items = vec![RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(replacement_history.clone()), window_number: Some(42), + first_window_id: Some(first_window_id.to_string()), + previous_window_id: Some(previous_window_id.to_string()), window_id: Some(window_id.to_string()), })]; @@ -1667,6 +1671,8 @@ async fn reconstruct_history_uses_replacement_history_verbatim() { assert_eq!(reconstructed.history, replacement_history); assert_eq!(42, reconstructed.window_number); + assert_eq!(Some(first_window_id), reconstructed.first_window_id); + assert_eq!(Some(previous_window_id), reconstructed.previous_window_id); assert_eq!(Some(window_id), reconstructed.window_id); } @@ -3059,6 +3065,8 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio user_message("turn 1 user"), user_message("summary after compaction"), ]; + let first_window_id = Uuid::now_v7(); + let previous_window_id = Uuid::now_v7(); let compacted_window_id = Uuid::now_v7(); sess.persist_rollout_items(&[ @@ -3102,6 +3110,8 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio message: "summary after compaction".to_string(), replacement_history: Some(compacted_history.clone()), window_number: Some(7), + first_window_id: Some(first_window_id.to_string()), + previous_window_id: Some(previous_window_id.to_string()), window_id: Some(compacted_window_id.to_string()), }), RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { @@ -3152,7 +3162,14 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio .await; { let mut state = sess.state.lock().await; - state.restore_auto_compact_window(/*window_number*/ 99, Uuid::now_v7()); + state.restore_auto_compact_window( + /*window_number*/ 99, + AutoCompactWindowIds { + first_window_id: Uuid::now_v7(), + previous_window_id: Some(Uuid::now_v7()), + window_id: Uuid::now_v7(), + }, + ); } handlers::thread_rollback(&sess, "sub-1".to_string(), /*num_turns*/ 1).await; @@ -3162,8 +3179,12 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio assert_eq!(sess.clone_history().await.raw_items(), compacted_history); assert!(sess.reference_context_item().await.is_none()); assert_eq!( - sess.state.lock().await.auto_compact_window_id(), - compacted_window_id + sess.state.lock().await.auto_compact_window_ids(), + AutoCompactWindowIds { + first_window_id, + previous_window_id: Some(previous_window_id), + window_id: compacted_window_id, + } ); assert!(sess.current_window_id().await.ends_with(":7")); } @@ -9910,12 +9931,14 @@ async fn sample_rollout( let user_messages1 = collect_user_messages(&snapshot1); let rebuilt1 = compact::build_compacted_history(Vec::new(), &user_messages1, summary1); live_history.replace(rebuilt1); - let (window_number, window_id) = session.advance_auto_compact_window().await; + let (window_number, window_ids) = session.advance_auto_compact_window().await; rollout_items.push(RolloutItem::Compacted(CompactedItem { message: summary1.to_string(), replacement_history: None, window_number: Some(window_number), - window_id: Some(window_id), + first_window_id: Some(window_ids.first_window_id.to_string()), + previous_window_id: window_ids.previous_window_id.map(|id| id.to_string()), + window_id: Some(window_ids.window_id.to_string()), })); let user2 = ResponseItem::Message { @@ -9955,12 +9978,14 @@ async fn sample_rollout( let user_messages2 = collect_user_messages(&snapshot2); let rebuilt2 = compact::build_compacted_history(Vec::new(), &user_messages2, summary2); live_history.replace(rebuilt2); - let (window_number, window_id) = session.advance_auto_compact_window().await; + let (window_number, window_ids) = session.advance_auto_compact_window().await; rollout_items.push(RolloutItem::Compacted(CompactedItem { message: summary2.to_string(), replacement_history: None, window_number: Some(window_number), - window_id: Some(window_id), + first_window_id: Some(window_ids.first_window_id.to_string()), + previous_window_id: window_ids.previous_window_id.map(|id| id.to_string()), + window_id: Some(window_ids.window_id.to_string()), })); let user3 = ResponseItem::Message { diff --git a/codex-rs/core/src/state/auto_compact_window.rs b/codex-rs/core/src/state/auto_compact_window.rs index e1fde2750..ff5d7a2a2 100644 --- a/codex-rs/core/src/state/auto_compact_window.rs +++ b/codex-rs/core/src/state/auto_compact_window.rs @@ -1,6 +1,13 @@ use codex_protocol::protocol::TokenUsage; use uuid::Uuid; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct AutoCompactWindowIds { + pub(crate) first_window_id: Uuid, + pub(crate) previous_window_id: Option, + pub(crate) window_id: Uuid, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct AutoCompactWindowSnapshot { pub(crate) prefill_input_tokens: Option, @@ -15,7 +22,7 @@ enum AutoCompactWindowPrefill { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) struct AutoCompactWindow { window_number: u64, - window_id: Uuid, + ids: AutoCompactWindowIds, new_context_window_requested: bool, /// Absolute input-token baseline for the current compaction window. /// @@ -27,9 +34,14 @@ pub(super) struct AutoCompactWindow { impl AutoCompactWindow { pub(super) fn new() -> Self { + let window_id = Uuid::now_v7(); Self { window_number: 0, - window_id: Uuid::now_v7(), + ids: AutoCompactWindowIds { + first_window_id: window_id, + previous_window_id: None, + window_id, + }, new_context_window_requested: false, prefill_input_tokens: None, } @@ -43,20 +55,21 @@ impl AutoCompactWindow { self.window_number } - pub(super) fn window_id(&self) -> Uuid { - self.window_id + pub(super) fn ids(&self) -> AutoCompactWindowIds { + self.ids } - pub(super) fn restore(&mut self, window_number: u64, window_id: Uuid) { + pub(super) fn restore(&mut self, window_number: u64, ids: AutoCompactWindowIds) { self.window_number = window_number; - self.window_id = window_id; + self.ids = ids; } - pub(super) fn advance(&mut self) -> (u64, Uuid) { + pub(super) fn advance(&mut self) -> (u64, AutoCompactWindowIds) { self.window_number = self.window_number.saturating_add(1); - self.window_id = Uuid::now_v7(); + self.ids.previous_window_id = Some(self.ids.window_id); + self.ids.window_id = Uuid::now_v7(); self.new_context_window_requested = false; - (self.window_number, self.window_id) + (self.window_number, self.ids) } pub(super) fn request_new_context_window(&mut self) { @@ -118,21 +131,41 @@ mod tests { let mut window = AutoCompactWindow::new(); assert_eq!(window.window_number(), 0); - assert_eq!(window.window_id().get_version_num(), 7); + let initial_window_id = window.ids().window_id; + assert_eq!(initial_window_id.get_version_num(), 7); + assert_eq!( + window.ids(), + AutoCompactWindowIds { + first_window_id: initial_window_id, + previous_window_id: None, + window_id: initial_window_id, + } + ); + let first_window_id = initial_window_id; let restored_window_id = Uuid::now_v7(); - window.restore(/*window_number*/ 3, restored_window_id); + let restored_previous_window_id = Uuid::now_v7(); + window.restore( + /*window_number*/ 3, + AutoCompactWindowIds { + first_window_id, + previous_window_id: Some(restored_previous_window_id), + window_id: restored_window_id, + }, + ); assert_eq!(window.window_number(), 3); - assert_eq!(window.window_id(), restored_window_id); + assert_eq!(window.ids().window_id, restored_window_id); window.request_new_context_window(); assert!(window.take_new_context_window_request()); assert!(!window.take_new_context_window_request()); window.request_new_context_window(); - let (window_number, window_id) = window.advance(); + let (window_number, ids) = window.advance(); assert_eq!(window_number, 4); assert_eq!(window.window_number(), 4); - assert_eq!(window.window_id(), window_id); - assert_eq!(window_id.get_version_num(), 7); - assert_ne!(window_id, restored_window_id); + assert_eq!(window.ids(), ids); + assert_eq!(ids.first_window_id, first_window_id); + assert_eq!(ids.previous_window_id, Some(restored_window_id)); + assert_eq!(ids.window_id.get_version_num(), 7); + assert_ne!(ids.window_id, restored_window_id); assert!(!window.take_new_context_window_request()); assert_eq!( diff --git a/codex-rs/core/src/state/mod.rs b/codex-rs/core/src/state/mod.rs index c693221be..20ff78b54 100644 --- a/codex-rs/core/src/state/mod.rs +++ b/codex-rs/core/src/state/mod.rs @@ -5,6 +5,7 @@ mod session; mod turn; pub(crate) use additional_context::AdditionalContextStore; +pub(crate) use auto_compact_window::AutoCompactWindowIds; pub(crate) use auto_compact_window::AutoCompactWindowSnapshot; pub(crate) use service::SessionServices; pub(crate) use session::SessionState; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 83cef227f..824f5f636 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -6,10 +6,10 @@ use codex_sandboxing::policy_transforms::merge_permission_profiles; use std::collections::HashMap; use std::collections::HashSet; use std::collections::VecDeque; -use uuid::Uuid; use super::AdditionalContextStore; use super::auto_compact_window::AutoCompactWindow; +use super::auto_compact_window::AutoCompactWindowIds; use super::auto_compact_window::AutoCompactWindowSnapshot; use crate::context_manager::ContextManager; use crate::session::PreviousTurnSettings; @@ -152,15 +152,19 @@ impl SessionState { self.auto_compact_window.window_number() } - pub(crate) fn auto_compact_window_id(&self) -> Uuid { - self.auto_compact_window.window_id() + pub(crate) fn auto_compact_window_ids(&self) -> AutoCompactWindowIds { + self.auto_compact_window.ids() } - pub(crate) fn restore_auto_compact_window(&mut self, window_number: u64, window_id: Uuid) { - self.auto_compact_window.restore(window_number, window_id); + pub(crate) fn restore_auto_compact_window( + &mut self, + window_number: u64, + ids: AutoCompactWindowIds, + ) { + self.auto_compact_window.restore(window_number, ids); } - pub(crate) fn advance_auto_compact_window(&mut self) -> (u64, Uuid) { + pub(crate) fn advance_auto_compact_window(&mut self) -> (u64, AutoCompactWindowIds) { self.auto_compact_window.advance() } @@ -168,7 +172,9 @@ impl SessionState { self.auto_compact_window.request_new_context_window(); } - pub(crate) fn start_new_context_window_if_requested(&mut self) -> Option<(u64, Uuid)> { + pub(crate) fn start_new_context_window_if_requested( + &mut self, + ) -> Option<(u64, AutoCompactWindowIds)> { if !self.auto_compact_window.take_new_context_window_request() { return None; } diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap b/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap index 12bab10ec..3b933b4ed 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap @@ -1,5 +1,6 @@ --- source: core/tests/suite/token_budget.rs +assertion_line: 539 expression: snapshot --- Scenario: New context window tool installs fresh full context before the next follow-up request. @@ -8,7 +9,7 @@ Scenario: New context window tool installs fresh full context before the next fo 00:message/developer[3]: [01] [02] - [03] \nThread id .\nCurrent context window id .\nYou have 121600 tokens left in this context window.\n + [03] \nThread id .\nFirst context window id .\nPrevious context window id .\nCurrent context window id .\nYou have 121600 tokens left in this context window.\n 01:message/user:> 02:function_call/update_plan 03:function_call_output:Plan updated diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index db70d3f33..7c4dfa1be 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -35,6 +35,35 @@ fn token_budget_texts(request: &ResponsesRequest) -> Vec { .collect() } +fn token_budget_window_ids( + text: &str, + thread_id: codex_protocol::ThreadId, + tokens_left: i64, +) -> (String, Option, String) { + let captures = assert_regex_match( + &format!( + r"^\nThread id {thread_id}\.\nFirst context window id ([0-9a-f-]{{36}})\.\nPrevious context window id (none|[0-9a-f-]{{36}})\.\nCurrent context window id ([0-9a-f-]{{36}})\.\nYou have {tokens_left} tokens left in this context window\.\n$" + ), + text, + ); + let first_window_id = captures + .get(1) + .expect("first window id capture") + .as_str() + .to_string(); + let previous_window_id = captures + .get(2) + .expect("previous window id capture") + .as_str(); + let previous_window_id = (previous_window_id != "none").then(|| previous_window_id.to_string()); + let window_id = captures + .get(3) + .expect("window id capture") + .as_str() + .to_string(); + (first_window_id, previous_window_id, window_id) +} + fn tool_names(request: &ResponsesRequest) -> Vec { request .body_json() @@ -83,12 +112,13 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> let thread_id = test.session_configured.thread_id; let initial_token_budget = token_budget_texts(&requests[0]); assert_eq!(initial_token_budget.len(), 1); - assert_regex_match( - &format!( - r"^\nThread id {thread_id}\.\nCurrent context window id [0-9a-f-]{{36}}\.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window\.\n$" - ), + let (first_window_id, previous_window_id, window_id) = token_budget_window_ids( &initial_token_budget[0], + thread_id, + EFFECTIVE_CONTEXT_WINDOW, ); + assert_eq!(previous_window_id, None); + assert_eq!(first_window_id, window_id); assert_eq!( token_budget_texts(&requests[1]), initial_token_budget, @@ -147,12 +177,7 @@ async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> R let thread_id = test.session_configured.thread_id; let full_context = token_budget_texts(&requests[0]); assert_eq!(full_context.len(), 1); - assert_regex_match( - &format!( - r"^\nThread id {thread_id}\.\nCurrent context window id [0-9a-f-]{{36}}\.\nYou have 9500 tokens left in this context window\.\n$" - ), - &full_context[0], - ); + token_budget_window_ids(&full_context[0], thread_id, /*tokens_left*/ 9_500); let full_context = full_context[0].clone(); let threshold_25 = "\nYou have 7000 tokens left in this context window.\n" @@ -245,12 +270,7 @@ async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Resu .to_string(); let token_budgets = token_budget_texts(&requests[1]); assert_eq!(token_budgets.len(), 2); - assert_regex_match( - &format!( - r"^\nThread id {thread_id}\.\nCurrent context window id [0-9a-f-]{{36}}\.\nYou have 9500 tokens left in this context window\.\n$" - ), - &token_budgets[0], - ); + token_budget_window_ids(&token_budgets[0], thread_id, /*tokens_left*/ 9_500); assert_eq!(token_budgets[1], remaining_context); assert_eq!( requests[2].function_call_output_content_and_success(call_id), @@ -373,28 +393,30 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> { let thread_id = test.session_configured.thread_id; let initial_token_budget = token_budget_texts(&requests[0]); assert_eq!(initial_token_budget.len(), 1); - let initial_window_id = assert_regex_match( - &format!( - r"^\nThread id {thread_id}\.\nCurrent context window id ([0-9a-f-]{{36}})\.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window\.\n$" - ), - &initial_token_budget[0], - ) - .get(1) - .expect("window id capture") - .as_str() - .to_string(); + let (initial_first_window_id, initial_previous_window_id, initial_window_id) = + token_budget_window_ids( + &initial_token_budget[0], + thread_id, + EFFECTIVE_CONTEXT_WINDOW, + ); let post_compaction_token_budget = token_budget_texts(&requests[2]); assert_eq!(post_compaction_token_budget.len(), 1); - let post_compaction_window_id = assert_regex_match( - &format!( - r"^\nThread id {thread_id}\.\nCurrent context window id ([0-9a-f-]{{36}})\.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window\.\n$" - ), + let ( + post_compaction_first_window_id, + post_compaction_previous_window_id, + post_compaction_window_id, + ) = token_budget_window_ids( &post_compaction_token_budget[0], - ) - .get(1) - .expect("window id capture") - .as_str() - .to_string(); + thread_id, + EFFECTIVE_CONTEXT_WINDOW, + ); + assert_eq!(initial_previous_window_id, None); + assert_eq!(initial_first_window_id, initial_window_id); + assert_eq!(post_compaction_first_window_id, initial_first_window_id); + assert_eq!( + post_compaction_previous_window_id.as_deref(), + Some(initial_window_id.as_str()) + ); assert_ne!(post_compaction_window_id, initial_window_id); Ok(()) @@ -458,29 +480,24 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { let thread_id = test.session_configured.thread_id; let initial_token_budget = token_budget_texts(&requests[0]); assert_eq!(initial_token_budget.len(), 1); - let initial_window_id = assert_regex_match( - &format!( - r"^\nThread id {thread_id}\.\nCurrent context window id ([0-9a-f-]{{36}})\.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window\.\n$" - ), + let (initial_first_window_id, _, initial_window_id) = token_budget_window_ids( &initial_token_budget[0], - ) - .get(1) - .expect("window id capture") - .as_str() - .to_string(); + thread_id, + EFFECTIVE_CONTEXT_WINDOW, + ); let new_window_token_budget = token_budget_texts(&requests[2]); assert_eq!(new_window_token_budget.len(), 1); - let window_id = assert_regex_match( - &format!( - r"^\nThread id {thread_id}\.\nCurrent context window id ([0-9a-f-]{{36}})\.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window\.\n$" - ), + let (new_first_window_id, new_previous_window_id, new_window_id) = token_budget_window_ids( &new_window_token_budget[0], - ) - .get(1) - .expect("window id capture") - .as_str() - .to_string(); - assert_ne!(window_id, initial_window_id); + thread_id, + EFFECTIVE_CONTEXT_WINDOW, + ); + assert_eq!(new_first_window_id, initial_first_window_id); + assert_eq!( + new_previous_window_id.as_deref(), + Some(initial_window_id.as_str()) + ); + assert_ne!(new_window_id, initial_window_id); assert!( !requests[2].body_contains_text("request new context window"), "new_context should drop the prior window history before continuing the turn" @@ -496,7 +513,8 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { ); let snapshot = snapshot .replace(&thread_id.to_string(), "") - .replace(&window_id, ""); + .replace(&new_first_window_id, "") + .replace(&new_window_id, ""); insta::assert_snapshot!( "token_budget_new_context_window_tool_full_context", snapshot diff --git a/codex-rs/protocol/src/compacted_item.rs b/codex-rs/protocol/src/compacted_item.rs index 341537389..0f437e0bb 100644 --- a/codex-rs/protocol/src/compacted_item.rs +++ b/codex-rs/protocol/src/compacted_item.rs @@ -23,6 +23,8 @@ impl<'de> Deserialize<'de> for CompactedItem { message: serialized.message, replacement_history: serialized.replacement_history, window_number, + first_window_id: serialized.first_window_id, + previous_window_id: serialized.previous_window_id, window_id, }) } @@ -36,6 +38,10 @@ struct SerializedCompactedItem { #[serde(default)] window_number: Option, #[serde(default)] + first_window_id: Option, + #[serde(default)] + previous_window_id: Option, + #[serde(default)] window_id: Option, } @@ -59,6 +65,8 @@ mod tests { message: "summary".to_string(), replacement_history: None, window_number: Some(3), + first_window_id: Some("019b3f6e-0000-7000-8000-000000000001".to_string()), + previous_window_id: Some("019b3f6e-0000-7000-8000-000000000002".to_string()), window_id: Some("019b3f6e-7a10-7cc3-8b6e-1d09e2f7a001".to_string()), }; @@ -67,6 +75,8 @@ mod tests { json!({ "message": "summary", "window_number": 3, + "first_window_id": "019b3f6e-0000-7000-8000-000000000001", + "previous_window_id": "019b3f6e-0000-7000-8000-000000000002", "window_id": "019b3f6e-7a10-7cc3-8b6e-1d09e2f7a001", }) ); @@ -86,6 +96,8 @@ mod tests { message: "summary".to_string(), replacement_history: None, window_number: Some(3), + first_window_id: None, + previous_window_id: None, window_id: None, } ); diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 1fe732c4f..081649b7a 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2996,6 +2996,12 @@ pub struct CompactedItem { /// Monotonic position of this context window within the thread. #[serde(default, skip_serializing_if = "Option::is_none")] pub window_number: Option, + /// UUIDv7 identity of the first context window in this thread's window chain. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub first_window_id: Option, + /// UUIDv7 identity of the context window immediately before this one. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub previous_window_id: Option, /// UUIDv7 identity of this context window. #[serde(default, skip_serializing_if = "Option::is_none")] pub window_id: Option, diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index 446743955..e6771e37c 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -155,6 +155,8 @@ fn builder_from_items_falls_back_to_filename() { message: "noop".to_string(), replacement_history: None, window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, })]; diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index b45bb2d63..0336006fb 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -484,6 +484,8 @@ mod tests { message: "compacted".to_string(), replacement_history: None, window_number: None, + first_window_id: None, + previous_window_id: None, window_id: None, });