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
Unverified
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,