diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 0fb1cd3bd2..cd5fc70394 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -9,6 +9,8 @@ use crate::event_mapping::is_contextual_user_message_content; use crate::session::turn_context::TurnContext; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_context_fragments::set_annotated_content; +use codex_context_fragments::to_annotated_content; use codex_extension_api::ConversationHistorySnapshot; use codex_history::CodexHarnessMetadata; use codex_history::ResponseItemEnvelope; @@ -345,17 +347,19 @@ impl ContextManager { { retained_items.retain_mut(|item| { if item.turn_id() == Some(first_turn_id) - && let ResponseItem::Message { role, content, .. } = &mut item.item - && role == "developer" + && matches!(&item.item, ResponseItem::Message { role, .. } if role == "developer") { + let Some(mut content) = to_annotated_content(&mut item.item) else { + return false; + }; content.retain(|content| { !matches!( - content, + content.content(), ContentItem::InputText { text } if ModelSwitchInstructions::matches_text(text) ) }); - !content.is_empty() + !content.is_empty() && set_annotated_content(&mut item.item, content).is_some() } else { true } diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index a1088310bc..5a0c6e2602 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -11,6 +11,7 @@ use codex_protocol::AgentPath; use codex_protocol::ResponseItemId; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; +use codex_protocol::models::ContentItemKind; use codex_protocol::models::DEFAULT_IMAGE_DETAIL; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; @@ -1255,6 +1256,67 @@ fn drop_last_n_user_turns_trims_context_updates_above_rolled_back_turn() { ); } +#[test] +fn drop_last_n_user_turns_preserves_annotations_for_surviving_developer_fragments() { + let turn_id = "rolled-back-turn"; + let model_switch = ModelSwitchInstructions::new("switched model instructions").render(); + let developer_message = ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ + ContentItem::InputText { + text: "persistent developer instructions".to_string(), + }, + ContentItem::InputText { text: model_switch }, + ContentItem::InputText { + text: "persistent environment instructions".to_string(), + }, + ], + phase: None, + internal_chat_message_metadata_passthrough: Some(InternalChatMessageMetadataPassthrough { + turn_id: Some(turn_id.to_string()), + content_item_kinds: Some(vec![ + ContentItemKind("generic.developer_instructions".to_string()), + ContentItemKind("model_switch.instructions".to_string()), + ContentItemKind("environments.instructions".to_string()), + ]), + ..Default::default() + }), + }; + let mut user_message = user_input_text_msg("turn to roll back"); + user_message.set_turn_id_if_missing(turn_id); + let mut history = create_history_with_items(vec![developer_message, user_message]); + + history.drop_last_n_user_turns(/*num_turns*/ 1); + + assert_eq!( + raw_items(&history), + vec![ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ + ContentItem::InputText { + text: "persistent developer instructions".to_string(), + }, + ContentItem::InputText { + text: "persistent environment instructions".to_string(), + }, + ], + phase: None, + internal_chat_message_metadata_passthrough: Some( + InternalChatMessageMetadataPassthrough { + turn_id: Some(turn_id.to_string()), + content_item_kinds: Some(vec![ + ContentItemKind("generic.developer_instructions".to_string()), + ContentItemKind("environments.instructions".to_string()), + ]), + ..Default::default() + }, + ), + }], + ); +} + #[test] fn drop_last_n_user_turns_trims_saved_prefix_update_above_rolled_back_turn() { let items = vec![ diff --git a/codex-rs/core/tests/suite/model_switching.rs b/codex-rs/core/tests/suite/model_switching.rs index 0eaa0748ae..4bd09d11ae 100644 --- a/codex-rs/core/tests/suite/model_switching.rs +++ b/codex-rs/core/tests/suite/model_switching.rs @@ -345,6 +345,20 @@ async fn rollback_first_turn_model_change_removes_its_instructions( let request = &response_mock.requests()[1]; assert_eq!(request.body_json()["model"], followup_model); + let misaligned_messages = request + .inputs_of_type("message") + .into_iter() + .filter(|message| { + message["internal_chat_message_metadata_passthrough"]["content_item_kinds"] + .as_array() + .is_some_and(|kinds| { + message["content"] + .as_array() + .is_none_or(|content| content.len() != kinds.len()) + }) + }) + .collect::>(); + assert_eq!(misaligned_messages, Vec::::new()); let model_switch_count = request .message_input_texts("developer") .iter()