openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-08 20:40:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Preserve context annotations in merged messages (#40184)

Browse Source 
## 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
This commit is contained in:
pakrym-oai
2026-08-23 04:41:11 +00:00
committed by copyberry
parent 422239eb4b
commit c4b434aae2
8 changed files with 199 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/context/world_state/multi_agent_mode_tests.rs
+14 -3
View File
@@ -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),
ContextualUserFragment::into(expected_mode),
vec![ContentItem::InputText {
text: current_role.render(),
}],
vec![ContentItem::InputText {
text: expected_mode.render(),
}],
],
);
}
codex-rs/core/src/context_manager/updates.rs
+55 -44
View File
@@ -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> {
build_text_message("user", text_sections)
}
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
pub(crate) fn build_rendered_message(fragments: Vec<RenderedFragment>) -> Option<ResponseItem> {
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<String>) -> Option<ResponseItem> {
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<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
+47 -24
View File
@@ -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 contextual_user_sections = Vec::<String>::with_capacity(2);
let mut separate_developer_sections = Vec::<String>::new();
let mut developer_sections = Vec::<RenderedFragment>::with_capacity(8);
let mut contextual_user_sections = Vec::<RenderedFragment>::with_capacity(2);
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()),
"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 {
codex-rs/core/src/session/tests.rs
+13 -1
View File
@@ -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()),