github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-21 09:54:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Retain client developer messages across context compaction (#38445)

Browse Source 
## Why

Client-authored developer instructions should remain available after a context
window is compacted when `retain_client_developer_messages` is enabled.

## What changed

- Preserve annotated client-authored developer messages during remote compaction
  v2 and local token-budget context resets.
- Apply the existing retained-message token budget when carrying those messages
  into the new context window.

## Testing

- Cover enabled and disabled retention for remote compaction v2, manual
  token-budget compaction, and mid-turn automatic compaction.

GitOrigin-RevId: 4cb3782cb86ef46458413a17a3255d0016fd675f
This commit is contained in:
rka-oai
2026-08-13 22:21:18 +00:00
committed by copyberry
parent 3ba52d6075
commit bff03ecce5
4 changed files with 304 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/compact_remote_v2.rs
+138 -37
View File
@@ -8,10 +8,11 @@ use crate::compact::CompactedHistoryMetadata;
use crate::compact::CompactionAnalyticsAttempt;
use crate::compact::CompactionAnalyticsDetails;
use crate::compact::InitialContextInjection;
use crate::compact::build_compaction_initial_context;
use crate::compact::compaction_status_from_result;
use crate::compact::insert_initial_context_before_last_real_user_or_summary;
use crate::compact_model_fallback::record_model_fallback;
use crate::compact_model_fallback::should_retry_with_current_model;
use crate::compact_remote::process_annotated_compacted_history;
use crate::compact_remote::should_keep_compacted_history_item;
use crate::compact_remote_history::HistoryItemGroup;
use crate::compact_remote_history::history_item_groups;
@@ -32,6 +33,7 @@ use codex_analytics::CompactionImplementation;
use codex_analytics::CompactionPhase;
use codex_analytics::CompactionReason;
use codex_analytics::CompactionTrigger;
use codex_features::Feature;
use codex_history::CodexHarnessMetadata;
use codex_history::ResponseItemEnvelope;
use codex_protocol::error::CodexErr;
@@ -60,7 +62,7 @@ use attempt::run_remote_compact_v2_attempt;
// Mirror the current /responses/compact retained-message default while the
// server-side path remains the reference implementation.
const RETAINED_MESSAGE_TOKEN_BUDGET: usize = 64_000;
pub(crate) const RETAINED_MESSAGE_TOKEN_BUDGET: usize = 64_000;
const MAX_RETAINED_AGENT_MESSAGE_TOKENS: i64 = 10_000;
// Compact attempts can run much longer than normal turns, so keep the per-transport
// retry budget smaller than the general Responses stream retry budget.
@@ -292,16 +294,18 @@ async fn run_remote_compact_task_inner_impl(
analytics_details.cached_input_tokens = Some(token_usage.cached_input_tokens);
analytics_details.cache_write_input_tokens = Some(token_usage.cache_write_input_tokens);
}
let (compacted_history, retained_images) =
build_v2_compacted_history(prompt_input, prompt_input_metadata, compaction_output);
let (compacted_history, retained_images) = build_v2_compacted_history(
prompt_input,
prompt_input_metadata,
compaction_output,
sess.enabled(Feature::RetainClientDeveloperMessages),
);
analytics_details.retained_image_count = Some(retained_images);
let (new_window_number, new_window_ids) = sess.advance_auto_compact_window().await;
let (new_history, world_state_baseline) = process_annotated_compacted_history(
sess.as_ref(),
compacted_history,
&initial_context_injection,
)
.await;
let (initial_context, world_state_baseline) =
build_compaction_initial_context(sess.as_ref(), &initial_context_injection).await;
let new_history =
insert_initial_context_before_last_real_user_or_summary(compacted_history, initial_context);
let reference_context_item = match initial_context_injection {
InitialContextInjection::DoNotInject => None,
@@ -456,6 +460,7 @@ fn build_v2_compacted_history(
prompt_input: Vec<ResponseItem>,
prompt_input_metadata: Vec<Option<CodexHarnessMetadata>>,
compaction_output: ResponseItem,
retain_client_developer_messages: bool,
) -> (Vec<ResponseItemEnvelope>, usize) {
debug_assert_eq!(prompt_input.len(), prompt_input_metadata.len());
let prompt_input = prompt_input
@@ -463,9 +468,13 @@ fn build_v2_compacted_history(
.zip(prompt_input_metadata)
.map(|(item, metadata)| ResponseItemEnvelope { item, metadata })
.collect::<Vec<_>>();
let retained = history_item_groups(prompt_input)
let retained = v2_history_item_groups(prompt_input)
.filter(|group| is_retained_for_remote_compaction_v2(&group.source.item))
.filter(|group| should_keep_compacted_history_item(&group.source.item))
.filter(|group| {
should_keep_compacted_history_item(&group.source.item)
|| (retain_client_developer_messages
&& is_client_authored_developer_message(&group.source))
})
.flat_map(HistoryItemGroup::into_items)
.collect::<Vec<_>>();
let mut retained =
@@ -478,6 +487,28 @@ fn build_v2_compacted_history(
(retained, retained_image_count)
}
pub(crate) fn is_client_authored_developer_message(item: &ResponseItemEnvelope) -> bool {
item.metadata
.as_ref()
.is_some_and(|metadata| metadata.client_authored)
&& matches!(&item.item, ResponseItem::Message { role, .. } if role == "developer")
}
fn v2_history_item_groups(
items: Vec<ResponseItemEnvelope>,
) -> impl Iterator<Item = HistoryItemGroup<ResponseItemEnvelope>> {
history_item_groups(items).flat_map(|mut group| {
let client_message = group
.attached_notice
.take_if(|item| is_client_authored_developer_message(item))
.map(|source| HistoryItemGroup {
source,
attached_notice: None,
});
std::iter::once(group).chain(client_message)
})
}
fn is_retained_for_remote_compaction_v2(item: &ResponseItem) -> bool {
if let ResponseItem::AgentMessage { content, .. } = item {
let is_completion = matches!(
@@ -507,13 +538,13 @@ fn retained_input_image_count(item: &ResponseItem) -> usize {
.count()
}
fn truncate_retained_messages_for_remote_compaction(
pub(crate) fn truncate_retained_messages_for_remote_compaction(
items: Vec<ResponseItemEnvelope>,
max_tokens: usize,
) -> Vec<ResponseItemEnvelope> {
let mut remaining = max_tokens;
let mut truncated_reversed = Vec::with_capacity(items.len());
for group in history_item_groups(items)
for group in v2_history_item_groups(items)
.collect::<Vec<_>>()
.into_iter()
.rev()
@@ -522,25 +553,58 @@ fn truncate_retained_messages_for_remote_compaction(
continue;
}
let client_developer = is_client_authored_developer_message(&group.source);
let notice_tokens = group
.attached_notice
.as_ref()
.map_or(0, |notice| message_text_token_count(&notice.item).max(1));
let token_count = message_text_token_count(&group.source.item)
.max(1)
.saturating_add(notice_tokens);
let source_tokens = if client_developer {
usize::try_from(estimate_item_token_count(&group.source.item)).unwrap_or(usize::MAX)
} else {
message_text_token_count(&group.source.item).max(1)
};
let token_count = source_tokens.saturating_add(notice_tokens);
if token_count <= remaining {
if let Some(notice) = group.attached_notice {
truncated_reversed.push(notice);
}
truncated_reversed.push(group.source);
remaining = remaining.saturating_sub(token_count);
} else if remaining > notice_tokens
&& let Some(truncated_item) = truncate_message_text_to_token_budget(
group.source,
/*max_tokens*/ remaining - notice_tokens,
)
{
} else if remaining > notice_tokens {
let available_tokens = remaining - notice_tokens;
let text_budget = if client_developer {
available_tokens.saturating_sub(
source_tokens.saturating_sub(message_text_token_count(&group.source.item)),
)
} else {
available_tokens
};
let Some(mut truncated_item) =
truncate_message_text_to_token_budget(group.source, text_budget)
else {
continue;
};
if client_developer {
let item_tokens = usize::try_from(estimate_item_token_count(&truncated_item.item))
.unwrap_or(usize::MAX);
if item_tokens > available_tokens {
let adjusted_budget = text_budget
.saturating_sub(item_tokens - available_tokens)
.saturating_sub(1);
let Some(adjusted) =
truncate_message_text_to_token_budget(truncated_item, adjusted_budget)
else {
continue;
};
if usize::try_from(estimate_item_token_count(&adjusted.item))
.unwrap_or(usize::MAX)
> available_tokens
{
continue;
}
truncated_item = adjusted;
}
}
if let Some(notice) = group.attached_notice {
truncated_reversed.push(notice);
}
@@ -655,7 +719,9 @@ mod tests {
output: ResponseItem,
) -> (Vec<ResponseItemEnvelope>, usize) {
let metadata = vec![None; input.len()];
build_v2_compacted_history(input, metadata, output)
build_v2_compacted_history(
input, metadata, output, /*retain_client_developer_messages*/ false,
)
}
fn annotated(items: Vec<ResponseItem>) -> Vec<ResponseItemEnvelope> {
@@ -730,28 +796,63 @@ mod tests {
#[test]
fn build_v2_compacted_history_preserves_retained_metadata_sidecar() {
let retained = message("user", "keep", /*phase*/ None);
let filtered = message("developer", "drop", /*phase*/ None);
let generated_notice = message(
"developer",
"<image_resize_notice>generated</image_resize_notice>",
/*phase*/ None,
);
let harness = message("developer", "drop", /*phase*/ None);
let client = message(
"developer",
"<image_resize_notice>client</image_resize_notice>",
/*phase*/ None,
);
let output = ResponseItem::Compaction {
id: None,
encrypted_content: "new".to_string(),
internal_chat_message_metadata_passthrough: None,
};
let (history, _) = build_v2_compacted_history(
vec![filtered, retained.clone()],
vec![None, Some(CodexHarnessMetadata::default())],
output.clone(),
);
assert_eq!(
history,
vec![
for enabled in [false, true] {
let (history, _) = build_v2_compacted_history(
vec![
harness.clone(),
client.clone(),
retained.clone(),
generated_notice.clone(),
],
vec![
None,
Some(CodexHarnessMetadata {
client_authored: true,
}),
Some(CodexHarnessMetadata::default()),
None,
],
output.clone(),
enabled,
);
let mut expected = vec![
ResponseItemEnvelope {
item: retained,
item: retained.clone(),
metadata: Some(CodexHarnessMetadata::default()),
},
ResponseItemEnvelope::new(output),
]
);
ResponseItemEnvelope::new(generated_notice.clone()),
ResponseItemEnvelope::new(output.clone()),
];
if enabled {
expected.insert(
0,
ResponseItemEnvelope {
item: client.clone(),
metadata: Some(CodexHarnessMetadata {
client_authored: true,
}),
},
);
}
assert_eq!(history, expected);
}
}
#[test]
codex-rs/core/src/session/mod.rs
+18
View File
@@ -3695,6 +3695,23 @@ impl Session {
world_state: Arc<WorldState>,
) -> u64 {
let turn_context = step_context.turn.as_ref();
let retained_client_developer_messages =
if self.enabled(Feature::RetainClientDeveloperMessages) {
let history = self.clone_history().await;
crate::compact_remote_v2::truncate_retained_messages_for_remote_compaction(
history
.annotated_items()
.iter()
.filter(|item| {
crate::compact_remote_v2::is_client_authored_developer_message(item)
})
.cloned()
.collect(),
crate::compact_remote_v2::RETAINED_MESSAGE_TOKEN_BUDGET,
)
} else {
Vec::new()
};
let window = {
let mut state = self.state.lock().await;
state.start_new_context_window()
@@ -3705,6 +3722,7 @@ impl Session {
.await
.into_iter()
.map(ResponseItemEnvelope::new)
.chain(retained_client_developer_messages)
.collect();
let turn_context_item = turn_context.to_turn_context_item();
self.replace_compacted_history(
codex-rs/core/tests/suite/compact_remote.rs
+97
View File
@@ -69,6 +69,7 @@ use core_test_support::wait_for_event_with_timeout;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
use test_case::test_case;
use tokio::time::Duration;
use wiremock::ResponseTemplate;
@@ -571,6 +572,102 @@ async fn remote_compact_v2_retains_metadata_from_resumed_history() -> Result<()>
Ok(())
}
#[test_case(false; "feature_disabled")]
#[test_case(true; "feature_enabled")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn remote_compact_v2_retains_only_client_developer_messages_when_enabled(
enabled: bool,
) -> Result<()> {
skip_if_no_network!(Ok(()));
let harness = TestCodexHarness::with_auto_env_builder(
test_codex()
.with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
.with_config(move |config| {
config
.features
.enable(Feature::RemoteCompactionV2)
.expect("remote compaction v2 should be configurable");
if enabled {
config
.features
.enable(Feature::RetainClientDeveloperMessages)
.expect("client developer retention should be configurable");
}
}),
)
.await?;
let developer = |text: &str| ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: text.to_string(),
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
};
let codex = &harness.test().codex;
let rollout_path = codex.rollout_path().context("rollout path")?;
codex
.inject_response_items(vec![developer("INJECTED_CLIENT_DEVELOPER")])
.await?;
let response_mock = responses::mount_sse_sequence(
harness.server(),
vec![
sse(vec![responses::ev_completed("before-compact")]),
sse(vec![
json!({
"type": "response.output_item.done",
"item": {
"type": "compaction",
"encrypted_content": "CLIENT_RETENTION_SUMMARY",
},
}),
responses::ev_completed("compact"),
]),
sse(vec![responses::ev_completed("after-compact")]),
],
)
.await;
harness.test().submit_turn("before compact").await?;
codex.submit(Op::Compact).await?;
wait_for_turn_complete(codex).await;
harness.test().submit_turn("after compact").await?;
let requests = response_mock.requests();
let follow_up = requests.last().context("follow-up request")?;
assert_eq!(
follow_up.body_contains_text("INJECTED_CLIENT_DEVELOPER"),
enabled
);
requests
.iter()
.for_each(assert_compact_request_omits_harness_metadata);
codex.shutdown_and_wait().await?;
let replacement_history = fs::read_to_string(&rollout_path)?
.lines()
.filter_map(|line| serde_json::from_str::<RolloutLine>(line).ok())
.filter_map(|line| match line.item {
RolloutItem::Compacted(compacted) => compacted.replacement_history,
_ => None,
})
.next_back()
.context("remote compaction should persist a checkpoint")?;
let retained_client_developers = replacement_history
.iter()
.filter(|item| {
item.metadata
.as_ref()
.is_some_and(|metadata| metadata.client_authored)
})
.count();
assert_eq!(retained_client_developers, usize::from(enabled));
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn amazon_bedrock_automatic_compaction_uses_v1_endpoint_when_v2_is_enabled() -> Result<()> {
skip_if_no_network!(Ok(()));
codex-rs/core/tests/suite/token_budget.rs
+51 -2
View File
@@ -56,6 +56,7 @@ use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use test_case::test_case;
const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000;
const AUTO_COMPACT_FALLBACK_PROMPT: &str = "Save the important state before rollover.";
@@ -951,8 +952,12 @@ async fn get_context_remaining_returns_unknown_when_threshold_is_unbounded() ->
Ok(())
}
#[test_case(false; "token_budget_only")]
#[test_case(true; "with_client_developer_retention")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
async fn token_budget_context_uses_new_window_after_compaction(
retain_client_developer_messages: bool,
) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
@@ -982,10 +987,25 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
if retain_client_developer_messages {
config
.features
.enable(Feature::RetainClientDeveloperMessages)
.expect("test config should allow client developer retention");
}
})
.build(&server)
.await?;
if retain_client_developer_messages {
test.codex
.inject_response_items(vec![serde_json::from_value(json!({
"type": "message",
"role": "developer",
"content": [{"type": "input_text", "text": "CLIENT_DEVELOPER_INSTRUCTIONS"}]
}))?])
.await?;
}
test.submit_turn("before compact").await?;
test.codex.submit(Op::Compact).await?;
assert_context_compaction_item_lifecycle(&test.codex).await;
@@ -1025,6 +1045,11 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
!requests[1].body_contains_text("assistant before compact"),
"token budget compaction should drop prior assistant messages"
);
assert_eq!(
requests[1].body_contains_text("CLIENT_DEVELOPER_INSTRUCTIONS"),
retain_client_developer_messages,
"token budget compaction should retain client-authored developer messages when enabled"
);
assert!(
requests[1].body_contains_text("after compact"),
"follow-up should still include the new turn input"
@@ -1082,8 +1107,12 @@ async fn token_budget_compaction_runs_compact_hooks() -> Result<()> {
Ok(())
}
#[test_case(false; "token_budget_only")]
#[test_case(true; "with_client_developer_retention")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up() -> Result<()> {
async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up(
retain_client_developer_messages: bool,
) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
@@ -1129,10 +1158,25 @@ async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up()
.features
.enable(Feature::TokenBudget)
.expect("test config should allow token budget");
if retain_client_developer_messages {
config
.features
.enable(Feature::RetainClientDeveloperMessages)
.expect("test config should allow client developer retention");
}
})
.build(&server)
.await?;
if retain_client_developer_messages {
test.codex
.inject_response_items(vec![serde_json::from_value(json!({
"type": "message",
"role": "developer",
"content": [{"type": "input_text", "text": "MID_TURN_CLIENT_INSTRUCTIONS"}]
}))?])
.await?;
}
test.submit_turn("trigger mid-turn auto compaction").await?;
let requests = responses.requests();
@@ -1179,6 +1223,11 @@ async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up()
!requests[1].body_contains_text("trigger mid-turn auto compaction"),
"fresh token-budget windows should drop prior user messages"
);
assert_eq!(
requests[1].body_contains_text("MID_TURN_CLIENT_INSTRUCTIONS"),
retain_client_developer_messages,
"mid-turn token-budget compaction should retain client-authored developer messages when enabled"
);
assert_ne!(
follow_up_window_id, initial_window_id,
"mid-turn token-budget auto-compaction should reset the context window"