Fix compaction context reinjection and model baselines (#12252)

## Summary
- move regular-turn context diff/full-context persistence into
`run_turn` so pre-turn compaction runs before incoming context updates
are recorded
- after successful pre-turn compaction, rely on a cleared
`reference_context_item` to trigger full context reinjection on the
follow-up regular turn (manual `/compact` keeps replacement history
summary-only and also clears the baseline)
- preserve `<model_switch>` when full context is reinjected, and inject
it *before* the rest of the full-context items
- scope `reference_context_item` and `previous_model` to regular user
turns only so standalone tasks (`/compact`, shell, review, undo) cannot
suppress future reinjection or `<model_switch>` behavior
- make context-diff persistence + `reference_context_item` updates
explicit in the regular-turn path, with clearer docs/comments around the
invariant
- stop persisting local `/compact` `RolloutItem::TurnContext` snapshots
(only regular turns persist `TurnContextItem` now)
- simplify resume/fork previous-model/reference-baseline hydration by
looking up the last surviving turn context from rollout lifecycle
events, including rollback and compaction-crossing handling
- remove the legacy fallback that guessed from bare `TurnContext`
rollouts without lifecycle events
- update compaction/remote-compaction/model-visible snapshots and
compact test assertions (including remote compaction mock response
shape)

## Why
We were persisting incoming context items before spawning the regular
turn task, which let pre-turn compaction requests accidentally include
incoming context diffs without the new user message. Fixing that exposed
follow-on baseline issues around `/compact`, resume/fork, and standalone
tasks that could cause duplicate context injection or suppress
`<model_switch>` instructions.

This PR re-centers the invariants around regular turns:
- regular turns persist model-visible context diffs/full reinjection and
update the `reference_context_item`
- standalone tasks do not advance those regular-turn baselines
- compaction clears the baseline when replacement history may have
stripped the referenced context diffs

## Follow-ups (TODOs left in code)
- `TODO(ccunningham)`: fix rollback/backtracking baseline handling more
comprehensively
- `TODO(ccunningham)`: include pending incoming context items in
pre-turn compaction threshold estimation
- `TODO(ccunningham)`: inject updated personality spec alongside
`<model_switch>` so some model-switch paths can avoid forced full
reinjection
- `TODO(ccunningham)`: review task turn lifecycle
(`TurnStarted`/`TurnComplete`) behavior and emit task-start context
diffs for task types that should have them (excluding `/compact`)

## Validation
- `just fmt`
- CI should cover the updated compaction/resume/model-visible snapshot
expectations and rollout-hydration behavior
- I did **not** rerun the full local test suite after the latest
resume-lookup / rollout-persistence simplifications
This commit is contained in:
Charley Cunningham
2026-02-20 23:13:08 -08:00
committed by GitHub
parent 264fc444b6
commit bb0ac5be70
31 changed files with 1289 additions and 1206 deletions
+582 -131
View File
@@ -18,6 +18,7 @@ use crate::analytics_client::build_track_events_context;
use crate::apps::render_apps_section;
use crate::commit_attribution::commit_message_trailer_instruction;
use crate::compact;
use crate::compact::InitialContextInjection;
use crate::compact::run_inline_auto_compact_task;
use crate::compact::should_use_remote_compact_task;
use crate::compact_remote::run_inline_remote_auto_compact_task;
@@ -1599,7 +1600,7 @@ impl Session {
self.record_conversation_items(&turn_context, &items).await;
{
let mut state = self.state.lock().await;
state.initial_context_seeded = true;
state.set_reference_context_item(Some(turn_context.to_turn_context_item()));
}
self.set_previous_model(None).await;
// Ensure initial items are visible to immediate readers (e.g., tests, forks).
@@ -1609,19 +1610,26 @@ impl Session {
let rollout_items = resumed_history.history;
let restored_tool_selection =
Self::extract_mcp_tool_selection_from_rollout(&rollout_items);
let previous_model = Self::last_rollout_model_name(&rollout_items)
.map(std::string::ToString::to_string);
let (previous_regular_turn_context_item, crossed_compaction_after_turn) =
Self::last_rollout_regular_turn_context_lookup(&rollout_items);
let previous_model =
previous_regular_turn_context_item.map(|ctx| ctx.model.clone());
let curr = turn_context.model_info.slug.as_str();
let reference_context_item = if !crossed_compaction_after_turn {
previous_regular_turn_context_item.cloned()
} else {
// Keep the baseline empty when compaction may have stripped the referenced
// context diffs so the first resumed regular turn fully reinjects context.
None
};
{
let mut state = self.state.lock().await;
state.initial_context_seeded = false;
state.set_reference_context_item(reference_context_item);
}
self.set_previous_model(previous_model).await;
self.set_previous_model(previous_model.clone()).await;
// If resuming, warn when the last recorded model differs from the current one.
let curr = turn_context.model_info.slug.as_str();
if let Some(prev) =
Self::last_rollout_model_name(&rollout_items).filter(|p| *p != curr)
{
if let Some(prev) = previous_model.as_deref().filter(|p| *p != curr) {
warn!("resuming session with different model: previous={prev}, current={curr}");
self.send_event(
&turn_context,
@@ -1661,8 +1669,10 @@ impl Session {
InitialHistory::Forked(rollout_items) => {
let restored_tool_selection =
Self::extract_mcp_tool_selection_from_rollout(&rollout_items);
let previous_model = Self::last_rollout_model_name(&rollout_items)
.map(std::string::ToString::to_string);
let (previous_regular_turn_context_item, _) =
Self::last_rollout_regular_turn_context_lookup(&rollout_items);
let previous_model =
previous_regular_turn_context_item.map(|ctx| ctx.model.clone());
self.set_previous_model(previous_model).await;
// Always add response items to conversation history
@@ -1695,7 +1705,7 @@ impl Session {
.await;
{
let mut state = self.state.lock().await;
state.initial_context_seeded = true;
state.set_reference_context_item(Some(turn_context.to_turn_context_item()));
}
// Forked threads should remain file-backed immediately after startup.
@@ -1707,14 +1717,148 @@ impl Session {
}
}
fn last_rollout_model_name(rollout_items: &[RolloutItem]) -> Option<&str> {
rollout_items.iter().rev().find_map(|it| {
if let RolloutItem::TurnContext(ctx) = it {
Some(ctx.model.as_str())
} else {
None
/// Returns `(last_turn_context_item, crossed_compaction_after_turn)` from the
/// rollback-adjusted rollout view.
///
/// This relies on the invariant that only regular turns persist `TurnContextItem`.
/// `ThreadRolledBack` markers are applied so resume/fork uses the post-rollback history view.
///
/// Returns `(None, false)` when no persisted `TurnContextItem` can be found.
///
/// Older/minimal rollouts may only contain `RolloutItem::TurnContext` entries without turn
/// lifecycle events. In that case we fall back to the last `TurnContextItem` (plus whether a
/// later `Compacted` item appears in rollout order).
// TODO(ccunningham): Simplify this lookup by sharing rollout traversal/rollback application
// with `reconstruct_history_from_rollout` so resume/fork baseline hydration does not need a
// second bespoke rollout scan.
fn last_rollout_regular_turn_context_lookup(
rollout_items: &[RolloutItem],
) -> (Option<&TurnContextItem>, bool) {
// Reverse scan over rollout items. `ThreadRolledBack(num_turns)` is naturally handled by
// skipping the next `num_turns` completed turn spans we encounter while walking backward.
//
// "Active turn" here means: we have seen `TurnComplete`/`TurnAborted` and are currently
// scanning backward through that completed turn until its matching `TurnStarted`.
let mut turns_to_skip_due_to_rollback = 0usize;
let mut saw_surviving_compaction_after_candidate = false;
let mut saw_turn_lifecycle_event = false;
let mut active_turn_id: Option<&str> = None;
let mut active_turn_saw_user_message = false;
let mut active_turn_context: Option<&TurnContextItem> = None;
let mut active_turn_contains_compaction = false;
for item in rollout_items.iter().rev() {
match item {
RolloutItem::EventMsg(EventMsg::ThreadRolledBack(rollback)) => {
// Rollbacks count completed turns, not `TurnContextItem`s. We must continue
// ignoring all items inside each skipped turn until we reach its
// corresponding `TurnStarted`.
let num_turns = usize::try_from(rollback.num_turns).unwrap_or(usize::MAX);
turns_to_skip_due_to_rollback =
turns_to_skip_due_to_rollback.saturating_add(num_turns);
}
RolloutItem::EventMsg(EventMsg::TurnComplete(event)) => {
saw_turn_lifecycle_event = true;
// Enter the reverse "turn span" for this completed turn.
active_turn_id = Some(event.turn_id.as_str());
active_turn_saw_user_message = false;
active_turn_context = None;
active_turn_contains_compaction = false;
}
RolloutItem::EventMsg(EventMsg::TurnAborted(event)) => {
saw_turn_lifecycle_event = true;
// Same reverse-turn handling as `TurnComplete`. Some aborted turns may not
// have a turn id; in that case we cannot match `TurnContextItem`s to them.
active_turn_id = event.turn_id.as_deref();
active_turn_saw_user_message = false;
active_turn_context = None;
active_turn_contains_compaction = false;
}
RolloutItem::EventMsg(EventMsg::UserMessage(_)) => {
if active_turn_id.is_some() {
active_turn_saw_user_message = true;
}
}
RolloutItem::EventMsg(EventMsg::TurnStarted(event)) => {
saw_turn_lifecycle_event = true;
if active_turn_id == Some(event.turn_id.as_str()) {
let active_turn_is_rolled_back =
active_turn_saw_user_message && turns_to_skip_due_to_rollback > 0;
if active_turn_is_rolled_back {
// `ThreadRolledBack(num_turns)` counts user turns, so only consume a
// skip once we've confirmed this reverse-scanned turn span contains a
// user message. Standalone task turns must not consume rollback skips.
turns_to_skip_due_to_rollback -= 1;
}
if !active_turn_is_rolled_back {
if let Some(context_item) = active_turn_context {
return (
Some(context_item),
saw_surviving_compaction_after_candidate,
);
}
// No `TurnContextItem` in this surviving turn; keep scanning older
// turns, but remember if this turn compacted so the eventual
// candidate reports "compaction happened after it".
if active_turn_contains_compaction {
saw_surviving_compaction_after_candidate = true;
}
}
active_turn_id = None;
active_turn_saw_user_message = false;
active_turn_context = None;
active_turn_contains_compaction = false;
}
}
RolloutItem::TurnContext(ctx) => {
// Capture the latest turn context seen in this reverse-scanned turn span. If
// the turn later proves to be rolled back, we discard it when we hit the
// matching `TurnStarted`. Older rollouts may have lifecycle events but omit
// `TurnContextItem.turn_id`; accept those as belonging to the active turn
// span for resume/fork hydration.
if let Some(active_id) = active_turn_id
&& ctx
.turn_id
.as_deref()
.is_none_or(|turn_id| turn_id == active_id)
{
// Reverse scan sees the latest `TurnContextItem` for the turn first.
active_turn_context.get_or_insert(ctx);
}
}
RolloutItem::Compacted(_) => {
if active_turn_id.is_some() {
// Compaction inside the currently scanned turn is only "after" the
// eventual candidate if this turn has no `TurnContextItem` and we keep
// scanning into older turns.
active_turn_contains_compaction = true;
} else {
saw_surviving_compaction_after_candidate = true;
}
}
_ => {}
}
})
}
// Legacy/minimal rollouts may only persist `TurnContextItem`/`Compacted` without turn
// lifecycle events. Fall back to the last `TurnContextItem` in rollout order so
// resume/fork can still hydrate `previous_model` and detect compaction-after-baseline.
if !saw_turn_lifecycle_event {
let mut saw_compaction_after_last_turn_context = false;
for item in rollout_items.iter().rev() {
match item {
RolloutItem::Compacted(_) => {
saw_compaction_after_last_turn_context = true;
}
RolloutItem::TurnContext(ctx) => {
return (Some(ctx), saw_compaction_after_last_turn_context);
}
_ => {}
}
}
}
(None, false)
}
fn last_token_info_from_rollout(rollout_items: &[RolloutItem]) -> Option<TokenUsageInfo> {
@@ -2033,33 +2177,21 @@ impl Session {
.await
}
pub(crate) fn is_model_switch_developer_message(item: &ResponseItem) -> bool {
let ResponseItem::Message { role, content, .. } = item else {
return false;
};
role == "developer"
&& content.iter().any(|content_item| {
matches!(
content_item,
ContentItem::InputText { text } if text.starts_with("<model_switch>")
)
})
}
fn build_settings_update_items(
&self,
previous_context: Option<&TurnContextItem>,
resumed_model: Option<&str>,
reference_context_item: Option<&TurnContextItem>,
previous_user_turn_model: Option<&str>,
current_context: &TurnContext,
) -> Vec<ResponseItem> {
// TODO: Make context updates a pure diff of persisted previous/current TurnContextItem
// state so replay/backtracking is deterministic. Runtime inputs that affect model-visible
// context (shell, exec policy, feature gates, resumed model bridge) should be persisted
// context (shell, exec policy, feature gates, previous-model bridge) should be persisted
// state or explicit non-state replay events.
let shell = self.user_shell();
let exec_policy = self.services.exec_policy.current();
crate::context_manager::updates::build_settings_update_items(
previous_context,
resumed_model,
reference_context_item,
previous_user_turn_model,
current_context,
shell.as_ref(),
exec_policy.as_ref(),
@@ -2465,15 +2597,6 @@ impl Session {
history.raw_items().to_vec()
}
pub(crate) async fn process_compacted_history(
&self,
turn_context: &TurnContext,
compacted_history: Vec<ResponseItem>,
) -> Vec<ResponseItem> {
let initial_context = self.build_initial_context(turn_context).await;
compact::process_compacted_history(compacted_history, &initial_context)
}
/// Append ResponseItems to the in-memory conversation history only.
pub(crate) async fn record_into_history(
&self,
@@ -2542,24 +2665,13 @@ impl Session {
true
}
pub(crate) async fn replace_history(&self, items: Vec<ResponseItem>) {
pub(crate) async fn replace_history(
&self,
items: Vec<ResponseItem>,
reference_context_item: Option<TurnContextItem>,
) {
let mut state = self.state.lock().await;
state.replace_history(items);
}
pub(crate) async fn seed_initial_context_if_needed(&self, turn_context: &TurnContext) {
{
let mut state = self.state.lock().await;
if state.initial_context_seeded {
return;
}
state.initial_context_seeded = true;
}
let initial_context = self.build_initial_context(turn_context).await;
self.record_conversation_items(turn_context, &initial_context)
.await;
self.flush_rollout().await;
state.replace_history(items, reference_context_item);
}
async fn persist_rollout_response_items(&self, items: &[ResponseItem]) {
@@ -2693,14 +2805,61 @@ impl Session {
state.clone_history()
}
pub(crate) async fn previous_context_item(&self) -> Option<TurnContextItem> {
pub(crate) async fn reference_context_item(&self) -> Option<TurnContextItem> {
let state = self.state.lock().await;
state.previous_context_item()
state.reference_context_item()
}
pub(crate) async fn set_previous_context_item(&self, item: Option<TurnContextItem>) {
/// Persist the latest turn context snapshot and emit any required model-visible context updates.
///
/// When the reference snapshot is missing, this injects full initial context. Otherwise, it
/// emits only settings diff items.
///
/// If full context is injected and a model switch occurred, this prepends the
/// `<model_switch>` developer message so model-specific instructions are not lost.
///
/// Invariant: this is the only runtime path that writes a non-`None`
/// `reference_context_item`. Non-regular tasks intentionally do not update that
/// baseline; `reference_context_item` tracks the latest regular model turn.
pub(crate) async fn record_context_updates_and_set_reference_context_item(
&self,
turn_context: &TurnContext,
previous_user_turn_model: Option<&str>,
) {
let reference_context_item = self.reference_context_item().await;
let should_inject_full_context = reference_context_item.is_none();
let context_items = if should_inject_full_context {
let mut initial_context = self.build_initial_context(turn_context).await;
// Full reinjection bypasses the settings-diff path, so add the model-switch
// instruction explicitly when needed. Keep it before the rest of full context so
// model-specific guidance is read first.
if let Some(model_switch_item) =
crate::context_manager::updates::build_model_instructions_update_item(
previous_user_turn_model,
turn_context,
)
{
// TODO(ccunningham): When a model switch changes the effective personality
// instructions, inject the updated personality spec alongside <model_switch>
// here so resume/model-switch paths can avoid forcing full reinjection.
initial_context.insert(0, model_switch_item);
}
initial_context
} else {
// Steady-state path: append only context diffs to minimize token overhead.
self.build_settings_update_items(
reference_context_item.as_ref(),
previous_user_turn_model,
turn_context,
)
};
if !context_items.is_empty() {
self.record_conversation_items(turn_context, &context_items)
.await;
}
let mut state = self.state.lock().await;
state.set_previous_context_item(item);
state.set_reference_context_item(Some(turn_context.to_turn_context_item()));
}
pub(crate) async fn update_token_usage_info(
@@ -3168,11 +3327,6 @@ impl Session {
}
async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiver<Submission>) {
// Seed with context in case there is an OverrideTurnContext first.
let initial_context = sess.new_default_turn().await;
sess.set_previous_context_item(Some(initial_context.to_turn_context_item()))
.await;
// To break out of this loop, send Op::Shutdown.
while let Ok(sub) = rx_sub.recv().await {
debug!(?sub, "Submission");
@@ -3486,26 +3640,11 @@ mod handlers {
// Attempt to inject input into current task.
if let Err(SteerInputError::NoActiveTurn(items)) = sess.steer_input(items, None).await {
sess.seed_initial_context_if_needed(&current_context).await;
let previous_model = sess.previous_model().await;
let previous_context_item = sess.previous_context_item().await;
let update_items = sess.build_settings_update_items(
previous_context_item.as_ref(),
previous_model.as_deref(),
&current_context,
);
if !update_items.is_empty() {
sess.record_conversation_items(&current_context, &update_items)
.await;
}
sess.refresh_mcp_servers_if_requested(&current_context)
.await;
let regular_task = sess.take_startup_regular_task().await.unwrap_or_default();
sess.spawn_task(Arc::clone(&current_context), items, regular_task)
.await;
sess.set_previous_context_item(Some(current_context.to_turn_context_item()))
.await;
}
}
@@ -3534,8 +3673,6 @@ mod handlers {
UserShellCommandTask::new(command),
)
.await;
sess.set_previous_context_item(Some(turn_context.to_turn_context_item()))
.await;
}
pub async fn resolve_elicitation(
@@ -3965,15 +4102,20 @@ mod handlers {
}
let turn_context = sess.new_default_turn_with_sub_id(sub_id).await;
sess.set_previous_model(Some(turn_context.model_info.slug.clone()))
.await;
let mut history = sess.clone_history().await;
// TODO(ccunningham): Fix rollback/backtracking baseline handling.
// We clear `reference_context_item` here, but should restore the
// post-rollback baseline from the surviving history/rollout instead.
// Truncating history should also invalidate/recompute `previous_model`
// so the next regular turn replays any dropped model-switch
// instructions.
history.drop_last_n_user_turns(num_turns);
// Replace with the raw items. We don't want to replace with a normalized
// version of the history.
sess.replace_history(history.raw_items().to_vec()).await;
sess.replace_history(history.raw_items().to_vec(), None)
.await;
sess.recompute_token_usage(turn_context.as_ref()).await;
sess.send_event_raw_flushed(Event {
@@ -4247,6 +4389,9 @@ async fn spawn_review_thread(
}];
let tc = Arc::new(review_turn_context);
tc.turn_metadata_state.spawn_git_enrichment_task();
// TODO(ccunningham): Review turns currently rely on `spawn_task` for TurnComplete but do not
// emit a parent TurnStarted. Consider giving review a full parent turn lifecycle
// (TurnStarted + TurnComplete) for consistency with other standalone tasks.
sess.spawn_task(tc.clone(), input, ReviewTask::new()).await;
// Announce entering review mode so UIs can switch modes.
@@ -4346,6 +4491,10 @@ pub(crate) async fn run_turn(
collaboration_mode_kind: turn_context.collaboration_mode.mode,
});
sess.send_event(&turn_context, event).await;
// TODO(ccunningham): Pre-turn compaction runs before context updates and the
// new user message are recorded. Estimate pending incoming items (context
// diffs/full reinjection + user input) and trigger compaction preemptively
// when they would push the thread over the compaction threshold.
if run_pre_sampling_compact(&sess, &turn_context)
.await
.is_err()
@@ -4354,6 +4503,13 @@ pub(crate) async fn run_turn(
return None;
}
let previous_model = sess.previous_model().await;
sess.record_context_updates_and_set_reference_context_item(
turn_context.as_ref(),
previous_model.as_deref(),
)
.await;
let skills_outcome = Some(
sess.services
.skills_manager
@@ -4465,6 +4621,11 @@ pub(crate) async fn run_turn(
let response_item: ResponseItem = initial_input_for_turn.clone().into();
sess.record_user_prompt_and_emit_turn_item(turn_context.as_ref(), &input, response_item)
.await;
// Track the previous-model baseline from the regular user-turn path only so
// standalone tasks (compact/shell/review/undo) cannot suppress future
// `<model_switch>` injections.
sess.set_previous_model(Some(turn_context.model_info.slug.clone()))
.await;
if !skill_items.is_empty() {
sess.record_conversation_items(&turn_context, &skill_items)
@@ -4568,7 +4729,14 @@ pub(crate) async fn run_turn(
// as long as compaction works well in getting us way below the token limit, we shouldn't worry about being in an infinite loop.
if token_limit_reached && needs_follow_up {
if run_auto_compact(&sess, &turn_context).await.is_err() {
if run_auto_compact(
&sess,
&turn_context,
InitialContextInjection::BeforeLastUserMessage,
)
.await
.is_err()
{
return None;
}
continue;
@@ -4688,7 +4856,7 @@ async fn run_pre_sampling_compact(
.unwrap_or(i64::MAX);
// Compact if the total usage tokens are greater than the auto compact limit
if total_usage_tokens >= auto_compact_limit {
run_auto_compact(sess, turn_context).await?;
run_auto_compact(sess, turn_context, InitialContextInjection::DoNotInject).await?;
}
Ok(())
}
@@ -4696,47 +4864,67 @@ async fn run_pre_sampling_compact(
/// Runs pre-sampling compaction against the previous model when switching to a smaller
/// context-window model.
///
/// Returns `Ok(())` when compaction either completed successfully or was skipped because the
/// model/context-window preconditions were not met. Returns `Err(_)` only when compaction was
/// attempted and failed.
/// Returns `Ok(true)` when compaction ran successfully, `Ok(false)` when compaction was skipped
/// because the model/context-window preconditions were not met, and `Err(_)` only when compaction
/// was attempted and failed.
async fn maybe_run_previous_model_inline_compact(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
total_usage_tokens: i64,
) -> CodexResult<()> {
) -> CodexResult<bool> {
let Some(previous_model) = sess.previous_model().await else {
return Ok(());
return Ok(false);
};
let previous_turn_context = Arc::new(
let previous_model_turn_context = Arc::new(
turn_context
.with_model(previous_model, &sess.services.models_manager)
.await,
);
let Some(old_context_window) = previous_turn_context.model_context_window() else {
return Ok(());
let Some(old_context_window) = previous_model_turn_context.model_context_window() else {
return Ok(false);
};
let Some(new_context_window) = turn_context.model_context_window() else {
return Ok(());
return Ok(false);
};
let new_auto_compact_limit = turn_context
.model_info
.auto_compact_token_limit()
.unwrap_or(i64::MAX);
let should_run = total_usage_tokens > new_auto_compact_limit
&& previous_turn_context.model_info.slug != turn_context.model_info.slug
&& previous_model_turn_context.model_info.slug != turn_context.model_info.slug
&& old_context_window > new_context_window;
if should_run {
run_auto_compact(sess, &previous_turn_context).await?;
run_auto_compact(
sess,
&previous_model_turn_context,
InitialContextInjection::DoNotInject,
)
.await?;
return Ok(true);
}
Ok(())
Ok(false)
}
async fn run_auto_compact(sess: &Arc<Session>, turn_context: &Arc<TurnContext>) -> CodexResult<()> {
async fn run_auto_compact(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
initial_context_injection: InitialContextInjection,
) -> CodexResult<()> {
if should_use_remote_compact_task(&turn_context.provider) {
run_inline_remote_auto_compact_task(Arc::clone(sess), Arc::clone(turn_context)).await?;
run_inline_remote_auto_compact_task(
Arc::clone(sess),
Arc::clone(turn_context),
initial_context_injection,
)
.await?;
} else {
run_inline_auto_compact_task(Arc::clone(sess), Arc::clone(turn_context)).await?;
run_inline_auto_compact_task(
Arc::clone(sess),
Arc::clone(turn_context),
initial_context_injection,
)
.await?;
}
Ok(())
}
@@ -5525,6 +5713,9 @@ async fn try_run_sampling_request(
prompt: &Prompt,
cancellation_token: CancellationToken,
) -> CodexResult<SamplingRequestResult> {
// Persist one TurnContext marker per sampling request (not just per user turn) so rollout
// analysis can reconstruct API-turn boundaries. `run_turn` persists model-visible context
// diffs/full reinjection earlier in the same regular turn before reaching this path.
let rollout_item = RolloutItem::TurnContext(turn_context.to_turn_context_item());
feedback_tags!(
@@ -6458,7 +6649,7 @@ mod tests {
async fn record_initial_history_resumed_hydrates_previous_model() {
let (session, turn_context) = make_session_and_context().await;
let previous_model = "previous-rollout-model";
let rollout_items = vec![RolloutItem::TurnContext(TurnContextItem {
let previous_context_item = TurnContextItem {
turn_id: Some(turn_context.sub_id.clone()),
cwd: turn_context.cwd.clone(),
approval_policy: turn_context.approval_policy,
@@ -6473,7 +6664,8 @@ mod tests {
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy.into()),
})];
};
let rollout_items = vec![RolloutItem::TurnContext(previous_context_item)];
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
@@ -6490,7 +6682,182 @@ mod tests {
}
#[tokio::test]
async fn resumed_history_seeds_initial_context_on_first_turn_only() {
async fn record_initial_history_resumed_hydrates_previous_model_from_lifecycle_turn_with_missing_turn_context_id()
{
let (session, turn_context) = make_session_and_context().await;
let previous_model = "previous-rollout-model";
let mut previous_context_item = TurnContextItem {
turn_id: Some(turn_context.sub_id.clone()),
cwd: turn_context.cwd.clone(),
approval_policy: turn_context.approval_policy,
sandbox_policy: turn_context.sandbox_policy.clone(),
network: None,
model: previous_model.to_string(),
personality: turn_context.personality,
collaboration_mode: Some(turn_context.collaboration_mode.clone()),
effort: turn_context.reasoning_effort,
summary: turn_context.reasoning_summary,
user_instructions: None,
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy.into()),
};
let turn_id = previous_context_item
.turn_id
.clone()
.expect("turn context should have turn_id");
previous_context_item.turn_id = None;
let rollout_items = vec![
RolloutItem::EventMsg(EventMsg::TurnStarted(
codex_protocol::protocol::TurnStartedEvent {
turn_id: turn_id.clone(),
model_context_window: Some(128_000),
collaboration_mode_kind: ModeKind::Default,
},
)),
RolloutItem::EventMsg(EventMsg::UserMessage(
codex_protocol::protocol::UserMessageEvent {
message: "seed".to_string(),
images: None,
local_images: Vec::new(),
text_elements: Vec::new(),
},
)),
RolloutItem::TurnContext(previous_context_item),
RolloutItem::EventMsg(EventMsg::TurnComplete(
codex_protocol::protocol::TurnCompleteEvent {
turn_id,
last_agent_message: None,
},
)),
];
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
conversation_id: ThreadId::default(),
history: rollout_items,
rollout_path: PathBuf::from("/tmp/resume.jsonl"),
}))
.await;
assert_eq!(
session.previous_model().await,
Some(previous_model.to_string())
);
}
#[tokio::test]
async fn record_initial_history_resumed_rollback_skips_only_user_turns() {
let (session, turn_context) = make_session_and_context().await;
let previous_context_item = turn_context.to_turn_context_item();
let user_turn_id = previous_context_item
.turn_id
.clone()
.expect("turn context should have turn_id");
let standalone_turn_id = "standalone-task-turn".to_string();
let rollout_items = vec![
RolloutItem::EventMsg(EventMsg::TurnStarted(
codex_protocol::protocol::TurnStartedEvent {
turn_id: user_turn_id.clone(),
model_context_window: Some(128_000),
collaboration_mode_kind: ModeKind::Default,
},
)),
RolloutItem::EventMsg(EventMsg::UserMessage(
codex_protocol::protocol::UserMessageEvent {
message: "seed".to_string(),
images: None,
local_images: Vec::new(),
text_elements: Vec::new(),
},
)),
RolloutItem::TurnContext(previous_context_item),
RolloutItem::EventMsg(EventMsg::TurnComplete(
codex_protocol::protocol::TurnCompleteEvent {
turn_id: user_turn_id,
last_agent_message: None,
},
)),
// Standalone task turn (no UserMessage) should not consume rollback skips.
RolloutItem::EventMsg(EventMsg::TurnStarted(
codex_protocol::protocol::TurnStartedEvent {
turn_id: standalone_turn_id.clone(),
model_context_window: Some(128_000),
collaboration_mode_kind: ModeKind::Default,
},
)),
RolloutItem::EventMsg(EventMsg::TurnComplete(
codex_protocol::protocol::TurnCompleteEvent {
turn_id: standalone_turn_id,
last_agent_message: None,
},
)),
RolloutItem::EventMsg(EventMsg::ThreadRolledBack(
codex_protocol::protocol::ThreadRolledBackEvent { num_turns: 1 },
)),
];
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
conversation_id: ThreadId::default(),
history: rollout_items,
rollout_path: PathBuf::from("/tmp/resume.jsonl"),
}))
.await;
assert_eq!(session.previous_model().await, None);
assert!(session.reference_context_item().await.is_none());
}
#[tokio::test]
async fn record_initial_history_resumed_seeds_reference_context_item_without_compaction() {
let (session, turn_context) = make_session_and_context().await;
let previous_context_item = turn_context.to_turn_context_item();
let rollout_items = vec![RolloutItem::TurnContext(previous_context_item.clone())];
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
conversation_id: ThreadId::default(),
history: rollout_items,
rollout_path: PathBuf::from("/tmp/resume.jsonl"),
}))
.await;
assert_eq!(
serde_json::to_value(session.reference_context_item().await)
.expect("serialize seeded reference context item"),
serde_json::to_value(Some(previous_context_item))
.expect("serialize expected reference context item")
);
}
#[tokio::test]
async fn record_initial_history_resumed_does_not_seed_reference_context_item_after_compaction()
{
let (session, turn_context) = make_session_and_context().await;
let previous_context_item = turn_context.to_turn_context_item();
let rollout_items = vec![
RolloutItem::TurnContext(previous_context_item),
RolloutItem::Compacted(CompactedItem {
message: String::new(),
replacement_history: Some(Vec::new()),
}),
];
session
.record_initial_history(InitialHistory::Resumed(ResumedHistory {
conversation_id: ThreadId::default(),
history: rollout_items,
rollout_path: PathBuf::from("/tmp/resume.jsonl"),
}))
.await;
assert!(session.reference_context_item().await.is_none());
}
#[tokio::test]
async fn resumed_history_injects_initial_context_on_first_context_update_only() {
let (session, turn_context) = make_session_and_context().await;
let (rollout_items, mut expected) = sample_rollout(&session, &turn_context).await;
@@ -6505,12 +6872,16 @@ mod tests {
let history_before_seed = session.state.lock().await.clone_history();
assert_eq!(expected, history_before_seed.raw_items());
session.seed_initial_context_if_needed(&turn_context).await;
session
.record_context_updates_and_set_reference_context_item(&turn_context, None)
.await;
expected.extend(session.build_initial_context(&turn_context).await);
let history_after_seed = session.clone_history().await;
assert_eq!(expected, history_after_seed.raw_items());
session.seed_initial_context_if_needed(&turn_context).await;
session
.record_context_updates_and_set_reference_context_item(&turn_context, None)
.await;
let history_after_second_seed = session.clone_history().await;
assert_eq!(expected, history_after_second_seed.raw_items());
}
@@ -6677,7 +7048,7 @@ mod tests {
async fn record_initial_history_forked_hydrates_previous_model() {
let (session, turn_context) = make_session_and_context().await;
let previous_model = "forked-rollout-model";
let rollout_items = vec![RolloutItem::TurnContext(TurnContextItem {
let previous_context_item = TurnContextItem {
turn_id: Some(turn_context.sub_id.clone()),
cwd: turn_context.cwd.clone(),
approval_policy: turn_context.approval_policy,
@@ -6692,7 +7063,8 @@ mod tests {
developer_instructions: None,
final_output_json_schema: None,
truncation_policy: Some(turn_context.truncation_policy.into()),
})];
};
let rollout_items = vec![RolloutItem::TurnContext(previous_context_item)];
session
.record_initial_history(InitialHistory::Forked(rollout_items))
@@ -6755,6 +7127,8 @@ mod tests {
},
];
sess.record_into_history(&turn_2, tc.as_ref()).await;
sess.set_previous_model(Some("previous-regular-model".to_string()))
.await;
handlers::thread_rollback(&sess, "sub-1".to_string(), 1).await;
@@ -6769,7 +7143,7 @@ mod tests {
assert_eq!(expected, history.raw_items());
assert_eq!(
sess.previous_model().await,
Some(tc.model_info.slug.clone())
Some("previous-regular-model".to_string())
);
}
@@ -7426,8 +7800,7 @@ mod tests {
session_configuration.session_source.clone(),
);
let mut state = SessionState::new(session_configuration.clone());
mark_state_initial_context_seeded(&mut state);
let state = SessionState::new(session_configuration.clone());
let skills_manager = Arc::new(SkillsManager::new(config.codex_home.clone()));
let network_approval = Arc::new(NetworkApprovalService::default());
@@ -7583,8 +7956,7 @@ mod tests {
session_configuration.session_source.clone(),
);
let mut state = SessionState::new(session_configuration.clone());
mark_state_initial_context_seeded(&mut state);
let state = SessionState::new(session_configuration.clone());
let skills_manager = Arc::new(SkillsManager::new(config.codex_home.clone()));
let network_approval = Arc::new(NetworkApprovalService::default());
@@ -7671,10 +8043,6 @@ mod tests {
(session, turn_context, rx_event)
}
fn mark_state_initial_context_seeded(state: &mut SessionState) {
state.initial_context_seeded = true;
}
#[tokio::test]
async fn refresh_mcp_servers_is_deferred_until_next_turn() {
let (session, turn_context) = make_session_and_context().await;
@@ -7746,7 +8114,7 @@ mod tests {
}
#[tokio::test]
async fn spawn_task_hydrates_previous_model() {
async fn spawn_task_does_not_update_previous_model_for_non_run_turn_tasks() {
let (sess, tc, _rx) = make_session_and_context_with_rx().await;
sess.set_previous_model(None).await;
let input = vec![UserInput::Text {
@@ -7765,10 +8133,7 @@ mod tests {
.await;
sess.abort_all_tasks(TurnAbortReason::Interrupted).await;
assert_eq!(
sess.previous_model().await,
Some(tc.model_info.slug.clone())
);
assert_eq!(sess.previous_model().await, None);
}
#[tokio::test]
@@ -7806,9 +8171,9 @@ mod tests {
.expect("rebuild config layer stack with network requirements");
current_context.config = Arc::new(config);
let previous_context_item = previous_context.to_turn_context_item();
let reference_context_item = previous_context.to_turn_context_item();
let update_items = session.build_settings_update_items(
Some(&previous_context_item),
Some(&reference_context_item),
None,
&current_context,
);
@@ -7830,6 +8195,92 @@ mod tests {
assert!(environment_update.contains("<denied>blocked.example.com</denied>"));
}
#[tokio::test]
async fn record_context_updates_and_set_reference_context_item_injects_full_context_when_baseline_missing()
{
let (session, turn_context) = make_session_and_context().await;
session
.record_context_updates_and_set_reference_context_item(&turn_context, None)
.await;
let history = session.clone_history().await;
let initial_context = session.build_initial_context(&turn_context).await;
assert_eq!(history.raw_items().to_vec(), initial_context);
let current_context = session.reference_context_item().await;
assert_eq!(
serde_json::to_value(current_context).expect("serialize current context item"),
serde_json::to_value(Some(turn_context.to_turn_context_item()))
.expect("serialize expected context item")
);
}
#[tokio::test]
async fn record_context_updates_and_set_reference_context_item_reinjects_full_context_after_clear()
{
let (session, turn_context) = make_session_and_context().await;
let compacted_summary = ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: format!("{}\nsummary", crate::compact::SUMMARY_PREFIX),
}],
end_turn: None,
phase: None,
};
session
.record_into_history(std::slice::from_ref(&compacted_summary), &turn_context)
.await;
session
.record_context_updates_and_set_reference_context_item(&turn_context, None)
.await;
{
let mut state = session.state.lock().await;
state.set_reference_context_item(None);
}
session
.replace_history(vec![compacted_summary.clone()], None)
.await;
session
.record_context_updates_and_set_reference_context_item(&turn_context, None)
.await;
let history = session.clone_history().await;
let mut expected_history = vec![compacted_summary];
expected_history.extend(session.build_initial_context(&turn_context).await);
assert_eq!(history.raw_items().to_vec(), expected_history);
}
#[tokio::test]
async fn run_user_shell_command_does_not_set_reference_context_item() {
let (session, _turn_context, rx) = make_session_and_context_with_rx().await;
{
let mut state = session.state.lock().await;
state.set_reference_context_item(None);
}
handlers::run_user_shell_command(&session, "sub-id".to_string(), "echo shell".to_string())
.await;
let deadline = StdDuration::from_secs(5);
let start = std::time::Instant::now();
loop {
let remaining = deadline.saturating_sub(start.elapsed());
let evt = tokio::time::timeout(remaining, rx.recv())
.await
.expect("timeout waiting for event")
.expect("event");
if matches!(evt.msg, EventMsg::TurnComplete(_)) {
break;
}
}
assert!(
session.reference_context_item().await.is_none(),
"standalone shell tasks should not mutate previous context"
);
}
#[derive(Clone, Copy)]
struct NeverEndingTask {
kind: TaskKind,
+265 -285
View File
@@ -32,10 +32,37 @@ pub const SUMMARIZATION_PROMPT: &str = include_str!("../templates/compact/prompt
pub const SUMMARY_PREFIX: &str = include_str!("../templates/compact/summary_prefix.md");
const COMPACT_USER_MESSAGE_MAX_TOKENS: usize = 20_000;
/// Controls whether compaction replacement history must include initial context.
///
/// Pre-turn/manual compaction variants use `DoNotInject`: they replace history with a summary and
/// clear `reference_context_item`, so the next regular turn will fully reinject initial context
/// after compaction.
///
/// Mid-turn compaction must use `BeforeLastUserMessage` because the model is trained to see the
/// compaction summary as the last item in history after mid-turn compaction; we therefore inject
/// initial context into the replacement history just above the last real user message.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum InitialContextInjection {
BeforeLastUserMessage,
DoNotInject,
}
pub(crate) fn should_use_remote_compact_task(provider: &ModelProviderInfo) -> bool {
provider.is_openai()
}
fn is_model_switch_developer_message(item: &ResponseItem) -> bool {
match item {
ResponseItem::Message { role, content, .. } if role == "developer" => {
matches!(
content.as_slice(),
[ContentItem::InputText { text }] if text.starts_with("<model_switch>\n")
)
}
_ => false,
}
}
pub(crate) fn extract_trailing_model_switch_update_for_compaction_request(
history: &mut ContextManager,
) -> Option<ResponseItem> {
@@ -49,7 +76,7 @@ pub(crate) fn extract_trailing_model_switch_update_for_compaction_request(
.rev()
.find_map(|(i, item)| {
let is_trailing = last_user_turn_boundary_index.is_none_or(|boundary| i > boundary);
if is_trailing && Session::is_model_switch_developer_message(item) {
if is_trailing && is_model_switch_developer_message(item) {
Some(i)
} else {
None
@@ -64,6 +91,7 @@ pub(crate) fn extract_trailing_model_switch_update_for_compaction_request(
pub(crate) async fn run_inline_auto_compact_task(
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
initial_context_injection: InitialContextInjection,
) -> CodexResult<()> {
let prompt = turn_context.compact_prompt().to_string();
let input = vec![UserInput::Text {
@@ -72,7 +100,7 @@ pub(crate) async fn run_inline_auto_compact_task(
text_elements: Vec::new(),
}];
run_compact_task_inner(sess, turn_context, input).await?;
run_compact_task_inner(sess, turn_context, input, initial_context_injection).await?;
Ok(())
}
@@ -87,13 +115,20 @@ pub(crate) async fn run_compact_task(
collaboration_mode_kind: turn_context.collaboration_mode.mode,
});
sess.send_event(&turn_context, start_event).await;
run_compact_task_inner(sess.clone(), turn_context, input).await
run_compact_task_inner(
sess.clone(),
turn_context,
input,
InitialContextInjection::DoNotInject,
)
.await
}
async fn run_compact_task_inner(
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
input: Vec<UserInput>,
initial_context_injection: InitialContextInjection,
) -> CodexResult<()> {
let compaction_item = TurnItem::ContextCompaction(ContextCompactionItem::new());
sess.emit_turn_item_started(&turn_context, &compaction_item)
@@ -118,9 +153,6 @@ async fn run_compact_task_inner(
// Reuse one client session so turn-scoped state (sticky routing, websocket append tracking)
// survives retries within this compact turn.
let rollout_item = RolloutItem::TurnContext(turn_context.to_turn_context_item());
sess.persist_rollout_items(&[rollout_item]).await;
loop {
// Clone is required because of the loop
let turn_input = history
@@ -202,8 +234,16 @@ async fn run_compact_task_inner(
let summary_text = format!("{SUMMARY_PREFIX}\n{summary_suffix}");
let user_messages = collect_user_messages(history_items);
let initial_context = sess.build_initial_context(turn_context.as_ref()).await;
let mut new_history = build_compacted_history(initial_context, &user_messages, &summary_text);
let mut new_history = build_compacted_history(Vec::new(), &user_messages, &summary_text);
if matches!(
initial_context_injection,
InitialContextInjection::BeforeLastUserMessage
) {
let initial_context = sess.build_initial_context(turn_context.as_ref()).await;
new_history =
insert_initial_context_before_last_real_user_or_summary(new_history, initial_context);
}
// Reattach the stripped model-switch update only after successful compaction so the model
// still sees the switch instructions on the next real sampling request.
if let Some(model_switch_item) = stripped_model_switch_item {
@@ -215,12 +255,17 @@ async fn run_compact_task_inner(
.cloned()
.collect();
new_history.extend(ghost_snapshots);
sess.replace_history(new_history).await;
let reference_context_item = match initial_context_injection {
InitialContextInjection::DoNotInject => None,
InitialContextInjection::BeforeLastUserMessage => Some(turn_context.to_turn_context_item()),
};
sess.replace_history(new_history.clone(), reference_context_item)
.await;
sess.recompute_token_usage(&turn_context).await;
let rollout_item = RolloutItem::Compacted(CompactedItem {
message: summary_text.clone(),
replacement_history: None,
replacement_history: Some(new_history),
});
sess.persist_rollout_items(&[rollout_item]).await;
@@ -272,24 +317,50 @@ pub(crate) fn is_summary_message(message: &str) -> bool {
message.starts_with(format!("{SUMMARY_PREFIX}\n").as_str())
}
pub(crate) fn process_compacted_history(
/// Inserts canonical initial context into compacted replacement history at the
/// model-expected boundary.
///
/// Placement rules:
/// - Prefer immediately before the last real user message.
/// - If no real user messages remain, insert before the compaction summary so
/// the summary stays last.
/// - If there are no user messages, insert before the last compaction item so
/// that item remains last (remote compaction may return only compaction items).
/// - If there are no user messages or compaction items, append the context.
pub(crate) fn insert_initial_context_before_last_real_user_or_summary(
mut compacted_history: Vec<ResponseItem>,
initial_context: &[ResponseItem],
initial_context: Vec<ResponseItem>,
) -> Vec<ResponseItem> {
compacted_history.retain(should_keep_compacted_history_item);
let initial_context = initial_context.to_vec();
let mut last_user_or_summary_index = None;
let mut last_real_user_index = None;
for (i, item) in compacted_history.iter().enumerate().rev() {
let Some(TurnItem::UserMessage(user)) = crate::event_mapping::parse_turn_item(item) else {
continue;
};
// Compaction summaries are encoded as user messages, so track both:
// the last real user message (preferred insertion point) and the last
// user-message-like item (fallback summary insertion point).
last_user_or_summary_index.get_or_insert(i);
if !is_summary_message(&user.message()) {
last_real_user_index = Some(i);
break;
}
}
let last_compaction_index = compacted_history
.iter()
.enumerate()
.rev()
.find_map(|(i, item)| matches!(item, ResponseItem::Compaction { .. }).then_some(i));
let insertion_index = last_real_user_index
.or(last_user_or_summary_index)
.or(last_compaction_index);
// Re-inject canonical context from the current session since we stripped it
// from the pre-compaction history. Keep it right before the last user
// message so older user messages remain earlier in the transcript.
if let Some(last_user_index) = compacted_history.iter().rposition(|item| {
matches!(
crate::event_mapping::parse_turn_item(item),
Some(TurnItem::UserMessage(_))
)
}) {
compacted_history.splice(last_user_index..last_user_index, initial_context);
// from the pre-compaction history. Prefer placing it before the last real
// user message; if there is no real user message left, place it before the
// summary or compaction item so the compaction item remains last.
if let Some(insertion_index) = insertion_index {
compacted_history.splice(insertion_index..insertion_index, initial_context);
} else {
compacted_history.extend(initial_context);
}
@@ -297,30 +368,6 @@ pub(crate) fn process_compacted_history(
compacted_history
}
/// Returns whether an item from remote compaction output should be preserved.
///
/// Called while processing the model-provided compacted transcript, before we
/// append fresh canonical context from the current session.
///
/// We drop:
/// - `developer` messages because remote output can include stale/duplicated
/// instruction content.
/// - non-user-content `user` messages (session prefix/instruction wrappers),
/// keeping only real user messages as parsed by `parse_turn_item`.
///
/// This intentionally keeps `user`-role warnings and compaction-generated
/// summary messages because they parse as `TurnItem::UserMessage`.
fn should_keep_compacted_history_item(item: &ResponseItem) -> bool {
match item {
ResponseItem::Message { role, .. } if role == "developer" => false,
ResponseItem::Message { role, .. } if role == "user" => matches!(
crate::event_mapping::parse_turn_item(item),
Some(TurnItem::UserMessage(_))
),
_ => true,
}
}
pub(crate) fn build_compacted_history(
initial_context: Vec<ResponseItem>,
user_messages: &[String],
@@ -442,6 +489,21 @@ mod tests {
use super::*;
use pretty_assertions::assert_eq;
async fn process_compacted_history_with_test_session(
compacted_history: Vec<ResponseItem>,
) -> (Vec<ResponseItem>, Vec<ResponseItem>) {
let (session, turn_context) = crate::codex::make_session_and_context().await;
let initial_context = session.build_initial_context(&turn_context).await;
let refreshed = crate::compact_remote::process_compacted_history(
&session,
&turn_context,
compacted_history,
InitialContextInjection::BeforeLastUserMessage,
)
.await;
(refreshed, initial_context)
}
#[test]
fn content_items_to_text_joins_non_empty_segments() {
let items = vec![
@@ -514,7 +576,7 @@ mod tests {
history
.raw_items()
.iter()
.all(|item| !Session::is_model_switch_developer_message(item))
.all(|item| !is_model_switch_developer_message(item))
);
}
@@ -569,7 +631,7 @@ mod tests {
history
.raw_items()
.iter()
.any(Session::is_model_switch_developer_message)
.any(is_model_switch_developer_message)
);
}
@@ -708,8 +770,8 @@ do things
assert_eq!(summary, summary_text);
}
#[test]
fn process_compacted_history_replaces_developer_messages() {
#[tokio::test]
async fn process_compacted_history_replaces_developer_messages() {
let compacted_history = vec![
ResponseItem::Message {
id: None,
@@ -739,88 +801,22 @@ do things
phase: None,
},
];
let initial_context = vec![
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: r#"<environment_context>
<cwd>/tmp</cwd>
<shell>zsh</shell>
</environment_context>"#
.to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh personality".to_string(),
}],
end_turn: None,
phase: None,
},
];
let refreshed = process_compacted_history(compacted_history, &initial_context);
let expected = vec![
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: r#"<environment_context>
<cwd>/tmp</cwd>
<shell>zsh</shell>
</environment_context>"#
.to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh personality".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
},
];
let (refreshed, mut expected) =
process_compacted_history_with_test_session(compacted_history).await;
expected.push(ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
});
assert_eq!(refreshed, expected);
}
#[test]
fn process_compacted_history_reinjects_full_initial_context() {
#[tokio::test]
async fn process_compacted_history_reinjects_full_initial_context() {
let compacted_history = vec![ResponseItem::Message {
id: None,
role: "user".to_string(),
@@ -830,124 +826,22 @@ do things
end_turn: None,
phase: None,
}];
let initial_context = vec![
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: r#"# AGENTS.md instructions for /repo
<INSTRUCTIONS>
keep me updated
</INSTRUCTIONS>"#
.to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: r#"<environment_context>
<cwd>/repo</cwd>
<shell>zsh</shell>
</environment_context>"#
.to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: r#"<turn_aborted>
<turn_id>turn-1</turn_id>
<reason>interrupted</reason>
</turn_aborted>"#
.to_string(),
}],
end_turn: None,
phase: None,
},
];
let refreshed = process_compacted_history(compacted_history, &initial_context);
let expected = vec![
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: r#"# AGENTS.md instructions for /repo
<INSTRUCTIONS>
keep me updated
</INSTRUCTIONS>"#
.to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: r#"<environment_context>
<cwd>/repo</cwd>
<shell>zsh</shell>
</environment_context>"#
.to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: r#"<turn_aborted>
<turn_id>turn-1</turn_id>
<reason>interrupted</reason>
</turn_aborted>"#
.to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
},
];
let (refreshed, mut expected) =
process_compacted_history_with_test_session(compacted_history).await;
expected.push(ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
});
assert_eq!(refreshed, expected);
}
#[test]
fn process_compacted_history_drops_non_user_content_messages() {
#[tokio::test]
async fn process_compacted_history_drops_non_user_content_messages() {
let compacted_history = vec![
ResponseItem::Message {
id: None,
@@ -1008,42 +902,22 @@ keep me updated
phase: None,
},
];
let initial_context = vec![ResponseItem::Message {
let (refreshed, mut expected) =
process_compacted_history_with_test_session(compacted_history).await;
expected.push(ResponseItem::Message {
id: None,
role: "developer".to_string(),
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "fresh developer instructions".to_string(),
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
}];
let refreshed = process_compacted_history(compacted_history, &initial_context);
let expected = vec![
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh developer instructions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "summary".to_string(),
}],
end_turn: None,
phase: None,
},
];
});
assert_eq!(refreshed, expected);
}
#[test]
fn process_compacted_history_inserts_context_before_last_real_user_message_only() {
#[tokio::test]
async fn process_compacted_history_inserts_context_before_last_real_user_message_only() {
let compacted_history = vec![
ResponseItem::Message {
id: None,
@@ -1073,18 +947,10 @@ keep me updated
phase: None,
},
];
let initial_context = vec![ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
}];
let refreshed = process_compacted_history(compacted_history, &initial_context);
let expected = vec![
let (refreshed, initial_context) =
process_compacted_history_with_test_session(compacted_history).await;
let mut expected = vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
@@ -1103,6 +969,75 @@ keep me updated
end_turn: None,
phase: None,
},
];
expected.extend(initial_context);
expected.push(ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "latest user".to_string(),
}],
end_turn: None,
phase: None,
});
assert_eq!(refreshed, expected);
}
#[test]
fn insert_initial_context_before_last_real_user_or_summary_keeps_summary_last() {
let compacted_history = vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "older user".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "latest user".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: format!("{SUMMARY_PREFIX}\nsummary text"),
}],
end_turn: None,
phase: None,
},
];
let initial_context = vec![ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
}];
let refreshed = insert_initial_context_before_last_real_user_or_summary(
compacted_history,
initial_context,
);
let expected = vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "older user".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "developer".to_string(),
@@ -1121,6 +1056,51 @@ keep me updated
end_turn: None,
phase: None,
},
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: format!("{SUMMARY_PREFIX}\nsummary text"),
}],
end_turn: None,
phase: None,
},
];
assert_eq!(refreshed, expected);
}
#[test]
fn insert_initial_context_before_last_real_user_or_summary_keeps_compaction_last() {
let compacted_history = vec![ResponseItem::Compaction {
encrypted_content: "encrypted".to_string(),
}];
let initial_context = vec![ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
}];
let refreshed = insert_initial_context_before_last_real_user_or_summary(
compacted_history,
initial_context,
);
let expected = vec![
ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: "fresh permissions".to_string(),
}],
end_turn: None,
phase: None,
},
ResponseItem::Compaction {
encrypted_content: "encrypted".to_string(),
},
];
assert_eq!(refreshed, expected);
}
+82 -7
View File
@@ -3,7 +3,9 @@ use std::sync::Arc;
use crate::Prompt;
use crate::codex::Session;
use crate::codex::TurnContext;
use crate::compact::InitialContextInjection;
use crate::compact::extract_trailing_model_switch_update_for_compaction_request;
use crate::compact::insert_initial_context_before_last_real_user_or_summary;
use crate::context_manager::ContextManager;
use crate::context_manager::TotalTokenUsageBreakdown;
use crate::context_manager::estimate_response_item_model_visible_bytes;
@@ -25,8 +27,9 @@ use tracing::info;
pub(crate) async fn run_inline_remote_auto_compact_task(
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
initial_context_injection: InitialContextInjection,
) -> CodexResult<()> {
run_remote_compact_task_inner(&sess, &turn_context).await?;
run_remote_compact_task_inner(&sess, &turn_context, initial_context_injection).await?;
Ok(())
}
@@ -41,14 +44,17 @@ pub(crate) async fn run_remote_compact_task(
});
sess.send_event(&turn_context, start_event).await;
run_remote_compact_task_inner(&sess, &turn_context).await
run_remote_compact_task_inner(&sess, &turn_context, InitialContextInjection::DoNotInject).await
}
async fn run_remote_compact_task_inner(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
initial_context_injection: InitialContextInjection,
) -> CodexResult<()> {
if let Err(err) = run_remote_compact_task_inner_impl(sess, turn_context).await {
if let Err(err) =
run_remote_compact_task_inner_impl(sess, turn_context, initial_context_injection).await
{
let event = EventMsg::Error(
err.to_error_event(Some("Error running remote compact task".to_string())),
);
@@ -61,6 +67,7 @@ async fn run_remote_compact_task_inner(
async fn run_remote_compact_task_inner_impl(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
initial_context_injection: InitialContextInjection,
) -> CodexResult<()> {
let compaction_item = TurnItem::ContextCompaction(ContextCompactionItem::new());
sess.emit_turn_item_started(turn_context, &compaction_item)
@@ -122,9 +129,13 @@ async fn run_remote_compact_task_inner_impl(
Err(err)
})
.await?;
new_history = sess
.process_compacted_history(turn_context, new_history)
.await;
new_history = process_compacted_history(
sess.as_ref(),
turn_context.as_ref(),
new_history,
initial_context_injection,
)
.await;
// Reattach the stripped model-switch update only after successful compaction so the model
// still sees the switch instructions on the next real sampling request.
if let Some(model_switch_item) = stripped_model_switch_item {
@@ -134,7 +145,12 @@ async fn run_remote_compact_task_inner_impl(
if !ghost_snapshots.is_empty() {
new_history.extend(ghost_snapshots);
}
sess.replace_history(new_history.clone()).await;
let reference_context_item = match initial_context_injection {
InitialContextInjection::DoNotInject => None,
InitialContextInjection::BeforeLastUserMessage => Some(turn_context.to_turn_context_item()),
};
sess.replace_history(new_history.clone(), reference_context_item)
.await;
sess.recompute_token_usage(turn_context).await;
let compacted_item = CompactedItem {
@@ -149,6 +165,65 @@ async fn run_remote_compact_task_inner_impl(
Ok(())
}
pub(crate) async fn process_compacted_history(
sess: &Session,
turn_context: &TurnContext,
mut compacted_history: Vec<ResponseItem>,
initial_context_injection: InitialContextInjection,
) -> Vec<ResponseItem> {
// Mid-turn compaction is the only path that must inject initial context above the last user
// message in the replacement history. Pre-turn compaction instead injects context after the
// compaction item, but mid-turn compaction keeps the compaction item last for model training.
let initial_context = if matches!(
initial_context_injection,
InitialContextInjection::BeforeLastUserMessage
) {
sess.build_initial_context(turn_context).await
} else {
Vec::new()
};
compacted_history.retain(should_keep_compacted_history_item);
insert_initial_context_before_last_real_user_or_summary(compacted_history, initial_context)
}
/// Returns whether an item from remote compaction output should be preserved.
///
/// Called while processing the model-provided compacted transcript, before we
/// append fresh canonical context from the current session.
///
/// We drop:
/// - `developer` messages because remote output can include stale/duplicated
/// instruction content.
/// - non-user-content `user` messages (session prefix/instruction wrappers),
/// keeping only real user messages as parsed by `parse_turn_item`.
///
/// This intentionally keeps:
/// - `assistant` messages (future remote compaction models may emit them)
/// - `user`-role warnings and compaction-generated summary messages because
/// they parse as `TurnItem::UserMessage`.
fn should_keep_compacted_history_item(item: &ResponseItem) -> bool {
match item {
ResponseItem::Message { role, .. } if role == "developer" => false,
ResponseItem::Message { role, .. } if role == "user" => matches!(
crate::event_mapping::parse_turn_item(item),
Some(TurnItem::UserMessage(_))
),
ResponseItem::Message { role, .. } if role == "assistant" => true,
ResponseItem::Message { .. } => false,
ResponseItem::Compaction { .. } => true,
ResponseItem::Reasoning { .. }
| ResponseItem::LocalShellCall { .. }
| ResponseItem::FunctionCall { .. }
| ResponseItem::FunctionCallOutput { .. }
| ResponseItem::CustomToolCall { .. }
| ResponseItem::CustomToolCallOutput { .. }
| ResponseItem::WebSearchCall { .. }
| ResponseItem::GhostSnapshot { .. }
| ResponseItem::Other => false,
}
}
#[derive(Debug)]
struct CompactRequestLogData {
failing_compaction_request_model_visible_bytes: i64,
+11 -8
View File
@@ -27,12 +27,15 @@ pub(crate) struct ContextManager {
/// The oldest items are at the beginning of the vector.
items: Vec<ResponseItem>,
token_info: Option<TokenUsageInfo>,
/// Previous turn context snapshot used for diffing context and producing
/// model-visible settings update items.
/// Reference context snapshot used for diffing and producing model-visible
/// settings update items.
///
/// This is the baseline for the next regular model turn, and may already
/// match the current turn after context updates are persisted.
///
/// When this is `None`, settings diffing treats the next turn as having no
/// baseline and emits a full reinjection of context state.
previous_context_item: Option<TurnContextItem>,
reference_context_item: Option<TurnContextItem>,
}
#[derive(Debug, Clone, Copy, Default)]
@@ -48,7 +51,7 @@ impl ContextManager {
Self {
items: Vec::new(),
token_info: TokenUsageInfo::new_or_append(&None, &None, None),
previous_context_item: None,
reference_context_item: None,
}
}
@@ -60,12 +63,12 @@ impl ContextManager {
self.token_info = info;
}
pub(crate) fn set_previous_context_item(&mut self, item: Option<TurnContextItem>) {
self.previous_context_item = item;
pub(crate) fn set_reference_context_item(&mut self, item: Option<TurnContextItem>) {
self.reference_context_item = item;
}
pub(crate) fn previous_context_item(&self) -> Option<TurnContextItem> {
self.previous_context_item.clone()
pub(crate) fn reference_context_item(&self) -> Option<TurnContextItem> {
self.reference_context_item.clone()
}
pub(crate) fn set_token_usage_full(&mut self, context_window: i64) {
+10 -9
View File
@@ -97,11 +97,10 @@ pub(crate) fn personality_message_for(
}
pub(crate) fn build_model_instructions_update_item(
previous: Option<&TurnContextItem>,
resumed_model: Option<&str>,
previous_user_turn_model: Option<&str>,
next: &TurnContext,
) -> Option<ResponseItem> {
let previous_model = resumed_model.or_else(|| previous.map(|prev| prev.model.as_str()))?;
let previous_model = previous_user_turn_model?;
if previous_model == next.model_info.slug {
return None;
}
@@ -116,7 +115,7 @@ pub(crate) fn build_model_instructions_update_item(
pub(crate) fn build_settings_update_items(
previous: Option<&TurnContextItem>,
resumed_model: Option<&str>,
previous_user_turn_model: Option<&str>,
next: &TurnContext,
shell: &Shell,
exec_policy: &Policy,
@@ -124,6 +123,13 @@ pub(crate) fn build_settings_update_items(
) -> Vec<ResponseItem> {
let mut update_items = Vec::new();
// Keep model-switch instructions first so model-specific guidance is read before
// any other context diffs on this turn.
if let Some(model_instructions_item) =
build_model_instructions_update_item(previous_user_turn_model, next)
{
update_items.push(model_instructions_item);
}
if let Some(env_item) = build_environment_update_item(previous, next, shell) {
update_items.push(env_item);
}
@@ -133,11 +139,6 @@ pub(crate) fn build_settings_update_items(
if let Some(collaboration_mode_item) = build_collaboration_mode_update_item(previous, next) {
update_items.push(collaboration_mode_item);
}
if let Some(model_instructions_item) =
build_model_instructions_update_item(previous, resumed_model, next)
{
update_items.push(model_instructions_item);
}
if let Some(personality_item) =
build_personality_update_item(previous, next, personality_feature_enabled)
{
+14 -12
View File
@@ -21,12 +21,9 @@ pub(crate) struct SessionState {
pub(crate) server_reasoning_included: bool,
pub(crate) dependency_env: HashMap<String, String>,
pub(crate) mcp_dependency_prompted: HashSet<String>,
/// Whether the session's initial context has been seeded into history.
///
/// TODO(owen): This is a temporary solution to avoid updating a thread's updated_at
/// timestamp when resuming a session. Remove this once SQLite is in place.
pub(crate) initial_context_seeded: bool,
/// Previous model seen by the session, used for model-switch handling on task start.
/// Model used by the latest regular user turn, used for model-switch handling
/// on subsequent regular turns (including full-context reinjection after
/// resume or `/compact`).
previous_model: Option<String>,
/// Startup regular task pre-created during session initialization.
pub(crate) startup_regular_task: Option<RegularTask>,
@@ -45,7 +42,6 @@ impl SessionState {
server_reasoning_included: false,
dependency_env: HashMap::new(),
mcp_dependency_prompted: HashSet::new(),
initial_context_seeded: false,
previous_model: None,
startup_regular_task: None,
active_mcp_tool_selection: None,
@@ -73,20 +69,26 @@ impl SessionState {
self.history.clone()
}
pub(crate) fn replace_history(&mut self, items: Vec<ResponseItem>) {
pub(crate) fn replace_history(
&mut self,
items: Vec<ResponseItem>,
reference_context_item: Option<TurnContextItem>,
) {
self.history.replace(items);
self.history
.set_reference_context_item(reference_context_item);
}
pub(crate) fn set_token_info(&mut self, info: Option<TokenUsageInfo>) {
self.history.set_token_info(info);
}
pub(crate) fn set_previous_context_item(&mut self, item: Option<TurnContextItem>) {
self.history.set_previous_context_item(item);
pub(crate) fn set_reference_context_item(&mut self, item: Option<TurnContextItem>) {
self.history.set_reference_context_item(item);
}
pub(crate) fn previous_context_item(&self) -> Option<TurnContextItem> {
self.history.previous_context_item()
pub(crate) fn reference_context_item(&self) -> Option<TurnContextItem> {
self.history.reference_context_item()
}
// Token/rate limit helpers
+4 -5
View File
@@ -25,22 +25,21 @@ impl SessionTask for CompactTask {
_cancellation_token: CancellationToken,
) -> Option<String> {
let session = session.clone_session();
if crate::compact::should_use_remote_compact_task(&ctx.provider) {
let _ = if crate::compact::should_use_remote_compact_task(&ctx.provider) {
let _ = session.services.otel_manager.counter(
"codex.task.compact",
1,
&[("type", "remote")],
);
let _ = crate::compact_remote::run_remote_compact_task(session, ctx).await;
crate::compact_remote::run_remote_compact_task(session.clone(), ctx).await
} else {
let _ = session.services.otel_manager.counter(
"codex.task.compact",
1,
&[("type", "local")],
);
let _ = crate::compact::run_compact_task(session, ctx, input).await;
}
crate::compact::run_compact_task(session.clone(), ctx, input).await
};
None
}
}
-10
View File
@@ -121,8 +121,6 @@ impl Session {
) {
self.abort_all_tasks(TurnAbortReason::Replaced).await;
self.clear_connector_selection().await;
self.seed_initial_context_if_needed(turn_context.as_ref())
.await;
let task: Arc<dyn SessionTask> = Arc::new(task);
let task_kind = task.kind();
@@ -140,7 +138,6 @@ impl Session {
tokio::spawn(
async move {
let ctx_for_finish = Arc::clone(&ctx);
let model_slug = ctx_for_finish.model_info.slug.clone();
let last_agent_message = task_for_run
.run(
Arc::clone(&session_ctx),
@@ -151,9 +148,6 @@ impl Session {
.await;
let sess = session_ctx.clone_session();
sess.flush_rollout().await;
// Update previous model before TurnComplete is emitted so
// immediately following turns observe the correct switch state.
sess.set_previous_model(Some(model_slug)).await;
if !task_cancellation_token.is_cancelled() {
// Emit completion uniformly from spawn site so all tasks share the same lifecycle.
sess.on_task_finished(Arc::clone(&ctx_for_finish), last_agent_message)
@@ -278,10 +272,6 @@ impl Session {
task.handle.abort();
// Set previous model even when interrupted so model-switch handling stays correct.
self.set_previous_model(Some(task.turn_context.model_info.slug.clone()))
.await;
let session_ctx = Arc::new(SessionTaskContext::new(Arc::clone(self)));
session_task
.abort(session_ctx, Arc::clone(&task.turn_context))
+2 -1
View File
@@ -101,7 +101,8 @@ impl SessionTask for UndoTask {
match restore_result {
Ok(Ok(())) => {
items.remove(idx);
sess.replace_history(items).await;
let reference_context_item = sess.reference_context_item().await;
sess.replace_history(items, reference_context_item).await;
let short_id: String = commit_id.chars().take(7).collect();
info!(commit_id = commit_id, "Undo restored ghost snapshot");
completed.success = true;
+4
View File
@@ -101,6 +101,10 @@ pub(crate) async fn execute_user_shell_command(
// Auxiliary mode runs within an existing active turn. That turn already
// emitted TurnStarted, so emitting another TurnStarted here would create
// duplicate turn lifecycle events and confuse clients.
// TODO(ccunningham): After TurnStarted, emit model-visible turn context diffs for
// standalone lifecycle tasks (for example /shell, and review once it emits TurnStarted).
// `/compact` is an intentional exception because compaction requests should not include
// freshly reinjected context before the summary/replacement history is applied.
let event = EventMsg::TurnStarted(TurnStartedEvent {
turn_id: turn_context.sub_id.clone(),
model_context_window: turn_context.model_context_window(),