Add per-item retention policies for injected history (#61344)

## What changed

Allow `thread/inject_items` to accept wrapped items with `retention: "default"` or `retention: "always"`, while preserving support for raw response items. Persist explicit policies separately from legacy client-authored metadata.

Carry explicitly retained items through local and remote compaction, context-window changes, and child-agent creation, subject to retention budgets. Preserve history order and keep retained tool exchanges together; require matching calls and outputs to be adjacent in the injection request and both marked `always`.

Update Quinn dependencies and record exceptions for additional Hickory DNSSEC advisories where DNSSEC is disabled.

## Testing

Update injection tests for explicit retention metadata and extend the compaction scenario to cover a retained tool exchange across repeated compaction and restart.

GitOrigin-RevId: 1195a4ed62dc162a62e58bc2ad733262fbabea27
This commit is contained in:
Vivian Fang
2026-11-13 18:46:13 +00:00
committed by copyberry
parent d17887a8d1
commit c33002259e
29 changed files with 578 additions and 300 deletions
+52 -4
View File
@@ -1,6 +1,8 @@
//! Loads stored agent context and prepares the history a child inherits when it forks.
use crate::agent::control::LocalAgentRuntime;
use crate::agent::types::SpawnAgentOptions;
use crate::compact_remote_v2::retained_client_items;
use crate::config::Config;
use crate::context::BaseInstructionsFragment;
use crate::context::ContextualUserFragment;
@@ -22,6 +24,7 @@ use crate::thread_manager::ThreadManagerState;
use crate::thread_manager::fork_history_from_snapshot;
use codex_context_fragments::set_annotated_content;
use codex_context_fragments::to_annotated_content;
use codex_history::CodexHarnessMetadata;
use codex_history::InitialHistory;
use codex_history::ResponseItemEnvelope;
use codex_history::RolloutItem;
@@ -39,8 +42,51 @@ use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::ThreadHistoryMode;
use codex_thread_store::LoadThreadHistoryParams;
impl LocalAgentRuntime {
pub(crate) async fn retained_items_for_fresh_child(
&self,
initial_history: &InitialHistory,
session_source: &SessionSource,
) -> Vec<ResponseItemEnvelope> {
// Full-history forks already carry their retained items.
let (InitialHistory::New, Some(parent_thread_id)) =
(initial_history, session_source.parent_thread_id())
else {
return Vec::new();
};
let Ok(manager) = self.upgrade() else {
return Vec::new();
};
let Ok(parent) = manager.get_thread(parent_thread_id).await else {
return Vec::new();
};
let history = parent.session.clone_history().await;
let mut items: Vec<_> = retained_client_items(
history.annotated_items(),
&history.file_audio_durations,
ResponseItemEnvelope::is_always_retained,
)
.into_values()
.collect();
for item in &mut items {
let parent_metadata = item.metadata.take().unwrap_or_default();
// Keep item metadata needed for retention and replay. Checkpoints belong to the child.
item.metadata = Some(CodexHarnessMetadata {
retain_across_windows: parent_metadata.retain_across_windows,
retained_source: parent_metadata.retained_source,
history_truncation_token_limit: parent_metadata.history_truncation_token_limit,
image_replay_sources: parent_metadata.image_replay_sources,
inherited_user_message: true,
..Default::default()
});
}
items
}
}
fn keep_forked_rollout_item(item: &RolloutItem, preserve_context_baselines: bool) -> bool {
match item {
RolloutItem::ResponseItem(envelope) if envelope.is_always_retained() => true,
RolloutItem::ResponseItem(envelope) => match &envelope.item {
ResponseItem::Message { role, phase, .. } => match role.as_str() {
"system" | "developer" | "user" => true,
@@ -256,9 +302,11 @@ pub(super) async fn prepare_forked_history(
forked_rollout_items = items;
}
let retain_forked_item = |envelope: &mut ResponseItemEnvelope, replaced: &mut bool| {
if multi_agent_version == MultiAgentVersion::V2
&& (preserve_fork_prefix
|| matches!(&envelope.item, ResponseItem::Message { role, .. } if role == "user" || role == "assistant"))
let explicitly_retained = envelope.is_always_retained();
if explicitly_retained
|| (multi_agent_version == MultiAgentVersion::V2
&& (preserve_fork_prefix
|| matches!(&envelope.item, ResponseItem::Message { role, .. } if role == "user" || role == "assistant")))
{
// Resume must not recapture copied parent messages as local authorization.
envelope
@@ -273,7 +321,7 @@ pub(super) async fn prepare_forked_history(
// Assistant and tool positions belong to the parent counter, not the child.
metadata.user_input_order = None;
}
if preserve_fork_prefix {
if preserve_fork_prefix || explicitly_retained {
return true;
}
let response_item = &mut envelope.item;
+39 -23
View File
@@ -62,6 +62,8 @@ use codex_protocol::protocol::TurnEnvironmentRequests;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_protocol::protocol::TurnEnvironmentSelections;
use codex_protocol::protocol::W3cTraceContext;
use codex_protocol::turn_input::InjectedResponseItem;
use codex_protocol::turn_input::ItemRetention;
use codex_protocol::turn_input::RecoverTurnRequest;
use codex_protocol::turn_input::StartIfIdleSubmission;
use codex_protocol::turn_input::SteerSubmission;
@@ -1032,7 +1034,15 @@ impl CodexThread {
/// Record raw Responses API items without starting a new turn.
pub async fn inject_response_items(&self, items: Vec<ResponseItem>) -> CodexResult<()> {
self.inject_response_items_for_turn(items).await?;
self.inject_response_items_with_retention(items).await
}
/// Inject Responses items with explicit, persisted retention policies.
pub async fn inject_response_items_with_retention(
&self,
items: Vec<impl Into<InjectedResponseItem>>,
) -> CodexResult<()> {
self.record_injected_items(items).await?;
self.checkpoint_preparation().await?;
Ok(())
}
@@ -1058,37 +1068,43 @@ impl CodexThread {
pub async fn inject_response_items_for_turn(
&self,
items: Vec<ResponseItem>,
) -> CodexResult<()> {
self.record_injected_items(items).await
}
async fn record_injected_items(
&self,
items: Vec<impl Into<InjectedResponseItem>>,
) -> CodexResult<()> {
if items.is_empty() {
return Err(CodexErr::InvalidRequest(
"items must not be empty".to_string(),
));
}
items.iter().try_for_each(validate_injected_response_item)?;
let had_reference_context = self.session.reference_context_item().await.is_some();
let mut turn_context = if had_reference_context {
self.session.new_inject_items_context().await
} else {
self.session.new_default_turn().await
};
if self.session.reference_context_item().await.is_none() {
// Compaction can clear the reference while the recording context is built.
// Initial context must capture a step with a complete skills snapshot.
if had_reference_context {
turn_context = self.session.new_default_turn().await;
let items: Vec<InjectedResponseItem> = items.into_iter().map(Into::into).collect();
let mut indexed_items = items.iter().enumerate();
while let Some((index, item)) = indexed_items.next() {
validate_injected_response_item(&item.item)?;
if item.retention != ItemRetention::Always {
continue;
}
let needs_call = crate::tools::output_call_id(&item.item).is_some()
&& !matches!(&item.item, ResponseItem::ToolSearchOutput { execution, .. } if execution == "server");
if (crate::tools::input_call_id(&item.item).is_some() || needs_call)
&& !indexed_items.next().is_some_and(|(_, output)| {
output.retention == ItemRetention::Always
&& crate::tools::is_matching_tool_output(&item.item, &output.item)
})
{
return Err(CodexErr::InvalidRequest(format!(
"items[{index}]: retained tool calls and outputs must be adjacent in the same request, call first, with both marked always"
)));
}
// This history-only API runs without run_turn, so it owns its initial step.
let step_context = self
.session
.capture_step_context(Arc::clone(&turn_context), &CancellationToken::new())
.await?;
self.session
.record_context_updates_and_set_reference_context_item(step_context.as_ref())
.await?;
}
let turn_context = self.session.history_injection_context().await?;
self.session
.inject_client_response_items(items, turn_context.as_ref())
.inject_client_response_items(items, &turn_context)
.await;
Ok(())
}
+38 -26
View File
@@ -7,12 +7,14 @@ use crate::Prompt;
use crate::client::ModelClientSession;
use crate::client_common::ResponseEvent;
use crate::compact_pending_input::is_retained_pending_mail;
use crate::compact_remote_v2::retained_client_items;
use crate::context::BaseInstructionsFragment;
use crate::context::CompactionSummary;
use crate::context::ContextualUserFragment;
use crate::context::world_state::WorldState;
use crate::context::world_state::WorldStateSnapshot;
use crate::context::world_state::split_prefix_updates;
use crate::context_manager::ContextManager;
use crate::context_manager::updates::merge_world_state_updates;
use crate::hook_runtime::PostCompactHookOutcome;
use crate::hook_runtime::PreCompactHookOutcome;
@@ -20,6 +22,7 @@ use crate::hook_runtime::run_post_compact_hooks;
use crate::hook_runtime::run_pre_compact_hooks;
use crate::responses_metadata::CodexResponsesMetadata;
use crate::responses_metadata::CompactionTurnMetadata;
use crate::session::FileAudioDurations;
use crate::session::RequestEffortUsage;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
@@ -379,7 +382,6 @@ async fn run_compact_task_inner_impl(
};
let history_snapshot = sess.clone_history().await;
let history_items = history_snapshot.annotated_items();
let summary_suffix = if matches!(compaction_metadata.phase(), CompactionPhase::PostTurn) {
get_last_assistant_message_from_turn(compaction_response.output.iter())
.filter(|summary| !summary.trim().is_empty())
@@ -392,9 +394,8 @@ async fn run_compact_task_inner_impl(
get_last_assistant_message_from_turn(history_snapshot.raw_items()).unwrap_or_default()
};
let summary_text = format!("{SUMMARY_PREFIX}\n{summary_suffix}");
let user_messages = collect_annotated_user_messages(history_items);
let mut new_history = build_compacted_history(Vec::new(), &user_messages, &summary_text);
let mut new_history = build_compacted_history(&history_snapshot, &summary_text);
if let Some(summary_item) = new_history.last_mut() {
// This replacement history skips `record_conversation_items`; only the appended summary
// belongs to this compaction turn.
@@ -407,6 +408,9 @@ async fn run_compact_task_inner_impl(
let summary = new_history
.pop()
.ok_or_else(|| CodexErr::Fatal("local compaction summary is missing".to_string()))?;
// A client may also have explicitly retained this mail. Use the final
// compactor input for pending mail so each delivery appears once.
new_history.retain(|item| !is_retained_pending_mail(item, first));
new_history.extend(
history
.annotated_items()
@@ -606,13 +610,13 @@ pub(crate) fn collect_user_messages(items: &[ResponseItem]) -> Vec<CompactedUser
.collect()
}
pub(crate) fn collect_annotated_user_messages(
fn annotated_user_messages(
items: &[ResponseItemEnvelope],
) -> Vec<CompactedUserMessage<'_>> {
items
.iter()
.filter_map(|envelope| compacted_user_message(&envelope.item, envelope.metadata.as_ref()))
.collect()
) -> impl DoubleEndedIterator<Item = (usize, CompactedUserMessage<'_>)> {
items.iter().enumerate().filter_map(|(position, envelope)| {
let message = compacted_user_message(&envelope.item, envelope.metadata.as_ref())?;
Some((position, message))
})
}
fn compacted_user_message<'a>(
@@ -714,28 +718,33 @@ pub(crate) fn insert_initial_context_before_last_real_user_or_summary(
}
pub(crate) fn build_compacted_history(
initial_context: Vec<ResponseItemEnvelope>,
user_messages: &[CompactedUserMessage<'_>],
history: &ContextManager,
summary_text: &str,
) -> Vec<ResponseItemEnvelope> {
build_compacted_history_with_limit(
initial_context,
user_messages,
history.annotated_items(),
summary_text,
COMPACT_USER_MESSAGE_MAX_TOKENS,
&history.file_audio_durations,
)
}
fn build_compacted_history_with_limit(
mut history: Vec<ResponseItemEnvelope>,
user_messages: &[CompactedUserMessage<'_>],
items: &[ResponseItemEnvelope],
summary_text: &str,
max_tokens: usize,
audio_durations: &FileAudioDurations,
) -> Vec<ResponseItemEnvelope> {
let mut selected_messages = Vec::new();
// An injected user message can qualify under both budgets. Its position
// identifies the occurrence even if another message has the same client ID.
let mut selected = retained_client_items(
items,
audio_durations,
ResponseItemEnvelope::is_always_retained,
);
if max_tokens > 0 {
let mut remaining = max_tokens;
for message in user_messages.iter().rev() {
for (position, message) in annotated_user_messages(items).rev() {
if remaining == 0 {
break;
}
@@ -749,6 +758,10 @@ fn build_compacted_history_with_limit(
else {
continue;
};
if selected.contains_key(&position) {
remaining = remaining.saturating_sub(tokens);
continue;
}
let mut harness_metadata = message.harness_metadata.cloned();
let preserve_content = tokens <= remaining
&& content
@@ -806,18 +819,17 @@ fn build_compacted_history_with_limit(
{
metadata.mark_retained_sources_incomplete();
}
selected_messages.push(ResponseItemEnvelope {
item,
metadata: harness_metadata,
});
if tokens > remaining {
break;
}
selected.insert(
position,
ResponseItemEnvelope {
item,
metadata: harness_metadata,
},
);
remaining = remaining.saturating_sub(tokens);
}
selected_messages.reverse();
}
history.extend(selected_messages);
let mut history: Vec<_> = selected.into_values().collect();
let summary_text = if summary_text.is_empty() {
"(no summary available)".to_string()
+78 -59
View File
@@ -1,13 +1,13 @@
use std::borrow::Borrow;
use std::collections::BTreeMap;
use crate::compact_retained_tool_outputs::is_retained_tool_output;
use crate::compact_remote_v2::should_retain_client_item;
use crate::context::ContextualUserFragment;
use crate::context::ImageResizeNotice;
use crate::context_manager::ContextManager;
use crate::context_manager::estimate_item_token_count_with_audio;
use crate::session::FileAudioDurations;
use crate::session::turn_context::TurnContext;
use codex_features::Feature;
use crate::tools::is_matching_tool_output;
use codex_history::ResponseItemEnvelope;
use codex_protocol::models::BaseInstructions;
use codex_protocol::models::ContentItem;
@@ -19,51 +19,85 @@ use codex_utils_output_truncation::approx_token_count;
const CONTEXT_WINDOW_TRUNCATED_OUTPUT_MESSAGE: &str =
"Output exceeded the available model context and was truncated";
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct HistoryItemGroup<T> {
pub(crate) source: T,
pub(crate) attached_notice: Option<T>,
pub(crate) struct HistoryItemGroup<'a> {
pub(crate) source_position: usize,
pub(crate) source: &'a ResponseItemEnvelope,
pub(crate) attached_notice: Option<&'a ResponseItemEnvelope>,
pub(crate) tool_call: Option<&'a ResponseItemEnvelope>,
}
impl<T: Borrow<ResponseItem>> HistoryItemGroup<T> {
pub(crate) fn into_items(self) -> impl Iterator<Item = T> {
std::iter::once(self.source).chain(self.attached_notice)
impl HistoryItemGroup<'_> {
pub(crate) fn keep(
&self,
selected: &mut BTreeMap<usize, ResponseItemEnvelope>,
source: ResponseItemEnvelope,
) {
if let Some(call) = self.tool_call {
selected.insert(self.source_position - 1, call.clone());
}
selected.insert(self.source_position, source);
if let Some(notice) = self.attached_notice {
selected.insert(self.source_position + 1, notice.clone());
}
}
pub(crate) fn estimated_token_count(&self, audio_durations: &FileAudioDurations) -> i128 {
let source_tokens = i128::from(estimate_item_token_count_with_audio(
self.source.borrow(),
audio_durations,
));
let notice_tokens = self.attached_notice.as_ref().map_or(0, |notice| {
i128::from(estimate_item_token_count_with_audio(
notice.borrow(),
audio_durations,
))
});
source_tokens.saturating_add(notice_tokens)
self.tool_call
.into_iter()
.chain(std::iter::once(self.source))
.chain(self.attached_notice)
.map(|item| i128::from(estimate_item_token_count_with_audio(item, audio_durations)))
.fold(0, i128::saturating_add)
}
}
pub(crate) fn history_item_groups<I>(items: I) -> impl Iterator<Item = HistoryItemGroup<I::Item>>
where
I: IntoIterator,
I::Item: Borrow<ResponseItem>,
{
let mut items = items.into_iter().peekable();
pub(crate) fn newest_history_item_groups(
mut items: &[ResponseItemEnvelope],
) -> impl Iterator<Item = HistoryItemGroup<'_>> {
std::iter::from_fn(move || {
let source = items.next()?;
let attached_notice = items.next_if(|notice| is_attached_notice(notice.borrow()));
let (last, preceding) = items.split_last()?;
let (source, attached_notice, rest) = match preceding.split_last() {
Some((source, rest)) if is_attached_notice(source, last) => (source, Some(last), rest),
_ => (last, None, preceding),
};
let source_position = rest.len();
// Injection requires explicitly retained exchanges to be adjacent in the request.
let (tool_call, rest) = match rest.split_last() {
Some((call, rest))
if source.is_always_retained()
&& call.is_always_retained()
&& is_matching_tool_output(&call.item, &source.item) =>
{
(Some(call), rest)
}
_ => (None, rest),
};
items = rest;
Some(HistoryItemGroup {
source_position,
source,
attached_notice,
tool_call,
})
})
}
fn is_attached_notice(notice: &ResponseItem) -> bool {
fn is_attached_notice(source: &ResponseItemEnvelope, notice: &ResponseItemEnvelope) -> bool {
let can_have_resize_notice = match &source.item {
ResponseItem::Message { role, .. } => role == "user",
ResponseItem::FunctionCallOutput { .. } | ResponseItem::CustomToolCallOutput { .. } => true,
_ => false,
};
// A client may send the same text as a resize notice with its own policy.
if !can_have_resize_notice
|| notice.metadata.as_ref().is_some_and(|metadata| {
metadata.client_authored || metadata.retain_across_windows.is_some()
})
{
return false;
}
matches!(
notice,
&notice.item,
ResponseItem::Message { role, content, .. }
if role == "developer"
&& matches!(
@@ -89,38 +123,24 @@ pub(crate) fn trim_function_call_history_to_fit_context_window(
i128::try_from(approx_token_count(&base_instructions.text)).unwrap_or(i128::MAX);
let audio_durations = &history.file_audio_durations;
let original_items = history.annotated_items();
let mut estimated_tokens = history_item_groups(original_items.iter().map(|item| &item.item))
.map(|group| group.estimated_token_count(audio_durations))
let mut estimated_tokens = original_items
.iter()
.map(|item| i128::from(estimate_item_token_count_with_audio(item, audio_durations)))
.fold(base_tokens, i128::saturating_add);
let initial_estimated_tokens = i64::try_from(estimated_tokens).unwrap_or(i64::MAX);
if initial_estimated_tokens <= context_window {
return (0, 0);
}
let mut rewritten_items = Vec::new();
let mut consumed_items: usize = 0;
for group in history_item_groups(original_items.iter().map(|item| &item.item))
.collect::<Vec<_>>()
.into_iter()
.rev()
{
let mut retained_len = original_items.len();
for group in newest_history_item_groups(original_items) {
if i64::try_from(estimated_tokens).unwrap_or(i64::MAX) <= context_window {
break;
}
let group_item_count = 1 + usize::from(group.attached_notice.is_some());
let source_index = original_items
.len()
.saturating_sub(consumed_items.saturating_add(group_item_count));
let Some(rewritten_item) = original_items.get(source_index).and_then(|envelope| {
if turn_context
.config
.features
.enabled(Feature::RetainClientToolOutputs)
&& is_retained_tool_output(envelope)
{
// Enabled retained outputs form the same boundary as user/developer messages.
None
} else {
rewritten_output_for_context_window(envelope)
}
}) else {
if should_retain_client_item(group.source, &turn_context.config.features) {
break;
}
let Some(rewritten_item) = rewritten_output_for_context_window(group.source) else {
break;
};
estimated_tokens = estimated_tokens
@@ -129,13 +149,12 @@ pub(crate) fn trim_function_call_history_to_fit_context_window(
&rewritten_item.item,
audio_durations,
)));
consumed_items += group_item_count;
retained_len = group.source_position;
rewritten_items.push(rewritten_item);
}
let rewritten_outputs = rewritten_items.len();
if rewritten_outputs > 0 {
let retained_len = original_items.len() - consumed_items;
let mut items = original_items[..retained_len].to_vec();
items.extend(rewritten_items.into_iter().rev());
history.replace_annotated(items);
+80 -63
View File
@@ -1,4 +1,5 @@
use codex_login::AuthOwner;
use std::collections::BTreeMap;
use std::sync::Arc;
use std::sync::OnceLock;
@@ -18,8 +19,7 @@ use crate::compact_model_fallback::record_model_fallback;
use crate::compact_model_fallback::should_retry_with_current_model;
use crate::compact_pending_input::MAX_RETAINED_AGENT_MESSAGE_TOKENS;
use crate::compact_pending_input::is_retained_pending_mail;
use crate::compact_remote_history::HistoryItemGroup;
use crate::compact_remote_history::history_item_groups;
use crate::compact_remote_history::newest_history_item_groups;
use crate::context::world_state::WorldState;
use crate::context_manager::estimate_item_token_count;
use crate::context_manager::estimate_item_token_count_with_audio;
@@ -510,20 +510,18 @@ fn build_v2_compacted_history(
&prompt_input,
);
let pending_sequence = crate::compact_pending_input::pending_mail_sequence(&prompt_input);
let retained = v2_history_item_groups(prompt_input)
.filter(|group| {
is_retained_for_remote_compaction_v2(&group.source, features)
|| pending_sequence
.is_some_and(|first| is_retained_pending_mail(&group.source, first))
})
.flat_map(HistoryItemGroup::into_items)
.collect::<Vec<_>>();
let retained = truncate_retained_messages(
retained,
let retained = select_retained_items(
&prompt_input,
|item| {
is_retained_for_remote_compaction_v2(item, features)
|| pending_sequence.is_some_and(|first| is_retained_pending_mail(item, first))
},
RETAINED_MESSAGE_TOKEN_BUDGET,
image_budget,
audio_durations,
);
)
.into_values()
.collect::<Vec<_>>();
let retained_image_count = retained
.iter()
.map(|envelope| retained_input_image_count(&envelope.item))
@@ -540,6 +538,20 @@ fn build_v2_compacted_history(
(retained, retained_image_count)
}
pub(crate) fn retained_client_items(
items: &[ResponseItemEnvelope],
audio_durations: &FileAudioDurations,
should_retain: impl Fn(&ResponseItemEnvelope) -> bool,
) -> BTreeMap<usize, ResponseItemEnvelope> {
select_retained_items(
items,
should_retain,
RETAINED_MESSAGE_TOKEN_BUDGET,
RetainedImageBudget::Disabled,
audio_durations,
)
}
pub(crate) fn is_client_authored_developer_message(item: &ResponseItemEnvelope) -> bool {
item.metadata
.as_ref()
@@ -547,26 +559,20 @@ pub(crate) fn is_client_authored_developer_message(item: &ResponseItemEnvelope)
&& 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)
})
pub(crate) fn should_retain_client_item(item: &ResponseItemEnvelope, features: &Features) -> bool {
if let Some(retain) = item.metadata.as_ref().and_then(|m| m.retain_across_windows) {
return retain;
}
(features.enabled(Feature::RetainClientDeveloperMessages)
&& is_client_authored_developer_message(item))
|| (features.enabled(Feature::RetainClientToolOutputs) && is_retained_tool_output(item))
}
fn is_retained_for_remote_compaction_v2(
envelope: &ResponseItemEnvelope,
features: &Features,
) -> bool {
if features.enabled(Feature::RetainClientToolOutputs) && is_retained_tool_output(envelope) {
if should_retain_client_item(envelope, features) {
return true;
}
let item = &envelope.item;
@@ -605,10 +611,6 @@ fn is_retained_for_remote_compaction_v2(
crate::event_mapping::parse_turn_item(item),
Some(TurnItem::UserMessage(_) | TurnItem::HookPrompt(_))
),
"developer" => {
features.enabled(Feature::RetainClientDeveloperMessages)
&& is_client_authored_developer_message(envelope)
}
_ => false,
}
}
@@ -651,19 +653,32 @@ fn truncate_retained_messages(
image_budget: RetainedImageBudget,
audio_durations: &FileAudioDurations,
) -> Vec<ResponseItemEnvelope> {
select_retained_items(&items, |_| true, max_tokens, image_budget, audio_durations)
.into_values()
.collect()
}
/// Select newest first, then return the items at their original positions.
fn select_retained_items(
items: &[ResponseItemEnvelope],
should_retain: impl Fn(&ResponseItemEnvelope) -> bool,
max_tokens: usize,
image_budget: RetainedImageBudget,
audio_durations: &FileAudioDurations,
) -> BTreeMap<usize, ResponseItemEnvelope> {
let mut remaining = max_tokens;
let mut truncated_reversed = Vec::with_capacity(items.len());
for group in v2_history_item_groups(items)
.collect::<Vec<_>>()
.into_iter()
.rev()
{
let mut selected = BTreeMap::new();
for group in newest_history_item_groups(items) {
if remaining == 0 {
break;
}
if !should_retain(group.source) {
continue;
}
let client_developer = is_client_authored_developer_message(&group.source);
let image_budget = if client_developer {
let charge_full_item =
group.source.is_always_retained() || is_client_authored_developer_message(group.source);
let image_budget = if charge_full_item {
RetainedImageBudget::Disabled
} else {
image_budget
@@ -677,7 +692,7 @@ fn truncate_retained_messages(
// Charged media is truncated newest-first as a message boundary. Each
// modality still defines its own grouping and token cost in `media`.
let requires_media_boundary_truncation = has_budgeted_image || has_file_audio;
let notice_tokens = group.attached_notice.as_ref().map_or(0, |notice| {
let notice_tokens = group.attached_notice.map_or(0, |notice| {
retained_message_token_count(
&notice.item,
RetainedImageBudget::Disabled,
@@ -685,11 +700,10 @@ fn truncate_retained_messages(
)
.max(1)
});
// Client-authored developer messages already charge non-text content via
// the serialized estimate. Preserve their text-only boundary correction.
// Explicitly retained items and client developer messages also pay for their media.
let content_tokens =
retained_message_token_count(&group.source.item, image_budget, audio_durations);
let source_tokens = if client_developer {
let source_tokens = if charge_full_item {
usize::try_from(estimate_item_token_count_with_audio(
&group.source.item,
audio_durations,
@@ -698,19 +712,23 @@ fn truncate_retained_messages(
} else {
content_tokens.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);
let token_count = if group.tool_call.is_some() {
usize::try_from(group.estimated_token_count(audio_durations)).unwrap_or(usize::MAX)
} else {
source_tokens.saturating_add(notice_tokens)
};
let source = if token_count <= remaining {
remaining = remaining.saturating_sub(token_count);
} else if is_retained_tool_output(&group.source) {
// Keep outputs whole and stop before older items to preserve newest-first ordering.
group.source.clone()
} else if is_retained_tool_output(group.source)
|| (group.source.is_always_retained()
&& !matches!(group.source.item, ResponseItem::Message { .. }))
{
// Structured items such as tool results must stay whole.
break;
} else if remaining > notice_tokens {
let available_tokens = remaining - notice_tokens;
let content_budget = if client_developer {
let content_budget = if charge_full_item {
available_tokens.saturating_sub(source_tokens.saturating_sub(content_tokens))
} else {
available_tokens
@@ -722,18 +740,18 @@ fn truncate_retained_messages(
}
let truncated_item = if requires_media_boundary_truncation {
media::truncate_message_to_token_budget(
group.source,
group.source.clone(),
content_budget,
image_budget,
audio_durations,
)
} else {
truncate_message_text_to_token_budget(group.source, content_budget)
truncate_message_text_to_token_budget(group.source.clone(), content_budget)
};
let Some(mut truncated_item) = truncated_item else {
continue;
};
if client_developer {
if charge_full_item {
let item_tokens = usize::try_from(estimate_item_token_count_with_audio(
&truncated_item.item,
audio_durations,
@@ -768,17 +786,16 @@ fn truncate_retained_messages(
truncated_item = adjusted;
}
}
if let Some(notice) = group.attached_notice {
truncated_reversed.push(notice);
}
truncated_reversed.push(truncated_item);
remaining = 0;
truncated_item
} else if requires_media_boundary_truncation {
remaining = 0;
}
break;
} else {
continue;
};
group.keep(&mut selected, source);
}
truncated_reversed.reverse();
truncated_reversed
selected
}
fn retained_message_token_count(
@@ -1,5 +1,4 @@
//! Retains explicitly marked, named function outputs that have no matching call.
//! Outputs share the normal history token budget and are retained whole or dropped.
//! Identifies standalone retained tool outputs in model history.
use codex_history::ResponseItemEnvelope;
use codex_protocol::models::ResponseItem;
@@ -8,7 +7,7 @@ pub(crate) fn is_retained_tool_output(envelope: &ResponseItemEnvelope) -> bool {
envelope
.metadata
.as_ref()
.is_some_and(|metadata| metadata.client_authored)
.is_some_and(|m| m.retain_across_windows.unwrap_or(m.client_authored))
&& matches!(
envelope.item,
ResponseItem::FunctionCallOutput {
+18 -13
View File
@@ -272,6 +272,15 @@ fn compacted_user_message<'a>(text: &str, original: &'a ResponseItem) -> Compact
}
}
fn compact_items(items: &[ResponseItemEnvelope], summary: &str) -> Vec<ResponseItemEnvelope> {
build_compacted_history_with_limit(
items,
summary,
COMPACT_USER_MESSAGE_MAX_TOKENS,
&FileAudioDurations::default(),
)
}
#[tokio::test]
async fn extension_window_prefix_precedes_retained_turns_and_is_not_repeated() {
use codex_extension_api::PreviousWorldStateSection;
@@ -429,7 +438,9 @@ fn collect_annotated_user_messages_extracts_user_text_only(
ResponseItemEnvelope::new(ResponseItem::Other),
];
let collected = collect_annotated_user_messages(&items);
let collected: Vec<_> = annotated_user_messages(&items)
.map(|(_, message)| message)
.collect();
if !preserve_content {
item = user_message(expected_text);
@@ -458,7 +469,7 @@ fn collect_annotated_user_messages_extracts_user_text_only(
}]
);
assert_eq!(
build_compacted_history(Vec::new(), &collected, "summary"),
compact_items(&items, "summary"),
vec![
expected,
ResponseItemEnvelope::new(ContextualUserFragment::into(CompactionSummary::new(
@@ -576,10 +587,10 @@ fn build_token_limited_compacted_history_truncates_overlong_user_messages(
.set_id(Some(ResponseItemId::with_suffix("msg", "long-user")));
original.metadata = Some(CodexHarnessMetadata::default());
let history = super::build_compacted_history_with_limit(
Vec::new(),
&collect_annotated_user_messages(std::slice::from_ref(&original)),
std::slice::from_ref(&original),
"SUMMARY",
max_tokens,
&FileAudioDurations::default(),
);
assert_eq!(history.len(), 2);
@@ -653,12 +664,10 @@ fn build_token_limited_compacted_history_truncates_overlong_user_messages(
#[test]
fn build_token_limited_compacted_history_appends_summary_message() {
let initial_context: Vec<ResponseItemEnvelope> = Vec::new();
let original = user_message("first user message");
let user_messages = collect_user_messages(std::slice::from_ref(&original));
let items = annotated(vec![user_message("first user message")]);
let summary_text = "summary text";
let history = build_compacted_history(initial_context, &user_messages, summary_text);
let history = compact_items(&items, summary_text);
assert!(
!history.is_empty(),
"expected compacted history to include summary"
@@ -704,11 +713,7 @@ fn build_compacted_history_preserves_user_message_passthrough_metadata(attribute
item: original,
metadata: Some(CodexHarnessMetadata::default()),
};
let history = build_compacted_history(
Vec::new(),
&collect_annotated_user_messages(std::slice::from_ref(&original)),
"summary text",
);
let history = compact_items(std::slice::from_ref(&original), "summary text");
assert_eq!(
history,
+49 -4
View File
@@ -7,6 +7,8 @@ use codex_protocol::error::CodexErr;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ModelInfo;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
/// Uploaded files need the admission checks on typed turn input before entering history.
pub(crate) fn validate_injected_response_item(item: &ResponseItem) -> CodexResult<()> {
@@ -19,6 +21,45 @@ pub(crate) fn validate_injected_response_item(item: &ResponseItem) -> CodexResul
}
impl Session {
/// History injection can precede the first turn. Capture full initial context
/// only when the thread has no reference context, including after compaction.
pub(crate) async fn history_injection_context(
self: &Arc<Self>,
) -> CodexResult<Arc<TurnContext>> {
if self.reference_context_item().await.is_some() {
let turn = self.new_inject_items_context().await;
// Compaction can clear the baseline while we build this context.
// Rebuilding initial context requires a full skills snapshot.
if self.reference_context_item().await.is_some() {
return Ok(turn);
}
}
let turn = self.new_default_turn().await;
if self.reference_context_item().await.is_none() {
let step = self
.capture_step_context(Arc::clone(&turn), &CancellationToken::new())
.await?;
self.record_context_updates_and_set_reference_context_item(step.as_ref())
.await?;
}
Ok(turn)
}
pub(crate) async fn record_inherited_items(
self: &Arc<Self>,
items: Vec<ResponseItemEnvelope>,
) -> CodexResult<()> {
if items.is_empty() {
return Ok(());
}
let turn = self.history_injection_context().await?;
// The parent already prepared the images and their dimension notices.
self.record_prepared_conversation_items(&turn, turn.model_info(), items, Vec::new())
.await;
self.flush_rollout().await?;
Ok(())
}
/// Returns the input if there is no active turn to inject into.
#[expect(
clippy::await_holding_invalid_type,
@@ -76,19 +117,19 @@ impl Session {
Ok(())
}
/// Preserves trusted client provenance while items wait for an active turn.
/// Applies client developer-message retention before queueing or recording items.
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state updates must remain atomic"
)]
pub(crate) async fn inject_client_response_items(
&self,
items: Vec<ResponseItem>,
items: Vec<impl Into<ResponseItemEnvelope>>,
turn_context: &TurnContext,
) {
let items = items
.into_iter()
.map(ResponseItemEnvelope::new)
.map(Into::into)
.map(|item| self.annotate_client_response_item(item))
.collect::<Vec<_>>();
let mut active = self.active_turn.lock().await;
@@ -120,7 +161,11 @@ impl Session {
if self.enabled(Feature::RetainClientDeveloperMessages)
&& matches!(&envelope.item, ResponseItem::Message { role, .. } if role == "developer")
{
envelope.metadata.get_or_insert_default().client_authored = true;
envelope
.metadata
.get_or_insert_default()
.retain_across_windows
.get_or_insert(true);
}
envelope
}
+20 -22
View File
@@ -225,8 +225,6 @@ use uuid::Uuid;
use crate::client::ModelClient;
use crate::codex_thread::CodexThreadSettingsOverrides;
use crate::codex_thread::ThreadConfigSnapshot;
#[cfg(test)]
use crate::compact::collect_user_messages;
use crate::config::Config;
use crate::config::ConstraintResult;
use crate::config::PermissionProfileSnapshot;
@@ -689,7 +687,15 @@ impl Session {
})
}
async fn spawn_internal(args: SessionSpawnArgs) -> CodexResult<(Arc<Self>, SessionIo)> {
// Construct the startup future outside the caller's poll frame.
#[inline(never)]
fn spawn_internal(
args: SessionSpawnArgs,
) -> BoxFuture<'static, CodexResult<(Arc<Self>, SessionIo)>> {
Box::pin(Self::initialize(args))
}
async fn initialize(args: SessionSpawnArgs) -> CodexResult<(Arc<Self>, SessionIo)> {
args.config
.config_layer_stack
.ensure_config_server_applied()?;
@@ -5472,24 +5478,16 @@ impl Session {
let turn_context = step_context.turn.as_ref();
let history = self.clone_history().await;
let input_goal_ids = crate::context::UserGoalUpdate::message_ids(history.raw_items());
let retained_client_items =
crate::compact_remote_v2::truncate_retained_messages_for_remote_compaction(
history
.annotated_items()
.iter()
.filter(|item| {
(self.enabled(Feature::RetainClientToolOutputs)
&& crate::compact_retained_tool_outputs::is_retained_tool_output(item))
|| (self.enabled(Feature::RetainClientDeveloperMessages)
&& crate::compact_remote_v2::is_client_authored_developer_message(
item,
))
})
.cloned()
.collect(),
crate::compact_remote_v2::RETAINED_MESSAGE_TOKEN_BUDGET,
&history.file_audio_durations,
);
let retained_client_items = crate::compact_remote_v2::retained_client_items(
history.annotated_items(),
&history.file_audio_durations,
|item| {
crate::compact_remote_v2::should_retain_client_item(
item,
&turn_context.config.features,
)
},
);
// Resolve the roster before allocating a new context window.
let rendered = world_state.render_full_with_host_roster().await;
let (window_number, window_ids) = self.state.lock().await.start_new_context_window();
@@ -5499,7 +5497,7 @@ impl Session {
crate::context_manager::updates::merge_world_state_updates(context_updates)
.into_iter()
.map(ResponseItemEnvelope::new)
.chain(retained_client_items)
.chain(retained_client_items.into_values())
.collect();
let turn_context_item = step_context.to_turn_context_item();
self.replace_compacted_history(
@@ -561,13 +561,8 @@ impl ContextManager {
// prompt shape.
// TODO(ccunningham): if we drop support for None replacement_history compaction items,
// we can get rid of this second loop entirely and just build `history` directly in the first loop.
let user_messages =
compact::collect_annotated_user_messages(history.annotated_items());
let rebuilt = compact::build_compacted_history(
Vec::new(),
&user_messages,
&compacted.message,
);
let rebuilt =
compact::build_compacted_history(&history, &compacted.message);
let retained_context = history.retained_context().clone();
history.replace_annotated(rebuilt);
history.restore_retained_context(Some(&retained_context));
+5
View File
@@ -1017,6 +1017,10 @@ impl Session {
SessionId::from(thread_id)
}
});
let inherited_items = agent_control
.runtime()
.retained_items_for_fresh_child(&initial_history, &session_source)
.await;
let has_provided_agent_control =
matches!(&agent_control, AgentControlInit::Provided { .. });
let (agent_control, local_agent_runtime): (Arc<dyn AgentControl>, _) = match agent_control {
@@ -2129,6 +2133,7 @@ impl Session {
// record_initial_history can emit events. We record only after the SessionConfiguredEvent is emitted.
Box::pin(sess.record_initial_history(initial_history)).await;
Box::pin(sess.record_inherited_items(inherited_items)).await?;
if restore_child_window {
sess.state.lock().await.restore_auto_compact_window(
/*window_number*/ 0,
+2 -6
View File
@@ -13433,9 +13433,7 @@ async fn sample_rollout(
rollout_items.push(RolloutItem::ResponseItem(assistant1.clone().into()));
let summary1 = "summary one";
let snapshot1 = raw_history_items(&live_history);
let user_messages1 = collect_user_messages(&snapshot1);
let rebuilt1 = compact::build_compacted_history(Vec::new(), &user_messages1, summary1);
let rebuilt1 = compact::build_compacted_history(&live_history, summary1);
live_history.replace_annotated(rebuilt1);
let (window_number, window_ids) = session.advance_auto_compact_window().await;
rollout_items.push(RolloutItem::Compacted(Box::new(CompactedItem {
@@ -13468,9 +13466,7 @@ async fn sample_rollout(
rollout_items.push(RolloutItem::ResponseItem(assistant2.clone().into()));
let summary2 = "summary two";
let snapshot2 = raw_history_items(&live_history);
let user_messages2 = collect_user_messages(&snapshot2);
let rebuilt2 = compact::build_compacted_history(Vec::new(), &user_messages2, summary2);
let rebuilt2 = compact::build_compacted_history(&live_history, summary2);
live_history.replace_annotated(rebuilt2);
let (window_number, window_ids) = session.advance_auto_compact_window().await;
rollout_items.push(RolloutItem::Compacted(Box::new(CompactedItem {
+8 -10
View File
@@ -1127,7 +1127,7 @@ async fn pending_turn_input(
SubmittedTurnInput::AnnotatedResponseItem(input) => input.annotations.retain,
_ => false,
};
let mut input = match input {
match input {
SubmittedTurnInput::UserInput { content, client_id } => TurnInput::UserInput {
content,
client_id,
@@ -1143,7 +1143,7 @@ async fn pending_turn_input(
ResponseItem::FunctionCallOutput { call_id: None, .. }
) =>
{
let metadata = if let Some(fragment) = sender_messages {
let mut metadata = if let Some(fragment) = sender_messages {
let mut messages = SenderUserMessages {
receiver_turn_id: turn_id.to_owned(),
receiver_message_id: fragment.delivery.clone(),
@@ -1158,6 +1158,12 @@ async fn pending_turn_input(
} else {
None
};
if retain
&& session.enabled(Feature::RetainClientToolOutputs)
&& matches!(item, ResponseItem::FunctionCallOutput { name: Some(_), .. })
{
metadata.get_or_insert_default().retain_across_windows = Some(true);
}
TurnInput::FunctionCallOutput(ResponseItemEnvelope { item, metadata })
}
SubmittedTurnInput::ResponseItem(item)
@@ -1167,13 +1173,5 @@ async fn pending_turn_input(
SubmittedTurnInput::InterAgentCommunication(communication) => {
TurnInput::InterAgentCommunication(communication.into())
}
};
if retain
&& let TurnInput::ResponseItem(envelope) | TurnInput::FunctionCallOutput(envelope) =
&mut input
&& matches!(envelope.item, ResponseItem::FunctionCallOutput { .. })
{
envelope.metadata.get_or_insert_default().client_authored = true;
}
input
}
+22 -2
View File
@@ -766,7 +766,7 @@ fn recorded_call(call: &ToolCall) -> (ExecutedToolCall, usize) {
(ExecutedToolCall::new(name, arguments), original_bytes)
}
fn input_call_id(item: &ResponseItem) -> Option<&str> {
pub(crate) fn input_call_id(item: &ResponseItem) -> Option<&str> {
match item {
ResponseItem::FunctionCall { call_id, .. }
| ResponseItem::CustomToolCall { call_id, .. } => Some(call_id),
@@ -776,7 +776,7 @@ fn input_call_id(item: &ResponseItem) -> Option<&str> {
}
}
fn output_call_id(item: &ResponseItem) -> Option<&str> {
pub(crate) fn output_call_id(item: &ResponseItem) -> Option<&str> {
match item {
ResponseItem::FunctionCallOutput { call_id, .. }
| ResponseItem::ToolSearchOutput { call_id, .. } => call_id.as_deref(),
@@ -784,3 +784,23 @@ fn output_call_id(item: &ResponseItem) -> Option<&str> {
_ => None,
}
}
pub(crate) fn is_matching_tool_output(call: &ResponseItem, output: &ResponseItem) -> bool {
let Some(call_id) = input_call_id(call) else {
return false;
};
output_call_id(output) == Some(call_id)
&& matches!(
(call, output),
(
ResponseItem::FunctionCall { .. } | ResponseItem::LocalShellCall { .. },
ResponseItem::FunctionCallOutput { .. }
) | (
ResponseItem::CustomToolCall { .. },
ResponseItem::CustomToolCallOutput { .. }
) | (
ResponseItem::ToolSearchCall { .. },
ResponseItem::ToolSearchOutput { .. }
)
)
}
+3
View File
@@ -41,6 +41,9 @@ use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::formatted_truncate_text;
use codex_utils_output_truncation::truncate_text;
pub(crate) use executed_tool_calls::ExecutedToolCalls;
pub(crate) use executed_tool_calls::input_call_id;
pub(crate) use executed_tool_calls::is_matching_tool_output;
pub(crate) use executed_tool_calls::output_call_id;
pub use multi_agent_tool::MULTI_AGENT_V2_NAMESPACE_DESCRIPTION;
pub use router::ToolRouter;