diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index 66e0cf27ce..e087240ccd 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -78,6 +78,16 @@ fn finalize_active_segment<'a>( return; } + // Full world-state snapshots are persisted after installing initial context. They still + // establish a baseline when a child fork removes the parent turn's agent message. Do not + // count these context-only segments as user turns for rollback, or use a snapshot from + // before the segment's latest compaction. + let has_context_baseline = active_segment.counts_as_user_turn + || active_segment + .world_state_replay + .iter() + .take_while(|item| !matches!(item, RolloutItem::Compacted(_))) + .any(|item| matches!(item, RolloutItem::WorldState(state) if state.full)); world_state_replay.extend(active_segment.world_state_replay); // A surviving replacement-history checkpoint is a complete history base. Once we @@ -92,15 +102,15 @@ fn finalize_active_segment<'a>( *window = active_segment.window; } - // `previous_turn_settings` come from the newest surviving user turn that established them. - if previous_turn_settings.is_none() && active_segment.counts_as_user_turn { + // Restore settings from the newest surviving context baseline. + if previous_turn_settings.is_none() && has_context_baseline { *previous_turn_settings = active_segment.previous_turn_settings; } - // `reference_context_item` comes from the newest surviving user turn baseline, or + // `reference_context_item` comes from the newest surviving context baseline, or // from a surviving compaction that explicitly cleared that baseline. if matches!(reference_context_item, TurnReferenceContextItem::NeverSet) - && (active_segment.counts_as_user_turn + && (has_context_baseline || matches!( active_segment.reference_context_item, TurnReferenceContextItem::Cleared diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index 4f8c6ceebd..4816b92d99 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -24,6 +24,7 @@ use pretty_assertions::assert_eq; use serde_json::json; use std::collections::BTreeMap; use std::path::PathBuf; +use test_case::test_case; use uuid::Uuid; macro_rules! object { @@ -181,8 +182,16 @@ async fn record_initial_history_ignores_security_risk_scores() { ); } +#[derive(Clone, Copy)] +enum BaselineTurnInput { + UserMessage, + RemovedByFork, +} + +#[test_case(BaselineTurnInput::UserMessage; "user turn")] +#[test_case(BaselineTurnInput::RemovedByFork; "fork removed the task message")] #[tokio::test] -async fn record_initial_history_restores_world_state_baseline() { +async fn record_initial_history_restores_world_state_baseline(input: BaselineTurnInput) { let (session, turn_context) = make_session_and_context().await; let turn_context = Arc::new(turn_context); let world_state = build_world_state_from_turn_context(&session, &turn_context).await; @@ -200,8 +209,19 @@ async fn record_initial_history_restores_world_state_baseline() { world_state_items.push(RolloutItem::WorldState(WorldStateItem::full( world_state.snapshot().into_object(), ))); + let context_item = turn_context.to_turn_context_item(); + let rollout_items = match input { + BaselineTurnInput::UserMessage => { + completed_user_turn_rollout(context_item.clone(), world_state_items) + } + BaselineTurnInput::RemovedByFork => { + world_state_items.push(RolloutItem::TurnContext(context_item.clone())); + world_state_items + } + }; + // Exercise the persisted representation, not just an in-memory fork. let rollout_items = - completed_user_turn_rollout(turn_context.to_turn_context_item(), world_state_items); + serde_json::from_value(serde_json::to_value(rollout_items).unwrap()).unwrap(); session .record_initial_history(InitialHistory::Resumed(ResumedHistory { @@ -210,6 +230,20 @@ async fn record_initial_history_restores_world_state_baseline() { rollout_path: Some(PathBuf::from("/tmp/resume.jsonl")), })) .await; + assert_eq!( + ( + session.previous_turn_settings().await, + serde_json::to_value(session.reference_context_item().await).unwrap(), + ), + ( + Some(PreviousTurnSettings { + model: context_item.model.clone(), + comp_hash: context_item.comp_hash.clone(), + realtime_active: context_item.realtime_active, + }), + serde_json::to_value(Some(context_item)).unwrap(), + ) + ); let step_context = StepContext::for_test(Arc::clone(&turn_context)); session .record_context_updates_and_set_reference_context_item(&step_context) @@ -674,6 +708,11 @@ async fn reconstruct_history_rollback_skips_non_user_turns_for_history_and_metad }, )), RolloutItem::ResponseItem(standalone_assistant.into()), + RolloutItem::WorldState(WorldStateItem::full(object!({}))), + RolloutItem::TurnContext(TurnContextItem { + turn_id: Some(standalone_turn_id.clone()), + ..first_context_item.clone() + }), RolloutItem::EventMsg(EventMsg::TurnComplete( codex_protocol::protocol::TurnCompleteEvent { turn_id: standalone_turn_id, @@ -1053,11 +1092,42 @@ async fn record_initial_history_resumed_rollback_drops_incomplete_user_turn_comp ); } +#[derive(Clone, Copy)] +enum MissingContextBaseline { + BareTurnContext, + WorldStatePatch, + CompactedSnapshot, +} + +#[test_case(MissingContextBaseline::BareTurnContext; "bare turn context")] +#[test_case(MissingContextBaseline::WorldStatePatch; "patch without a full snapshot")] +#[test_case(MissingContextBaseline::CompactedSnapshot; "full snapshot before compaction")] #[tokio::test] -async fn record_initial_history_resumed_bare_turn_context_does_not_seed_reference_context_item() { +async fn record_initial_history_requires_surviving_full_snapshot_without_user_turn( + baseline: MissingContextBaseline, +) { let (session, turn_context) = make_session_and_context().await; - let previous_context_item = turn_context.to_turn_context_item(); - let rollout_items = vec![RolloutItem::TurnContext(previous_context_item.clone())]; + let mut rollout_items = match baseline { + MissingContextBaseline::BareTurnContext => Vec::new(), + MissingContextBaseline::WorldStatePatch => { + vec![RolloutItem::WorldState(WorldStateItem::patch(object!({})))] + } + MissingContextBaseline::CompactedSnapshot => vec![ + RolloutItem::WorldState(WorldStateItem::full(object!({}))), + RolloutItem::Compacted(CompactedItem { + message: String::new(), + replacement_history: Some(Vec::new()), + mcp_resource_origins: None, + window_number: None, + first_window_id: None, + previous_window_id: None, + window_id: None, + }), + ], + }; + rollout_items.push(RolloutItem::TurnContext( + turn_context.to_turn_context_item(), + )); session .record_initial_history(InitialHistory::Resumed(ResumedHistory { diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index 5e0f4c8f1c..f7bef441ad 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -33,6 +33,8 @@ use core_test_support::responses::assert_parent_turn; use core_test_support::responses::assert_root_turn; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_completed_with_tokens; +use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_function_call_with_namespace; use core_test_support::responses::ev_response_created; use core_test_support::responses::ev_tool_search_call; @@ -1122,6 +1124,248 @@ async fn spawned_child_receives_forked_parent_context( Ok(()) } +#[derive(Clone, Copy, Debug)] +enum GrandchildParentContext { + FullHistory, + LastTurn, + NoHistory, + Compacted, +} + +#[test_case(GrandchildParentContext::FullHistory, ThreadHistoryMode::Legacy; "legacy full history")] +#[test_case(GrandchildParentContext::LastTurn, ThreadHistoryMode::Legacy; "legacy last turn")] +#[test_case(GrandchildParentContext::NoHistory, ThreadHistoryMode::Legacy; "legacy no history")] +#[test_case(GrandchildParentContext::FullHistory, ThreadHistoryMode::Paginated; "paginated full history")] +#[test_case(GrandchildParentContext::LastTurn, ThreadHistoryMode::Paginated; "paginated last turn")] +#[test_case(GrandchildParentContext::NoHistory, ThreadHistoryMode::Paginated; "paginated no history")] +#[test_case(GrandchildParentContext::Compacted, ThreadHistoryMode::Legacy; "legacy full history after compaction")] +#[test_case(GrandchildParentContext::Compacted, ThreadHistoryMode::Paginated; "paginated full history after compaction")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn grandchild_full_fork_preserves_context_baseline( + parent_context: GrandchildParentContext, + history_mode: ThreadHistoryMode, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + const ROOT_PROMPT: &str = "root: delegate the context check"; + const CHILD_TASK: &str = "child: delegate the context check"; + const GRANDCHILD_TASK: &str = "grandchild: check inherited context"; + const ROOT_CALL: &str = "root-context-baseline-spawn"; + const CHILD_CALL: &str = "child-context-baseline-spawn"; + const INSTRUCTIONS: &str = "UNIQUE_CONTEXT_BASELINE_DEVELOPER_INSTRUCTIONS"; + const COMPACT_PROMPT: &str = "CONTEXT_BASELINE_COMPACTION_PROMPT"; + const COMPACT_SUMMARY: &str = "CONTEXT_BASELINE_COMPACTION_SUMMARY"; + const PRELUDE_CALL: &str = "context-baseline-prelude-call"; + + let server = start_mock_server().await; + let (parent_fork_turns, compact_parent) = match parent_context { + GrandchildParentContext::FullHistory => ("all", false), + GrandchildParentContext::LastTurn => ("1", false), + GrandchildParentContext::NoHistory => ("none", false), + GrandchildParentContext::Compacted => ("all", true), + }; + let root_spawn_args = serde_json::to_string(&json!({ + "task_name": "child", + "message": CHILD_TASK, + "fork_turns": parent_fork_turns, + }))?; + let root_log = mount_sse_once_match( + &server, + |req: &wiremock::Request| { + body_contains(req, ROOT_PROMPT) + && !body_contains(req, CHILD_TASK) + && !body_contains(req, GRANDCHILD_TASK) + && !body_contains(req, ROOT_CALL) + }, + sse(vec![ + ev_response_created("baseline-root"), + ev_function_call_with_namespace( + ROOT_CALL, + MULTI_AGENT_V2_NAMESPACE, + "spawn_agent", + &root_spawn_args, + ), + ev_completed("baseline-root"), + ]), + ) + .await; + let child_spawn_args = serde_json::to_string(&json!({ + "task_name": "grandchild", + "message": GRANDCHILD_TASK, + "fork_turns": "all", + }))?; + if compact_parent { + mount_sse_once_match( + &server, + |req: &wiremock::Request| { + body_contains(req, CHILD_TASK) + && !body_contains(req, ROOT_CALL) + && !body_contains(req, PRELUDE_CALL) + && !body_contains(req, COMPACT_SUMMARY) + }, + sse(vec![ + ev_response_created("baseline-prelude"), + ev_function_call( + PRELUDE_CALL, + "update_plan", + r#"{"plan":[{"step":"Check inherited context","status":"in_progress"}]}"#, + ), + ev_completed_with_tokens("baseline-prelude", /*total_tokens*/ 250_000), + ]), + ) + .await; + mount_sse_once_match( + &server, + |req: &wiremock::Request| body_contains(req, COMPACT_PROMPT), + sse(vec![ + ev_response_created("baseline-compaction"), + ev_assistant_message("baseline-summary", COMPACT_SUMMARY), + ev_completed("baseline-compaction"), + ]), + ) + .await; + } + let child_log = mount_sse_once_match( + &server, + move |req: &wiremock::Request| { + body_contains( + req, + if compact_parent { + COMPACT_SUMMARY + } else { + CHILD_TASK + }, + ) && !body_contains(req, GRANDCHILD_TASK) + && !body_contains(req, ROOT_CALL) + && !body_contains(req, CHILD_CALL) + && !body_contains(req, COMPACT_PROMPT) + }, + sse(vec![ + ev_response_created("baseline-child"), + ev_function_call_with_namespace( + CHILD_CALL, + MULTI_AGENT_V2_NAMESPACE, + "spawn_agent", + &child_spawn_args, + ), + ev_completed("baseline-child"), + ]), + ) + .await; + let grandchild_log = mount_sse_once_match( + &server, + |req: &wiremock::Request| { + body_contains(req, GRANDCHILD_TASK) && !body_contains(req, CHILD_CALL) + }, + sse(vec![ + ev_response_created("baseline-grandchild"), + ev_assistant_message("baseline-grandchild-answer", "done"), + ev_completed("baseline-grandchild"), + ]), + ) + .await; + let _parent_followups = mount_sse_sequence( + &server, + vec![ + sse(vec![ev_completed("baseline-parent-finished-1")]), + sse(vec![ev_completed("baseline-parent-finished-2")]), + ], + ) + .await; + let test = test_codex() + .with_history_mode(history_mode) + .with_config(move |config| { + config + .features + .enable(Feature::Collab) + .expect("test config should allow feature update"); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + config.model = Some(V2_DEFAULT_MODEL.to_string()); + config.agent_default_subagent_model = Some(V2_DEFAULT_MODEL.to_string()); + config.developer_instructions = Some(INSTRUCTIONS.to_string()); + if compact_parent { + // Use local compaction so the test controls the replacement history. + config.model_provider.name = "test-provider".to_string(); + config + .features + .disable(Feature::RemoteCompactionV2) + .expect("test config should allow feature update"); + config.compact_prompt = Some(COMPACT_PROMPT.to_string()); + config.model_auto_compact_token_limit = Some(200_000); + config.model_context_window = Some(1_000_000); + } + }) + .build_with_auto_env(&server) + .await?; + + test.submit_turn(ROOT_PROMPT).await?; + let root_request = root_log.single_request(); + let mut descendant_requests = Vec::new(); + for (mock, agent_name) in [ + (&child_log, "/root/child"), + (&grandchild_log, "/root/child/grandchild"), + ] { + let request = timeout(Duration::from_secs(/*secs*/ 10), async { + loop { + let request = mock.requests().into_iter().find(|request| { + request.body_json()["client_metadata"]["x-codex-turn-metadata"] + .as_str() + .and_then(|text| serde_json::from_str::(text).ok()) + .is_some_and(|metadata| metadata["agent_name"] == agent_name) + && (!compact_parent || request.body_contains_text(COMPACT_SUMMARY)) + }); + if let Some(request) = request { + break request; + } + sleep(Duration::from_millis(/*millis*/ 10)).await; + } + }) + .await?; + let thread_id = ThreadId::from_string( + request.body_json()["client_metadata"]["thread_id"] + .as_str() + .expect("descendant thread id"), + )?; + let thread = test.thread_manager.get_thread(thread_id).await?; + timeout(Duration::from_secs(/*secs*/ 10), async { + while !matches!(thread.agent_status().await, AgentStatus::Completed(_)) { + sleep(Duration::from_millis(/*millis*/ 10)).await; + } + }) + .await?; + descendant_requests.push(request); + } + let context_counts = [ + &root_request, + &descendant_requests[0], + &descendant_requests[1], + ] + .map(|request| { + ( + request + .message_input_texts("developer") + .iter() + .filter(|text| text.contains(INSTRUCTIONS)) + .count(), + request + .message_input_texts("user") + .iter() + .filter(|text| text.contains("")) + .count(), + ) + }); + assert_eq!( + context_counts, + [(1, 1); 3], + "Initial context should appear once per agent: {parent_context:?}, {history_mode:?}" + ); + assert!(!descendant_requests[1].body_contains_text(CHILD_TASK)); + Ok(()) +} + #[derive(Clone, Copy)] enum FullHistoryV2ModelSelection { ConfiguredDefault,