core: pass step environments to turn input extensions

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

@@ -171,8 +171,13 @@ pub(crate) async fn run_turn(
         .record_context_updates_and_set_reference_context_item(first_step_context.as_ref())
         .await;
 
-    let Some((injection_items, explicitly_enabled_connectors)) =
-        build_skills_and_plugins(&sess, turn_context.as_ref(), &input, &cancellation_token).await
+    let Some((injection_items, explicitly_enabled_connectors)) = build_skills_and_plugins(
+        &sess,
+        first_step_context.as_ref(),
+        &input,
+        &cancellation_token,
+    )
+    .await
     else {
         return Ok(None);
     };
@@ -504,10 +509,11 @@ async fn run_hooks_and_record_inputs(
 #[instrument(level = "trace", skip_all)]
 async fn build_skills_and_plugins(
     sess: &Arc<Session>,
-    turn_context: &TurnContext,
+    step_context: &StepContext,
     input: &[TurnInput],
     cancellation_token: &CancellationToken,
 ) -> Option<(Vec<ResponseItem>, HashSet<String>)> {
+    let turn_context = step_context.turn.as_ref();
     // Guardian input embeds the parent transcript as untrusted evidence. Do not interpret skill or
     // plugin mentions from that generated prompt as requests to inject additional instructions.
     if crate::guardian::is_guardian_reviewer_source(&turn_context.session_source) {
@@ -575,7 +581,7 @@ async fn build_skills_and_plugins(
     let skills_outcome = turn_context.turn_skills.snapshot.outcome();
     let connector_slug_counts = build_connector_slug_counts(&available_connectors);
     let extension_injection_items =
-        build_extension_turn_input_items(sess, turn_context, &user_input, cancellation_token)
+        build_extension_turn_input_items(sess, step_context, &user_input, cancellation_token)
             .await?;
     let skill_name_counts_lower =
         build_skill_name_counts(&skills_outcome.skills, &skills_outcome.disabled_paths).1;
@@ -678,7 +684,7 @@ async fn build_skills_and_plugins(
 )]
 async fn build_extension_turn_input_items(
     sess: &Arc<Session>,
-    turn_context: &TurnContext,
+    step_context: &StepContext,
     user_input: &[UserInput],
     cancellation_token: &CancellationToken,
 ) -> Option<Vec<ResponseItem>> {
@@ -687,7 +693,7 @@ async fn build_extension_turn_input_items(
         return Some(Vec::new());
     }
 
-    let environments = turn_context
+    let environments = step_context
         .environments
         .turn_environments
         .iter()
@@ -703,6 +709,7 @@ async fn build_extension_turn_input_items(
         })
         .collect::<Vec<_>>();
 
+    let turn_context = step_context.turn.as_ref();
     let input = TurnInputContext {
         turn_id: turn_context.sub_id.to_string(),
         user_input: user_input.to_vec(),

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

@@ -1,12 +1,35 @@
 use super::*;
 use codex_extension_api::ExtensionData;
+use codex_extension_api::TurnInputContext;
+use codex_extension_api::TurnInputContributor;
 use codex_extension_api::TurnItemContributor;
 use codex_protocol::items::AgentMessageContent;
 use pretty_assertions::assert_eq;
 use std::sync::Arc;
+use std::sync::Mutex;
+use tokio_util::sync::CancellationToken;
 
 struct RewriteAgentMessageContributor;
 
+struct CaptureTurnInputContributor {
+    input: Arc<Mutex<Option<TurnInputContext>>>,
+}
+
+impl TurnInputContributor for CaptureTurnInputContributor {
+    fn contribute<'a>(
+        &'a self,
+        input: TurnInputContext,
+        _session_store: &'a ExtensionData,
+        _thread_store: &'a ExtensionData,
+        _turn_store: &'a ExtensionData,
+    ) -> codex_extension_api::ExtensionFuture<'a, Vec<Box<dyn ContextualUserFragment + Send>>> {
+        Box::pin(async move {
+            *self.input.lock().expect("capture input lock") = Some(input);
+            Vec::new()
+        })
+    }
+}
+
 impl TurnItemContributor for RewriteAgentMessageContributor {
     fn contribute<'a>(
         &'a self,
@@ -65,3 +88,57 @@ async fn plan_mode_uses_contributed_turn_item_for_last_agent_message() {
         Some("plan contributed assistant text")
     );
 }
+
+#[tokio::test]
+async fn turn_input_contributors_receive_step_environments() {
+    let (mut session, mut turn_context) = crate::session::tests::make_session_and_context().await;
+    let step_environments = turn_context.environments.clone();
+    turn_context.environments = Default::default();
+    let turn_context = Arc::new(turn_context);
+    let step_context = StepContext::new(
+        Arc::clone(&turn_context),
+        step_environments.clone(),
+        /*loaded_agents_md*/ None,
+    );
+    let captured_input = Arc::new(Mutex::new(None));
+    let mut builder = codex_extension_api::ExtensionRegistryBuilder::new();
+    builder.turn_input_contributor(Arc::new(CaptureTurnInputContributor {
+        input: Arc::clone(&captured_input),
+    }));
+    session.services.extensions = Arc::new(builder.build());
+
+    build_extension_turn_input_items(
+        &Arc::new(session),
+        &step_context,
+        &[],
+        &CancellationToken::new(),
+    )
+    .await
+    .expect("contributor should complete");
+
+    let input = captured_input
+        .lock()
+        .expect("capture input lock")
+        .take()
+        .expect("contributor should receive input");
+    let [environment] = input.environments.as_slice() else {
+        panic!("expected one contributed environment");
+    };
+    let expected_environment = step_environments.primary().expect("step environment");
+    let expected_cwd = expected_environment
+        .cwd()
+        .to_abs_path()
+        .expect("test cwd should be host native");
+    assert_eq!(
+        (
+            environment.environment_id.as_str(),
+            environment.cwd.as_path(),
+            environment.is_primary,
+        ),
+        (
+            expected_environment.environment_id.as_str(),
+            expected_cwd.as_path(),
+            true,
+        )
+    );
+}