From 7a55e553169dfe965dae20eeca160bbc07558842 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 22 Jun 2026 18:31:47 -0700 Subject: [PATCH] core: persist initial context window metadata Record the initial AutoCompactWindowIds in session_meta.context_window so JSONL tailers can observe the first context window before any compaction item exists. Thread creation now passes the same initial IDs to rollout persistence and SessionState, and resume reconstruction uses the session_meta value as a fallback behind compacted window metadata. Verification: - just test -p codex-rollout recorder_materializes_on_flush_with_pending_items - just test -p codex-core reconstruct_history_restores_initial_window_from_session_meta - just test -p codex-core reconstruct_history_prefers_compacted_window_over_session_meta - just test -p codex-thread-store - just test -p codex-state - just test -p codex-app-server thread_read_returns_summary_without_turns - just fix -p codex-protocol -p codex-rollout -p codex-thread-store -p codex-state -p codex-core -p codex-app-server - just fix -p codex-rollout - just fix -p codex-thread-store --- .../external_agent_session_import.rs | 1 + codex-rs/app-server/tests/common/rollout.rs | 2 + .../tests/suite/conversation_summary.rs | 1 + .../tests/suite/v2/remote_thread_store.rs | 1 + .../app-server/tests/suite/v2/thread_read.rs | 1 + .../tests/suite/v2/thread_resume.rs | 1 + .../tests/suite/v2/thread_unarchive.rs | 1 + .../context/world_state/environment_tests.rs | 2 +- .../core/src/personality_migration_tests.rs | 1 + .../src/session/rollout_reconstruction.rs | 28 ++++++- .../session/rollout_reconstruction_tests.rs | 83 +++++++++++++++++++ codex-rs/core/src/session/session.rs | 10 ++- codex-rs/core/src/session/tests.rs | 2 + .../core/src/state/auto_compact_window.rs | 33 ++++++-- codex-rs/core/src/state/session.rs | 13 ++- .../core/tests/suite/personality_migration.rs | 2 + codex-rs/core/tests/suite/sqlite_state.rs | 1 + codex-rs/protocol/src/protocol.rs | 17 ++++ codex-rs/rollout/src/compression_tests.rs | 1 + codex-rs/rollout/src/metadata_tests.rs | 3 + codex-rs/rollout/src/recorder.rs | 16 ++++ codex-rs/rollout/src/recorder_tests.rs | 13 ++- codex-rs/rollout/src/session_index_tests.rs | 1 + codex-rs/rollout/src/state_db_tests.rs | 1 + codex-rs/rollout/src/tests.rs | 1 + codex-rs/state/src/extract.rs | 2 + codex-rs/state/src/runtime/threads.rs | 2 + codex-rs/thread-store/src/in_memory.rs | 2 + .../thread-store/src/local/create_thread.rs | 3 +- codex-rs/thread-store/src/local/mod.rs | 1 + codex-rs/thread-store/src/types.rs | 4 + 31 files changed, 238 insertions(+), 12 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs index 0e17cb5737..a51bbf4e30 100644 --- a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs +++ b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs @@ -215,6 +215,7 @@ impl ExternalAgentSessionImporter { }, dynamic_tools: Vec::new(), multi_agent_version: Some(MultiAgentVersion::V1), + initial_context_window: None, metadata: ThreadPersistenceMetadata { cwd: Some(cwd.clone()), model_provider: model_provider.clone(), diff --git a/codex-rs/app-server/tests/common/rollout.rs b/codex-rs/app-server/tests/common/rollout.rs index 7b389e8fc1..ec01dacc47 100644 --- a/codex-rs/app-server/tests/common/rollout.rs +++ b/codex-rs/app-server/tests/common/rollout.rs @@ -202,6 +202,7 @@ fn create_fake_rollout_with_source_and_parent_thread_id( dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }; let payload = serde_json::to_value(SessionMetaLine { meta, @@ -289,6 +290,7 @@ pub fn create_fake_rollout_with_text_elements( dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }; let payload = serde_json::to_value(SessionMetaLine { meta, diff --git a/codex-rs/app-server/tests/suite/conversation_summary.rs b/codex-rs/app-server/tests/suite/conversation_summary.rs index 17f73db3ed..64f7818a66 100644 --- a/codex-rs/app-server/tests/suite/conversation_summary.rs +++ b/codex-rs/app-server/tests/suite/conversation_summary.rs @@ -131,6 +131,7 @@ async fn get_conversation_summary_by_thread_id_reads_pathless_store_thread() -> base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_context_window: None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index 6deec70ffe..e546a55395 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -157,6 +157,7 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_context_window: None, metadata: ThreadPersistenceMetadata { cwd: Some(codex_home.path().to_path_buf()), model_provider: "mock_provider".to_string(), diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index 58990aaf04..c71f69d8cc 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -1367,6 +1367,7 @@ async fn seed_pathless_store_thread( base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_context_window: None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index ea7a1a9d46..a67282a82d 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -2072,6 +2072,7 @@ stream_max_retries = 0 dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }; std::fs::write( &rollout_path, diff --git a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs index 46a989545c..7c13234aa8 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs @@ -218,6 +218,7 @@ async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> { base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_context_window: None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), diff --git a/codex-rs/core/src/context/world_state/environment_tests.rs b/codex-rs/core/src/context/world_state/environment_tests.rs index 77a3adde17..961ee5ed73 100644 --- a/codex-rs/core/src/context/world_state/environment_tests.rs +++ b/codex-rs/core/src/context/world_state/environment_tests.rs @@ -251,6 +251,6 @@ fn user_message(text: &str) -> ResponseItem { text: text.to_string(), }], phase: None, - metadata: None, + internal_chat_message_metadata_passthrough: None, } } diff --git a/codex-rs/core/src/personality_migration_tests.rs b/codex-rs/core/src/personality_migration_tests.rs index 57f06317a1..dd0da327c4 100644 --- a/codex-rs/core/src/personality_migration_tests.rs +++ b/codex-rs/core/src/personality_migration_tests.rs @@ -61,6 +61,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: None, }; diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index bf534c1e8a..d9a9e71cb8 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 codex_protocol::protocol::SessionContextWindow; use uuid::Uuid; // Return value of `Session::reconstruct_history_from_rollout`, bundling the rebuilt history with @@ -113,6 +114,14 @@ impl Session { // stopping once a surviving replacement-history checkpoint and the required resume metadata // are both known; then replay only the buffered surviving tail forward to preserve exact // history semantics. + let initial_window = rollout_items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(session_meta) => session_meta + .meta + .context_window + .as_ref() + .and_then(reconstructed_window_from_session_context_window), + _ => None, + }); let mut base_replacement_history: Option<&[ResponseItem]> = None; let mut previous_turn_settings = None; let mut reference_context_item = TurnReferenceContextItem::NeverSet; @@ -348,7 +357,7 @@ impl Session { reference_context_item }; - let window = window.unwrap_or(ReconstructedWindow { + let window = window.or(initial_window).unwrap_or(ReconstructedWindow { number: fallback_window_number, first_id: None, previous_id: None, @@ -371,3 +380,20 @@ fn parse_uuid_v7(value: &str) -> Option { .ok() .filter(|uuid| uuid.get_version_num() == 7) } + +fn reconstructed_window_from_session_context_window( + context_window: &SessionContextWindow, +) -> Option { + let first_id = parse_uuid_v7(&context_window.first_window_id)?; + let previous_id = match context_window.previous_window_id.as_deref() { + Some(previous_window_id) => Some(parse_uuid_v7(previous_window_id)?), + None => None, + }; + let id = parse_uuid_v7(&context_window.window_id)?; + Some(ReconstructedWindow { + number: context_window.window_number, + first_id: Some(first_id), + previous_id, + id: Some(id), + }) +} diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index b28bd37c79..8b353d0133 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -9,8 +9,12 @@ use codex_protocol::protocol::CompactedItem; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::ResumedHistory; +use codex_protocol::protocol::SessionContextWindow; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; use pretty_assertions::assert_eq; use std::path::PathBuf; +use uuid::Uuid; fn user_message(text: &str) -> ResponseItem { ResponseItem::Message { @@ -907,6 +911,85 @@ async fn record_initial_history_resumed_does_not_seed_reference_context_item_aft assert!(session.reference_context_item().await.is_none()); } +#[tokio::test] +async fn reconstruct_history_restores_initial_window_from_session_meta() { + let (session, turn_context) = make_session_and_context().await; + let thread_id = ThreadId::default(); + let initial_window_id = Uuid::now_v7(); + let rollout_items = vec![RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + session_id: thread_id.into(), + id: thread_id, + context_window: Some(SessionContextWindow { + window_number: 0, + first_window_id: initial_window_id.to_string(), + previous_window_id: None, + window_id: initial_window_id.to_string(), + }), + ..SessionMeta::default() + }, + git: None, + })]; + + let reconstructed = session + .reconstruct_history_from_rollout(&turn_context, &rollout_items) + .await; + + assert_eq!(reconstructed.window_number, 0); + assert_eq!(reconstructed.first_window_id, Some(initial_window_id)); + assert_eq!(reconstructed.previous_window_id, None); + assert_eq!(reconstructed.window_id, Some(initial_window_id)); +} + +#[tokio::test] +async fn reconstruct_history_prefers_compacted_window_over_session_meta() { + let (session, turn_context) = make_session_and_context().await; + let thread_id = ThreadId::default(); + let initial_window_id = Uuid::now_v7(); + let compacted_first_window_id = Uuid::now_v7(); + let compacted_previous_window_id = Uuid::now_v7(); + let compacted_window_id = Uuid::now_v7(); + let rollout_items = vec![ + RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + session_id: thread_id.into(), + id: thread_id, + context_window: Some(SessionContextWindow { + window_number: 0, + first_window_id: initial_window_id.to_string(), + previous_window_id: None, + window_id: initial_window_id.to_string(), + }), + ..SessionMeta::default() + }, + git: None, + }), + RolloutItem::Compacted(CompactedItem { + message: String::new(), + replacement_history: Some(Vec::new()), + window_number: Some(2), + first_window_id: Some(compacted_first_window_id.to_string()), + previous_window_id: Some(compacted_previous_window_id.to_string()), + window_id: Some(compacted_window_id.to_string()), + }), + ]; + + let reconstructed = session + .reconstruct_history_from_rollout(&turn_context, &rollout_items) + .await; + + assert_eq!(reconstructed.window_number, 2); + assert_eq!( + reconstructed.first_window_id, + Some(compacted_first_window_id) + ); + assert_eq!( + reconstructed.previous_window_id, + Some(compacted_previous_window_id) + ); + assert_eq!(reconstructed.window_id, Some(compacted_window_id)); +} + #[tokio::test] async fn reconstruct_history_legacy_compaction_without_replacement_history_does_not_inject_current_initial_context() { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index b4740bea99..334082d4c7 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -546,6 +546,7 @@ impl Session { SessionId::from(thread_id) } }); + let initial_auto_compact_window_ids = AutoCompactWindowIds::new_initial(); let agent_control = agent_control.with_session_id( session_id, config @@ -585,6 +586,10 @@ impl Session { }, dynamic_tools: session_configuration.dynamic_tools.clone(), multi_agent_version: initial_multi_agent_version, + initial_context_window: Some( + initial_auto_compact_window_ids + .to_session_context_window(/*window_number*/ 0), + ), metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), @@ -898,7 +903,10 @@ impl Session { session_configuration.thread_name = thread_name.clone(); validate_config_lock_if_configured(&session_configuration).await?; export_config_lock_if_configured(&session_configuration, thread_id).await?; - let state = SessionState::new(session_configuration.clone()); + let state = SessionState::new_with_auto_compact_window_ids( + session_configuration.clone(), + initial_auto_compact_window_ids, + ); let managed_network_requirements_configured = config .config_layer_stack .requirements_toml() diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index ff71775aa5..f05b1be671 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3760,6 +3760,7 @@ async fn attach_thread_persistence(session: &mut Session) -> PathBuf { base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_context_window: None, metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), @@ -6609,6 +6610,7 @@ async fn shutdown_complete_does_not_append_to_thread_store_after_shutdown() { base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_context_window: None, metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), model_provider: config.model_provider_id.clone(), diff --git a/codex-rs/core/src/state/auto_compact_window.rs b/codex-rs/core/src/state/auto_compact_window.rs index 419e85be04..6ba8ff14c4 100644 --- a/codex-rs/core/src/state/auto_compact_window.rs +++ b/codex-rs/core/src/state/auto_compact_window.rs @@ -1,3 +1,4 @@ +use codex_protocol::protocol::SessionContextWindow; use codex_protocol::protocol::TokenUsage; use uuid::Uuid; @@ -8,6 +9,26 @@ pub(crate) struct AutoCompactWindowIds { pub(crate) window_id: Uuid, } +impl AutoCompactWindowIds { + pub(crate) fn new_initial() -> Self { + let window_id = Uuid::now_v7(); + Self { + first_window_id: window_id, + previous_window_id: None, + window_id, + } + } + + pub(crate) fn to_session_context_window(self, window_number: u64) -> SessionContextWindow { + SessionContextWindow { + window_number, + first_window_id: self.first_window_id.to_string(), + previous_window_id: self.previous_window_id.map(|id| id.to_string()), + window_id: self.window_id.to_string(), + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct AutoCompactWindowSnapshot { pub(crate) prefill_input_tokens: Option, @@ -34,15 +55,15 @@ pub(super) struct AutoCompactWindow { } impl AutoCompactWindow { + #[cfg(test)] pub(super) fn new() -> Self { - let window_id = Uuid::now_v7(); + Self::new_with_ids(AutoCompactWindowIds::new_initial()) + } + + pub(super) fn new_with_ids(ids: AutoCompactWindowIds) -> Self { Self { window_number: 0, - ids: AutoCompactWindowIds { - first_window_id: window_id, - previous_window_id: None, - window_id, - }, + ids, new_context_window_requested: false, prefill_input_tokens: None, token_budget_reminder_delivered: false, diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 6444dc8cf3..5215594867 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -47,7 +47,18 @@ pub(crate) struct SessionState { impl SessionState { /// Create a new session state mirroring previous `State::default()` semantics. + #[cfg(test)] pub(crate) fn new(session_configuration: SessionConfiguration) -> Self { + Self::new_with_auto_compact_window_ids( + session_configuration, + AutoCompactWindowIds::new_initial(), + ) + } + + pub(crate) fn new_with_auto_compact_window_ids( + session_configuration: SessionConfiguration, + auto_compact_window_ids: AutoCompactWindowIds, + ) -> Self { let history = ContextManager::new(); Self { session_configuration, @@ -57,7 +68,7 @@ impl SessionState { mcp_dependency_prompted: HashSet::new(), additional_context: AdditionalContextStore::default(), previous_turn_settings: None, - auto_compact_window: AutoCompactWindow::new(), + auto_compact_window: AutoCompactWindow::new_with_ids(auto_compact_window_ids), startup_prewarm: None, current_time_reminder: CurrentTimeReminderState::default(), active_connector_selection: HashSet::new(), diff --git a/codex-rs/core/tests/suite/personality_migration.rs b/codex-rs/core/tests/suite/personality_migration.rs index bbab961817..ca1d148728 100644 --- a/codex-rs/core/tests/suite/personality_migration.rs +++ b/codex-rs/core/tests/suite/personality_migration.rs @@ -77,6 +77,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: None, }; @@ -128,6 +129,7 @@ async fn write_rollout_with_meta_only(dir: &Path, thread_id: ThreadId) -> io::Re dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: None, }; diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index 36a51f80c3..ddfc8ef210 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -374,6 +374,7 @@ async fn backfill_scans_existing_rollouts() -> Result<()> { dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: None, }; diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 600aa475bd..9b49710f8d 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2905,6 +2905,19 @@ pub enum MultiAgentVersion { V2, } +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS)] +pub struct SessionContextWindow { + /// Monotonic position of this context window within the thread. + pub window_number: u64, + /// UUIDv7 identity of the first context window in this thread's window chain. + pub first_window_id: String, + /// 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. + pub window_id: String, +} + /// SessionMeta contains session-level data that doesn't correspond to a specific turn. /// /// NOTE: There used to be an `instructions` field here, which stored user_instructions, but we @@ -2951,6 +2964,9 @@ pub struct SessionMeta { pub memory_mode: Option, #[serde(skip_serializing_if = "Option::is_none")] pub multi_agent_version: Option, + /// Initial context-window identity for consumers that tail rollout JSONL before compaction. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub context_window: Option, } impl Default for SessionMeta { @@ -2975,6 +2991,7 @@ impl Default for SessionMeta { dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, } } } diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index 60f47d078f..a4db9a965c 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -474,6 +474,7 @@ fn write_rollout(path: &std::path::Path, thread_id: ThreadId, message: &str) -> dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: None, }; diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index 5f923f16d0..0744e7780e 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -51,6 +51,7 @@ async fn extract_metadata_from_rollout_uses_session_meta() { dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }; let session_meta_line = SessionMetaLine { meta: session_meta, @@ -107,6 +108,7 @@ async fn extract_metadata_from_rollout_returns_latest_memory_mode() { dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }; let polluted_meta = SessionMeta { memory_mode: Some("polluted".to_string()), @@ -375,6 +377,7 @@ fn write_rollout_in_sessions_with_cwd( dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }; let session_meta_line = SessionMetaLine { meta: session_meta, diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 084acb0d73..1ce0d816f4 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -57,6 +57,7 @@ use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::ResumedHistory; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::SessionContextWindow; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; @@ -91,6 +92,7 @@ pub enum RolloutRecorderParams { base_instructions: BaseInstructions, dynamic_tools: Vec, multi_agent_version: Option, + context_window: Option>, }, Resume { path: PathBuf, @@ -178,6 +180,7 @@ impl RolloutRecorderParams { base_instructions, dynamic_tools, multi_agent_version: None, + context_window: None, } } @@ -202,6 +205,17 @@ impl RolloutRecorderParams { self } + pub fn with_context_window(mut self, context_window: Option) -> Self { + if let Self::Create { + context_window: window, + .. + } = &mut self + { + *window = context_window.map(Box::new); + } + self + } + pub fn resume(path: PathBuf) -> Self { Self::Resume { path } } @@ -715,6 +729,7 @@ impl RolloutRecorder { base_instructions, dynamic_tools, multi_agent_version, + context_window, } => { let log_file_info = precompute_log_file_info(config, conversation_id)?; let path = log_file_info.path.clone(); @@ -752,6 +767,7 @@ impl RolloutRecorder { }, memory_mode: (!config.generate_memories()).then_some("disabled".to_string()), multi_agent_version, + context_window: context_window.map(|context_window| *context_window), }; (None, Some(log_file_info), path, Some(session_meta)) diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 738fde8992..31770bfc00 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -12,6 +12,7 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::protocol::SessionContextWindow; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; @@ -103,6 +104,7 @@ async fn state_db_init_backfills_before_returning() -> anyhow::Result<()> { dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: None, }; @@ -375,6 +377,13 @@ async fn recorder_materializes_on_flush_with_pending_items() -> std::io::Result< let config = test_config(home.path()); let session_id = SessionId::default(); let thread_id = ThreadId::new(); + let initial_window_id = Uuid::now_v7().to_string(); + let context_window = SessionContextWindow { + window_number: 0, + first_window_id: initial_window_id.clone(), + previous_window_id: None, + window_id: initial_window_id, + }; let recorder = RolloutRecorder::new( &config, RolloutRecorderParams::new( @@ -386,7 +395,8 @@ async fn recorder_materializes_on_flush_with_pending_items() -> std::io::Result< BaseInstructions::default(), Vec::new(), ) - .with_session_id(session_id), + .with_session_id(session_id) + .with_context_window(Some(context_window.clone())), ) .await?; @@ -437,6 +447,7 @@ async fn recorder_materializes_on_flush_with_pending_items() -> std::io::Result< panic!("expected session metadata in rollout"); }; assert_eq!(session_meta.meta.session_id, session_id); + assert_eq!(session_meta.meta.context_window, Some(context_window)); let buffered_idx = text .find("buffered-event") .expect("buffered event in rollout"); diff --git a/codex-rs/rollout/src/session_index_tests.rs b/codex-rs/rollout/src/session_index_tests.rs index 33e4ae4a59..7f57c2ba08 100644 --- a/codex-rs/rollout/src/session_index_tests.rs +++ b/codex-rs/rollout/src/session_index_tests.rs @@ -43,6 +43,7 @@ fn write_rollout_with_metadata(path: &Path, thread_id: ThreadId) -> std::io::Res dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: None, }), diff --git a/codex-rs/rollout/src/state_db_tests.rs b/codex-rs/rollout/src/state_db_tests.rs index 4b2ae4f25c..5712b65cee 100644 --- a/codex-rs/rollout/src/state_db_tests.rs +++ b/codex-rs/rollout/src/state_db_tests.rs @@ -176,6 +176,7 @@ fn write_rollout_with_user_message( dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: None, }), diff --git a/codex-rs/rollout/src/tests.rs b/codex-rs/rollout/src/tests.rs index 983fff824c..e6a3c728ef 100644 --- a/codex-rs/rollout/src/tests.rs +++ b/codex-rs/rollout/src/tests.rs @@ -1290,6 +1290,7 @@ async fn test_updated_at_uses_file_mtime() -> Result<()> { dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: None, }), diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index aa1dd04036..5b8388bf83 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -341,6 +341,7 @@ mod tests { dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: None, }), @@ -532,6 +533,7 @@ mod tests { dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: None, }), diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index a448dce56b..e4181d47db 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -1975,6 +1975,7 @@ mod tests { dynamic_tools: None, memory_mode: Some("polluted".to_string()), multi_agent_version: None, + context_window: None, }, git: None, })]; @@ -2037,6 +2038,7 @@ mod tests { dynamic_tools: None, memory_mode: None, multi_agent_version: None, + context_window: None, }, git: Some(GitInfo { commit_hash: Some(codex_git_utils::GitSha::new("rollout-sha")), diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index a93e3644c0..ab6b29699e 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -120,6 +120,7 @@ mod tests { base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_context_window: None, metadata: ThreadPersistenceMetadata { cwd: None, model_provider: "test-provider".to_string(), @@ -248,6 +249,7 @@ impl InMemoryThreadStore { memory_mode: matches!(params.metadata.memory_mode, ThreadMemoryMode::Disabled) .then_some("disabled".to_string()), multi_agent_version: params.multi_agent_version, + context_window: params.initial_context_window.clone(), ..SessionMeta::default() }; state diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs index 51a96b0815..859f3e313b 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -37,7 +37,8 @@ pub(super) async fn create_thread( params.dynamic_tools, ) .with_session_id(params.session_id) - .with_multi_agent_version(params.multi_agent_version), + .with_multi_agent_version(params.multi_agent_version) + .with_context_window(params.initial_context_window), ) .await .map_err(|err| ThreadStoreError::Internal { diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index b96e05cfd5..db9a85d2c9 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -1133,6 +1133,7 @@ mod tests { base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), multi_agent_version: None, + initial_context_window: None, metadata: thread_metadata(), } } diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 74412bb9cd..c7ac603190 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -12,6 +12,7 @@ use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::GitInfo; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::SessionContextWindow; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadMemoryMode as MemoryMode; use codex_protocol::protocol::ThreadSource; @@ -84,6 +85,9 @@ pub struct CreateThreadParams { pub dynamic_tools: Vec, /// Multi-agent runtime selected when the thread was created. pub multi_agent_version: Option, + /// Initial context-window identity captured when the thread was created. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub initial_context_window: Option, /// Metadata captured for the newly created thread. pub metadata: ThreadPersistenceMetadata, }