diff --git a/codex-rs/core/src/agent/control/runtime.rs b/codex-rs/core/src/agent/control/runtime.rs index 08a8ad2efc..fa05eafb59 100644 --- a/codex-rs/core/src/agent/control/runtime.rs +++ b/codex-rs/core/src/agent/control/runtime.rs @@ -9,6 +9,7 @@ use crate::agent::registry::AgentRegistry; use crate::config::RolloutBudgetConfig; use crate::rollout_budget::RolloutBudget; use crate::session::session::Session; +use crate::state::ContextWindowIdGenerator; use crate::thread_manager::AgentTreeShutdownFailure; use crate::thread_manager::AgentTreeShutdownFailureReason; use crate::thread_manager::AgentTreeShutdownReport; @@ -195,6 +196,7 @@ pub(crate) struct LocalAgentRuntime { /// Captured at construction so delegates retain their manager's allocation policy. pub(super) thread_id_generator: ThreadIdGenerator, pub(crate) item_id_generator: ItemIdGenerator, + pub(crate) context_window_id_generator: ContextWindowIdGenerator, pub(super) agent_execution_limiter: Arc, /// Session-scoped state shared by the root thread and every cloned sub-agent control handle. pub(super) rollout_budget: Arc, @@ -221,6 +223,7 @@ impl LocalAgentRuntime { manager, thread_id_generator, item_id_generator: Session::new_item_id_generator(), + context_window_id_generator: Arc::new(codex_protocol::ContextWindowId::new_uuid_v7), registry: Arc::default(), residency: Arc::default(), mailboxes: Arc::default(), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 7677e92296..33391ccb6b 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1057,7 +1057,10 @@ impl Session { SessionId::from(thread_id) } }); - let initial_auto_compact_window_ids = AutoCompactWindowIds::new_initial(); + let context_window_id_generator = + Arc::clone(&agent_control.runtime().context_window_id_generator); + let initial_auto_compact_window_ids = + AutoCompactWindowIds::new_initial(context_window_id_generator()); let restore_child_window = matches!(&initial_history, InitialHistory::Forked(_)) && session_configuration.session_source.is_non_root_agent() && config.features.enabled(Feature::TokenBudget); @@ -1633,6 +1636,7 @@ impl Session { let mut state = SessionState::new_with_auto_compact_window_ids( session_configuration.clone(), initial_auto_compact_window_ids.clone(), + context_window_id_generator, ContextManager::for_session( session_configuration.root_capabilities_restricted, &config.features, diff --git a/codex-rs/core/src/state/auto_compact_window.rs b/codex-rs/core/src/state/auto_compact_window.rs index 600df8c0cd..f72ace9fe8 100644 --- a/codex-rs/core/src/state/auto_compact_window.rs +++ b/codex-rs/core/src/state/auto_compact_window.rs @@ -1,5 +1,9 @@ use codex_protocol::ContextWindowId; use codex_protocol::protocol::TokenUsage; +use std::sync::Arc; + +/// Allocates context window IDs. Clones share the same allocator. +pub(crate) type ContextWindowIdGenerator = Arc ContextWindowId + Send + Sync>; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct AutoCompactWindowIds { @@ -9,8 +13,7 @@ pub(crate) struct AutoCompactWindowIds { } impl AutoCompactWindowIds { - pub(crate) fn new_initial() -> Self { - let window_id = ContextWindowId::new_uuid_v7(); + pub(crate) fn new_initial(window_id: ContextWindowId) -> Self { Self { first_window_id: window_id.clone(), previous_window_id: None, @@ -74,12 +77,9 @@ impl AutoCompactWindow { self.ids = ids; } - pub(super) fn advance(&mut self) -> (u64, AutoCompactWindowIds) { + pub(super) fn advance(&mut self, window_id: ContextWindowId) -> (u64, AutoCompactWindowIds) { self.window_number = self.window_number.saturating_add(1); - self.ids.previous_window_id = Some(std::mem::replace( - &mut self.ids.window_id, - ContextWindowId::new_uuid_v7(), - )); + self.ids.previous_window_id = Some(std::mem::replace(&mut self.ids.window_id, window_id)); self.new_context_window_requested = false; self.token_budget_reminder_delivered = false; self.auto_compact_fallback_delivered = false; @@ -150,7 +150,9 @@ mod tests { #[test] fn tracks_prefill_and_window_boundaries() { - let mut window = AutoCompactWindow::new_with_ids(AutoCompactWindowIds::new_initial()); + let mut window = AutoCompactWindow::new_with_ids(AutoCompactWindowIds::new_initial( + ContextWindowId::new_uuid_v7(), + )); assert_eq!(window.window_number(), 0); let initial_window_id = window.ids().window_id.clone(); @@ -189,18 +191,14 @@ mod tests { assert!(window.take_new_context_window_request()); assert!(!window.take_new_context_window_request()); window.request_new_context_window(); - let (window_number, ids) = window.advance(); + let next_window_id = ContextWindowId::new_uuid_v7(); + let (window_number, ids) = window.advance(next_window_id.clone()); assert_eq!(window_number, 4); assert_eq!(window.window_number(), 4); assert_eq!(window.ids(), &ids); assert_eq!(ids.first_window_id, first_window_id); - assert_eq!(ids.previous_window_id, Some(restored_window_id.clone())); - assert_eq!( - uuid::Uuid::parse_str(&ids.window_id.to_string()) - .unwrap() - .get_version_num(), - 7 - ); + assert_eq!(ids.previous_window_id.as_ref(), Some(&restored_window_id)); + assert_eq!(ids.window_id, next_window_id); assert_ne!(ids.window_id, restored_window_id); assert!(!window.take_new_context_window_request()); assert!(window.claim_token_budget_reminder()); diff --git a/codex-rs/core/src/state/mod.rs b/codex-rs/core/src/state/mod.rs index 66df9927d0..8b59c9af89 100644 --- a/codex-rs/core/src/state/mod.rs +++ b/codex-rs/core/src/state/mod.rs @@ -10,6 +10,7 @@ pub(crate) use additional_context::AdditionalContextFormat; pub(crate) use additional_context::AdditionalContextStore; pub(crate) use auto_compact_window::AutoCompactWindowIds; pub(crate) use auto_compact_window::AutoCompactWindowSnapshot; +pub(crate) use auto_compact_window::ContextWindowIdGenerator; pub(crate) use service::SessionServices; pub(crate) use session::ReasoningEffortPin; pub(crate) use session::SessionState; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index a9853cf486..b7ad2a4d85 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -14,6 +14,7 @@ use super::AdditionalContextStore; use super::auto_compact_window::AutoCompactWindow; use super::auto_compact_window::AutoCompactWindowIds; use super::auto_compact_window::AutoCompactWindowSnapshot; +use super::auto_compact_window::ContextWindowIdGenerator; use crate::TurnStartOptions; use crate::context_manager::ContextManager; use crate::context_manager::HistoryReplacement; @@ -97,6 +98,7 @@ pub(crate) struct SessionState { pub(crate) turn_attribution: Option, /// Runtime accounting state for the active auto-compaction window. auto_compact_window: AutoCompactWindow, + context_window_id_generator: ContextWindowIdGenerator, pub(crate) async_compaction: crate::session::async_compaction::AsyncCompactionState, /// Original request effort for the current model while configuration updates remain active. pub(crate) reasoning_effort_pin: ReasoningEffortPin, @@ -117,9 +119,12 @@ impl SessionState { /// Create a new session state mirroring previous `State::default()` semantics. #[cfg(test)] pub(crate) fn new(session_configuration: SessionConfiguration) -> Self { + let context_window_id_generator: ContextWindowIdGenerator = + Arc::new(codex_protocol::ContextWindowId::new_uuid_v7); Self::new_with_auto_compact_window_ids( session_configuration, - AutoCompactWindowIds::new_initial(), + AutoCompactWindowIds::new_initial(context_window_id_generator()), + context_window_id_generator, ContextManager::default(), ) } @@ -127,6 +132,7 @@ impl SessionState { pub(crate) fn new_with_auto_compact_window_ids( session_configuration: SessionConfiguration, auto_compact_window_ids: AutoCompactWindowIds, + context_window_id_generator: ContextWindowIdGenerator, history: ContextManager, ) -> Self { Self { @@ -145,6 +151,7 @@ impl SessionState { last_started_turn_id: None, turn_attribution: None, auto_compact_window: AutoCompactWindow::new_with_ids(auto_compact_window_ids), + context_window_id_generator, async_compaction: crate::session::async_compaction::AsyncCompactionState::default(), reasoning_effort_pin: ReasoningEffortPin::Unset, shutting_down: false, @@ -350,7 +357,8 @@ impl SessionState { } pub(crate) fn advance_auto_compact_window(&mut self) -> (u64, AutoCompactWindowIds) { - self.auto_compact_window.advance() + self.auto_compact_window + .advance((self.context_window_id_generator)()) } pub(crate) fn request_new_context_window(&mut self) { @@ -362,7 +370,7 @@ impl SessionState { } pub(crate) fn start_new_context_window(&mut self) -> (u64, AutoCompactWindowIds) { - let window = self.auto_compact_window.advance(); + let window = self.advance_auto_compact_window(); self.auto_compact_window.clear_prefill(); window }