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

@@ -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 {
+        RenderedFragment::new(
-            role: self.role(),
+            self.role(),
-            content: AnnotatedContent::input_text(self.render(), self.content_kind()),
+            AnnotatedContent::input_text(self.render(), self.content_kind()),
-        }
+        )
     }
 
     fn into(self) -> ResponseItem

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<_>>(),
         vec![
-            ContextualUserFragment::into(current_role),
+            vec![ContentItem::InputText {
-            ContextualUserFragment::into(expected_mode),
+                text: current_role.render(),
+            }],
+            vec![ContentItem::InputText {
+                text: expected_mode.render(),
+            }],
         ],
     );
 }

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<String>) -> Option<ResponseItem> {
-    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<String>) -> Option<ResponseItem> {
+pub(crate) fn build_rendered_message(fragments: Vec<RenderedFragment>) -> Option<ResponseItem> {
-    build_text_message("user", text_sections)
+    let role = fragments.first()?.role();
-}
+    debug_assert!(fragments.iter().all(|fragment| fragment.role() == role));
-
+    let (content, content_item_kinds): (Vec<_>, Vec<_>) = fragments
-pub(crate) fn merge_contextual_fragments(
-    fragments: Vec<Box<dyn ContextualUserFragment>>,
-) -> Vec<ResponseItem> {
-    let mut messages: Vec<(&str, MessageGroup, Vec<String>)> = 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
         .into_iter()
-        .filter_map(|(role, _, text_sections)| build_text_message(role, text_sections))
+        .map(|fragment| fragment.into_parts().1.into_parts())
-        .collect()
+        .unzip();
-}
-
-fn build_text_message(role: &str, text_sections: Vec<String>) -> Option<ResponseItem> {
-    if text_sections.is_empty() {
-        return None;
-    }
-
-    let content = text_sections
-        .into_iter()
-        .map(|text| ContentItem::InputText { text })
-        .collect();
 
     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<Box<dyn ContextualUserFragment>>,
+) -> Vec<ResponseItem> {
+    let mut messages: Vec<(&str, MessageGroup, Vec<RenderedFragment>)> =
+        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()
+}

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<ResponseItem> {
-        let mut developer_sections = Vec::<String>::with_capacity(8);
+        let mut developer_sections = Vec::<RenderedFragment>::with_capacity(8);
-        let mut contextual_user_sections = Vec::<String>::with_capacity(2);
+        let mut contextual_user_sections = Vec::<RenderedFragment>::with_capacity(2);
-        let mut separate_developer_sections = Vec::<String>::new();
+        let mut separate_developer_sections = Vec::<RenderedFragment>::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()),
+                "developer" => developer_sections.push(fragment.render_fragment()),
-                "user" => contextual_user_sections.push(fragment.render()),
+                "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 {
+        if let Some(initial_multi_agent_mode) = initial_multi_agent_mode
-            items.push(initial_multi_agent_mode.into_boxed_response_item());
+            && 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![
+                crate::context_manager::updates::build_rendered_message(vec![
-                    developer_instructions.to_string(),
+                    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 {
+        if let Some(managed_developer_instructions) = managed_developer_instructions
-            items.push(managed_developer_instructions.into_boxed_response_item());
+            && 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 {

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()),

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::<Vec<_>>()
+        .join("\n");
+    insta::assert_snapshot!(items, @r#"
+    message developer ["generic.developer_instructions","generic.permissions_instructions"]
+    message user ["environments.environment_context"]
+    message user null
+    "#);
+
+    Ok(())
+}

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;

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<PromptFragment> for RenderedFragment {
+    fn from(fragment: PromptFragment) -> Self {
+        Self::new(
+            "developer",
+            AnnotatedContent::input_text(fragment.text, fragment.content_kind),
+        )
+    }
+}