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 43200833df..86fc7067c3 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -4263,6 +4263,8 @@ fn read_compacted_resume_prefix_items( let file = std::fs::File::open(path)?; let reader = std::io::BufReader::new(file.take(byte_len)); let mut prefix_items = Vec::with_capacity(3); + let mut latest_metadata_candidates = Vec::new(); + let mut latest_metadata_before_checkpoint = Vec::new(); let mut saw_session_meta = false; let mut saw_user_event = false; let mut saw_user_response_item = false; @@ -4284,33 +4286,54 @@ fn read_compacted_resume_prefix_items( Ok(rollout_line) => rollout_line, Err(_) => return Ok(None), }; + let is_reconstruction_checkpoint = + rollout_item_is_reconstruction_checkpoint(&rollout_line.item); match &rollout_line.item { RolloutItem::SessionMeta(meta_line) if meta_line.meta.id == expected_thread_id && !saw_session_meta => { saw_session_meta = true; - prefix_items.push(rollout_line.item); + prefix_items.push(rollout_line.item.clone()); } RolloutItem::SessionMeta(_) => return Ok(None), RolloutItem::EventMsg(EventMsg::UserMessage(_)) if !saw_user_event => { saw_user_event = true; - prefix_items.push(rollout_line.item); + prefix_items.push(rollout_line.item.clone()); } RolloutItem::ResponseItem(codex_protocol::models::ResponseItem::Message { role, .. }) if role == "user" && !saw_user_response_item => { saw_user_response_item = true; - prefix_items.push(rollout_line.item); + prefix_items.push(rollout_line.item.clone()); + } + item if rollout_item_needs_full_resume_load(item) => { + return Ok(None); + } + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(_)) => return Ok(None), + RolloutItem::EventMsg(EventMsg::ThreadGoalUpdated(_)) => { + latest_metadata_candidates.retain(|item| { + !matches!(item, RolloutItem::EventMsg(EventMsg::ThreadGoalUpdated(_))) + }); + latest_metadata_candidates.push(rollout_line.item.clone()); + } + RolloutItem::EventMsg(EventMsg::TokenCount(_)) => { + latest_metadata_candidates + .retain(|item| !matches!(item, RolloutItem::EventMsg(EventMsg::TokenCount(_)))); + latest_metadata_candidates.push(rollout_line.item.clone()); } - item if rollout_item_needs_full_resume_load(item) => return Ok(None), _ => {} } - if saw_session_meta && saw_user_event && saw_user_response_item { - return Ok(Some(prefix_items)); + if is_reconstruction_checkpoint { + latest_metadata_before_checkpoint = latest_metadata_candidates.clone(); } } - Ok(saw_session_meta.then_some(prefix_items)) + if saw_session_meta { + prefix_items.extend(latest_metadata_before_checkpoint); + Ok(Some(prefix_items)) + } else { + Ok(None) + } } fn read_compacted_resume_suffix_items( @@ -4320,6 +4343,7 @@ fn read_compacted_resume_suffix_items( let mut suffix_items = Vec::new(); let mut found_checkpoint = false; let mut found_checkpoint_turn_context = false; + let mut found_checkpoint_user_boundary = false; let mut should_fallback = false; read_rollout_lines_from_end(path, byte_len, |line| { @@ -4335,12 +4359,29 @@ fn read_compacted_resume_suffix_items( return Ok(false); } }; + if found_checkpoint + && matches!( + &rollout_line.item, + RolloutItem::EventMsg( + EventMsg::ThreadRolledBack(_) + | EventMsg::ThreadGoalUpdated(_) + | EventMsg::TokenCount(_) + ) + ) + { + should_fallback = true; + return Ok(false); + } if rollout_item_needs_full_resume_load(&rollout_line.item) { should_fallback = true; return Ok(false); } let is_checkpoint = rollout_item_is_reconstruction_checkpoint(&rollout_line.item); let is_turn_context = matches!(&rollout_line.item, RolloutItem::TurnContext(_)); + let is_user_boundary = matches!( + &rollout_line.item, + RolloutItem::EventMsg(EventMsg::UserMessage(_)) + ); suffix_items.push(rollout_line.item); if is_checkpoint { found_checkpoint = true; @@ -4348,12 +4389,20 @@ fn read_compacted_resume_suffix_items( } if found_checkpoint && is_turn_context { found_checkpoint_turn_context = true; + return Ok(true); + } + if found_checkpoint && is_user_boundary { + found_checkpoint_user_boundary = true; return Ok(false); } Ok(true) })?; - if should_fallback || !found_checkpoint || !found_checkpoint_turn_context { + if should_fallback + || !found_checkpoint + || !found_checkpoint_turn_context + || !found_checkpoint_user_boundary + { return Ok(None); } @@ -4433,60 +4482,89 @@ fn rollout_item_needs_full_resume_load(item: &RolloutItem) -> bool { } } +const RECENT_NOT_LOADED_SUFFIX_SCAN_LIMIT_BYTES: usize = 2 * 1024 * 1024; + fn read_recent_not_loaded_rollout_items( path: &Path, page_size: usize, ) -> std::io::Result> { - let contents = std::fs::read_to_string(path)?; - if contents - .lines() - .find(|line| !line.trim().is_empty()) - .is_some_and(rollout_line_has_fork_parent_rollout_ref) - { + let metadata = std::fs::metadata(path)?; + if !metadata.is_file() || metadata.len() == 0 { + return Ok(None); + } + if rollout_head_has_fork_parent_rollout_ref(path)? { return Ok(None); } let desired_turns = page_size.saturating_add(1); - let mut record_index_after = contents - .lines() - .filter(|line| !line.trim().is_empty()) - .count(); - let mut suffix_start_index = record_index_after; let mut suffix_items = Vec::new(); - let mut turn_openings = 0usize; + let mut explicit_turn_starts = 0usize; + let mut scanned_bytes = 0usize; + let mut should_fallback = false; - for line in contents.lines().rev() { + read_rollout_lines_from_end(path, metadata.len(), |line| { let trimmed = line.trim(); if trimmed.is_empty() { - continue; + return Ok(true); + } + scanned_bytes = scanned_bytes.saturating_add(line.len()).saturating_add(1); + if scanned_bytes > RECENT_NOT_LOADED_SUFFIX_SCAN_LIMIT_BYTES + && explicit_turn_starts < desired_turns + { + should_fallback = true; + return Ok(false); } - record_index_after = record_index_after.saturating_sub(1); let rollout_line = match serde_json::from_str::(trimmed) { Ok(rollout_line) => rollout_line, - Err(_) => return Ok(None), + Err(_) => { + should_fallback = true; + return Ok(false); + } }; - if rollout_item_disables_recent_not_loaded_fast_path(&rollout_line.item) { - return Ok(None); + if !is_persisted_rollout_item(&rollout_line.item) + || rollout_item_disables_recent_not_loaded_fast_path(&rollout_line.item) + { + should_fallback = true; + return Ok(false); } - if rollout_item_opens_implicit_turn(&rollout_line.item) { - turn_openings = turn_openings.saturating_add(1); + if rollout_item_starts_explicit_turn(&rollout_line.item) { + explicit_turn_starts = explicit_turn_starts.saturating_add(1); } - suffix_start_index = record_index_after; suffix_items.push(rollout_line.item); - if turn_openings >= desired_turns { - break; + if explicit_turn_starts >= desired_turns { + return Ok(false); } + Ok(true) + })?; + + if should_fallback || explicit_turn_starts == 0 { + return Ok(None); } suffix_items.reverse(); Ok(Some(RecentRolloutItems { - initial_rollout_index: suffix_start_index, + initial_rollout_index: 0, items: suffix_items, })) } -fn rollout_item_opens_implicit_turn(item: &RolloutItem) -> bool { - matches!(item, RolloutItem::EventMsg(EventMsg::UserMessage(_))) +fn rollout_head_has_fork_parent_rollout_ref(path: &Path) -> std::io::Result { + use std::io::BufRead; + + let file = std::fs::File::open(path)?; + let reader = std::io::BufReader::new(file); + for line in reader.lines() { + let line = line?; + let trimmed = line.trim(); + if !trimmed.is_empty() { + return Ok(rollout_line_has_fork_parent_rollout_ref(trimmed)); + } + } + Ok(false) +} + +fn rollout_item_starts_explicit_turn(item: &RolloutItem) -> bool { + matches!(item, RolloutItem::EventMsg(EventMsg::TurnStarted(_))) } fn rollout_item_disables_recent_not_loaded_fast_path(item: &RolloutItem) -> bool { diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 383c56ec6f..db71136ca5 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -416,6 +416,34 @@ mod thread_processor_behavior_tests { }) } + fn compacted_resume_test_token_count() -> RolloutItem { + RolloutItem::EventMsg(EventMsg::TokenCount( + codex_protocol::protocol::TokenCountEvent { + info: None, + rate_limits: None, + }, + )) + } + + fn compacted_resume_test_goal_update(thread_id: ThreadId) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::ThreadGoalUpdated( + codex_protocol::protocol::ThreadGoalUpdatedEvent { + thread_id, + turn_id: None, + goal: codex_protocol::protocol::ThreadGoal { + thread_id, + objective: "finish the task".to_string(), + status: codex_protocol::protocol::ThreadGoalStatus::Active, + token_budget: None, + tokens_used: 0, + time_used_seconds: 0, + created_at: 1, + updated_at: 1, + }, + }, + )) + } + fn compacted_resume_test_rollout_contents(items: Vec) -> Result { let mut contents = String::new(); for item in items { @@ -462,20 +490,24 @@ mod thread_processor_behavior_tests { let items = read_compacted_resume_rollout_items(path.as_path(), thread_id, byte_len)? .expect("compacted resume items"); - assert_eq!(items.len(), 6); + assert_eq!(items.len(), 7); assert!(matches!(items[0], RolloutItem::SessionMeta(_))); assert!(matches!( items[1], RolloutItem::EventMsg(EventMsg::UserMessage(_)) )); assert!(matches!(items[2], RolloutItem::ResponseItem(_))); - assert!(matches!(items[3], RolloutItem::TurnContext(_))); + assert!(matches!( + items[3], + RolloutItem::EventMsg(EventMsg::UserMessage(_)) + )); + assert!(matches!(items[4], RolloutItem::TurnContext(_))); assert_eq!( - compacted_resume_test_replacement_text(&items[4]), + compacted_resume_test_replacement_text(&items[5]), Some("latest checkpoint") ); assert!(matches!( - &items[5], + &items[6], RolloutItem::EventMsg(EventMsg::AgentMessage(event)) if event.message == large_tail_message )); @@ -489,6 +521,67 @@ mod thread_processor_behavior_tests { Ok(()) } + #[test] + fn compacted_resume_rollout_items_carries_metadata_before_checkpoint() -> Result<()> { + let temp_dir = TempDir::new()?; + let path = temp_dir.path().join("rollout.jsonl"); + let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?; + let contents = compacted_resume_test_rollout_contents(vec![ + compacted_resume_test_session(thread_id), + compacted_resume_test_user_event("first user event"), + RolloutItem::ResponseItem(compacted_resume_test_message("user", "first user response")), + compacted_resume_test_token_count(), + compacted_resume_test_goal_update(thread_id), + compacted_resume_test_user_event("middle user event"), + compacted_resume_test_turn_context(), + compacted_resume_test_compaction("latest checkpoint", /*window_id*/ 2), + compacted_resume_test_agent_event("tail"), + ])?; + std::fs::write(&path, contents)?; + + let byte_len = std::fs::metadata(path.as_path())?.len(); + let items = read_compacted_resume_rollout_items(path.as_path(), thread_id, byte_len)? + .expect("compacted resume items"); + + assert!( + items + .iter() + .any(|item| matches!(item, RolloutItem::EventMsg(EventMsg::TokenCount(_)))) + ); + assert!( + items + .iter() + .any(|item| matches!(item, RolloutItem::EventMsg(EventMsg::ThreadGoalUpdated(_)))) + ); + Ok(()) + } + + #[test] + fn compacted_resume_rollout_items_falls_back_on_rollback_before_checkpoint() -> Result<()> { + let temp_dir = TempDir::new()?; + let path = temp_dir.path().join("rollout.jsonl"); + let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?; + let contents = compacted_resume_test_rollout_contents(vec![ + compacted_resume_test_session(thread_id), + compacted_resume_test_user_event("first user event"), + RolloutItem::ResponseItem(compacted_resume_test_message("user", "first user response")), + compacted_resume_test_user_event("middle user event"), + compacted_resume_test_turn_context(), + RolloutItem::EventMsg(EventMsg::ThreadRolledBack( + codex_protocol::protocol::ThreadRolledBackEvent { num_turns: 1 }, + )), + compacted_resume_test_compaction("rolled back checkpoint", /*window_id*/ 2), + compacted_resume_test_agent_event("tail"), + ])?; + std::fs::write(&path, contents)?; + + let byte_len = std::fs::metadata(path.as_path())?.len(); + assert!( + read_compacted_resume_rollout_items(path.as_path(), thread_id, byte_len)?.is_none() + ); + Ok(()) + } + #[test] fn compacted_resume_rollout_items_falls_back_on_malformed_tail_line() -> Result<()> { let temp_dir = TempDir::new()?; @@ -562,6 +655,12 @@ mod thread_processor_behavior_tests { let recent = read_recent_not_loaded_rollout_items(path.as_path(), /*page_size*/ 1)? .expect("recent notLoaded suffix"); + assert!( + recent + .items + .iter() + .any(|item| { matches!(item, RolloutItem::EventMsg(EventMsg::TurnStarted(_))) }) + ); assert!( recent .items @@ -588,6 +687,50 @@ mod thread_processor_behavior_tests { Ok(()) } + #[test] + fn recent_not_loaded_rollout_items_falls_back_without_explicit_turn_starts() -> Result<()> { + let temp_dir = TempDir::new()?; + let path = temp_dir.path().join("rollout.jsonl"); + let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?; + let contents = compacted_resume_test_rollout_contents(vec![ + compacted_resume_test_session(thread_id), + compacted_resume_test_user_event("first user event"), + compacted_resume_test_agent_event("first agent event"), + ])?; + std::fs::write(&path, contents)?; + + assert!(read_recent_not_loaded_rollout_items(path.as_path(), /*page_size*/ 1)?.is_none()); + Ok(()) + } + + #[test] + fn recent_not_loaded_rollout_items_falls_back_on_non_persisted_suffix_item() -> Result<()> { + let temp_dir = TempDir::new()?; + let path = temp_dir.path().join("rollout.jsonl"); + let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?; + let contents = compacted_resume_test_rollout_contents(vec![ + compacted_resume_test_session(thread_id), + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: "turn-a".to_string(), + trace_id: None, + started_at: None, + model_context_window: None, + collaboration_mode_kind: Default::default(), + }, + )), + compacted_resume_test_user_event("first user event"), + RolloutItem::ResponseItem(codex_protocol::models::ResponseItem::CompactionTrigger { + metadata: None, + }), + compacted_resume_test_agent_event("first agent event"), + ])?; + std::fs::write(&path, contents)?; + + assert!(read_recent_not_loaded_rollout_items(path.as_path(), /*page_size*/ 1)?.is_none()); + Ok(()) + } + #[test] fn validate_dynamic_tools_rejects_empty_namespace() { let tools = vec![dynamic_tool(