From 5c12034e424f84af3c1e69f063203b19918fa374 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Thu, 18 Jun 2026 17:00:49 -0700 Subject: [PATCH] core: add UUIDv7 context window IDs (#28953) ## Why The token-budget context currently identifies a context window by its thread-local sequence number. A UUIDv7 gives the model a stable opaque identity that remains fixed for a window and rotates when compaction or `new_context` starts the next one. ## What changed - Preserve the existing monotonic value as `window_number` and add a UUIDv7 `window_id` to `CompactedItem`. - Generate and rotate the UUID with auto-compaction window state, persist it alongside the number, and reconstruct it on resume and rollback. - Accept legacy compacted rollout records where the numeric `window_id` represented the window number. - Use the UUID only in token-budget context; existing request headers and metadata continue using `thread_id:window_number`. ## Testing - `just test -p codex-protocol compacted_item::tests` - `just test -p codex-core token_budget` --- .../src/protocol/thread_history.rs | 1 + codex-rs/core/src/agent/control_tests.rs | 1 + codex-rs/core/src/compact.rs | 3 +- codex-rs/core/src/compact_remote.rs | 3 +- codex-rs/core/src/compact_remote_v2.rs | 3 +- .../core/src/context/token_budget_context.rs | 7 +- codex-rs/core/src/session/mod.rs | 22 ++-- .../src/session/rollout_reconstruction.rs | 48 ++++++-- .../session/rollout_reconstruction_tests.rs | 8 ++ codex-rs/core/src/session/tests.rs | 34 ++++-- .../core/src/state/auto_compact_window.rs | 40 +++++-- codex-rs/core/src/state/session.rs | 21 ++-- ..._new_context_window_tool_full_context.snap | 2 +- codex-rs/core/tests/suite/token_budget.rs | 108 +++++++++++++----- codex-rs/protocol/src/compacted_item.rs | 94 +++++++++++++++ codex-rs/protocol/src/lib.rs | 1 + codex-rs/protocol/src/protocol.rs | 8 +- codex-rs/rollout/src/metadata_tests.rs | 1 + .../thread-store/src/thread_metadata_sync.rs | 1 + 19 files changed, 319 insertions(+), 87 deletions(-) create mode 100644 codex-rs/protocol/src/compacted_item.rs 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 6b27e72d4..30221d944 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -3340,6 +3340,7 @@ mod tests { RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: None, + window_number: 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 2d4b221a9..fdd9b0837 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1105,6 +1105,7 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(replacement_history), + window_number: 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 9d79138e8..d7d2f57bf 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_id = sess.advance_auto_compact_window_id().await; + let (window_number, window_id) = sess.advance_auto_compact_window().await; if matches!( initial_context_injection, @@ -319,6 +319,7 @@ async fn run_compact_task_inner_impl( let compacted_item = CompactedItem { message: summary_text.clone(), replacement_history: Some(new_history.clone()), + window_number: Some(window_number), window_id: Some(window_id), }; sess.replace_compacted_history(new_history, reference_context_item, compacted_item) diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index c88d8062d..1eebc2317 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_id = sess.advance_auto_compact_window_id().await; + let (new_window_number, new_window_id) = sess.advance_auto_compact_window().await; new_history = process_compacted_history( sess.as_ref(), turn_context.as_ref(), @@ -274,6 +274,7 @@ async fn run_remote_compact_task_inner_impl( let compacted_item = CompactedItem { message: String::new(), replacement_history: Some(new_history.clone()), + window_number: Some(new_window_number), window_id: Some(new_window_id), }; // Install is the semantic boundary where the compact endpoint's output becomes live diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index f9c2ddc71..099b9c30b 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_id = sess.advance_auto_compact_window_id().await; + let (new_window_number, new_window_id) = sess.advance_auto_compact_window().await; let new_history = process_compacted_history( sess.as_ref(), turn_context.as_ref(), @@ -308,6 +308,7 @@ async fn run_remote_compact_task_inner_impl( let compacted_item = CompactedItem { message: String::new(), replacement_history: Some(new_history.clone()), + window_number: Some(new_window_number), window_id: Some(new_window_id), }; compaction_trace.record_installed(&CompactionCheckpointTracePayload { diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index 37b355d51..b21b61735 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -1,15 +1,16 @@ use super::ContextualUserFragment; use codex_protocol::ThreadId; +use uuid::Uuid; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct TokenBudgetContext { thread_id: ThreadId, - window_id: u64, + window_id: Uuid, tokens_left: i64, } impl TokenBudgetContext { - pub(crate) fn new(thread_id: ThreadId, window_id: u64, tokens_left: i64) -> Self { + pub(crate) fn new(thread_id: ThreadId, window_id: Uuid, tokens_left: i64) -> Self { Self { thread_id, window_id, @@ -36,7 +37,7 @@ impl ContextualUserFragment for TokenBudgetContext { let window_id = self.window_id; let tokens_left = self.tokens_left; format!( - "Thread id {thread_id}.\nCurrent context window {window_id}.\nYou have {tokens_left} tokens left in this context window." + "Thread id {thread_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 5d579f14a..d721a162d 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1352,6 +1352,7 @@ impl Session { mut history, previous_turn_settings, reference_context_item, + window_number, window_id, } = self .reconstruct_history_from_rollout(turn_context, rollout_items) @@ -1370,7 +1371,8 @@ impl Session { { let mut state = self.state.lock().await; state.replace_history(history, reference_context_item); - state.set_auto_compact_window_id(window_id); + let window_id = window_id.unwrap_or_else(|| state.auto_compact_window_id()); + state.restore_auto_compact_window(window_number, window_id); state.set_previous_turn_settings(previous_turn_settings.clone()); } let prefix_tokens = if matches!( @@ -3231,13 +3233,14 @@ impl Session { pub(crate) async fn current_window_id(&self) -> String { let state = self.state.lock().await; let thread_id = self.thread_id; - let window_id = state.auto_compact_window_id(); - format!("{thread_id}:{window_id}") + let window_number = state.auto_compact_window_number(); + format!("{thread_id}:{window_number}") } - pub(crate) async fn advance_auto_compact_window_id(&self) -> u64 { + pub(crate) async fn advance_auto_compact_window(&self) -> (u64, String) { let mut state = self.state.lock().await; - state.advance_auto_compact_window_id() + let (window_number, window_id) = state.advance_auto_compact_window(); + (window_number, window_id.to_string()) } pub(crate) async fn request_new_context_window(&self) { @@ -3249,11 +3252,11 @@ impl Session { &self, turn_context: &TurnContext, ) -> Option { - let window_id = { + let window = { let mut state = self.state.lock().await; state.start_new_context_window_if_requested() }; - let window_id = window_id?; + let (window_number, window_id) = 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; @@ -3265,7 +3268,8 @@ impl Session { RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(replacement_history), - window_id: Some(window_id), + window_number: Some(window_number), + window_id: Some(window_id.to_string()), }), RolloutItem::TurnContext(turn_context_item), ]) @@ -3275,7 +3279,7 @@ impl Session { state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact); } self.recompute_token_usage(turn_context).await; - Some(window_id) + Some(window_number) } pub(crate) async fn reference_context_item(&self) -> Option { diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index ab4838b2f..22d0990e3 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -1,5 +1,6 @@ use super::*; use crate::context_manager::is_user_turn_boundary; +use uuid::Uuid; // Return value of `Session::reconstruct_history_from_rollout`, bundling the rebuilt history with // the resume/fork hydration metadata derived from the same replay. @@ -8,7 +9,14 @@ pub(super) struct RolloutReconstruction { pub(super) history: Vec, pub(super) previous_turn_settings: Option, pub(super) reference_context_item: Option, - pub(super) window_id: u64, + pub(super) window_number: u64, + pub(super) window_id: Option, +} + +#[derive(Debug, Clone, Copy)] +struct ReconstructedWindow { + number: u64, + id: Option, } #[derive(Debug, Default)] @@ -34,7 +42,7 @@ struct ActiveReplaySegment<'a> { previous_turn_settings: Option, reference_context_item: TurnReferenceContextItem, base_replacement_history: Option<&'a [ResponseItem]>, - window_id: Option, + window: Option, } fn turn_ids_are_compatible(active_turn_id: Option<&str>, item_turn_id: Option<&str>) -> bool { @@ -47,7 +55,7 @@ fn finalize_active_segment<'a>( base_replacement_history: &mut Option<&'a [ResponseItem]>, previous_turn_settings: &mut Option, reference_context_item: &mut TurnReferenceContextItem, - window_id: &mut Option, + window: &mut Option, pending_rollback_turns: &mut usize, ) { // Thread rollback drops the newest surviving real user-message boundaries. In replay, that @@ -68,8 +76,8 @@ fn finalize_active_segment<'a>( *base_replacement_history = Some(segment_base_replacement_history); } - if window_id.is_none() { - *window_id = active_segment.window_id; + if window.is_none() { + *window = active_segment.window; } // `previous_turn_settings` come from the newest surviving user turn that established them. @@ -104,7 +112,7 @@ impl Session { let mut base_replacement_history: Option<&[ResponseItem]> = None; let mut previous_turn_settings = None; let mut reference_context_item = TurnReferenceContextItem::NeverSet; - let mut window_id = None; + let mut window = None; // Rollback is "drop the newest N user turns". While scanning in reverse, that becomes // "skip the next N user-turn segments we finalize". let mut pending_rollback_turns = 0usize; @@ -120,8 +128,13 @@ impl Session { RolloutItem::Compacted(compacted) => { let active_segment = active_segment.get_or_insert_with(ActiveReplaySegment::default); - if active_segment.window_id.is_none() { - active_segment.window_id = compacted.window_id; + if active_segment.window.is_none() + && let Some(window_number) = compacted.window_number + { + active_segment.window = Some(ReconstructedWindow { + number: window_number, + id: compacted.window_id.as_deref().and_then(parse_uuid_v7), + }); } // Looking backward, compaction clears any older baseline unless a newer // `TurnContextItem` in this same segment has already re-established it. @@ -210,7 +223,7 @@ impl Session { &mut base_replacement_history, &mut previous_turn_settings, &mut reference_context_item, - &mut window_id, + &mut window, &mut pending_rollback_turns, ); } @@ -245,12 +258,12 @@ impl Session { &mut base_replacement_history, &mut previous_turn_settings, &mut reference_context_item, - &mut window_id, + &mut window, &mut pending_rollback_turns, ); } - let fallback_window_id = u64::try_from( + let fallback_window_number = u64::try_from( rollout_items .iter() .filter(|item| matches!(item, RolloutItem::Compacted(_))) @@ -326,11 +339,22 @@ impl Session { reference_context_item }; + let window = window.unwrap_or(ReconstructedWindow { + number: fallback_window_number, + id: None, + }); RolloutReconstruction { history: history.raw_items().to_vec(), previous_turn_settings, reference_context_item, - window_id: window_id.unwrap_or(fallback_window_id), + window_number: window.number, + window_id: window.id, } } } + +fn parse_uuid_v7(value: &str) -> Option { + Uuid::parse_str(value) + .ok() + .filter(|uuid| uuid.get_version_num() == 7) +} diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index 084553f15..0b551ce72 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -826,6 +826,7 @@ async fn record_initial_history_resumed_rollback_drops_incomplete_user_turn_comp RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_number: None, window_id: None, }), RolloutItem::EventMsg(EventMsg::ThreadRolledBack( @@ -883,6 +884,7 @@ async fn record_initial_history_resumed_does_not_seed_reference_context_item_aft RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_number: None, window_id: None, }), ]; @@ -909,6 +911,7 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_does_ RolloutItem::Compacted(CompactedItem { message: "legacy summary".to_string(), replacement_history: None, + window_number: None, window_id: None, }), ]; @@ -941,6 +944,7 @@ async fn reconstruct_history_legacy_compaction_without_replacement_history_clear RolloutItem::Compacted(CompactedItem { message: "legacy summary".to_string(), replacement_history: None, + window_number: None, window_id: None, }), RolloutItem::EventMsg(EventMsg::TurnStarted( @@ -1035,6 +1039,7 @@ async fn record_initial_history_resumed_turn_context_after_compaction_reestablis RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_number: None, window_id: None, }), RolloutItem::TurnContext(previous_context_item), @@ -1185,6 +1190,7 @@ async fn record_initial_history_resumed_aborted_turn_without_id_clears_active_tu RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_number: None, window_id: None, }), ]; @@ -1419,6 +1425,7 @@ async fn record_initial_history_resumed_trailing_incomplete_turn_compaction_clea RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_number: None, window_id: None, }), ]; @@ -1583,6 +1590,7 @@ async fn record_initial_history_resumed_replaced_incomplete_compacted_turn_clear RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(Vec::new()), + window_number: 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 93faae548..67dd6becf 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -168,6 +168,7 @@ use tokio::sync::Semaphore; use tokio::time::sleep; use tokio::time::timeout; use tracing_opentelemetry::OpenTelemetrySpanExt; +use uuid::Uuid; use codex_protocol::mcp::CallToolResult as McpCallToolResult; use pretty_assertions::assert_eq; @@ -1592,7 +1593,13 @@ async fn reconstruct_history_matches_live_compactions() { .await; assert_eq!(expected, reconstructed.history); - assert_eq!(2, reconstructed.window_id); + assert_eq!(2, reconstructed.window_number); + assert_eq!( + reconstructed + .window_id + .map(|window_id| window_id.get_version_num()), + Some(7) + ); } #[tokio::test] @@ -1622,10 +1629,12 @@ async fn reconstruct_history_uses_replacement_history_verbatim() { metadata: None, }, ]; + let window_id = Uuid::now_v7(); let rollout_items = vec![RolloutItem::Compacted(CompactedItem { message: String::new(), replacement_history: Some(replacement_history.clone()), - window_id: Some(42), + window_number: Some(42), + window_id: Some(window_id.to_string()), })]; let reconstructed = session @@ -1633,7 +1642,8 @@ async fn reconstruct_history_uses_replacement_history_verbatim() { .await; assert_eq!(reconstructed.history, replacement_history); - assert_eq!(42, reconstructed.window_id); + assert_eq!(42, reconstructed.window_number); + assert_eq!(Some(window_id), reconstructed.window_id); } #[tokio::test] @@ -3014,6 +3024,7 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio user_message("turn 1 user"), user_message("summary after compaction"), ]; + let compacted_window_id = Uuid::now_v7(); sess.persist_rollout_items(&[ RolloutItem::EventMsg(EventMsg::TurnStarted( @@ -3055,7 +3066,8 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio RolloutItem::Compacted(CompactedItem { message: "summary after compaction".to_string(), replacement_history: Some(compacted_history.clone()), - window_id: Some(7), + window_number: Some(7), + window_id: Some(compacted_window_id.to_string()), }), RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { turn_id: compact_turn_id, @@ -3105,7 +3117,7 @@ async fn thread_rollback_restores_cleared_reference_context_item_after_compactio .await; { let mut state = sess.state.lock().await; - state.set_auto_compact_window_id(/*window_id*/ 99); + state.restore_auto_compact_window(/*window_number*/ 99, Uuid::now_v7()); } handlers::thread_rollback(&sess, "sub-1".to_string(), /*num_turns*/ 1).await; @@ -3114,6 +3126,10 @@ 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 + ); assert!(sess.current_window_id().await.ends_with(":7")); } @@ -9846,10 +9862,12 @@ 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; rollout_items.push(RolloutItem::Compacted(CompactedItem { message: summary1.to_string(), replacement_history: None, - window_id: None, + window_number: Some(window_number), + window_id: Some(window_id), })); let user2 = ResponseItem::Message { @@ -9889,10 +9907,12 @@ 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; rollout_items.push(RolloutItem::Compacted(CompactedItem { message: summary2.to_string(), replacement_history: None, - window_id: None, + window_number: Some(window_number), + window_id: Some(window_id), })); 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 1e3eceada..e1fde2750 100644 --- a/codex-rs/core/src/state/auto_compact_window.rs +++ b/codex-rs/core/src/state/auto_compact_window.rs @@ -1,4 +1,5 @@ use codex_protocol::protocol::TokenUsage; +use uuid::Uuid; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct AutoCompactWindowSnapshot { @@ -13,7 +14,8 @@ enum AutoCompactWindowPrefill { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) struct AutoCompactWindow { - window_id: u64, + window_number: u64, + window_id: Uuid, new_context_window_requested: bool, /// Absolute input-token baseline for the current compaction window. /// @@ -26,7 +28,8 @@ pub(super) struct AutoCompactWindow { impl AutoCompactWindow { pub(super) fn new() -> Self { Self { - window_id: 0, + window_number: 0, + window_id: Uuid::now_v7(), new_context_window_requested: false, prefill_input_tokens: None, } @@ -36,18 +39,24 @@ impl AutoCompactWindow { self.prefill_input_tokens = None; } - pub(super) fn window_id(&self) -> u64 { + pub(super) fn window_number(&self) -> u64 { + self.window_number + } + + pub(super) fn window_id(&self) -> Uuid { self.window_id } - pub(super) fn set_window_id(&mut self, window_id: u64) { + pub(super) fn restore(&mut self, window_number: u64, window_id: Uuid) { + self.window_number = window_number; self.window_id = window_id; } - pub(super) fn advance_window_id(&mut self) -> u64 { - self.window_id = self.window_id.saturating_add(1); + pub(super) fn advance(&mut self) -> (u64, Uuid) { + self.window_number = self.window_number.saturating_add(1); + self.window_id = Uuid::now_v7(); self.new_context_window_requested = false; - self.window_id + (self.window_number, self.window_id) } pub(super) fn request_new_context_window(&mut self) { @@ -108,15 +117,22 @@ mod tests { fn tracks_prefill_and_window_boundaries() { let mut window = AutoCompactWindow::new(); - assert_eq!(window.window_id(), 0); - window.set_window_id(/*window_id*/ 3); - assert_eq!(window.window_id(), 3); + assert_eq!(window.window_number(), 0); + assert_eq!(window.window_id().get_version_num(), 7); + let restored_window_id = Uuid::now_v7(); + window.restore(/*window_number*/ 3, restored_window_id); + assert_eq!(window.window_number(), 3); + assert_eq!(window.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(); - assert_eq!(window.advance_window_id(), 4); - assert_eq!(window.window_id(), 4); + let (window_number, window_id) = 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!(!window.take_new_context_window_request()); assert_eq!( diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 2bd9da5ce..83cef227f 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -6,6 +6,7 @@ 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; @@ -147,30 +148,34 @@ impl SessionState { self.auto_compact_window.snapshot() } - pub(crate) fn auto_compact_window_id(&self) -> u64 { + pub(crate) fn auto_compact_window_number(&self) -> u64 { + self.auto_compact_window.window_number() + } + + pub(crate) fn auto_compact_window_id(&self) -> Uuid { self.auto_compact_window.window_id() } - pub(crate) fn set_auto_compact_window_id(&mut self, window_id: u64) { - self.auto_compact_window.set_window_id(window_id); + 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 advance_auto_compact_window_id(&mut self) -> u64 { - self.auto_compact_window.advance_window_id() + pub(crate) fn advance_auto_compact_window(&mut self) -> (u64, Uuid) { + self.auto_compact_window.advance() } pub(crate) fn request_new_context_window(&mut self) { self.auto_compact_window.request_new_context_window(); } - pub(crate) fn start_new_context_window_if_requested(&mut self) -> Option { + pub(crate) fn start_new_context_window_if_requested(&mut self) -> Option<(u64, Uuid)> { if !self.auto_compact_window.take_new_context_window_request() { return None; } - let window_id = self.auto_compact_window.advance_window_id(); + let window = self.auto_compact_window.advance(); self.auto_compact_window.clear_prefill(); - Some(window_id) + Some(window) } pub(crate) fn token_info(&self) -> Option { 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 e3dddb45c..12bab10ec 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 @@ -8,7 +8,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 1.\nYou have 121600 tokens left in this context window.\n + [03] \nThread 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 c64251e24..db70d3f33 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -4,6 +4,7 @@ use codex_model_provider_info::built_in_model_providers; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use core_test_support::PathBufExt; +use core_test_support::assert_regex_match; use core_test_support::context_snapshot; use core_test_support::context_snapshot::ContextSnapshotOptions; use core_test_support::responses::ResponsesRequest; @@ -80,17 +81,17 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> assert_eq!(requests.len(), 2); let thread_id = test.session_configured.thread_id; - let expected = vec![format!( - "\nThread id {thread_id}.\nCurrent context window 0.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n" - )]; - assert_eq!( - token_budget_texts(&requests[0]), - expected, - "initial full context should report context window 0" + 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$" + ), + &initial_token_budget[0], ); assert_eq!( token_budget_texts(&requests[1]), - expected, + initial_token_budget, "steady-state context update should not advance the context window" ); @@ -144,9 +145,15 @@ async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> R assert_eq!(requests.len(), 5); let thread_id = test.session_configured.thread_id; - let full_context = format!( - "\nThread id {thread_id}.\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n" + 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], ); + let full_context = full_context[0].clone(); let threshold_25 = "\nYou have 7000 tokens left in this context window.\n" .to_string(); @@ -233,16 +240,18 @@ async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Resu ); let thread_id = test.session_configured.thread_id; - let full_context = format!( - "\nThread id {thread_id}.\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n" - ); let remaining_context = "\nYou have 7000 tokens left in this context window.\n" .to_string(); - assert_eq!( - token_budget_texts(&requests[1]), - vec![full_context, remaining_context.clone()] + 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], ); + assert_eq!(token_budgets[1], remaining_context); assert_eq!( requests[2].function_call_output_content_and_success(call_id), Some((Some(remaining_context), None)) @@ -362,13 +371,31 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> { assert_eq!(requests.len(), 3); let thread_id = test.session_configured.thread_id; - assert_eq!( - token_budget_texts(&requests[2]), - vec![format!( - "\nThread id {thread_id}.\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n" - )], - "post-compaction full context should report context window 1" - ); + 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 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$" + ), + &post_compaction_token_budget[0], + ) + .get(1) + .expect("window id capture") + .as_str() + .to_string(); + assert_ne!(post_compaction_window_id, initial_window_id); Ok(()) } @@ -429,12 +456,31 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { "new_context should be exposed when token budget is enabled" ); let thread_id = test.session_configured.thread_id; - assert_eq!( - token_budget_texts(&requests[2]), - vec![format!( - "\nThread id {thread_id}.\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n" - )] - ); + 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 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$" + ), + &new_window_token_budget[0], + ) + .get(1) + .expect("window id capture") + .as_str() + .to_string(); + assert_ne!(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" @@ -448,7 +494,9 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { &[("Final Follow-Up Request", &requests[2])], &ContextSnapshotOptions::default(), ); - let snapshot = snapshot.replace(&thread_id.to_string(), ""); + let snapshot = snapshot + .replace(&thread_id.to_string(), "") + .replace(&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 new file mode 100644 index 000000000..341537389 --- /dev/null +++ b/codex-rs/protocol/src/compacted_item.rs @@ -0,0 +1,94 @@ +use crate::models::ResponseItem; +use crate::protocol::CompactedItem; +use serde::Deserialize; + +// Before `window_number` was introduced, the numeric window number was serialized as +// `window_id`. Accept that shape so existing rollouts remain resumable. +impl<'de> Deserialize<'de> for CompactedItem { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let serialized = SerializedCompactedItem::deserialize(deserializer)?; + let mut window_number = serialized.window_number; + let window_id = match serialized.window_id { + Some(SerializedWindowId::Id(window_id)) => Some(window_id), + Some(SerializedWindowId::LegacyWindowNumber(legacy_window_number)) => { + window_number.get_or_insert(legacy_window_number); + None + } + None => None, + }; + Ok(Self { + message: serialized.message, + replacement_history: serialized.replacement_history, + window_number, + window_id, + }) + } +} + +#[derive(Deserialize)] +struct SerializedCompactedItem { + message: String, + #[serde(default)] + replacement_history: Option>, + #[serde(default)] + window_number: Option, + #[serde(default)] + window_id: Option, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum SerializedWindowId { + Id(String), + LegacyWindowNumber(u64), +} + +#[cfg(test)] +mod tests { + use super::*; + use anyhow::Result; + use pretty_assertions::assert_eq; + use serde_json::json; + + #[test] + fn serializes_window_number_and_id() -> Result<()> { + let item = CompactedItem { + message: "summary".to_string(), + replacement_history: None, + window_number: Some(3), + window_id: Some("019b3f6e-7a10-7cc3-8b6e-1d09e2f7a001".to_string()), + }; + + assert_eq!( + serde_json::to_value(item)?, + json!({ + "message": "summary", + "window_number": 3, + "window_id": "019b3f6e-7a10-7cc3-8b6e-1d09e2f7a001", + }) + ); + Ok(()) + } + + #[test] + fn migrates_legacy_numeric_window_id() -> Result<()> { + let item = serde_json::from_value::(json!({ + "message": "summary", + "window_id": 3, + }))?; + + assert_eq!( + item, + CompactedItem { + message: "summary".to_string(), + replacement_history: None, + window_number: Some(3), + window_id: None, + } + ); + Ok(()) + } +} diff --git a/codex-rs/protocol/src/lib.rs b/codex-rs/protocol/src/lib.rs index 3be16d08b..c50e95d37 100644 --- a/codex-rs/protocol/src/lib.rs +++ b/codex-rs/protocol/src/lib.rs @@ -10,6 +10,7 @@ pub use thread_id::ThreadId; pub use tool_name::ToolName; pub mod approvals; pub mod capabilities; +mod compacted_item; pub mod config_types; pub mod dynamic_tools; pub mod error; diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 258150aa8..285114d8b 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2971,13 +2971,17 @@ pub enum RolloutItem { EventMsg(EventMsg), } -#[derive(Serialize, Deserialize, Clone, Debug, JsonSchema, TS)] +#[derive(Serialize, Clone, Debug, PartialEq, JsonSchema, TS)] pub struct CompactedItem { pub message: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub replacement_history: Option>, + /// Monotonic position of this context window within the thread. #[serde(default, skip_serializing_if = "Option::is_none")] - pub window_id: Option, + pub window_number: Option, + /// UUIDv7 identity of this context window. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub window_id: Option, } impl From for ResponseItem { diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index 01e56f569..446743955 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -154,6 +154,7 @@ fn builder_from_items_falls_back_to_filename() { let items = vec![RolloutItem::Compacted(CompactedItem { message: "noop".to_string(), replacement_history: None, + window_number: 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 b2953d7ef..b45bb2d63 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -483,6 +483,7 @@ mod tests { let item = RolloutItem::Compacted(CompactedItem { message: "compacted".to_string(), replacement_history: None, + window_number: None, window_id: None, });