From c4b434aae27d5e589c9129f0b7ea312ccc8cbacd Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Sun, 23 Aug 2026 04:35:53 +0000 Subject: [PATCH] Preserve context annotations in merged messages (#40184) ## What changed - Carry each contextual fragment's `ContentItemKind` through initial context assembly and world-state updates. - Attach the classifications to merged response messages through `content_item_kinds`, preserving their order alongside the message content. - Convert extension prompt fragments into annotated developer fragments and classify built-in developer and guardian instructions. ## Testing - Add an integration test covering roles and content annotations in the first request. - Update world-state and persisted-history tests to account for annotation metadata. GitOrigin-RevId: 71c48dad5b17bffd47ba03263e3eb6d5e74fbf4b --- codex-rs/context-fragments/src/fragment.rs | 13 ++- .../world_state/multi_agent_mode_tests.rs | 17 +++- codex-rs/core/src/context_manager/updates.rs | 99 ++++++++++--------- codex-rs/core/src/session/mod.rs | 71 ++++++++----- codex-rs/core/src/session/tests.rs | 14 ++- .../core/tests/suite/context_annotations.rs | 49 +++++++++ codex-rs/core/tests/suite/mod.rs | 1 + .../extension-api/src/contributors/prompt.rs | 11 +++ 8 files changed, 199 insertions(+), 76 deletions(-) create mode 100644 codex-rs/core/tests/suite/context_annotations.rs diff --git a/codex-rs/context-fragments/src/fragment.rs b/codex-rs/context-fragments/src/fragment.rs index 2ee321ef3a..d6840dc1fe 100644 --- a/codex-rs/context-fragments/src/fragment.rs +++ b/codex-rs/context-fragments/src/fragment.rs @@ -12,6 +12,11 @@ pub struct RenderedFragment { } impl RenderedFragment { + /// Creates a rendered fragment without separating its role and annotated content. + pub fn new(role: &'static str, content: AnnotatedContent) -> Self { + Self { role, content } + } + /// Returns the response role associated with this fragment. pub fn role(&self) -> &'static str { self.role @@ -76,10 +81,10 @@ pub trait ContextualUserFragment { /// Renders the role, model-visible content, and classification together. fn render_fragment(&self) -> RenderedFragment { - RenderedFragment { - role: self.role(), - content: AnnotatedContent::input_text(self.render(), self.content_kind()), - } + RenderedFragment::new( + self.role(), + AnnotatedContent::input_text(self.render(), self.content_kind()), + ) } fn into(self) -> ResponseItem diff --git a/codex-rs/core/src/context/world_state/multi_agent_mode_tests.rs b/codex-rs/core/src/context/world_state/multi_agent_mode_tests.rs index 49c67a43f1..fa73781bc3 100644 --- a/codex-rs/core/src/context/world_state/multi_agent_mode_tests.rs +++ b/codex-rs/core/src/context/world_state/multi_agent_mode_tests.rs @@ -2,6 +2,7 @@ use super::super::test_support::render_section_cases; use super::*; use crate::context::MultiAgentRoleInstructions; use crate::context::world_state::WorldState; +use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_utils_output_truncation::approx_token_count; @@ -112,10 +113,20 @@ fn catalog_role_updates_remain_separate_from_active_mode() { let expected_mode = MultiAgentModeInstructions::from_mode(MultiAgentMode::Proactive) .expect("proactive mode should render"); assert_eq!( - updates, + updates + .into_iter() + .map(|item| match item { + ResponseItem::Message { content, .. } => content, + _ => panic!("expected world-state message"), + }) + .collect::>(), vec![ - ContextualUserFragment::into(current_role), - ContextualUserFragment::into(expected_mode), + vec![ContentItem::InputText { + text: current_role.render(), + }], + vec![ContentItem::InputText { + text: expected_mode.render(), + }], ], ); } diff --git a/codex-rs/core/src/context_manager/updates.rs b/codex-rs/core/src/context_manager/updates.rs index e07f7bd24e..4cd567c5a0 100644 --- a/codex-rs/core/src/context_manager/updates.rs +++ b/codex-rs/core/src/context_manager/updates.rs @@ -1,5 +1,8 @@ use crate::context::ContextualUserFragment; -use codex_protocol::models::ContentItem; +use codex_context_fragments::AnnotatedContent; +use codex_context_fragments::RenderedFragment; +use codex_protocol::models::ContentItemKind; +use codex_protocol::models::InternalChatMessageMetadataPassthrough; use codex_protocol::models::ResponseItem; #[derive(Clone, Copy, PartialEq, Eq)] @@ -9,57 +12,65 @@ enum MessageGroup { } pub(crate) fn build_developer_update_item(text_sections: Vec) -> Option { - build_text_message("developer", text_sections) + build_rendered_message( + text_sections + .into_iter() + .map(|text| { + RenderedFragment::new( + "developer", + AnnotatedContent::input_text(text, ContentItemKind("unknown".to_string())), + ) + }) + .collect(), + ) } -pub(crate) fn build_contextual_user_message(text_sections: Vec) -> Option { - build_text_message("user", text_sections) -} - -pub(crate) fn merge_contextual_fragments( - fragments: Vec>, -) -> Vec { - let mut messages: Vec<(&str, MessageGroup, Vec)> = Vec::with_capacity(fragments.len()); - for fragment in fragments { - let role = fragment.role(); - let group = if fragment.requires_separate_message() { - MessageGroup::Standalone - } else { - MessageGroup::Mergeable - }; - let text = fragment.render(); - match messages.last_mut() { - Some((previous_role, previous_group, text_sections)) - if *previous_role == role - && *previous_group == MessageGroup::Mergeable - && group == MessageGroup::Mergeable => - { - text_sections.push(text); - } - _ => messages.push((role, group, vec![text])), - } - } - messages +pub(crate) fn build_rendered_message(fragments: Vec) -> Option { + let role = fragments.first()?.role(); + debug_assert!(fragments.iter().all(|fragment| fragment.role() == role)); + let (content, content_item_kinds): (Vec<_>, Vec<_>) = fragments .into_iter() - .filter_map(|(role, _, text_sections)| build_text_message(role, text_sections)) - .collect() -} - -fn build_text_message(role: &str, text_sections: Vec) -> Option { - if text_sections.is_empty() { - return None; - } - - let content = text_sections - .into_iter() - .map(|text| ContentItem::InputText { text }) - .collect(); + .map(|fragment| fragment.into_parts().1.into_parts()) + .unzip(); Some(ResponseItem::Message { id: None, role: role.to_string(), content, phase: None, - internal_chat_message_metadata_passthrough: None, + internal_chat_message_metadata_passthrough: Some(InternalChatMessageMetadataPassthrough { + content_item_kinds: Some(content_item_kinds), + ..Default::default() + }), }) } + +pub(crate) fn merge_contextual_fragments( + fragments: Vec>, +) -> Vec { + let mut messages: Vec<(&str, MessageGroup, Vec)> = + Vec::with_capacity(fragments.len()); + for fragment in fragments { + let group = if fragment.requires_separate_message() { + MessageGroup::Standalone + } else { + MessageGroup::Mergeable + }; + let rendered = fragment.render_fragment(); + let role = rendered.role(); + match messages.last_mut() { + Some((previous_role, previous_group, rendered_fragments)) + if *previous_role == role + && *previous_group == MessageGroup::Mergeable + && group == MessageGroup::Mergeable => + { + rendered_fragments.push(rendered); + } + _ => messages.push((role, group, vec![rendered])), + } + } + messages + .into_iter() + .filter_map(|(_, _, fragments)| build_rendered_message(fragments)) + .collect() +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 28567c880f..8c120ec3fd 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -56,6 +56,8 @@ use codex_analytics::SubAgentThreadStartedInput; use codex_analytics::TurnCodexErrorFact; use codex_async_utils::OrCancelExt; use codex_connectors::connector_runtime_context_key; +use codex_context_fragments::AnnotatedContent; +use codex_context_fragments::RenderedFragment; use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; use codex_execpolicy::prefix_rule_migration; @@ -104,6 +106,7 @@ use codex_protocol::models::ActivePermissionProfile; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::models::BaseInstructions; use codex_protocol::models::BaseInstructionsProvenance; +use codex_protocol::models::ContentItemKind; use codex_protocol::models::PermissionProfile; use codex_protocol::models::SandboxEnforcement; use codex_protocol::openai_models::ModelInfo; @@ -3479,11 +3482,11 @@ impl Session { }) .await { - developer_sections.push(fragment.text().to_string()); + developer_sections.push(fragment.into()); } } - crate::context_manager::updates::build_developer_update_item(developer_sections) + crate::context_manager::updates::build_rendered_message(developer_sections) .into_iter() .collect() } @@ -3493,9 +3496,9 @@ impl Session { turn_context: &TurnContext, world_state: &WorldState, ) -> Vec { - let mut developer_sections = Vec::::with_capacity(8); - let mut contextual_user_sections = Vec::::with_capacity(2); - let mut separate_developer_sections = Vec::::new(); + let mut developer_sections = Vec::::with_capacity(8); + let mut contextual_user_sections = Vec::::with_capacity(2); + let mut separate_developer_sections = Vec::::new(); let (session_source, auto_compact_window_ids) = { let state = self.state.lock().await; ( @@ -3511,7 +3514,13 @@ impl Session { && let Some(developer_instructions) = turn_context.developer_instructions.as_deref() && !developer_instructions.is_empty() { - developer_sections.push(developer_instructions.to_string()); + developer_sections.push(RenderedFragment::new( + "developer", + AnnotatedContent::input_text( + developer_instructions, + ContentItemKind("generic.developer_instructions".to_string()), + ), + )); } let loaded_plugins = self .services @@ -3543,7 +3552,7 @@ impl Session { .as_deref() .and_then(RecommendedPluginsInstructions::from_plugins) { - contextual_user_sections.push(recommended_plugins.render()); + contextual_user_sections.push(recommended_plugins.render_fragment()); } let context_contributors = self.services.extensions.context_contributors().to_vec(); for contributor in &context_contributors { @@ -3554,7 +3563,7 @@ impl Session { ) .await { - developer_sections.push(fragment.text().to_string()); + developer_sections.push(fragment.into()); } } for contributor in &context_contributors { @@ -3569,7 +3578,7 @@ impl Session { }) .await { - developer_sections.push(fragment.text().to_string()); + developer_sections.push(fragment.into()); } } // This is full-context metadata. Steady-state context diffs should not re-emit it. @@ -3614,7 +3623,7 @@ impl Session { auto_compact_window_ids.window_id, mcp_result, ) - .render(), + .render_fragment(), ); } // Render the active mode after the usage hint so it can override that hint. @@ -3626,7 +3635,7 @@ impl Session { if fragment.markers().0 == ModelSwitchInstructions::type_markers().0 => { // New-model instructions must precede the rest of the developer context. - developer_sections.insert(0, fragment.render()); + developer_sections.insert(0, fragment.render_fragment()); } "developer" if fragment.markers().0 == MULTI_AGENT_MODE_OPEN_TAG => { initial_multi_agent_mode = Some(fragment); @@ -3639,37 +3648,41 @@ impl Session { "developer" if fragment.markers().0 == MultiAgentRoleInstructions::type_markers().0 => { - separate_developer_sections.push(fragment.render()); + separate_developer_sections.push(fragment.render_fragment()); } "developer" if fragment.requires_separate_message() && fragment.markers().0.is_empty() => { - separate_developer_sections.push(fragment.render()); + separate_developer_sections.push(fragment.render_fragment()); } - "developer" => developer_sections.push(fragment.render()), - "user" => contextual_user_sections.push(fragment.render()), + "developer" => developer_sections.push(fragment.render_fragment()), + "user" => contextual_user_sections.push(fragment.render_fragment()), _ => {} } } let mut items = Vec::with_capacity(4); if let Some(developer_message) = - crate::context_manager::updates::build_developer_update_item(developer_sections) + crate::context_manager::updates::build_rendered_message(developer_sections) { items.push(developer_message); } for section in separate_developer_sections { if let Some(developer_message) = - crate::context_manager::updates::build_developer_update_item(vec![section]) + crate::context_manager::updates::build_rendered_message(vec![section]) { items.push(developer_message); } } - if let Some(initial_multi_agent_mode) = initial_multi_agent_mode { - items.push(initial_multi_agent_mode.into_boxed_response_item()); + if let Some(initial_multi_agent_mode) = initial_multi_agent_mode + && let Some(message) = crate::context_manager::updates::build_rendered_message(vec![ + initial_multi_agent_mode.render_fragment(), + ]) + { + items.push(message); } if let Some(contextual_user_message) = - crate::context_manager::updates::build_contextual_user_message(contextual_user_sections) + crate::context_manager::updates::build_rendered_message(contextual_user_sections) { items.push(contextual_user_message); } @@ -3679,14 +3692,24 @@ impl Session { && let Some(developer_instructions) = turn_context.developer_instructions.as_deref() && !developer_instructions.is_empty() && let Some(guardian_developer_message) = - crate::context_manager::updates::build_developer_update_item(vec![ - developer_instructions.to_string(), + crate::context_manager::updates::build_rendered_message(vec![ + RenderedFragment::new( + "developer", + AnnotatedContent::input_text( + developer_instructions, + ContentItemKind("guardian.policy".to_string()), + ), + ), ]) { items.push(guardian_developer_message); } - if let Some(managed_developer_instructions) = managed_developer_instructions { - items.push(managed_developer_instructions.into_boxed_response_item()); + if let Some(managed_developer_instructions) = managed_developer_instructions + && let Some(message) = crate::context_manager::updates::build_rendered_message(vec![ + managed_developer_instructions.render_fragment(), + ]) + { + items.push(message); } // New context windows and compaction install these items directly into replacement history. for item in &mut items { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index a2334f863b..cadbb868ce 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -126,6 +126,7 @@ use codex_protocol::items::HookPromptFragment; use codex_protocol::items::build_hook_prompt_message; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; +use codex_protocol::models::ContentItemKind; use codex_protocol::models::InternalChatMessageMetadataPassthrough; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::AskForApproval; @@ -3323,7 +3324,18 @@ async fn record_initial_history_assigns_and_persists_id_for_forked_response_item "Subagent guidance.".to_string(), ]) .expect("developer message"); - let mut expected_item = response_item.clone(); + let mut expected_item = ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ContentItem::InputText { + text: "Subagent guidance.".to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: Some(InternalChatMessageMetadataPassthrough { + content_item_kinds: Some(vec![ContentItemKind("unknown".to_string())]), + ..Default::default() + }), + }; let response_item = ResponseItemEnvelope { item: response_item, metadata: Some(CodexHarnessMetadata::default()), diff --git a/codex-rs/core/tests/suite/context_annotations.rs b/codex-rs/core/tests/suite/context_annotations.rs new file mode 100644 index 0000000000..b0106b241b --- /dev/null +++ b/codex-rs/core/tests/suite/context_annotations.rs @@ -0,0 +1,49 @@ +use anyhow::Result; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_once; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn first_request_item_types_roles_and_content_annotations() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let response = mount_sse_once( + &server, + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + ) + .await; + let test = test_codex() + .with_config(|config| { + config.developer_instructions = Some("Keep world-state annotations aligned.".into()); + }) + .build_with_auto_env(&server) + .await?; + + test.submit_text_turn("inspect world state").await?; + + let items = response + .single_request() + .input() + .into_iter() + .map(|item| { + let item_type = item["type"].as_str().expect("response item type"); + let role = item["role"].as_str().unwrap_or("-"); + let content_annotations = + &item["internal_chat_message_metadata_passthrough"]["content_item_kinds"]; + format!("{item_type} {role} {content_annotations}") + }) + .collect::>() + .join("\n"); + insta::assert_snapshot!(items, @r#" + message developer ["generic.developer_instructions","generic.permissions_instructions"] + message user ["environments.environment_context"] + message user null + "#); + + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index fc63ac0acb..52fcac4bc2 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -58,6 +58,7 @@ mod compact; mod compact_remote; mod compact_remote_parity; mod compact_resume_fork; +mod context_annotations; mod current_time_reminder; mod cyber_exec_policy; mod deprecation_notice; diff --git a/codex-rs/ext/extension-api/src/contributors/prompt.rs b/codex-rs/ext/extension-api/src/contributors/prompt.rs index dd7c45e70e..cc4e261925 100644 --- a/codex-rs/ext/extension-api/src/contributors/prompt.rs +++ b/codex-rs/ext/extension-api/src/contributors/prompt.rs @@ -1,5 +1,7 @@ // All this file should be replaced by the existing fragment implementation ofc +use codex_context_fragments::AnnotatedContent; +use codex_context_fragments::RenderedFragment; use codex_protocol::models::ContentItemKind; #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] @@ -50,3 +52,12 @@ impl PromptFragment { &self.content_kind } } + +impl From for RenderedFragment { + fn from(fragment: PromptFragment) -> Self { + Self::new( + "developer", + AnnotatedContent::input_text(fragment.text, fragment.content_kind), + ) + } +}