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

@@ -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)
         {

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<Feature>,
+    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;