From 0d226929622ce177b56e35d09cf39dd001721466 Mon Sep 17 00:00:00 2001 From: Krish Chainani Date: Fri, 28 Aug 2026 21:47:14 +0000 Subject: [PATCH] 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/compact_remote.rs | 2 + codex-rs/core/src/compact_remote_v2.rs | 2 + codex-rs/core/src/session/mod.rs | 51 ++- .../core/src/session/step_activation_tests.rs | 35 +-- codex-rs/core/src/session/tests.rs | 296 ++++++++++++++++++ codex-rs/core/src/session/turn.rs | 2 +- codex-rs/core/src/state/turn.rs | 4 + 7 files changed, 356 insertions(+), 36 deletions(-) diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 932993d14c..ea1d5a372d 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -230,6 +230,8 @@ async fn run_remote_compact_task_inner_impl( if !should_retry_with_current_model(&error) { return Err(error); } + sess.set_last_known_step_context(fallback_step_context) + .await; let fallback_turn_context = &fallback_step_context.turn; let fallback_compaction_trace = sess.services.rollout_thread_trace.compaction_trace_context( diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index adc4b36c0d..e36d11387f 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -260,6 +260,8 @@ async fn run_remote_compact_task_inner_impl( if !should_retry_with_current_model(&error) { return Err(error); } + sess.set_last_known_step_context(fallback_step_context) + .await; let fallback_turn_context = &fallback_step_context.turn; let fallback_compaction_trace = sess.services.rollout_thread_trace.compaction_trace_context( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 790afd253f..996a18ab5f 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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) { + 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, turn_context: Arc, @@ -3329,6 +3345,35 @@ impl Session { turn_context: Arc, cancellation_token: &CancellationToken, required_servers: &[String], + ) -> CodexResult> { + 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, + turn_context: Arc, + cancellation_token: &CancellationToken, + ) -> CodexResult> { + 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, + turn_context: Arc, + cancellation_token: &CancellationToken, + required_servers: &[String], ) -> CodexResult> { // Capture once before asynchronous planning; all request consumers // retain this immutable settings version even if the turn is updated. diff --git a/codex-rs/core/src/session/step_activation_tests.rs b/codex-rs/core/src/session/step_activation_tests.rs index f0f8ed0d77..a584dd4987 100644 --- a/codex-rs/core/src/session/step_activation_tests.rs +++ b/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, -} - -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, - _session: Arc, - _turn: Arc, - _input: Vec, - cancellation_token: CancellationToken, - ) -> SessionTaskResult { - tokio::select! { - _ = cancellation_token.cancelled() => {}, - _ = self.finish.notified() => {}, - } - Ok(None) - } -} - struct ActivationFixture { session: Arc, turn: Arc, @@ -237,6 +206,7 @@ async fn activation_fixture(models: Vec) -> 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()), }, ) diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 3ec7f616ce..95566b13a3 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/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, +} + +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, + _session: Arc, + _turn: Arc, + _input: Vec, + 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, + Arc, + async_channel::Receiver, +) { + 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::>(), + 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<()>); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 8636a4bd8f..69b67a8eba 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/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) } diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index 056de27955..fdd1612015 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -21,6 +21,7 @@ use tokio::sync::oneshot; use crate::agent::control::AgentExecutionGuard; use crate::mcp_tool_call::McpToolApprovalMetadata; use crate::session::TurnInputQueue; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tasks::AnySessionTask; use codex_protocol::models::AdditionalPermissionProfile; @@ -100,6 +101,9 @@ pub(crate) struct TurnState { pub(crate) tool_calls: u64, pub(crate) has_memory_citation: bool, pub(crate) token_usage_at_turn_start: TokenUsage, + /// The last step captured for execution or selected from a speculative fallback. + /// Remains absent until a step is captured; standalone local compaction has no step. + pub(crate) last_known_step_context: Option>, } pub(crate) struct PendingRequestPermissions {