From 4f1a2bb5ffed8fd28518925c1d4085ead158e304 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 7 Sep 2026 18:17:37 +0000 Subject: [PATCH] Preserve fork runtime versions without loading full model context (#43545) ## Why A fork cutoff can exclude the `TurnContext` that records the source's multi-agent runtime version. Recovering that version should not require loading the source's full model context. ## What changed - Preserve `multi_agent_version` in fork session metadata before applying the cutoff. If absent, scan backward for the newest version-bearing `TurnContext`, respecting inherited segment boundaries and stopping once resolved. - Load the latest model context for paginated forks only when approval or permission settings need to be restored. ## Testing Add regression tests covering stored-version precedence, early termination without reading older segments, and version recovery within inherited segment cutoffs. GitOrigin-RevId: 00e66bcce6f4092805e273cde99bd50f304fb39e --- .../request_processors/thread_processor.rs | 1 + .../thread-store/src/local/model_context.rs | 41 +++++- .../src/local/model_context_tests.rs | 121 ++++++++++++++++++ 3 files changed, 162 insertions(+), 1 deletion(-) diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 78ea161c26..ee850c04a5 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -4944,6 +4944,7 @@ impl ThreadRequestProcessor { None }; let latest_context = if paginated_source + && needs_latest_settings && loaded_parent.is_none() && (last_turn_id.is_some() || before_turn_id.is_some()) { diff --git a/codex-rs/thread-store/src/local/model_context.rs b/codex-rs/thread-store/src/local/model_context.rs index 6d88ad53cb..c96883747d 100644 --- a/codex-rs/thread-store/src/local/model_context.rs +++ b/codex-rs/thread-store/src/local/model_context.rs @@ -1,3 +1,5 @@ +//! Reconstructs model context and preserves source runtime metadata across fork cutoffs. + use std::io; use codex_protocol::protocol::HistoryPosition; @@ -89,7 +91,7 @@ pub(super) async fn load_for_fork( .ok_or_else(|| ThreadStoreError::Internal { message: "fork lineage has no source segment".to_string(), })?; - let session_meta = codex_rollout::read_session_meta_line(source_path) + let mut session_meta = codex_rollout::read_session_meta_line(source_path) .await .map_err(|err| ThreadStoreError::Internal { message: format!( @@ -97,6 +99,43 @@ pub(super) async fn load_for_fork( source_path.display() ), })?; + if session_meta.meta.multi_agent_version.is_none() { + // Recover only the runtime version before applying the fork cutoff. Stop at the + // newest version-bearing context instead of retaining the source's full replay. + let source_lineage = lineage.clone(); + session_meta.meta.multi_agent_version = tokio::task::spawn_blocking(move || { + for segment in source_lineage.segments().iter().rev() { + let file = + codex_rollout::open_rollout_seekable_reader(segment.rollout_path.as_path())?; + let mut scanner = match segment.end.map(|end| end.end_byte_offset) { + Some(end_byte_offset) => ReverseJsonlScanner::new_at(file, end_byte_offset)?, + None => ReverseJsonlScanner::new(file)?, + }; + while let Some(outcome) = scanner.scan_next_rollout_line()? { + let ScanOutcome::Parsed(line) = outcome else { + continue; + }; + if let RolloutItem::TurnContext(context) = &line.item + && let Some(version) = context.multi_agent_version + { + return Ok(Some(version)); + } + // Ancestor metadata does not describe the immediate source's runtime. + if matches!(line.item, RolloutItem::SessionMeta(_)) { + break; + } + } + } + Ok::<_, io::Error>(None) + }) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to join fork runtime version scan: {err}"), + })? + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to read fork runtime version: {err}"), + })?; + } match history_base { Some(history_base) => { let lineage = lineage.truncate_at(history_base).await?; diff --git a/codex-rs/thread-store/src/local/model_context_tests.rs b/codex-rs/thread-store/src/local/model_context_tests.rs index 6c32c8ea07..d86f4ca652 100644 --- a/codex-rs/thread-store/src/local/model_context_tests.rs +++ b/codex-rs/thread-store/src/local/model_context_tests.rs @@ -14,6 +14,7 @@ use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::HistoryPosition; use codex_protocol::protocol::ItemCompletedEvent; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::TurnCompleteEvent; @@ -228,6 +229,126 @@ async fn fork_context_excludes_items_after_frozen_cutoff() { ); } +#[tokio::test] +async fn fork_version_stops_before_older_segments_once_resolved() { + for stored_version in [None, Some(MultiAgentVersion::Disabled)] { + let home = TempDir::new().expect("temp dir"); + let root_uuid = Uuid::from_u128(/*v*/ 3001); + let root_id = ThreadId::from_string(&root_uuid.to_string()).expect("root id"); + let root_path = write_ordinaled_paginated_rollout( + home.path(), + "2025-01-03T13-02-00", + root_uuid, + [user_message("older history")], + ); + let child_uuid = Uuid::from_u128(/*v*/ 3002); + let child_id = ThreadId::from_string(&child_uuid.to_string()).expect("child id"); + let child_path = + write_ordinaled_paginated_rollout(home.path(), "2025-01-03T13-02-01", child_uuid, []); + let mut source_meta = codex_rollout::read_session_meta_line(&child_path) + .await + .expect("read child metadata"); + source_meta.meta.multi_agent_version = stored_version; + source_meta.meta.history_base = Some(history_position( + &root_path, root_id, /*end_ordinal_exclusive*/ 2, + )); + let head = RolloutLine { + timestamp: "2025-01-03T13:02:01Z".to_string(), + ordinal: Some(0), + item: RolloutItem::SessionMeta(source_meta.clone()), + }; + std::fs::write( + &child_path, + format!( + "{}\n", + serde_json::to_string(&head).expect("serialize head") + ), + ) + .expect("write child metadata"); + let RolloutItem::TurnContext(mut context) = turn_context(home.path(), "version-turn") + else { + unreachable!(); + }; + context.multi_agent_version = Some(MultiAgentVersion::V2); + append_items( + &child_path, + [ + RolloutItem::TurnContext(context), + turn_context(home.path(), "unset-turn"), + ], + ); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let lineage = store + .resolve_rollout_lineage(child_id) + .await + .expect("resolve source lineage"); + // No checkpoint exists. A full model-context read would reach this missing segment. + std::fs::remove_file(root_path).expect("remove older segment after resolving lineage"); + + let context = load_for_fork(lineage, /*history_base*/ None) + .await + .expect("resolve version without reading older segments"); + source_meta.meta.multi_agent_version = stored_version.or(Some(MultiAgentVersion::V2)); + assert_eq!( + serde_json::to_value(context).expect("serialize fork context"), + serde_json::to_value(vec![RolloutItem::SessionMeta(source_meta)]) + .expect("serialize expected context") + ); + } +} + +#[tokio::test] +async fn fork_version_respects_inherited_segment_cutoffs() { + let home = TempDir::new().expect("temp dir"); + let root_uuid = Uuid::from_u128(/*v*/ 3003); + let root_id = ThreadId::from_string(&root_uuid.to_string()).expect("root id"); + let RolloutItem::TurnContext(mut inherited) = turn_context(home.path(), "inherited") else { + unreachable!(); + }; + inherited.multi_agent_version = Some(MultiAgentVersion::V2); + let mut excluded = inherited.clone(); + excluded.multi_agent_version = Some(MultiAgentVersion::V1); + let root_path = write_ordinaled_paginated_rollout( + home.path(), + "2025-01-03T13-02-02", + root_uuid, + [ + RolloutItem::TurnContext(inherited), + RolloutItem::TurnContext(excluded), + ], + ); + let child_uuid = Uuid::from_u128(/*v*/ 3004); + let child_id = ThreadId::from_string(&child_uuid.to_string()).expect("child id"); + let child_path = write_ordinaled_paginated_rollout( + home.path(), + "2025-01-03T13-02-03", + child_uuid, + [turn_context(home.path(), "unset-child")], + ); + set_history_base( + &child_path, + history_position(&root_path, root_id, /*end_ordinal_exclusive*/ 2), + ); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let lineage = store + .resolve_rollout_lineage(child_id) + .await + .expect("resolve source lineage"); + let mut source_meta = codex_rollout::read_session_meta_line(&child_path) + .await + .expect("read child metadata"); + + let context = load_for_fork(lineage, /*history_base*/ None) + .await + .expect("recover version from inherited prefix"); + source_meta.meta.multi_agent_version = Some(MultiAgentVersion::V2); + assert_eq!( + serde_json::to_value(context).expect("serialize fork context"), + serde_json::to_value(vec![RolloutItem::SessionMeta(source_meta)]) + .expect("serialize expected context") + ); +} + #[tokio::test] async fn loads_turn_metadata_across_an_older_checkpoint() { let home = TempDir::new().expect("temp dir");