diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index c513cbb55..c7af52ced 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1337,8 +1337,15 @@ impl Session { let _ = self.flush_rollout().await; } } - InitialHistory::Forked(rollout_items) => { + InitialHistory::Forked(mut rollout_items) => { let turn_context = self.new_default_turn().await; + if turn_context.config.features.enabled(Feature::ItemIds) { + for rollout_item in &mut rollout_items { + if let RolloutItem::ResponseItem(response_item) = rollout_item { + Self::assign_missing_response_item_id(response_item); + } + } + } self.apply_rollout_reconstruction(&turn_context, &rollout_items) .await; @@ -2723,31 +2730,35 @@ impl Session { } let mut items = items; for item in items.to_mut() { - if item.id().is_some() { - continue; - } - let prefix = match item { - ResponseItem::AdditionalTools { .. } => "at", - ResponseItem::Message { .. } => "msg", - ResponseItem::Reasoning { .. } => "rs", - ResponseItem::LocalShellCall { .. } => "lsh", - ResponseItem::FunctionCall { .. } => "fc", - ResponseItem::ToolSearchCall { .. } => "tsc", - ResponseItem::FunctionCallOutput { .. } => "fco", - ResponseItem::CustomToolCall { .. } => "ctc", - ResponseItem::CustomToolCallOutput { .. } => "ctco", - ResponseItem::ToolSearchOutput { .. } => "tso", - ResponseItem::WebSearchCall { .. } => "ws", - ResponseItem::ImageGenerationCall { .. } => "ig", - ResponseItem::Compaction { .. } | ResponseItem::ContextCompaction { .. } => "cmp", - ResponseItem::AgentMessage { .. } => "amsg", - ResponseItem::CompactionTrigger { .. } | ResponseItem::Other => continue, - }; - item.set_id(Some(format!("{prefix}_{}", Uuid::now_v7()))); + Self::assign_missing_response_item_id(item); } items } + fn assign_missing_response_item_id(item: &mut ResponseItem) { + if item.id().is_some() { + return; + } + let prefix = match item { + ResponseItem::AdditionalTools { .. } => "at", + ResponseItem::Message { .. } => "msg", + ResponseItem::Reasoning { .. } => "rs", + ResponseItem::LocalShellCall { .. } => "lsh", + ResponseItem::FunctionCall { .. } => "fc", + ResponseItem::ToolSearchCall { .. } => "tsc", + ResponseItem::FunctionCallOutput { .. } => "fco", + ResponseItem::CustomToolCall { .. } => "ctc", + ResponseItem::CustomToolCallOutput { .. } => "ctco", + ResponseItem::ToolSearchOutput { .. } => "tso", + ResponseItem::WebSearchCall { .. } => "ws", + ResponseItem::ImageGenerationCall { .. } => "ig", + ResponseItem::Compaction { .. } | ResponseItem::ContextCompaction { .. } => "cmp", + ResponseItem::AgentMessage { .. } => "amsg", + ResponseItem::CompactionTrigger { .. } | ResponseItem::Other => return, + }; + item.set_id(Some(format!("{prefix}_{}", Uuid::now_v7()))); + } + pub(crate) fn response_item_from_user_input(&self, input: Vec) -> ResponseItem { ResponseItem::from(ResponseInputItem::from_user_input( input, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 48e48b6f3..51f5694f8 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2736,6 +2736,60 @@ async fn start_new_context_window_assigns_and_persists_item_ids() { ); } +#[tokio::test] +async fn record_initial_history_assigns_and_persists_id_for_forked_response_item() { + let (mut session, _turn_context, _rx) = make_session_and_context_with_auth_and_config_and_rx( + CodexAuth::from_api_key("Test API Key"), + Vec::new(), + |config| { + let _ = config.features.enable(Feature::ItemIds); + }, + ) + .await; + let rollout_path = + attach_thread_persistence(Arc::get_mut(&mut session).expect("unique session")).await; + let response_item = crate::context_manager::updates::build_developer_update_item(vec![ + "Subagent guidance.".to_string(), + ]) + .expect("developer message"); + let mut expected_item = response_item.clone(); + + session + .record_initial_history(InitialHistory::Forked(vec![RolloutItem::ResponseItem( + response_item, + )])) + .await; + + let live_history = session.clone_history().await; + let [live_item] = live_history.raw_items() else { + panic!("expected one forked response item"); + }; + let live_item_id = live_item + .id() + .expect("forked response item should have an id") + .to_string(); + assert!(live_item_id.starts_with("msg_")); + expected_item.set_id(Some(live_item_id.clone())); + assert_eq!(live_history.raw_items(), &[expected_item]); + + session.flush_rollout().await.expect("rollout should flush"); + let InitialHistory::Resumed(resumed) = RolloutRecorder::get_rollout_history(&rollout_path) + .await + .expect("read rollout history") + else { + panic!("expected resumed rollout history"); + }; + let persisted_item_id = resumed.history.iter().find_map(|item| match item { + RolloutItem::ResponseItem(response_item) => response_item.id(), + RolloutItem::SessionMeta(_) + | RolloutItem::InterAgentCommunication(_) + | RolloutItem::Compacted(_) + | RolloutItem::TurnContext(_) + | RolloutItem::EventMsg(_) => None, + }); + assert_eq!(persisted_item_id, Some(live_item_id.as_str())); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn session_configured_reports_permission_profile_for_external_sandbox() -> anyhow::Result<()> {