diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index d2bcd38a87..ee7fd9168b 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -218,6 +218,7 @@ use codex_core::find_archived_thread_path_by_id_str; use codex_core::find_thread_name_by_id; use codex_core::find_thread_names_by_ids; use codex_core::find_thread_path_by_id_str; +use codex_core::materialize_rollout_items_for_replay; use codex_core::mcp::auth::discover_supported_scopes; use codex_core::mcp::auth::resolve_oauth_scopes; use codex_core::mcp::collect_mcp_snapshot; @@ -235,7 +236,6 @@ use codex_core::plugins::load_plugin_apps; use codex_core::plugins::load_plugin_mcp_servers; use codex_core::read_head_for_summary; use codex_core::read_session_meta_line; -use codex_core::resolve_fork_reference_rollout_path; use codex_core::rollout_date_parts; use codex_core::sandboxing::SandboxPermissions; use codex_core::state_db::StateDbHandle; @@ -8486,7 +8486,11 @@ pub(crate) async fn read_rollout_items_from_rollout( InitialHistory::Resumed(resumed) => resumed.history, }; - Ok(materialize_rollout_items_for_replay(codex_home_from_rollout_path(path), &items).await) + if let Some(codex_home) = codex_home_from_rollout_path(path) { + return Ok(materialize_rollout_items_for_replay(codex_home, &items).await); + } + + Ok(items) } fn extract_conversation_summary( @@ -8593,51 +8597,6 @@ fn preview_from_rollout_items(items: &[RolloutItem]) -> String { .unwrap_or_default() } -fn user_message_positions_in_rollout(items: &[RolloutItem]) -> Vec { - let mut user_positions = Vec::new(); - for (idx, item) in items.iter().enumerate() { - match item { - RolloutItem::ResponseItem(item) - if matches!( - codex_core::parse_turn_item(item), - Some(TurnItem::UserMessage(_)) - ) => - { - user_positions.push(idx); - } - RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => { - let num_turns = usize::try_from(rollback.num_turns).unwrap_or(usize::MAX); - let new_len = user_positions.len().saturating_sub(num_turns); - user_positions.truncate(new_len); - } - RolloutItem::ResponseItem(_) => {} - RolloutItem::SessionMeta(_) - | RolloutItem::ForkReference(_) - | RolloutItem::Compacted(_) - | RolloutItem::TurnContext(_) - | RolloutItem::EventMsg(_) => {} - } - } - user_positions -} - -fn truncate_rollout_before_nth_user_message_from_start( - items: &[RolloutItem], - n_from_start: usize, -) -> Vec { - if n_from_start == usize::MAX { - return items.to_vec(); - } - - let user_positions = user_message_positions_in_rollout(items); - if user_positions.len() <= n_from_start { - return Vec::new(); - } - - let cut_idx = user_positions[n_from_start]; - items[..cut_idx].to_vec() -} - fn codex_home_from_rollout_path(path: &Path) -> Option<&Path> { path.ancestors().find_map(|ancestor| { let name = ancestor.file_name().and_then(OsStr::to_str)?; @@ -8649,81 +8608,6 @@ fn codex_home_from_rollout_path(path: &Path) -> Option<&Path> { }) } -async fn materialize_rollout_items_for_replay( - codex_home: Option<&Path>, - rollout_items: &[RolloutItem], -) -> Vec { - const MAX_FORK_REFERENCE_DEPTH: usize = 8; - - let mut materialized = Vec::new(); - let mut stack: Vec<(Vec, usize, usize)> = vec![(rollout_items.to_vec(), 0, 0)]; - - while let Some((items, mut idx, depth)) = stack.pop() { - while idx < items.len() { - match &items[idx] { - RolloutItem::ForkReference(reference) => { - if depth >= MAX_FORK_REFERENCE_DEPTH { - warn!( - "skipping fork reference recursion at depth {} for {:?}", - depth, reference.rollout_path - ); - idx += 1; - continue; - } - - let resolved_rollout_path = if let Some(codex_home) = codex_home { - match resolve_fork_reference_rollout_path( - codex_home, - &reference.rollout_path, - ) - .await - { - Ok(path) => path, - Err(err) => { - warn!( - "failed to resolve fork reference rollout {:?}: {err}", - reference.rollout_path - ); - idx += 1; - continue; - } - } - } else { - reference.rollout_path.clone() - }; - let parent_history = match RolloutRecorder::get_rollout_history( - &resolved_rollout_path, - ) - .await - { - Ok(history) => history, - Err(err) => { - warn!( - "failed to load fork reference rollout {:?} (resolved from {:?}): {err}", - resolved_rollout_path, reference.rollout_path - ); - idx += 1; - continue; - } - }; - let parent_items = truncate_rollout_before_nth_user_message_from_start( - &parent_history.get_rollout_items(), - reference.nth_user_message, - ); - - stack.push((items, idx + 1, depth)); - stack.push((parent_items, 0, depth + 1)); - break; - } - item => materialized.push(item.clone()), - } - idx += 1; - } - } - - materialized -} - fn with_thread_spawn_agent_metadata( source: codex_protocol::protocol::SessionSource, agent_nickname: Option, diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index c04720ebd9..ef102f23ac 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -15,6 +15,8 @@ use crate::session_prefix::format_subagent_notification_message; use crate::shell_snapshot::ShellSnapshot; use crate::state_db; use crate::thread_manager::ThreadManagerState; +use crate::thread_rollout_truncation::fork_reference_user_message_boundary; +use crate::thread_rollout_truncation::materialize_rollout_items_for_replay; use crate::thread_rollout_truncation::truncate_rollout_to_last_n_fork_turns; use codex_features::Feature; use codex_protocol::AgentPath; @@ -314,20 +316,20 @@ impl AgentControl { .iter() .any(|item| matches!(item, RolloutItem::ForkReference(_))) { - forked_rollout_items = - crate::thread_rollout_truncation::materialize_rollout_items_for_replay( - config.codex_home.as_path(), - &forked_rollout_items, - ) - .await; + forked_rollout_items = materialize_rollout_items_for_replay( + config.codex_home.as_path(), + &forked_rollout_items, + ) + .await; } if let SpawnAgentForkMode::LastNTurns(last_n_turns) = fork_mode { forked_rollout_items = truncate_rollout_to_last_n_fork_turns(&forked_rollout_items, *last_n_turns); } + let fork_boundary = fork_reference_user_message_boundary(&forked_rollout_items); forked_rollout_items.push(RolloutItem::ForkReference(ForkReferenceItem { rollout_path: rollout_path.clone(), - nth_user_message: usize::MAX, + nth_user_message: fork_boundary, })); let mut output = @@ -1154,1250 +1156,3 @@ fn thread_spawn_depth(session_source: &SessionSource) -> Option { #[cfg(test)] #[path = "control_tests.rs"] mod tests; -// Keep this inline fork-reference test module disabled on the refreshed main API; -// branch coverage now comes from the package/integration tests that match current types. -#[cfg(any())] -mod fork_reference_tests { - use super::*; - use crate::CodexAuth; - use crate::CodexThread; - use crate::ThreadManager; - use crate::agent::agent_status_from_event; - use crate::config::AgentRoleConfig; - use crate::config::Config; - use crate::config::ConfigBuilder; - use crate::config_loader::LoaderOverrides; - use crate::contextual_user_message::SUBAGENT_NOTIFICATION_OPEN_TAG; - use codex_features::Feature; - use codex_protocol::config_types::ModeKind; - use codex_protocol::models::ContentItem; - use codex_protocol::models::ResponseItem; - use codex_protocol::protocol::ErrorEvent; - use codex_protocol::protocol::EventMsg; - use codex_protocol::protocol::SessionSource; - use codex_protocol::protocol::SubAgentSource; - use codex_protocol::protocol::TurnAbortReason; - use codex_protocol::protocol::TurnAbortedEvent; - use codex_protocol::protocol::TurnCompleteEvent; - use codex_protocol::protocol::TurnStartedEvent; - use pretty_assertions::assert_eq; - use tempfile::TempDir; - use tokio::time::Duration; - use tokio::time::sleep; - use tokio::time::timeout; - use toml::Value as TomlValue; - - async fn test_config_with_cli_overrides( - cli_overrides: Vec<(String, TomlValue)>, - ) -> (TempDir, Config) { - let home = TempDir::new().expect("create temp dir"); - let config = ConfigBuilder::default() - .codex_home(home.path().to_path_buf()) - .cli_overrides(cli_overrides) - .loader_overrides(LoaderOverrides { - #[cfg(target_os = "macos")] - managed_preferences_base64: Some(String::new()), - macos_managed_config_requirements_base64: Some(String::new()), - ..LoaderOverrides::default() - }) - .build() - .await - .expect("load default test config"); - (home, config) - } - - async fn test_config() -> (TempDir, Config) { - test_config_with_cli_overrides(Vec::new()).await - } - - fn text_input(text: &str) -> Op { - vec![UserInput::Text { - text: text.to_string(), - text_elements: Vec::new(), - }] - .into() - } - - struct AgentControlHarness { - _home: TempDir, - config: Config, - manager: ThreadManager, - control: AgentControl, - } - - impl AgentControlHarness { - async fn new() -> Self { - let (home, config) = test_config().await; - let manager = ThreadManager::with_models_provider_and_home_for_tests( - CodexAuth::from_api_key("dummy"), - config.model_provider.clone(), - config.codex_home.clone(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, - )), - ); - let control = manager.agent_control(); - Self { - _home: home, - config, - manager, - control, - } - } - - async fn start_thread(&self) -> (ThreadId, Arc) { - let new_thread = self - .manager - .start_thread(self.config.clone()) - .await - .expect("start thread"); - (new_thread.thread_id, new_thread.thread) - } - } - - fn has_subagent_notification(history_items: &[ResponseItem]) -> bool { - history_items.iter().any(|item| { - let ResponseItem::Message { role, content, .. } = item else { - return false; - }; - if role != "user" { - return false; - } - content.iter().any(|content_item| match content_item { - ContentItem::InputText { text } | ContentItem::OutputText { text } => { - text.contains(SUBAGENT_NOTIFICATION_OPEN_TAG) - } - ContentItem::InputImage { .. } => false, - }) - }) - } - - /// Returns true when any message item contains `needle` in a text span. - fn history_contains_text(history_items: &[ResponseItem], needle: &str) -> bool { - history_items.iter().any(|item| { - let ResponseItem::Message { content, .. } = item else { - return false; - }; - content.iter().any(|content_item| match content_item { - ContentItem::InputText { text } | ContentItem::OutputText { text } => { - text.contains(needle) - } - ContentItem::InputImage { .. } => false, - }) - }) - } - - async fn wait_for_subagent_notification(parent_thread: &Arc) -> bool { - let wait = async { - loop { - let history_items = parent_thread - .codex - .session - .clone_history() - .await - .raw_items() - .to_vec(); - if has_subagent_notification(&history_items) { - return true; - } - sleep(Duration::from_millis(25)).await; - } - }; - timeout(Duration::from_secs(2), wait).await.is_ok() - } - - #[tokio::test] - async fn send_input_errors_when_manager_dropped() { - let control = AgentControl::default(); - let err = control - .send_input( - ThreadId::new(), - vec![UserInput::Text { - text: "hello".to_string(), - text_elements: Vec::new(), - }], - ) - .await - .expect_err("send_input should fail without a manager"); - assert_eq!( - err.to_string(), - "unsupported operation: thread manager dropped" - ); - } - - #[tokio::test] - async fn get_status_returns_not_found_without_manager() { - let control = AgentControl::default(); - let got = control.get_status(ThreadId::new()).await; - assert_eq!(got, AgentStatus::NotFound); - } - - #[tokio::test] - async fn on_event_updates_status_from_task_started() { - let status = agent_status_from_event(&EventMsg::TurnStarted(TurnStartedEvent { - turn_id: "turn-1".to_string(), - model_context_window: None, - collaboration_mode_kind: ModeKind::Default, - })); - assert_eq!(status, Some(AgentStatus::Running)); - } - - #[tokio::test] - async fn on_event_updates_status_from_task_complete() { - let status = agent_status_from_event(&EventMsg::TurnComplete(TurnCompleteEvent { - turn_id: "turn-1".to_string(), - last_agent_message: Some("done".to_string()), - })); - let expected = AgentStatus::Completed(Some("done".to_string())); - assert_eq!(status, Some(expected)); - } - - #[tokio::test] - async fn on_event_updates_status_from_error() { - let status = agent_status_from_event(&EventMsg::Error(ErrorEvent { - message: "boom".to_string(), - codex_error_info: None, - })); - - let expected = AgentStatus::Errored("boom".to_string()); - assert_eq!(status, Some(expected)); - } - - #[tokio::test] - async fn on_event_updates_status_from_turn_aborted() { - let status = agent_status_from_event(&EventMsg::TurnAborted(TurnAbortedEvent { - turn_id: Some("turn-1".to_string()), - reason: TurnAbortReason::Interrupted, - })); - - let expected = AgentStatus::Errored("Interrupted".to_string()); - assert_eq!(status, Some(expected)); - } - - #[tokio::test] - async fn on_event_updates_status_from_shutdown_complete() { - let status = agent_status_from_event(&EventMsg::ShutdownComplete); - assert_eq!(status, Some(AgentStatus::Shutdown)); - } - - #[tokio::test] - async fn spawn_agent_errors_when_manager_dropped() { - let control = AgentControl::default(); - let (_home, config) = test_config().await; - let err = control - .spawn_agent(config, text_input("hello"), None) - .await - .expect_err("spawn_agent should fail without a manager"); - assert_eq!( - err.to_string(), - "unsupported operation: thread manager dropped" - ); - } - - #[tokio::test] - async fn resume_agent_errors_when_manager_dropped() { - let control = AgentControl::default(); - let (_home, config) = test_config().await; - let err = control - .resume_agent_from_rollout(config, ThreadId::new(), SessionSource::Exec) - .await - .expect_err("resume_agent should fail without a manager"); - assert_eq!( - err.to_string(), - "unsupported operation: thread manager dropped" - ); - } - - #[tokio::test] - async fn send_input_errors_when_thread_missing() { - let harness = AgentControlHarness::new().await; - let thread_id = ThreadId::new(); - let err = harness - .control - .send_input( - thread_id, - vec![UserInput::Text { - text: "hello".to_string(), - text_elements: Vec::new(), - }], - ) - .await - .expect_err("send_input should fail for missing thread"); - assert!(matches!(err, CodexErr::ThreadNotFound(id) if id == thread_id)); - } - - #[tokio::test] - async fn get_status_returns_not_found_for_missing_thread() { - let harness = AgentControlHarness::new().await; - let status = harness.control.get_status(ThreadId::new()).await; - assert_eq!(status, AgentStatus::NotFound); - } - - #[tokio::test] - async fn get_status_returns_pending_init_for_new_thread() { - let harness = AgentControlHarness::new().await; - let (thread_id, _) = harness.start_thread().await; - let status = harness.control.get_status(thread_id).await; - assert_eq!(status, AgentStatus::PendingInit); - } - - #[tokio::test] - async fn subscribe_status_errors_for_missing_thread() { - let harness = AgentControlHarness::new().await; - let thread_id = ThreadId::new(); - let err = harness - .control - .subscribe_status(thread_id) - .await - .expect_err("subscribe_status should fail for missing thread"); - assert!(matches!(err, CodexErr::ThreadNotFound(id) if id == thread_id)); - } - - #[tokio::test] - async fn subscribe_status_updates_on_shutdown() { - let harness = AgentControlHarness::new().await; - let (thread_id, thread) = harness.start_thread().await; - let mut status_rx = harness - .control - .subscribe_status(thread_id) - .await - .expect("subscribe_status should succeed"); - assert_eq!(status_rx.borrow().clone(), AgentStatus::PendingInit); - - let _ = thread - .submit(Op::Shutdown {}) - .await - .expect("shutdown should submit"); - - let _ = status_rx.changed().await; - assert_eq!(status_rx.borrow().clone(), AgentStatus::Shutdown); - } - - #[tokio::test] - async fn send_input_submits_user_message() { - let harness = AgentControlHarness::new().await; - let (thread_id, _thread) = harness.start_thread().await; - - let submission_id = harness - .control - .send_input( - thread_id, - vec![UserInput::Text { - text: "hello from tests".to_string(), - text_elements: Vec::new(), - }], - ) - .await - .expect("send_input should succeed"); - assert!(!submission_id.is_empty()); - let expected = ( - thread_id, - Op::UserInput { - items: vec![UserInput::Text { - text: "hello from tests".to_string(), - text_elements: Vec::new(), - }], - final_output_json_schema: None, - }, - ); - let captured = harness - .manager - .captured_ops() - .into_iter() - .find(|entry| *entry == expected); - assert_eq!(captured, Some(expected)); - } - - #[tokio::test] - async fn spawn_agent_creates_thread_and_sends_prompt() { - let harness = AgentControlHarness::new().await; - let thread_id = harness - .control - .spawn_agent(harness.config.clone(), text_input("spawned"), None) - .await - .expect("spawn_agent should succeed"); - let _thread = harness - .manager - .get_thread(thread_id) - .await - .expect("thread should be registered"); - let expected = ( - thread_id, - Op::UserInput { - items: vec![UserInput::Text { - text: "spawned".to_string(), - text_elements: Vec::new(), - }], - final_output_json_schema: None, - }, - ); - let captured = harness - .manager - .captured_ops() - .into_iter() - .find(|entry| *entry == expected); - assert_eq!(captured, Some(expected)); - } - - #[tokio::test] - async fn spawn_agent_can_fork_parent_thread_history() { - let harness = AgentControlHarness::new().await; - let (parent_thread_id, parent_thread) = harness.start_thread().await; - parent_thread - .inject_user_message_without_turn("parent seed context".to_string()) - .await; - let turn_context = parent_thread.codex.session.new_default_turn().await; - let parent_spawn_call_id = "spawn-call-history".to_string(); - let parent_spawn_call = ResponseItem::FunctionCall { - id: None, - name: "spawn_agent".to_string(), - arguments: "{}".to_string(), - call_id: parent_spawn_call_id.clone(), - namespace: None, - }; - parent_thread - .codex - .session - .record_conversation_items(turn_context.as_ref(), &[parent_spawn_call]) - .await; - parent_thread - .codex - .session - .ensure_rollout_materialized() - .await; - parent_thread.codex.session.flush_rollout().await; - - let child_thread_id = harness - .control - .spawn_agent_with_metadata( - harness.config.clone(), - text_input("child task"), - Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: None, - })), - SpawnAgentOptions { - fork_parent_spawn_call_id: Some(parent_spawn_call_id), - }, - ) - .await - .expect("forked spawn should succeed"); - - let child_thread = harness - .manager - .get_thread(child_thread_id) - .await - .expect("child thread should be registered"); - assert_ne!(child_thread_id, parent_thread_id); - let history = child_thread.codex.session.clone_history().await; - assert!(history_contains_text( - history.raw_items(), - "parent seed context" - )); - - let expected = ( - child_thread_id, - Op::UserInput { - items: vec![UserInput::Text { - text: "child task".to_string(), - text_elements: Vec::new(), - }], - final_output_json_schema: None, - }, - ); - let captured = harness - .manager - .captured_ops() - .into_iter() - .find(|entry| *entry == expected); - assert_eq!(captured, Some(expected)); - - let _ = harness - .control - .shutdown_live_agent(child_thread_id) - .await - .expect("child shutdown should submit"); - let _ = parent_thread - .submit(Op::Shutdown {}) - .await - .expect("parent shutdown should submit"); - } - - #[tokio::test] - async fn spawn_agent_fork_injects_output_for_parent_spawn_call() { - let harness = AgentControlHarness::new().await; - let (parent_thread_id, parent_thread) = harness.start_thread().await; - let turn_context = parent_thread.codex.session.new_default_turn().await; - let parent_spawn_call_id = "spawn-call-1".to_string(); - let parent_spawn_call = ResponseItem::FunctionCall { - id: None, - name: "spawn_agent".to_string(), - arguments: "{}".to_string(), - call_id: parent_spawn_call_id.clone(), - namespace: None, - }; - parent_thread - .codex - .session - .record_conversation_items(turn_context.as_ref(), &[parent_spawn_call]) - .await; - parent_thread - .codex - .session - .ensure_rollout_materialized() - .await; - parent_thread.codex.session.flush_rollout().await; - - let child_thread_id = harness - .control - .spawn_agent_with_metadata( - harness.config.clone(), - text_input("child task"), - Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: None, - })), - SpawnAgentOptions { - fork_parent_spawn_call_id: Some(parent_spawn_call_id.clone()), - }, - ) - .await - .expect("forked spawn should succeed"); - - let child_thread = harness - .manager - .get_thread(child_thread_id) - .await - .expect("child thread should be registered"); - let history = child_thread.codex.session.clone_history().await; - let injected_output = history.raw_items().iter().find_map(|item| match item { - ResponseItem::FunctionCallOutput { call_id, output } - if call_id == &parent_spawn_call_id => - { - Some(output) - } - _ => None, - }); - let injected_output = - injected_output.expect("forked child should contain synthetic tool output"); - assert_eq!( - injected_output.text_content(), - Some(FORKED_SPAWN_AGENT_OUTPUT_MESSAGE) - ); - assert_eq!(injected_output.success, Some(true)); - - let _ = harness - .control - .shutdown_live_agent(child_thread_id) - .await - .expect("child shutdown should submit"); - let _ = parent_thread - .submit(Op::Shutdown {}) - .await - .expect("parent shutdown should submit"); - } - - #[tokio::test] - async fn spawn_agent_fork_flushes_parent_rollout_before_loading_history() { - let harness = AgentControlHarness::new().await; - let (parent_thread_id, parent_thread) = harness.start_thread().await; - let turn_context = parent_thread.codex.session.new_default_turn().await; - let parent_spawn_call_id = "spawn-call-unflushed".to_string(); - let parent_spawn_call = ResponseItem::FunctionCall { - id: None, - name: "spawn_agent".to_string(), - arguments: "{}".to_string(), - call_id: parent_spawn_call_id.clone(), - namespace: None, - }; - parent_thread - .codex - .session - .record_conversation_items(turn_context.as_ref(), &[parent_spawn_call]) - .await; - - let child_thread_id = harness - .control - .spawn_agent_with_metadata( - harness.config.clone(), - text_input("child task"), - Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: None, - })), - SpawnAgentOptions { - fork_parent_spawn_call_id: Some(parent_spawn_call_id.clone()), - }, - ) - .await - .expect("forked spawn should flush parent rollout before loading history"); - - let child_thread = harness - .manager - .get_thread(child_thread_id) - .await - .expect("child thread should be registered"); - let history = child_thread.codex.session.clone_history().await; - - let mut parent_call_index = None; - let mut injected_output_index = None; - for (idx, item) in history.raw_items().iter().enumerate() { - match item { - ResponseItem::FunctionCall { call_id, .. } if call_id == &parent_spawn_call_id => { - parent_call_index = Some(idx); - } - ResponseItem::FunctionCallOutput { call_id, .. } - if call_id == &parent_spawn_call_id => - { - injected_output_index = Some(idx); - } - _ => {} - } - } - - let parent_call_index = - parent_call_index.expect("forked child should include the parent spawn_agent call"); - let injected_output_index = injected_output_index - .expect("forked child should include synthetic output for the parent spawn_agent call"); - assert!(parent_call_index < injected_output_index); - - let _ = harness - .control - .shutdown_live_agent(child_thread_id) - .await - .expect("child shutdown should submit"); - let _ = parent_thread - .submit(Op::Shutdown {}) - .await - .expect("parent shutdown should submit"); - } - - #[tokio::test] - async fn spawn_agent_fork_persists_fork_reference_instead_of_parent_history() { - let harness = AgentControlHarness::new().await; - let (parent_thread_id, parent_thread) = harness.start_thread().await; - parent_thread - .inject_user_message_without_turn("parent seed context".to_string()) - .await; - let turn_context = parent_thread.codex.session.new_default_turn().await; - let parent_spawn_call_id = "spawn-call-dedup".to_string(); - let parent_spawn_call = ResponseItem::FunctionCall { - id: None, - name: "spawn_agent".to_string(), - arguments: "{}".to_string(), - call_id: parent_spawn_call_id.clone(), - namespace: None, - }; - parent_thread - .codex - .session - .record_conversation_items(turn_context.as_ref(), &[parent_spawn_call]) - .await; - parent_thread - .codex - .session - .ensure_rollout_materialized() - .await; - parent_thread.codex.session.flush_rollout().await; - let parent_rollout_path = parent_thread - .rollout_path() - .expect("parent rollout path should be available"); - - let child_thread_id = harness - .control - .spawn_agent_with_metadata( - harness.config.clone(), - text_input("child task"), - Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: None, - })), - SpawnAgentOptions { - fork_parent_spawn_call_id: Some(parent_spawn_call_id), - }, - ) - .await - .expect("forked spawn should succeed"); - - let child_thread = harness - .manager - .get_thread(child_thread_id) - .await - .expect("child thread should be registered"); - let child_rollout_path = child_thread - .rollout_path() - .expect("child rollout path should be available"); - let InitialHistory::Resumed(resumed) = - RolloutRecorder::get_rollout_history(child_rollout_path.as_path()) - .await - .expect("child rollout should load") - else { - panic!("child rollout should include session metadata"); - }; - - assert!( - resumed.history.iter().any(|item| { - matches!( - item, - RolloutItem::ForkReference(ForkReferenceItem { - rollout_path, - nth_user_message, - }) if rollout_path == &parent_rollout_path && *nth_user_message == usize::MAX - ) - }), - "child rollout should persist a fork reference to the parent rollout" - ); - - let raw_response_items: Vec = resumed - .history - .iter() - .filter_map(|item| match item { - RolloutItem::ResponseItem(response_item) => Some(response_item.clone()), - RolloutItem::SessionMeta(_) - | RolloutItem::ForkReference(_) - | RolloutItem::Compacted(_) - | RolloutItem::TurnContext(_) - | RolloutItem::EventMsg(_) => None, - }) - .collect(); - assert!( - !history_contains_text(&raw_response_items, "parent seed context"), - "child rollout should not duplicate the parent's raw transcript" - ); - - let history = child_thread.codex.session.clone_history().await; - assert!(history_contains_text( - history.raw_items(), - "parent seed context" - )); - - let _ = harness - .control - .shutdown_live_agent(child_thread_id) - .await - .expect("child shutdown should submit"); - let _ = parent_thread - .submit(Op::Shutdown {}) - .await - .expect("parent shutdown should submit"); - } - - #[tokio::test] - async fn spawn_agent_respects_max_threads_limit() { - let max_threads = 1usize; - let (_home, config) = test_config_with_cli_overrides(vec![( - "agents.max_threads".to_string(), - TomlValue::Integer(max_threads as i64), - )]) - .await; - let manager = ThreadManager::with_models_provider_and_home_for_tests( - CodexAuth::from_api_key("dummy"), - config.model_provider.clone(), - config.codex_home.clone(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, - )), - ); - let control = manager.agent_control(); - - let _ = manager - .start_thread(config.clone()) - .await - .expect("start thread"); - - let first_agent_id = control - .spawn_agent(config.clone(), text_input("hello"), None) - .await - .expect("spawn_agent should succeed"); - - let err = control - .spawn_agent(config, text_input("hello again"), None) - .await - .expect_err("spawn_agent should respect max threads"); - let CodexErr::AgentLimitReached { - max_threads: seen_max_threads, - } = err - else { - panic!("expected CodexErr::AgentLimitReached"); - }; - assert_eq!(seen_max_threads, max_threads); - - let _ = control - .shutdown_live_agent(first_agent_id) - .await - .expect("shutdown agent"); - } - - #[tokio::test] - async fn spawn_agent_releases_slot_after_shutdown() { - let max_threads = 1usize; - let (_home, config) = test_config_with_cli_overrides(vec![( - "agents.max_threads".to_string(), - TomlValue::Integer(max_threads as i64), - )]) - .await; - let manager = ThreadManager::with_models_provider_and_home_for_tests( - CodexAuth::from_api_key("dummy"), - config.model_provider.clone(), - config.codex_home.clone(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, - )), - ); - let control = manager.agent_control(); - - let first_agent_id = control - .spawn_agent(config.clone(), text_input("hello"), None) - .await - .expect("spawn_agent should succeed"); - let _ = control - .shutdown_live_agent(first_agent_id) - .await - .expect("shutdown agent"); - - let second_agent_id = control - .spawn_agent(config.clone(), text_input("hello again"), None) - .await - .expect("spawn_agent should succeed after shutdown"); - let _ = control - .shutdown_live_agent(second_agent_id) - .await - .expect("shutdown agent"); - } - - #[tokio::test] - async fn spawn_agent_limit_shared_across_clones() { - let max_threads = 1usize; - let (_home, config) = test_config_with_cli_overrides(vec![( - "agents.max_threads".to_string(), - TomlValue::Integer(max_threads as i64), - )]) - .await; - let manager = ThreadManager::with_models_provider_and_home_for_tests( - CodexAuth::from_api_key("dummy"), - config.model_provider.clone(), - config.codex_home.clone(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, - )), - ); - let control = manager.agent_control(); - let cloned = control.clone(); - - let first_agent_id = cloned - .spawn_agent(config.clone(), text_input("hello"), None) - .await - .expect("spawn_agent should succeed"); - - let err = control - .spawn_agent(config, text_input("hello again"), None) - .await - .expect_err("spawn_agent should respect shared guard"); - let CodexErr::AgentLimitReached { max_threads } = err else { - panic!("expected CodexErr::AgentLimitReached"); - }; - assert_eq!(max_threads, 1); - - let _ = control - .shutdown_live_agent(first_agent_id) - .await - .expect("shutdown agent"); - } - - #[tokio::test] - async fn resume_agent_respects_max_threads_limit() { - let max_threads = 1usize; - let (_home, config) = test_config_with_cli_overrides(vec![( - "agents.max_threads".to_string(), - TomlValue::Integer(max_threads as i64), - )]) - .await; - let manager = ThreadManager::with_models_provider_and_home_for_tests( - CodexAuth::from_api_key("dummy"), - config.model_provider.clone(), - config.codex_home.clone(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, - )), - ); - let control = manager.agent_control(); - - let resumable_id = control - .spawn_agent(config.clone(), text_input("hello"), None) - .await - .expect("spawn_agent should succeed"); - let _ = control - .shutdown_live_agent(resumable_id) - .await - .expect("shutdown resumable thread"); - - let active_id = control - .spawn_agent(config.clone(), text_input("occupy"), None) - .await - .expect("spawn_agent should succeed for active slot"); - - let err = control - .resume_agent_from_rollout(config, resumable_id, SessionSource::Exec) - .await - .expect_err("resume should respect max threads"); - let CodexErr::AgentLimitReached { - max_threads: seen_max_threads, - } = err - else { - panic!("expected CodexErr::AgentLimitReached"); - }; - assert_eq!(seen_max_threads, max_threads); - - let _ = control - .shutdown_live_agent(active_id) - .await - .expect("shutdown active thread"); - } - - #[tokio::test] - async fn resume_agent_releases_slot_after_resume_failure() { - let max_threads = 1usize; - let (_home, config) = test_config_with_cli_overrides(vec![( - "agents.max_threads".to_string(), - TomlValue::Integer(max_threads as i64), - )]) - .await; - let manager = ThreadManager::with_models_provider_and_home_for_tests( - CodexAuth::from_api_key("dummy"), - config.model_provider.clone(), - config.codex_home.clone(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, - )), - ); - let control = manager.agent_control(); - - let _ = control - .resume_agent_from_rollout(config.clone(), ThreadId::new(), SessionSource::Exec) - .await - .expect_err("resume should fail for missing rollout path"); - - let resumed_id = control - .spawn_agent(config, text_input("hello"), None) - .await - .expect("spawn should succeed after failed resume"); - let _ = control - .shutdown_live_agent(resumed_id) - .await - .expect("shutdown resumed thread"); - } - - #[tokio::test] - async fn spawn_child_completion_notifies_parent_history() { - let harness = AgentControlHarness::new().await; - let (parent_thread_id, parent_thread) = harness.start_thread().await; - - let child_thread_id = harness - .control - .spawn_agent( - harness.config.clone(), - text_input("hello child"), - Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: Some("explorer".to_string()), - })), - ) - .await - .expect("child spawn should succeed"); - - let child_thread = harness - .manager - .get_thread(child_thread_id) - .await - .expect("child thread should exist"); - let _ = child_thread - .submit(Op::Shutdown {}) - .await - .expect("child shutdown should submit"); - - assert_eq!(wait_for_subagent_notification(&parent_thread).await, true); - } - - #[tokio::test] - async fn completion_watcher_notifies_parent_when_child_is_missing() { - let harness = AgentControlHarness::new().await; - let (parent_thread_id, parent_thread) = harness.start_thread().await; - let child_thread_id = ThreadId::new(); - - harness.control.maybe_start_completion_watcher( - child_thread_id, - Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: Some("explorer".to_string()), - })), - child_thread_id.to_string(), - None, - ); - - assert_eq!(wait_for_subagent_notification(&parent_thread).await, true); - - let history_items = parent_thread - .codex - .session - .clone_history() - .await - .raw_items() - .to_vec(); - assert_eq!( - history_contains_text( - &history_items, - &format!("\"agent_id\":\"{child_thread_id}\"") - ), - true - ); - assert_eq!( - history_contains_text(&history_items, "\"status\":\"not_found\""), - true - ); - } - - #[tokio::test] - async fn spawn_thread_subagent_gets_random_nickname_in_session_source() { - let harness = AgentControlHarness::new().await; - let (parent_thread_id, _parent_thread) = harness.start_thread().await; - - let child_thread_id = harness - .control - .spawn_agent( - harness.config.clone(), - text_input("hello child"), - Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: Some("explorer".to_string()), - })), - ) - .await - .expect("child spawn should succeed"); - - let child_thread = harness - .manager - .get_thread(child_thread_id) - .await - .expect("child thread should be registered"); - let snapshot = child_thread.config_snapshot().await; - - let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id: seen_parent_thread_id, - depth, - agent_nickname, - agent_role, - .. - }) = snapshot.session_source - else { - panic!("expected thread-spawn sub-agent source"); - }; - assert_eq!(seen_parent_thread_id, parent_thread_id); - assert_eq!(depth, 1); - assert!(agent_nickname.is_some()); - assert_eq!(agent_role, Some("explorer".to_string())); - } - - #[tokio::test] - async fn spawn_thread_subagent_uses_role_specific_nickname_candidates() { - let mut harness = AgentControlHarness::new().await; - harness.config.agent_roles.insert( - "researcher".to_string(), - AgentRoleConfig { - description: Some("Research role".to_string()), - config_file: None, - nickname_candidates: Some(vec!["Atlas".to_string()]), - }, - ); - let (parent_thread_id, _parent_thread) = harness.start_thread().await; - - let child_thread_id = harness - .control - .spawn_agent( - harness.config.clone(), - text_input("hello child"), - Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: Some("researcher".to_string()), - })), - ) - .await - .expect("child spawn should succeed"); - - let child_thread = harness - .manager - .get_thread(child_thread_id) - .await - .expect("child thread should be registered"); - let snapshot = child_thread.config_snapshot().await; - - let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { agent_nickname, .. }) = - snapshot.session_source - else { - panic!("expected thread-spawn sub-agent source"); - }; - assert_eq!(agent_nickname, Some("Atlas".to_string())); - } - - #[tokio::test] - async fn resume_thread_subagent_restores_stored_nickname_and_role() { - let (home, mut config) = test_config().await; - config - .features - .enable(Feature::Sqlite) - .expect("test config should allow sqlite"); - let manager = ThreadManager::with_models_provider_and_home_for_tests( - CodexAuth::from_api_key("dummy"), - config.model_provider.clone(), - config.codex_home.clone(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, - )), - ); - let control = manager.agent_control(); - let harness = AgentControlHarness { - _home: home, - config, - manager, - control, - }; - let (parent_thread_id, _parent_thread) = harness.start_thread().await; - - let child_thread_id = harness - .control - .spawn_agent( - harness.config.clone(), - text_input("hello child"), - Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: Some("explorer".to_string()), - })), - ) - .await - .expect("child spawn should succeed"); - - let child_thread = harness - .manager - .get_thread(child_thread_id) - .await - .expect("child thread should exist"); - let mut status_rx = harness - .control - .subscribe_status(child_thread_id) - .await - .expect("status subscription should succeed"); - if matches!(status_rx.borrow().clone(), AgentStatus::PendingInit) { - timeout(Duration::from_secs(5), async { - loop { - status_rx - .changed() - .await - .expect("child status should advance past pending init"); - if !matches!(status_rx.borrow().clone(), AgentStatus::PendingInit) { - break; - } - } - }) - .await - .expect("child should initialize before shutdown"); - } - let original_snapshot = child_thread.config_snapshot().await; - let original_nickname = original_snapshot - .session_source - .get_nickname() - .expect("spawned sub-agent should have a nickname"); - let state_db = child_thread - .state_db() - .expect("sqlite state db should be available for nickname resume test"); - timeout(Duration::from_secs(5), async { - loop { - if let Ok(Some(metadata)) = state_db.get_thread(child_thread_id).await - && metadata.agent_nickname.is_some() - && metadata.agent_role.as_deref() == Some("explorer") - { - break; - } - sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("child thread metadata should be persisted to sqlite before shutdown"); - - let _ = harness - .control - .shutdown_live_agent(child_thread_id) - .await - .expect("child shutdown should submit"); - - let resumed_thread_id = harness - .control - .resume_agent_from_rollout( - harness.config.clone(), - child_thread_id, - SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: None, - }), - ) - .await - .expect("resume should succeed"); - assert_eq!(resumed_thread_id, child_thread_id); - - let resumed_snapshot = harness - .manager - .get_thread(resumed_thread_id) - .await - .expect("resumed child thread should exist") - .config_snapshot() - .await; - let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id: resumed_parent_thread_id, - depth: resumed_depth, - agent_nickname: resumed_nickname, - agent_role: resumed_role, - .. - }) = resumed_snapshot.session_source - else { - panic!("expected thread-spawn sub-agent source"); - }; - assert_eq!(resumed_parent_thread_id, parent_thread_id); - assert_eq!(resumed_depth, 1); - assert_eq!(resumed_nickname, Some(original_nickname)); - assert_eq!(resumed_role, Some("explorer".to_string())); - - let _ = harness - .control - .shutdown_live_agent(resumed_thread_id) - .await - .expect("resumed child shutdown should submit"); - } -} diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 834ab2037b..aefb60eb5f 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -17,6 +17,7 @@ use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::ForkReferenceItem; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; @@ -906,14 +907,8 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { history.raw_items(), "old parent context" )); - assert!(!history_contains_text( - history.raw_items(), - "queued message" - )); - assert!(history_contains_text( - history.raw_items(), - "triggered context" - )); + assert!(!history_contains_text(history.raw_items(), "queued message")); + assert!(history_contains_text(history.raw_items(), "triggered context")); assert!(history_contains_text( history.raw_items(), "current parent task" @@ -930,6 +925,125 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { .expect("parent shutdown should submit"); } +#[tokio::test] +async fn spawn_agent_fork_snapshots_parent_boundary_for_persisted_fork_reference() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + parent_thread + .inject_user_message_without_turn("parent seed context".to_string()) + .await; + let turn_context = parent_thread.codex.session.new_default_turn().await; + let parent_spawn_call_id = "spawn-call-dedup".to_string(); + let parent_spawn_call = ResponseItem::FunctionCall { + id: None, + name: "spawn_agent".to_string(), + namespace: None, + arguments: "{}".to_string(), + call_id: parent_spawn_call_id.clone(), + }; + parent_thread + .codex + .session + .record_conversation_items(turn_context.as_ref(), &[parent_spawn_call]) + .await; + parent_thread + .codex + .session + .ensure_rollout_materialized() + .await; + parent_thread.codex.session.flush_rollout().await; + let parent_rollout_path = parent_thread + .rollout_path() + .expect("parent rollout path should be available"); + + let child_thread_id = harness + .control + .spawn_agent_with_metadata( + harness.config.clone(), + text_input("child task"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })), + SpawnAgentOptions { + fork_parent_spawn_call_id: Some(parent_spawn_call_id), + fork_mode: Some(SpawnAgentForkMode::FullHistory), + }, + ) + .await + .expect("forked spawn should succeed") + .thread_id; + + parent_thread + .inject_user_message_without_turn("parent late turn".to_string()) + .await; + parent_thread + .codex + .session + .ensure_rollout_materialized() + .await; + parent_thread.codex.session.flush_rollout().await; + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should be registered"); + let child_rollout_path = child_thread + .rollout_path() + .expect("child rollout path should be available"); + let InitialHistory::Resumed(resumed) = + RolloutRecorder::get_rollout_history(child_rollout_path.as_path()) + .await + .expect("child rollout should load") + else { + panic!("child rollout should include session metadata"); + }; + + assert!(resumed.history.iter().any(|item| { + matches!( + item, + RolloutItem::ForkReference(ForkReferenceItem { + rollout_path, + nth_user_message: 1, + }) if rollout_path == &parent_rollout_path + ) + })); + let materialized_child_rollout = crate::rollout::truncation::materialize_rollout_items_for_replay( + harness.config.codex_home.as_path(), + &resumed.history, + ) + .await; + let materialized_child_response_items: Vec = materialized_child_rollout + .iter() + .filter_map(|item| match item { + RolloutItem::ResponseItem(response_item) => Some(response_item.clone()), + _ => None, + }) + .collect(); + assert!(history_contains_text( + &materialized_child_response_items, + "parent seed context", + )); + assert!(!history_contains_text( + &materialized_child_response_items, + "parent late turn", + )); + + let _ = harness + .control + .shutdown_live_agent(child_thread_id) + .await + .expect("child shutdown should submit"); + let _ = parent_thread + .submit(Op::Shutdown {}) + .await + .expect("parent shutdown should submit"); +} + #[tokio::test] async fn spawn_agent_respects_max_threads_limit() { let max_threads = 1usize; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 8beb0dc3ab..6cf6602910 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -184,6 +184,7 @@ pub use rollout::policy::EventPersistenceMode; pub use rollout::resolve_fork_reference_rollout_path; pub use rollout::rollout_date_parts; pub use rollout::session_index::find_thread_names_by_ids; +pub use thread_rollout_truncation::materialize_rollout_items_for_replay; mod function_tool; mod state; mod tasks; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index e051a851da..3e86247115 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -619,7 +619,7 @@ impl ThreadManager { truncate_before_nth_user_message( config.codex_home.as_path(), history, - nth_user_message, + i64::try_from(nth_user_message).unwrap_or(i64::MAX), &snapshot_state, ) .await @@ -654,7 +654,7 @@ impl ThreadManager { .chain(std::iter::once(RolloutItem::ForkReference( ForkReferenceItem { rollout_path: path.clone(), - nth_user_message, + nth_user_message: i64::try_from(nth_user_message).unwrap_or(i64::MAX), }, ))) .collect(); @@ -954,7 +954,7 @@ impl ThreadManagerState { async fn truncate_before_nth_user_message( codex_home: &Path, history: InitialHistory, - n: usize, + n: i64, snapshot_state: &SnapshotTurnState, ) -> InitialHistory { let mut items: Vec = history.get_rollout_items(); @@ -965,7 +965,9 @@ async fn truncate_before_nth_user_message( items = truncation::materialize_rollout_items_for_replay(codex_home, &items).await; } let user_positions = truncation::user_message_positions_in_rollout(&items); - let rolled = if snapshot_state.ends_mid_turn && n >= user_positions.len() { + let rolled = if snapshot_state.ends_mid_turn + && usize::try_from(n).map_or(true, |n| n >= user_positions.len()) + { if let Some(cut_idx) = snapshot_state .active_turn_start_index .or_else(|| user_positions.last().copied()) @@ -1077,150 +1079,3 @@ fn append_interrupted_boundary(history: InitialHistory, turn_id: Option) #[cfg(test)] #[path = "thread_manager_tests.rs"] mod tests; -// Keep this inline fork-reference test module disabled on the refreshed main API; -// branch coverage now comes from the package/integration tests that match current types. -#[cfg(any())] -mod fork_reference_tests { - use super::*; - use crate::codex::make_session_and_context; - use codex_protocol::models::ContentItem; - use codex_protocol::models::ReasoningItemReasoningSummary; - use codex_protocol::models::ResponseItem; - use pretty_assertions::assert_eq; - - fn user_msg(text: &str) -> ResponseItem { - ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::OutputText { - text: text.to_string(), - }], - end_turn: None, - phase: None, - } - } - fn assistant_msg(text: &str) -> ResponseItem { - ResponseItem::Message { - id: None, - role: "assistant".to_string(), - content: vec![ContentItem::OutputText { - text: text.to_string(), - }], - end_turn: None, - phase: None, - } - } - - #[tokio::test] - async fn drops_from_last_user_only() { - let items = [ - user_msg("u1"), - assistant_msg("a1"), - assistant_msg("a2"), - user_msg("u2"), - assistant_msg("a3"), - ResponseItem::Reasoning { - id: "r1".to_string(), - summary: vec![ReasoningItemReasoningSummary::SummaryText { - text: "s".to_string(), - }], - content: None, - encrypted_content: None, - }, - ResponseItem::FunctionCall { - id: None, - call_id: "c1".to_string(), - name: "tool".to_string(), - namespace: None, - arguments: "{}".to_string(), - }, - assistant_msg("a4"), - ]; - - let initial: Vec = items - .iter() - .cloned() - .map(RolloutItem::ResponseItem) - .collect(); - let truncated = truncate_before_nth_user_message( - Path::new("/tmp"), - InitialHistory::Forked(initial), - 1, - &SnapshotTurnState { - ends_mid_turn: false, - active_turn_id: None, - active_turn_start_index: None, - }, - ) - .await; - let got_items = truncated.get_rollout_items(); - let expected_items = vec![ - RolloutItem::ResponseItem(items[0].clone()), - RolloutItem::ResponseItem(items[1].clone()), - RolloutItem::ResponseItem(items[2].clone()), - ]; - assert_eq!( - serde_json::to_value(&got_items).unwrap(), - serde_json::to_value(&expected_items).unwrap() - ); - - let initial2: Vec = items - .iter() - .cloned() - .map(RolloutItem::ResponseItem) - .collect(); - let truncated2 = truncate_before_nth_user_message( - Path::new("/tmp"), - InitialHistory::Forked(initial2), - 2, - &SnapshotTurnState { - ends_mid_turn: false, - active_turn_id: None, - active_turn_start_index: None, - }, - ) - .await; - assert!(matches!(truncated2, InitialHistory::New)); - } - - #[tokio::test] - async fn ignores_session_prefix_messages_when_truncating() { - let (session, turn_context) = make_session_and_context().await; - let mut items = session.build_initial_context(&turn_context).await; - items.push(user_msg("feature request")); - items.push(assistant_msg("ack")); - items.push(user_msg("second question")); - items.push(assistant_msg("answer")); - - let rollout_items: Vec = items - .iter() - .cloned() - .map(RolloutItem::ResponseItem) - .collect(); - - let truncated = truncate_before_nth_user_message( - Path::new("/tmp"), - InitialHistory::Forked(rollout_items), - 1, - &SnapshotTurnState { - ends_mid_turn: false, - active_turn_id: None, - active_turn_start_index: None, - }, - ) - .await; - let got_items = truncated.get_rollout_items(); - - let expected: Vec = vec![ - RolloutItem::ResponseItem(items[0].clone()), - RolloutItem::ResponseItem(items[1].clone()), - RolloutItem::ResponseItem(items[2].clone()), - RolloutItem::ResponseItem(items[3].clone()), - ]; - - assert_eq!( - serde_json::to_value(&got_items).unwrap(), - serde_json::to_value(&expected).unwrap() - ); - } -} diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index dda9fa9484..e9b6eca3e9 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -129,7 +129,7 @@ async fn out_of_range_truncation_drops_only_unfinished_suffix_mid_turn() { let truncated = truncate_before_nth_user_message( Path::new("/tmp"), InitialHistory::Forked(items.clone()), - usize::MAX, + /*n*/ -1, &SnapshotTurnState { ends_mid_turn: true, active_turn_id: None, @@ -190,7 +190,7 @@ async fn out_of_range_truncation_drops_pre_user_active_turn_prefix() { let truncated = truncate_before_nth_user_message( Path::new("/tmp"), InitialHistory::Forked(items.clone()), - usize::MAX, + /*n*/ -1, &snapshot_state, ) .await; diff --git a/codex-rs/core/src/thread_rollout_truncation.rs b/codex-rs/core/src/thread_rollout_truncation.rs index b116ccb469..f370e306e6 100644 --- a/codex-rs/core/src/thread_rollout_truncation.rs +++ b/codex-rs/core/src/thread_rollout_truncation.rs @@ -92,21 +92,30 @@ pub(crate) fn fork_turn_positions_in_rollout(items: &[RolloutItem]) -> Vec i64 { + i64::try_from(user_message_positions_in_rollout(items).len()).unwrap_or(i64::MAX) +} + /// Return a prefix of `items` obtained by cutting strictly before the nth user message. /// /// The boundary index is 0-based from the start of `items` (so `n_from_start = 0` returns /// a prefix that excludes the first user message and everything after it). /// -/// If `n_from_start` is `usize::MAX`, this returns the full rollout (no truncation). +/// If `n_from_start` is negative, this returns the full rollout (no truncation). /// If fewer than or equal to `n_from_start` user messages exist, this returns the full /// rollout unchanged. pub(crate) fn truncate_rollout_before_nth_user_message_from_start( items: &[RolloutItem], - n_from_start: usize, + n_from_start: i64, ) -> Vec { - if n_from_start == usize::MAX { + let Ok(n_from_start) = usize::try_from(n_from_start) else { return items.to_vec(); - } + }; let user_positions = user_message_positions_in_rollout(items); @@ -157,7 +166,12 @@ fn is_trigger_turn_boundary(item: &ResponseItem) -> bool { .is_some_and(|communication| communication.trigger_turn) } -pub(crate) async fn materialize_rollout_items_for_replay( +/// Expand `ForkReference` items into the referenced parent rollout slices they encode. +/// +/// This preserves child rollout compactness on disk while letting replay callers rebuild the +/// effective inherited transcript before reconstructing conversation history or deriving thread +/// summaries. +pub async fn materialize_rollout_items_for_replay( codex_home: &Path, rollout_items: &[RolloutItem], ) -> Vec { diff --git a/codex-rs/core/src/thread_rollout_truncation_tests.rs b/codex-rs/core/src/thread_rollout_truncation_tests.rs index d1380b7d45..3420454605 100644 --- a/codex-rs/core/src/thread_rollout_truncation_tests.rs +++ b/codex-rs/core/src/thread_rollout_truncation_tests.rs @@ -102,7 +102,8 @@ fn truncation_max_keeps_full_rollout() { RolloutItem::ResponseItem(user_msg("u2")), ]; - let truncated = truncate_rollout_before_nth_user_message_from_start(&rollout, usize::MAX); + let truncated = + truncate_rollout_before_nth_user_message_from_start(&rollout, /*n_from_start*/ -1); assert_eq!( serde_json::to_value(&truncated).unwrap(), diff --git a/codex-rs/core/tests/suite/fork_thread.rs b/codex-rs/core/tests/suite/fork_thread.rs index 5685aa100c..7227aaec24 100644 --- a/codex-rs/core/tests/suite/fork_thread.rs +++ b/codex-rs/core/tests/suite/fork_thread.rs @@ -2,6 +2,8 @@ use codex_core::ForkSnapshot; use codex_core::NewThread; use codex_core::parse_turn_item; use codex_protocol::items::TurnItem; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::protocol::RolloutItem; @@ -33,18 +35,54 @@ fn find_user_input_positions(items: &[RolloutItem]) -> Vec { fn truncate_before_nth_user_message( items: &[RolloutItem], - nth_user_message: usize, + nth_user_message: i64, ) -> Vec { - if nth_user_message == usize::MAX { + let Ok(nth_user_message) = usize::try_from(nth_user_message) else { return items.to_vec(); - } + }; let user_inputs = find_user_input_positions(items); let Some(cut_idx) = user_inputs.get(nth_user_message).copied() else { - return Vec::new(); + return items.to_vec(); }; items[..cut_idx].to_vec() } +fn test_user_message(text: &str) -> RolloutItem { + RolloutItem::ResponseItem(ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::OutputText { + text: text.to_string(), + }], + end_turn: None, + phase: None, + }) +} + +#[test] +fn truncate_before_nth_user_message_keeps_full_history_for_out_of_range_boundaries() { + let rollout_items = vec![test_user_message("u1"), test_user_message("u2")]; + + pretty_assertions::assert_eq!( + serde_json::to_value(truncate_before_nth_user_message( + &rollout_items, + /*nth_user_message*/ 2, + )) + .unwrap(), + serde_json::to_value(&rollout_items).unwrap(), + ); +} + +#[test] +fn truncate_before_nth_user_message_keeps_full_history_for_i64_max_boundaries() { + let rollout_items = vec![test_user_message("u1"), test_user_message("u2")]; + + pretty_assertions::assert_eq!( + serde_json::to_value(truncate_before_nth_user_message(&rollout_items, i64::MAX,)).unwrap(), + serde_json::to_value(&rollout_items).unwrap(), + ); +} + fn read_items_materialized(p: &std::path::Path) -> Vec { let text = std::fs::read_to_string(p).unwrap_or_else(|err| panic!("read rollout file {p:?}: {err}")); @@ -147,9 +185,8 @@ async fn fork_thread_twice_drops_to_first_message() { // GetHistory on fork1 flushed; the file is ready. let fork1_items = read_items_materialized(&fork1_path); - assert!(fork1_items.len() > expected_after_first.len()); pretty_assertions::assert_eq!( - serde_json::to_value(&fork1_items[..expected_after_first.len()]).unwrap(), + serde_json::to_value(&fork1_items).unwrap(), serde_json::to_value(&expected_after_first).unwrap() ); @@ -178,9 +215,8 @@ async fn fork_thread_twice_drops_to_first_message() { .unwrap_or(0); let expected_after_second: Vec = fork1_items[..cut_last_on_fork1].to_vec(); let fork2_items = read_items_materialized(&fork2_path); - assert!(fork2_items.len() > expected_after_second.len()); pretty_assertions::assert_eq!( - serde_json::to_value(&fork2_items[..expected_after_second.len()]).unwrap(), + serde_json::to_value(&fork2_items).unwrap(), serde_json::to_value(&expected_after_second).unwrap() ); } diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index a1d840056c..d2d55c63db 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2312,11 +2312,7 @@ impl InitialHistory { .iter() .filter_map(|ri| match ri { RolloutItem::EventMsg(ev) => Some(ev.clone()), - RolloutItem::SessionMeta(_) - | RolloutItem::ForkReference(_) - | RolloutItem::ResponseItem(_) - | RolloutItem::Compacted(_) - | RolloutItem::TurnContext(_) => None, + _ => None, }) .collect(), ), @@ -2325,11 +2321,7 @@ impl InitialHistory { .iter() .filter_map(|ri| match ri { RolloutItem::EventMsg(ev) => Some(ev.clone()), - RolloutItem::SessionMeta(_) - | RolloutItem::ForkReference(_) - | RolloutItem::ResponseItem(_) - | RolloutItem::Compacted(_) - | RolloutItem::TurnContext(_) => None, + _ => None, }) .collect(), ), @@ -2343,20 +2335,12 @@ impl InitialHistory { InitialHistory::Resumed(resumed) => { resumed.history.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.base_instructions.clone(), - RolloutItem::ForkReference(_) - | RolloutItem::ResponseItem(_) - | RolloutItem::Compacted(_) - | RolloutItem::TurnContext(_) - | RolloutItem::EventMsg(_) => None, + _ => None, }) } InitialHistory::Forked(items) => items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.base_instructions.clone(), - RolloutItem::ForkReference(_) - | RolloutItem::ResponseItem(_) - | RolloutItem::Compacted(_) - | RolloutItem::TurnContext(_) - | RolloutItem::EventMsg(_) => None, + _ => None, }), } } @@ -2367,20 +2351,12 @@ impl InitialHistory { InitialHistory::Resumed(resumed) => { resumed.history.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.dynamic_tools.clone(), - RolloutItem::ForkReference(_) - | RolloutItem::ResponseItem(_) - | RolloutItem::Compacted(_) - | RolloutItem::TurnContext(_) - | RolloutItem::EventMsg(_) => None, + _ => None, }) } InitialHistory::Forked(items) => items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.dynamic_tools.clone(), - RolloutItem::ForkReference(_) - | RolloutItem::ResponseItem(_) - | RolloutItem::Compacted(_) - | RolloutItem::TurnContext(_) - | RolloutItem::EventMsg(_) => None, + _ => None, }), } } @@ -2389,11 +2365,7 @@ impl InitialHistory { fn session_cwd_from_items(items: &[RolloutItem]) -> Option { items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => Some(meta_line.meta.cwd.clone()), - RolloutItem::ForkReference(_) - | RolloutItem::ResponseItem(_) - | RolloutItem::Compacted(_) - | RolloutItem::TurnContext(_) - | RolloutItem::EventMsg(_) => None, + _ => None, }) } @@ -2600,10 +2572,38 @@ pub struct SessionMetaLine { pub git: Option, } -#[derive(Serialize, Deserialize, Debug, Clone, JsonSchema, TS)] +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] pub struct ForkReferenceItem { pub rollout_path: PathBuf, - pub nth_user_message: usize, + #[serde( + deserialize_with = "deserialize_fork_reference_nth_user_message", + default + )] + pub nth_user_message: i64, +} + +fn deserialize_fork_reference_nth_user_message<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + let value = Value::deserialize(deserializer)?; + let Value::Number(number) = value else { + return Err(serde::de::Error::custom( + "expected integer fork reference boundary", + )); + }; + + if let Some(nth_user_message) = number.as_i64() { + return Ok(nth_user_message); + } + + if number.as_u64().is_some() { + return Ok(i64::MAX); + } + + Err(serde::de::Error::custom( + "expected integer fork reference boundary", + )) } #[derive(Serialize, Deserialize, Debug, Clone, JsonSchema, TS)] @@ -3727,6 +3727,23 @@ mod tests { ); } + #[test] + fn fork_reference_item_deserializes_legacy_usize_max_boundary() { + let item: ForkReferenceItem = serde_json::from_value(json!({ + "rollout_path": "/tmp/rollout.jsonl", + "nth_user_message": u64::MAX, + })) + .expect("legacy fork reference item should deserialize"); + + assert_eq!( + item, + ForkReferenceItem { + rollout_path: PathBuf::from("/tmp/rollout.jsonl"), + nth_user_message: i64::MAX, + } + ); + } + #[test] fn session_source_restriction_product_does_not_guess_subagent_products() { assert_eq!( diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 163c8a1ee8..f999d64402 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -7,6 +7,8 @@ use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; use codex_protocol::protocol::AgentMessageEvent; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::ForkReferenceItem; +use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::UserMessageEvent; @@ -62,6 +64,47 @@ fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io::Result std::io::Result<()> { + let home = TempDir::new().expect("temp dir"); + let parent_id = Uuid::new_v4(); + let child_id = Uuid::new_v4(); + let ts = "2025-01-03T12:00:00.000Z"; + let parent_rollout_path = write_session_file(home.path(), ts, parent_id)?; + let child_rollout_path = write_session_file(home.path(), ts, child_id)?; + + let mut file = fs::OpenOptions::new() + .append(true) + .open(&child_rollout_path)?; + let fork_reference_line = serde_json::json!({ + "timestamp": ts, + "type": "fork_reference", + "payload": { + "rollout_path": parent_rollout_path, + "nth_user_message": u64::MAX, + }, + }); + writeln!(file, "{fork_reference_line}")?; + + let history = RolloutRecorder::get_rollout_history(&child_rollout_path).await?; + let InitialHistory::Resumed(resumed) = history else { + panic!("expected resumed history"); + }; + + let loaded_fork_reference = resumed.history.last().and_then(|item| match item { + RolloutItem::ForkReference(fork_reference) => Some(fork_reference), + _ => None, + }); + assert_eq!( + loaded_fork_reference, + Some(&ForkReferenceItem { + rollout_path: parent_rollout_path, + nth_user_message: i64::MAX, + }), + ); + Ok(()) +} + #[tokio::test] async fn recorder_materializes_only_after_explicit_persist() -> std::io::Result<()> { let home = TempDir::new().expect("temp dir");