From 72be4508546f03b02d4da0a6c44775720e355d98 Mon Sep 17 00:00:00 2001 From: Max Johnson <162359438+maxj-oai@users.noreply.github.com> Date: Fri, 13 Nov 2026 04:47:53 +0000 Subject: [PATCH] Refactor session window access and async compaction test setup (#61263) ## What changed - Read response window metadata through `current_window_from_state` with an explicit session state lock, releasing it before assembling response metadata. - Name the token budget feature check used when restoring a forked child window. - Pass an explicit `TestEnv` into the async compaction test helper and build with `build_with_environment`. GitOrigin-RevId: 7ad852dbd65689e814f62cad645330ac5dac80f4 --- codex-rs/core/src/session/session.rs | 7 +++++-- codex-rs/core/tests/suite/scenarios.rs | 26 ++++++++++++++++++++------ 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index db260bfd21..98fa46b3ae 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -694,7 +694,9 @@ impl Session { step_context: &StepContext, request_kind: CodexResponsesRequestKind, ) -> CodexResponsesMetadata { - let (window_id, window_number, context_window_id) = self.current_window().await; + let state = self.state.lock().await; + let (window_id, window_number, context_window_id) = self.current_window_from_state(&state); + drop(state); let mut responses_metadata = step_context.turn.turn_metadata_state.to_responses_metadata( self.installation_id.clone(), window_id, @@ -911,9 +913,10 @@ impl Session { // writer has validated or refreshed the history snapshot. let new_thread_ids = (!matches!(initial_history, InitialHistory::Resumed(_))) .then(|| prepare_ids(&initial_history)); + let reset_fork_window = config.features.enabled(Feature::TokenBudget); let restore_child_window = matches!(&initial_history, InitialHistory::Forked(_)) && session_configuration.session_source.is_non_root_agent() - && config.features.enabled(Feature::TokenBudget); + && reset_fork_window; if let (InitialHistory::Forked(records), Some(ids)) = (&mut initial_history, &new_thread_ids) { diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index 2fb2dc43de..fada4eb881 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -103,9 +103,11 @@ use core_test_support::streaming_sse::StreamingSseChunk; use core_test_support::streaming_sse::start_streaming_sse_server; use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::TestCodexBuilder; +use core_test_support::test_codex::TestEnv; use core_test_support::test_codex::environment_config_for_selection; use core_test_support::test_codex::executor_path_uri; use core_test_support::test_codex::test_codex; +use core_test_support::test_codex::test_env; use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; use core_test_support::wait_for_mcp_server; @@ -3667,6 +3669,7 @@ async fn async_compaction_test( builder: TestCodexBuilder, responses_uri: &str, features: Vec, + environment: TestEnv, configure: impl Fn(&mut Config) + Send + Sync + 'static, ) -> Result<(wiremock::MockServer, TestCodex)> { let config_server = start_mock_server().await; @@ -3694,7 +3697,7 @@ async fn async_compaction_test( configure(config); }) .with_model("gpt-6-astra") - .build_with_auto_env(&config_server) + .build_with_environment(&config_server, environment) .await?; Ok((config_server, test)) } @@ -3910,14 +3913,19 @@ async fn astra_async_compaction_keeps_work_added_after_its_snapshot( extensions.prompt_contributor(reminders.clone()); builder = builder.with_extensions(Arc::new(extensions.build())); } - let (_config_server, test) = - async_compaction_test(builder, streaming.uri(), features, move |config| { + let (_config_server, test) = async_compaction_test( + builder, + streaming.uri(), + features, + test_env().await?, + move |config| { if over_budget { config.async_compaction_threshold_percent = 1; config.model_auto_compact_token_limit = Some(5_000); } - }) - .await?; + }, + ) + .await?; test.codex.ensure_rollout_materialized().await; let compaction_event = |event: &EventMsg| match event { EventMsg::ItemStarted(event) if matches!(event.item, TurnItem::ContextCompaction(_)) => { @@ -4177,6 +4185,7 @@ async fn astra_async_compaction_preserves_exec_catalog(change: ExecCatalogRefres Feature::IncrementalExecTools, Feature::ViewImage, ], + test_env().await?, move |config| { config .features @@ -4426,7 +4435,10 @@ async fn async_compaction_reuses_pending_work( const BACKGROUND_TOKENS: i64 = 3_001; skip_if_no_network!(Ok(())); - let over_budget = outcome == PendingCompactionOutcome::SuccessOverBudget; + let over_budget = match outcome { + PendingCompactionOutcome::SuccessOverBudget => true, + _ => false, + }; let retained_output = format!("LATE_OUTPUT {}", "x ".repeat(12_000)); let manual = matches!(trigger, CompactionTrigger::Manual); let has_background = outcome != PendingCompactionOutcome::Absent; @@ -4499,6 +4511,7 @@ async fn async_compaction_reuses_pending_work( extensions.thread_lifecycle_contributor(accounting.clone()); builder = builder.with_extensions(Arc::new(extensions.build())); } + let env = test_env().await?; let (_config_server, test) = async_compaction_test( builder, streaming.uri(), @@ -4507,6 +4520,7 @@ async fn async_compaction_reuses_pending_work( } else { vec![] }, + env, move |config| { if over_budget { config.async_compaction_threshold_percent = 1;