Retain the last selected step context for each turn (#41429)

## What changed

- Store the most recently captured execution `StepContext` in the active turn state.
- Keep speculative model-fallback captures from replacing that context, and update it only when the fallback is selected for remote compaction.
- Preserve the retained context when a turn completes or is interrupted.

## Testing

- Add coverage for completed and interrupted turns, successful primary compaction, fallback compaction, and interruption while a fallback is running.

GitOrigin-RevId: c7dee36ec8e88ba2f4fb4a9c0322e5c86a0caf0e

codex-rs/core/src/session/mod.rs

@@ -3305,12 +3305,28 @@ impl Session {
         Ok(world_state)
     }
 
-    /// Captures one request-scoped view of dynamic state.
+    /// Retains the step captured for execution.
+    pub(crate) async fn set_last_known_step_context(&self, step_context: &Arc<StepContext>) {
+        let turn_state = {
+            let active_turn = self.active_turn.lock().await;
+            active_turn.as_ref().and_then(|active_turn| {
+                active_turn
+                    .task
+                    .as_ref()
+                    .filter(|task| task.turn_context.sub_id == step_context.turn.sub_id)
+                    .map(|_| Arc::clone(&active_turn.turn_state))
+            })
+        };
+        if let Some(turn_state) = turn_state {
+            turn_state.lock().await.last_known_step_context = Some(Arc::clone(step_context));
+        }
+    }
+
+    /// Captures one request-scoped view of dynamic state and retains it for the active turn.
     ///
     /// This may refresh filesystem-derived state. Normal turns should call it only from
     /// `run_turn` and pass the result down; standalone request or history boundaries may capture
-    /// their own step.
-    #[tracing::instrument(name = "step_context.capture", level = "info", skip_all)]
+    /// their own step. Use speculative capture for a step that may not execute.
     pub(crate) async fn capture_step_context(
         self: &Arc<Self>,
         turn_context: Arc<TurnContext>,
@@ -3329,6 +3345,35 @@ impl Session {
         turn_context: Arc<TurnContext>,
         cancellation_token: &CancellationToken,
         required_servers: &[String],
+    ) -> CodexResult<Arc<StepContext>> {
+        let step_context = self
+            .capture_step_context_inner(turn_context, cancellation_token, required_servers)
+            .await?;
+        self.set_last_known_step_context(&step_context).await;
+        Ok(step_context)
+    }
+
+    /// Prepares a candidate step without replacing the active turn's retained context.
+    /// The caller must retain it explicitly if it is selected for execution.
+    async fn capture_speculative_step_context(
+        self: &Arc<Self>,
+        turn_context: Arc<TurnContext>,
+        cancellation_token: &CancellationToken,
+    ) -> CodexResult<Arc<StepContext>> {
+        self.capture_step_context_inner(
+            turn_context,
+            cancellation_token,
+            /*required_servers*/ &[],
+        )
+        .await
+    }
+
+    #[tracing::instrument(name = "step_context.capture", level = "info", skip_all)]
+    async fn capture_step_context_inner(
+        self: &Arc<Self>,
+        turn_context: Arc<TurnContext>,
+        cancellation_token: &CancellationToken,
+        required_servers: &[String],
     ) -> CodexResult<Arc<StepContext>> {
         // Capture once before asynchronous planning; all request consumers
         // retain this immutable settings version even if the turn is updated.

codex-rs/core/src/session/step_activation_tests.rs

@@ -1,15 +1,13 @@
 use super::*;
 use crate::guardian::BUNDLED_GUARDIAN_POLICY;
-use crate::session::TurnInput;
 use crate::session::handlers::submission_loop;
 use crate::session::step_context::StepContext;
 use crate::session::step_settings::StepSettings;
+use crate::session::tests::HeldStepTask;
 use crate::session::tests::make_session_and_context;
 use crate::session::tests::update_selected_settings_for_test;
 use crate::session::tests::update_turn_settings_for_test;
 use crate::state::TaskKind;
-use crate::tasks::SessionTask;
-use crate::tasks::SessionTaskResult;
 use codex_config::AutoReviewRequirementsToml;
 use codex_config::ConfigLayerStack;
 use codex_config::ConfigRequirements;
@@ -163,35 +161,6 @@ impl ModelsManager for GatedModelsManager {
     }
 }
 
-struct HeldStepTask {
-    finish: Arc<Notify>,
-}
-
-impl SessionTask for HeldStepTask {
-    fn kind(&self) -> TaskKind {
-        // Publication also works for a non-regular task that never samples.
-        TaskKind::Compact
-    }
-
-    fn span_name(&self) -> &'static str {
-        "session_task.step_activation_test"
-    }
-
-    async fn run(
-        self: Arc<Self>,
-        _session: Arc<Session>,
-        _turn: Arc<TurnContext>,
-        _input: Vec<TurnInput>,
-        cancellation_token: CancellationToken,
-    ) -> SessionTaskResult {
-        tokio::select! {
-            _ = cancellation_token.cancelled() => {},
-            _ = self.finish.notified() => {},
-        }
-        Ok(None)
-    }
-}
-
 struct ActivationFixture {
     session: Arc<Session>,
     turn: Arc<TurnContext>,
@@ -237,6 +206,7 @@ async fn activation_fixture(models: Vec<ModelInfo>) -> ActivationFixture {
             prepared,
             Vec::new(),
             HeldStepTask {
+                kind: TaskKind::Compact,
                 finish: Arc::clone(&finish),
             },
         )
@@ -607,6 +577,7 @@ async fn delayed_activation_does_not_retarget_a_task(change: TaskChangeDuringLoo
                     Arc::clone(&replacement),
                     Vec::new(),
                     HeldStepTask {
+                        kind: TaskKind::Compact,
                         finish: Arc::new(Notify::new()),
                     },
                 )

codex-rs/core/src/session/tests.rs

@@ -6,6 +6,7 @@ pub(crate) use super::step_settings::tests::update_selected_settings_for_test;
 use super::turn_context::TurnEnvironment;
 use super::*;
 use crate::agents_md_manager::AgentsMdManager;
+use crate::compact::InitialContextInjection;
 use crate::config::ConfigBuilder;
 use crate::config::ConfigOverrides;
 use crate::config::test_config;
@@ -24,6 +25,8 @@ use crate::shell_snapshot::ShellSnapshot;
 use crate::test_support::models_manager_with_provider;
 use crate::tools::format_exec_output_str;
 use crate::tools::registry::ToolRegistry;
+use codex_analytics::CompactionPhase;
+use codex_analytics::CompactionReason;
 use codex_config::ConfigLayerStack;
 use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID;
 use codex_config::LoaderOverrides;
@@ -49,6 +52,7 @@ use codex_login::CodexAuth;
 use codex_login::auth::AgentIdentityAuthPolicy;
 use codex_model_provider::create_model_provider;
 use codex_model_provider_info::ModelProviderInfo;
+use codex_model_provider_info::built_in_model_providers;
 use codex_models_manager::bundled_models_response;
 use codex_models_manager::model_info;
 use codex_models_manager::test_support::construct_model_info_offline_for_tests;
@@ -171,6 +175,7 @@ use core_test_support::PathExt;
 use core_test_support::context_snapshot;
 use core_test_support::context_snapshot::ContextSnapshotOptions;
 use core_test_support::context_snapshot::ContextSnapshotRenderMode;
+use core_test_support::responses;
 use core_test_support::responses::ev_assistant_message;
 use core_test_support::responses::ev_completed;
 use core_test_support::responses::ev_response_created;
@@ -180,6 +185,8 @@ use core_test_support::responses::start_mock_server;
 use core_test_support::responses::strip_metadata_from_items;
 use core_test_support::responses::strip_response_item_ids;
 use core_test_support::responses::strip_response_item_ids_from_json;
+use core_test_support::streaming_sse::StreamingSseChunk;
+use core_test_support::streaming_sse::start_streaming_sse_server;
 use core_test_support::test_codex::local;
 use core_test_support::test_codex::test_codex;
 use core_test_support::test_path_buf;
@@ -194,10 +201,13 @@ use opentelemetry_sdk::metrics::data::MetricData;
 use opentelemetry_sdk::metrics::data::ResourceMetrics;
 use std::path::Path;
 use std::time::Duration;
+use test_case::test_case;
+use tokio::sync::Notify;
 use tokio::sync::Semaphore;
 use tokio::time::sleep;
 use tokio::time::timeout;
 use tracing_opentelemetry::OpenTelemetrySpanExt;
+use wiremock::ResponseTemplate;
 
 use uuid::Uuid;
 
@@ -10847,6 +10857,292 @@ impl SessionTask for GuardianDeniedApprovalTask {
     }
 }
 
+pub(super) struct HeldStepTask {
+    pub(super) kind: TaskKind,
+    pub(super) finish: Arc<Notify>,
+}
+
+impl SessionTask for HeldStepTask {
+    fn kind(&self) -> TaskKind {
+        self.kind
+    }
+
+    fn span_name(&self) -> &'static str {
+        "session_task.step_activation_test"
+    }
+
+    async fn run(
+        self: Arc<Self>,
+        _session: Arc<Session>,
+        _turn: Arc<TurnContext>,
+        _input: Vec<TurnInput>,
+        cancellation_token: CancellationToken,
+    ) -> SessionTaskResult {
+        tokio::select! {
+            _ = cancellation_token.cancelled() => {},
+            _ = self.finish.notified() => {},
+        }
+        Ok(None)
+    }
+}
+
+#[test_case(TerminalEventKind::TurnComplete; "completion")]
+#[test_case(TerminalEventKind::TurnAborted; "interruption")]
+#[tokio::test]
+async fn finished_turn_retains_last_known_step_context(terminal: TerminalEventKind) {
+    let (session, turn, events) = make_session_and_context_with_rx().await;
+    let finish = Arc::new(Notify::new());
+    session
+        .spawn_task(
+            Arc::clone(&turn),
+            Vec::new(),
+            HeldStepTask {
+                kind: TaskKind::Regular,
+                finish: Arc::clone(&finish),
+            },
+        )
+        .await;
+    let expected = session
+        .capture_step_context(turn, &CancellationToken::new())
+        .await
+        .expect("capture executing step");
+    let state = {
+        let active = session.active_turn.lock().await;
+        Arc::clone(&active.as_ref().expect("active turn").turn_state)
+    };
+
+    match terminal {
+        TerminalEventKind::TurnComplete => finish.notify_one(),
+        TerminalEventKind::TurnAborted => {
+            session.abort_all_tasks(TurnAbortReason::Interrupted).await;
+        }
+    }
+    recv_terminal_event(&events, terminal).await;
+
+    assert!(session.active_turn.lock().await.is_none());
+    assert_eq!(
+        state
+            .lock()
+            .await
+            .last_known_step_context
+            .as_ref()
+            .map(Arc::as_ptr),
+        Some(Arc::as_ptr(&expected)),
+    );
+}
+
+#[derive(Clone, Copy)]
+enum FirstAttempt {
+    Succeeds,
+    Retries,
+}
+
+async fn make_remote_compaction_session(
+    server_uri: &str,
+) -> (
+    Arc<Session>,
+    Arc<TurnContext>,
+    async_channel::Receiver<Event>,
+) {
+    let mut provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone();
+    provider.base_url = Some(format!("{server_uri}/v1"));
+    provider.supports_websockets = false;
+    make_session_and_context_with_auth_and_config_and_rx(
+        CodexAuth::create_dummy_chatgpt_auth_for_testing(),
+        Vec::new(),
+        move |config| {
+            config.model = Some("gpt-5.2".to_string());
+            config.model_provider = provider;
+            let _ = config.features.enable(Feature::RemoteCompactionV2);
+            let _ = config.features.disable(Feature::TokenBudget);
+        },
+    )
+    .await
+}
+
+#[test_case(FirstAttempt::Succeeds; "primary succeeds")]
+#[test_case(FirstAttempt::Retries; "fallback executes")]
+#[tokio::test]
+async fn legacy_compaction_retains_only_the_selected_step(first_attempt: FirstAttempt) {
+    let server = responses::start_mock_server().await;
+    let (session, turn, events) = make_remote_compaction_session(&server.uri()).await;
+    session
+        .record_conversation_items(&turn, &[user_message("before compaction")])
+        .await;
+    session
+        .spawn_task(
+            Arc::clone(&turn),
+            Vec::new(),
+            NeverEndingTask {
+                kind: TaskKind::Regular,
+                listen_to_cancellation_token: true,
+            },
+        )
+        .await;
+    let primary_turn = Arc::new(
+        turn.with_model("gpt-5.4".to_string(), &session.services.models_manager)
+            .await,
+    );
+    let primary = session
+        .capture_step_context(primary_turn, &CancellationToken::new())
+        .await
+        .expect("capture primary step");
+    let fallback = session
+        .capture_speculative_step_context(turn, &CancellationToken::new())
+        .await
+        .expect("capture speculative fallback");
+    let state = {
+        let active = session.active_turn.lock().await;
+        Arc::clone(&active.as_ref().expect("active turn").turn_state)
+    };
+    assert_eq!(
+        state
+            .lock()
+            .await
+            .last_known_step_context
+            .as_ref()
+            .map(Arc::as_ptr),
+        Some(Arc::as_ptr(&primary)),
+    );
+
+    let success = ResponseTemplate::new(/*status*/ 200).set_body_json(json!({
+        "output": [{ "type": "compaction", "encrypted_content": "summary" }]
+    }));
+    let replies = match first_attempt {
+        FirstAttempt::Succeeds => vec![success],
+        FirstAttempt::Retries => vec![
+            ResponseTemplate::new(/*status*/ 400)
+                .set_body_json(json!({ "detail": "previous model unavailable" })),
+            success,
+        ],
+    };
+    let requests = responses::mount_compact_response_sequence(&server, replies).await;
+    crate::compact_remote::run_inline_remote_auto_compact_task(
+        Arc::clone(&session),
+        Arc::clone(&primary),
+        Some(Arc::clone(&fallback)),
+        Arc::new(OnceLock::new()),
+        InitialContextInjection::DoNotInject,
+        CompactionReason::ModelDownshift,
+        CompactionPhase::PreTurn,
+    )
+    .await
+    .expect("compaction succeeds");
+
+    let (expected, models) = match first_attempt {
+        FirstAttempt::Succeeds => (&primary, vec![json!("gpt-5.4")]),
+        FirstAttempt::Retries => (&fallback, vec![json!("gpt-5.4"), json!("gpt-5.2")]),
+    };
+    assert_eq!(
+        state
+            .lock()
+            .await
+            .last_known_step_context
+            .as_ref()
+            .map(Arc::as_ptr),
+        Some(Arc::as_ptr(expected)),
+    );
+    assert_eq!(
+        requests
+            .requests()
+            .iter()
+            .map(|request| request.body_json()["model"].clone())
+            .collect::<Vec<_>>(),
+        models,
+    );
+    session.abort_all_tasks(TurnAbortReason::Interrupted).await;
+    recv_terminal_event(&events, TerminalEventKind::TurnAborted).await;
+}
+
+#[tokio::test]
+async fn interrupting_compaction_fallback_retains_last_known_step_context() {
+    let (release_primary, primary_gate) = tokio::sync::oneshot::channel();
+    let (release_fallback, fallback_gate) = tokio::sync::oneshot::channel();
+    let (server, _) = start_streaming_sse_server(vec![
+        vec![StreamingSseChunk {
+            gate: Some(primary_gate),
+            body: responses::sse_failed(
+                "primary",
+                "context_length_exceeded",
+                "compact with the current model",
+            ),
+        }],
+        vec![StreamingSseChunk {
+            gate: Some(fallback_gate),
+            body: responses::sse_completed("fallback"),
+        }],
+    ])
+    .await;
+    let (session, mut turn, events) = make_remote_compaction_session(server.uri()).await;
+    update_turn_settings_for_test(
+        Arc::get_mut(&mut turn).expect("unshared turn"),
+        |settings| {
+            Arc::make_mut(&mut settings.model_info).comp_hash = Some("new".to_string());
+        },
+    );
+    session
+        .set_previous_turn_settings(Some(PreviousTurnSettings {
+            model: "gpt-5.4".to_string(),
+            comp_hash: Some("old".to_string()),
+            realtime_active: Some(turn.realtime_active),
+        }))
+        .await;
+    session
+        .record_conversation_items(&turn, &[user_message("before compaction")])
+        .await;
+    session
+        .spawn_task(turn, Vec::new(), crate::tasks::RegularTask::new())
+        .await;
+    let state = {
+        let active = session.active_turn.lock().await;
+        Arc::clone(&active.as_ref().expect("active turn").turn_state)
+    };
+
+    // The real turn loop has prepared both contexts before sending its first compact request.
+    timeout(
+        Duration::from_secs(/*secs*/ 10),
+        server.wait_for_request_count(/*count*/ 1),
+    )
+    .await
+    .expect("primary compaction request");
+    let primary = state
+        .lock()
+        .await
+        .last_known_step_context
+        .clone()
+        .expect("primary step");
+    assert_eq!(primary.settings.model_info.slug, "gpt-5.4");
+
+    release_primary.send(()).expect("release primary failure");
+    timeout(
+        Duration::from_secs(/*secs*/ 10),
+        server.wait_for_request_count(/*count*/ 2),
+    )
+    .await
+    .expect("fallback compaction request");
+    let fallback = state
+        .lock()
+        .await
+        .last_known_step_context
+        .clone()
+        .expect("fallback step");
+    assert_eq!(fallback.settings.model_info.slug, "gpt-5.2");
+
+    session.abort_all_tasks(TurnAbortReason::Interrupted).await;
+    recv_terminal_event(&events, TerminalEventKind::TurnAborted).await;
+    assert_eq!(
+        state
+            .lock()
+            .await
+            .last_known_step_context
+            .as_ref()
+            .map(Arc::as_ptr),
+        Some(Arc::as_ptr(&fallback)),
+    );
+    drop(release_fallback);
+    server.shutdown().await;
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn guardian_auto_review_emits_thread_idle_after_interrupt() {
     struct ThreadIdleRecorder(async_channel::Sender<()>);

codex-rs/core/src/session/turn.rs

@@ -1088,7 +1088,7 @@ async fn capture_current_model_fallback_step_context(
     {
         return Ok(None);
     }
-    sess.capture_step_context(Arc::clone(turn_context), cancellation_token)
+    sess.capture_speculative_step_context(Arc::clone(turn_context), cancellation_token)
         .await
         .map(Some)
 }