From d1209bddfce79e8f6cafd7f9212a2df6915aec5d Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Sat, 20 Jun 2026 13:15:49 -0700 Subject: [PATCH] core: add context window lineage IDs (#29256) ## Why The rendered `` fragment identifies the thread and current context window, but it does not expose enough lineage to identify the first window in the thread or the immediately preceding window. Those IDs also need to remain stable across compaction, resume, and rollback. ## What changed - Track first, previous, and current UUIDv7 context-window IDs in auto-compaction state. - Render `thread_id`, `first_window_id`, `previous_window_id`, and the current window ID in the full `` fragment. - Persist the first and previous window IDs in compacted rollout checkpoints and restore them during rollout reconstruction. - Preserve compatibility with older compacted records that do not contain the new optional fields. - Update focused state, rendering, reconstruction, rollback, and serialization coverage. ## Validation - `just test -p codex-core token_budget` - `just test -p codex-protocol compacted_item::tests` - `just test -p codex-core tracks_prefill_and_window_boundaries` - `just test -p codex-core reconstruct_history_uses_replacement_history_verbatim` - `just test -p codex-core thread_rollback_restores_cleared_reference_context_item_after_compaction` --- .../src/protocol/thread_history.rs | 2 + codex-rs/core/src/agent/control_tests.rs | 2 + codex-rs/core/src/compact.rs | 6 +- codex-rs/core/src/compact_remote.rs | 6 +- codex-rs/core/src/compact_remote_v2.rs | 6 +- .../core/src/context/token_budget_context.rs | 18 ++- codex-rs/core/src/session/mod.rs | 34 +++-- .../src/session/rollout_reconstruction.rs | 13 ++ .../session/rollout_reconstruction_tests.rs | 16 +++ codex-rs/core/src/session/tests.rs | 39 +++++- .../core/src/state/auto_compact_window.rs | 65 ++++++--- codex-rs/core/src/state/mod.rs | 1 + codex-rs/core/src/state/session.rs | 20 ++- ..._new_context_window_tool_full_context.snap | 3 +- codex-rs/core/tests/suite/token_budget.rs | 128 ++++++++++-------- codex-rs/protocol/src/compacted_item.rs | 12 ++ codex-rs/protocol/src/protocol.rs | 6 + codex-rs/rollout/src/metadata_tests.rs | 2 + .../thread-store/src/thread_metadata_sync.rs | 2 + 19 files changed, 277 insertions(+), 104 deletions(-) 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, });